tycho_simulation/evm/protocol/vm/
erc20_token.rs

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

tycho_simulation/evm/protocol/vm/
erc20_token.rs
