tycho_simulation/evm/protocol/vm/
erc20_token.rs

1use std::{collections::HashMap, str::FromStr};
2
3use alloy::primitives::{Address, U256};
4use lazy_static::lazy_static;
5
6use super::utils::get_storage_slot_index_at_key;
7use crate::evm::{ContractCompiler, SlotId};
8
9pub(crate) type Overwrites = HashMap<SlotId, U256>;
10
11// Storage slots constants for TokenProxy contract
12lazy_static! {
13    pub static ref IMPLEMENTATION_SLOT: SlotId =
14        U256::from_str("0x6677C72CDEB41ACAF2B17EC8A6E275C4205F27DBFE4DE34EBAF2E928A7E610DB")
15            .unwrap();
16    static ref BALANCES_MAPPING_POSITION: SlotId =
17        U256::from_str("0x474F5FD57EE674F7B6851BC6F07E751B49076DFB356356985B9DAF10E9ABC941")
18            .unwrap();
19    static ref HAS_CUSTOM_BALANCE_POSITION: SlotId =
20        U256::from_str("0x7EAD8EDE9DBB385B0664952C7462C9938A5821E6F78E859DA2E683216E99411B")
21            .unwrap();
22    static ref CUSTOM_APPROVAL_MAPPING_POSITION: SlotId =
23        U256::from_str("0x71A54E125991077003BEF7E7CA57369C919DAC6D2458895F1EAB4D03960F4AEB")
24            .unwrap();
25    static ref HAS_CUSTOM_APPROVAL_MAPPING_POSITION: SlotId =
26        U256::from_str("0x9F0C1BC0E9C3078F9AD5FC59C8606416B3FABCBD4C8353FED22937C66C866CE3")
27            .unwrap();
28    static ref CUSTOM_NAME_POSITION: SlotId =
29        U256::from_str("0xCC1E513FB5BDA80DC466AD9D44DF38805A8DEE4C82B3C6DF3D9B25D3D5355D1C")
30            .unwrap();
31    static ref CUSTOM_SYMBOL_POSITION: SlotId =
32        U256::from_str("0xDC17DD3380A9A034A702A2B2B1C6C25D39EBF0E89796E0D15E1E04D23E3BB221")
33            .unwrap();
34    static ref CUSTOM_DECIMALS_POSITION: SlotId =
35        U256::from_str("0xADD486B234562DE9AC745F036F538CDA2547EF6DBB4DA3FA1C017625F888A8E8")
36            .unwrap();
37    static ref CUSTOM_TOTAL_SUPPLY_POSITION: SlotId =
38        U256::from_str("0x6014AF1E8E9BB2844581B2FA9E5E3620181C3192EEFD3258319AEC23538DA9F5")
39            .unwrap();
40    static ref HAS_CUSTOM_METADATA_POSITION: SlotId =
41        U256::from_str("0x9F37243DE61714BE9CC00628D4B9BF9897AE670218AF52ADE6D192B4339D7616")
42            .unwrap();
43}
44
45pub(crate) struct TokenProxyOverwriteFactory {
46    token_address: Address,
47    overwrites: Overwrites,
48    compiler: ContractCompiler,
49}
50
51impl TokenProxyOverwriteFactory {
52    pub(crate) fn new(token_address: Address, proxy_address: Option<Address>) -> Self {
53        let mut instance = Self {
54            token_address,
55            overwrites: HashMap::new(),
56            compiler: ContractCompiler::Solidity,
57        };
58
59        if let Some(proxy_addr) = proxy_address {
60            instance.set_original_address(proxy_addr);
61        }
62
63        instance
64    }
65
66    /// Sets the original address of the token contract in the implementation slot.
67    pub(crate) fn set_original_address(&mut self, implementation: Address) {
68        self.overwrites
69            .insert(*IMPLEMENTATION_SLOT, U256::from_be_slice(implementation.as_slice()));
70    }
71
72    pub(crate) fn set_balance(&mut self, balance: U256, owner: Address) {
73        // Set the balance in the custom storage slot
74        let storage_index =
75            get_storage_slot_index_at_key(owner, *BALANCES_MAPPING_POSITION, self.compiler);
76        self.overwrites
77            .insert(storage_index, balance);
78
79        // Set the has_custom_balance flag to true
80        let has_balance_index =
81            get_storage_slot_index_at_key(owner, *HAS_CUSTOM_BALANCE_POSITION, self.compiler);
82        self.overwrites
83            .insert(has_balance_index, U256::from(1)); // true in Solidity
84    }
85
86    pub(crate) fn set_allowance(&mut self, allowance: U256, spender: Address, owner: Address) {
87        // Set the allowance in the custom storage slot
88        let owner_slot =
89            get_storage_slot_index_at_key(owner, *CUSTOM_APPROVAL_MAPPING_POSITION, self.compiler);
90        let storage_index = get_storage_slot_index_at_key(spender, owner_slot, self.compiler);
91        self.overwrites
92            .insert(storage_index, allowance);
93
94        // Set the has_custom_approval flag to true
95        let has_approval_index = get_storage_slot_index_at_key(
96            owner,
97            *HAS_CUSTOM_APPROVAL_MAPPING_POSITION,
98            self.compiler,
99        );
100        self.overwrites
101            .insert(has_approval_index, U256::from(1)); // true in Solidity
102    }
103
104    #[allow(dead_code)]
105    pub(crate) fn set_total_supply(&mut self, supply: U256) {
106        self.overwrites
107            .insert(*CUSTOM_TOTAL_SUPPLY_POSITION, supply);
108    }
109
110    /// Sets the has_custom_metadata flag for a given key
111    #[allow(dead_code)]
112    fn set_metadata_flag(&mut self, key: &str) {
113        let key_bytes = string_to_storage_bytes(key);
114        let mapping_slot_bytes: [u8; 32] = HAS_CUSTOM_METADATA_POSITION.to_be_bytes();
115        let has_metadata_index = self
116            .compiler
117            .compute_map_slot(&key_bytes, &mapping_slot_bytes);
118        self.overwrites
119            .insert(has_metadata_index, U256::from(1)); // true in Solidity
120    }
121
122    #[allow(dead_code)]
123    pub(crate) fn set_name(&mut self, name: &str) {
124        // Store the name value
125        let name_value = U256::from_be_bytes(string_to_storage_bytes(name));
126        self.overwrites
127            .insert(*CUSTOM_NAME_POSITION, name_value);
128
129        // Set the has_custom_metadata flag for name to true
130        self.set_metadata_flag("name");
131    }
132
133    #[allow(dead_code)]
134    pub(crate) fn set_symbol(&mut self, symbol: &str) {
135        // Store the symbol value
136        let symbol_value = U256::from_be_bytes(string_to_storage_bytes(symbol));
137        self.overwrites
138            .insert(*CUSTOM_SYMBOL_POSITION, symbol_value);
139
140        // Set the has_custom_metadata flag for symbol to true
141        self.set_metadata_flag("symbol");
142    }
143
144    #[allow(dead_code)]
145    pub(crate) fn set_decimals(&mut self, decimals: u8) {
146        self.overwrites
147            .insert(*CUSTOM_DECIMALS_POSITION, U256::from(decimals));
148
149        // Set the has_custom_metadata flag for decimals to true
150        self.set_metadata_flag("decimals");
151    }
152
153    pub(crate) fn get_overwrites(&self) -> HashMap<Address, Overwrites> {
154        let mut result = HashMap::new();
155        result.insert(self.token_address, self.overwrites.clone());
156        result
157    }
158}
159
160/// Converts a string to a 32-byte array for storage, truncating if necessary
161pub fn string_to_storage_bytes(s: &str) -> [u8; 32] {
162    let mut padded = [0u8; 32];
163    let len = s.len().min(31);
164    padded[..len].copy_from_slice(&s.as_bytes()[..len]);
165    padded[31] = (len * 2) as u8; // Length * 2 for short strings
166    padded
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    fn get_metadata_slot(key: &str) -> SlotId {
174        let key_bytes = string_to_storage_bytes(key);
175        let mapping_slot_bytes: [u8; 32] = HAS_CUSTOM_METADATA_POSITION.to_be_bytes();
176        ContractCompiler::Solidity.compute_map_slot(&key_bytes, &mapping_slot_bytes)
177    }
178
179    #[test]
180    fn test_token_proxy_factory_new() {
181        let token_address = Address::random();
182        let factory = TokenProxyOverwriteFactory::new(token_address, None);
183        assert_eq!(factory.token_address, token_address);
184        assert!(factory.overwrites.is_empty());
185    }
186
187    #[test]
188    fn test_token_proxy_factory_with_implementation() {
189        let token_address = Address::random();
190        let implementation = Address::random();
191        let factory = TokenProxyOverwriteFactory::new(token_address, Some(implementation));
192
193        // Check if implementation was set correctly
194        let mut expected_bytes = [0u8; 32];
195        expected_bytes[12..].copy_from_slice(implementation.as_slice());
196        let expected_value = U256::from_be_bytes(expected_bytes);
197
198        assert_eq!(factory.overwrites[&*IMPLEMENTATION_SLOT], expected_value);
199    }
200
201    #[test]
202    fn test_token_proxy_set_balance() {
203        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
204        let owner = Address::random();
205        let balance = U256::from(1000);
206
207        factory.set_balance(balance, owner);
208
209        // Check balance storage
210        let storage_index =
211            get_storage_slot_index_at_key(owner, *BALANCES_MAPPING_POSITION, factory.compiler);
212        assert_eq!(factory.overwrites[&storage_index], balance);
213
214        // Check has_custom_balance flag
215        let has_balance_index =
216            get_storage_slot_index_at_key(owner, *HAS_CUSTOM_BALANCE_POSITION, factory.compiler);
217        assert_eq!(factory.overwrites[&has_balance_index], U256::from(1));
218    }
219
220    #[test]
221    fn test_token_proxy_set_allowance() {
222        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
223        let owner = Address::random();
224        let spender = Address::random();
225        let allowance = U256::from(500);
226
227        factory.set_allowance(allowance, spender, owner);
228
229        // Check allowance storage
230        let owner_slot = get_storage_slot_index_at_key(
231            owner,
232            *CUSTOM_APPROVAL_MAPPING_POSITION,
233            factory.compiler,
234        );
235        let storage_index = get_storage_slot_index_at_key(spender, owner_slot, factory.compiler);
236        assert_eq!(factory.overwrites[&storage_index], allowance);
237
238        // Check has_custom_approval flag
239        let has_approval_index = get_storage_slot_index_at_key(
240            owner,
241            *HAS_CUSTOM_APPROVAL_MAPPING_POSITION,
242            factory.compiler,
243        );
244        assert_eq!(factory.overwrites[&has_approval_index], U256::from(1));
245    }
246
247    #[test]
248    fn test_token_proxy_set_total_supply() {
249        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
250        let supply = U256::from(1_000_000);
251
252        factory.set_total_supply(supply);
253
254        assert_eq!(factory.overwrites[&*CUSTOM_TOTAL_SUPPLY_POSITION], supply);
255    }
256
257    #[test]
258    fn test_token_proxy_set_name() {
259        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
260        let name = "Test Token";
261
262        factory.set_name(name);
263
264        // Check name storage
265        let mut expected_bytes = [0u8; 32];
266        let name_bytes = name.as_bytes();
267        expected_bytes[..name_bytes.len()].copy_from_slice(name_bytes);
268        expected_bytes[31] = (name_bytes.len() * 2) as u8; // Length * 2 for short strings
269        let expected_value = U256::from_be_bytes(expected_bytes);
270        assert_eq!(factory.overwrites[&*CUSTOM_NAME_POSITION], expected_value);
271
272        // Check has_custom_metadata flag
273        let has_metadata_index = get_metadata_slot("name");
274        assert_eq!(factory.overwrites[&has_metadata_index], U256::from(1));
275    }
276
277    #[test]
278    fn test_token_proxy_set_symbol() {
279        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
280        let symbol = "TEST";
281
282        factory.set_symbol(symbol);
283
284        // Check symbol storage
285        let mut expected_bytes = [0u8; 32];
286        let symbol_bytes = symbol.as_bytes();
287        expected_bytes[..symbol_bytes.len()].copy_from_slice(symbol_bytes);
288        expected_bytes[31] = (symbol_bytes.len() * 2) as u8; // Length * 2 for short strings
289        let expected_value = U256::from_be_bytes(expected_bytes);
290        assert_eq!(factory.overwrites[&*CUSTOM_SYMBOL_POSITION], expected_value);
291
292        // Check has_custom_metadata flag
293        let has_metadata_index = get_metadata_slot("symbol");
294        assert_eq!(factory.overwrites[&has_metadata_index], U256::from(1));
295    }
296
297    #[test]
298    fn test_token_proxy_set_decimals() {
299        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
300        let decimals = 18u8;
301
302        factory.set_decimals(decimals);
303
304        assert_eq!(factory.overwrites[&*CUSTOM_DECIMALS_POSITION], U256::from(decimals));
305
306        // Check has_custom_metadata flag
307        let has_metadata_index = get_metadata_slot("decimals");
308        assert_eq!(factory.overwrites[&has_metadata_index], U256::from(1));
309    }
310
311    #[test]
312    fn test_token_proxy_get_overwrites() {
313        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
314        let supply = U256::from(1_000_000);
315        factory.set_total_supply(supply);
316
317        let overwrites = factory.get_overwrites();
318
319        assert_eq!(overwrites.len(), 1);
320        assert!(overwrites.contains_key(&factory.token_address));
321        assert_eq!(overwrites[&factory.token_address][&*CUSTOM_TOTAL_SUPPLY_POSITION], supply);
322    }
323
324    #[test]
325    fn test_token_proxy_set_long_name_truncated() {
326        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
327        let name = "This is a very long token name that exceeds 31 bytes";
328
329        factory.set_name(name);
330
331        // Check name storage for truncated string
332        let mut expected_bytes = [0u8; 32];
333        expected_bytes[..31].copy_from_slice(&name.as_bytes()[..31]);
334        expected_bytes[31] = 62; // 31 * 2 for short strings
335        let expected_value = U256::from_be_bytes(expected_bytes);
336        assert_eq!(factory.overwrites[&*CUSTOM_NAME_POSITION], expected_value);
337
338        // Check has_custom_metadata flag
339        let has_metadata_index = get_metadata_slot("name");
340        assert_eq!(factory.overwrites[&has_metadata_index], U256::from(1));
341    }
342
343    #[test]
344    fn test_token_proxy_set_long_symbol_truncated() {
345        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
346        let symbol = "This is a very long token symbol that exceeds 31 bytes";
347
348        factory.set_symbol(symbol);
349
350        // Check symbol storage for truncated string
351        let mut expected_bytes = [0u8; 32];
352        expected_bytes[..31].copy_from_slice(&symbol.as_bytes()[..31]);
353        expected_bytes[31] = 62; // 31 * 2 for short strings
354        let expected_value = U256::from_be_bytes(expected_bytes);
355        assert_eq!(factory.overwrites[&*CUSTOM_SYMBOL_POSITION], expected_value);
356
357        // Check has_custom_metadata flag
358        let has_metadata_index = get_metadata_slot("symbol");
359        assert_eq!(factory.overwrites[&has_metadata_index], U256::from(1));
360    }
361
362    #[test]
363    fn test_string_to_storage_bytes() {
364        // Test short string
365        let short = "Test";
366        let bytes = string_to_storage_bytes(short);
367        assert_eq!(bytes[..4], short.as_bytes()[..4]);
368        assert_eq!(bytes[31], 8); // 4 * 2 for length
369
370        // Test long string (should be truncated)
371        let long = "This is a very long string that exceeds 31 bytes";
372        let bytes = string_to_storage_bytes(long);
373        assert_eq!(bytes[..31], long.as_bytes()[..31]);
374        assert_eq!(bytes[31], 62); // 31 * 2 for length
375    }
376
377    #[test]
378    fn test_set_metadata_flag() {
379        let mut factory = TokenProxyOverwriteFactory::new(Address::random(), None);
380
381        // Test setting metadata flag for a key
382        factory.set_metadata_flag("test_key");
383
384        // Verify the flag was set correctly
385        let key_bytes = string_to_storage_bytes("test_key");
386        let mapping_slot_bytes: [u8; 32] = HAS_CUSTOM_METADATA_POSITION.to_be_bytes();
387        let has_metadata_index =
388            ContractCompiler::Solidity.compute_map_slot(&key_bytes, &mapping_slot_bytes);
389        assert_eq!(factory.overwrites[&has_metadata_index], U256::from(1));
390    }
391}
tycho_simulation/evm/protocol/vm/erc20_token.rs

tycho_simulation/evm/protocol/vm/
erc20_token.rs
