From 7a80cd3a06937522b6b0bcc34680165a22e97455 Mon Sep 17 00:00:00 2001 From: TAMARA LIPOWSKI Date: Sat, 19 Oct 2024 13:19:00 -0400 Subject: [PATCH] feat: create ERC20OverwriteFactory utils --- src/protocol/vm/utils.rs | 167 ++++++++++++++++++++++++++++++++------- 1 file changed, 137 insertions(+), 30 deletions(-) diff --git a/src/protocol/vm/utils.rs b/src/protocol/vm/utils.rs index f5048868..cfaa0432 100644 --- a/src/protocol/vm/utils.rs +++ b/src/protocol/vm/utils.rs @@ -2,12 +2,11 @@ #![allow(dead_code)] use ethabi::{self, decode, ParamType}; -use std::path::PathBuf; use ethers::{ + abi::Abi, core::utils::keccak256, types::{Address, H256, U256}, }; -use ethers::abi::Abi; use hex::FromHex; use mini_moka::sync::Cache; use reqwest::{blocking::Client, StatusCode}; @@ -131,55 +130,60 @@ fn parse_solidity_error_message(data: &str) -> String { format!("Failed to decode: {}", data) } -pub fn get_storage_slot_at_key(key: Address, mapping_slot: usize) -> U256 { +pub fn get_storage_slot_at_key(key: Address, mapping_slot: H256) -> H256 { // Convert key to bytes let mut key_bytes = key.as_bytes().to_vec(); key_bytes.resize(32, 0); // Right pad with zeros // Convert mapping slot to bytes let mut mapping_slot_bytes = [0u8; 32]; - mapping_slot_bytes[32 - std::mem::size_of::()..] - .copy_from_slice(&mapping_slot.to_be_bytes()); + mapping_slot_bytes.copy_from_slice(mapping_slot.as_bytes()); // Concatenate key and mapping slot bytes, then hash let slot_bytes = keccak256([&key_bytes[..], &mapping_slot_bytes[..]].concat()); + ethers::types::TxHash(slot_bytes) +} - // Convert the hashed bytes to a U256 - U256::from_big_endian(&slot_bytes) +pub struct GethOverwrite { + pub state_diff: HashMap, // The formatted overwrites + pub code: String, // The bytecode as a string } pub struct ERC20OverwriteFactory { token_address: Address, overwrites: HashMap, - balance_slot: u64, - allowance_slot: u64, - total_supply_slot: u64, + balance_slot: H256, + allowance_slot: H256, + total_supply_slot: H256, } impl ERC20OverwriteFactory { - pub fn new(token_address: Address, token_slots: (u64, u64)) -> Self { + pub fn new(token_address: Address, token_slots: (H256, H256)) -> Self { ERC20OverwriteFactory { token_address, overwrites: HashMap::new(), balance_slot: token_slots.0, allowance_slot: token_slots.1, - total_supply_slot: 2, + total_supply_slot: H256::from_low_u64_be(2), } } pub fn set_balance(&mut self, balance: U256, owner: Address) { - let storage_index = get_storage_slot_at_key(owner, self.balance_slot as usize); - self.overwrites.insert(H256::from(storage_index), balance); + let storage_index = get_storage_slot_at_key(owner, self.balance_slot); + self.overwrites + .insert(storage_index, balance); } pub fn set_allowance(&mut self, allowance: U256, spender: Address, owner: Address) { - let owner_slot = get_storage_slot_at_key(owner, self.allowance_slot as usize); - let storage_index = get_storage_slot_at_key(spender, owner_slot.as_usize()); - self.overwrites.insert(H256::from(storage_index), allowance); + let owner_slot = get_storage_slot_at_key(owner, self.allowance_slot); + let storage_index = get_storage_slot_at_key(spender, owner_slot); + self.overwrites + .insert(storage_index, allowance); } pub fn set_total_supply(&mut self, supply: U256) { - self.overwrites.insert(H256::from_low_u64_be(self.total_supply_slot), supply); + self.overwrites + .insert(self.total_supply_slot, supply); } pub fn get_protosim_overwrites(&self) -> HashMap> { @@ -188,24 +192,32 @@ impl ERC20OverwriteFactory { result } - pub fn get_geth_overwrites(&self) -> HashMap> { + pub fn get_geth_overwrites(&self) -> HashMap { let mut formatted_overwrites = HashMap::new(); + for (key, val) in &self.overwrites { - let formatted_key = format!("0x{}", key.to_hex::()); - let formatted_val = format!("0x{:0>64}", val.to_hex::()); - formatted_overwrites.insert(formatted_key, formatted_val); + let hex_key = hex::encode(key.as_bytes()); + + let mut bytes = [0u8; 32]; + val.to_big_endian(&mut bytes); + let hex_val = format!("0x{:0>64}", hex::encode(bytes)); + + formatted_overwrites.insert(hex_key, hex_val); } - let code = match get_contract_bytecode(&PathBuf::from(ASSETS_FOLDER).join("ERC20.bin")) { - Ok(bytes) => format!("0x{}", hex::encode(bytes)), - Err(_) => "0x".to_string(), // Handle error appropriately in your actual code - }; + let erc20_abi_path = Path::new(file!()) + .parent() + .unwrap() + .join("assets") + .join("ERC20.abi"); + + let code = format!( + "0x{}", + hex::encode(get_contract_bytecode(erc20_abi_path.to_str().unwrap()).unwrap()) + ); let mut result = HashMap::new(); - let mut inner = HashMap::new(); - inner.insert("stateDiff".to_string(), formatted_overwrites); - inner.insert("code".to_string(), code); - result.insert(self.token_address, inner); + result.insert(self.token_address, GethOverwrite { state_diff: formatted_overwrites, code }); result } @@ -523,4 +535,99 @@ mod tests { let abi: Abi = result.expect("Failed to retrieve ERC20 ABI result"); assert!(!abi.functions.is_empty(), "The ERC20 ABI should contain functions."); } + + fn setup_factory() -> ERC20OverwriteFactory { + let token_address = Address::random(); + let balance_slot = H256::random(); + let allowance_slot = H256::random(); + ERC20OverwriteFactory::new(token_address, (balance_slot, allowance_slot)) + } + + #[test] + fn test_set_balance() { + let mut factory = setup_factory(); + let owner = Address::random(); + let balance = U256::from(1000); + + factory.set_balance(balance, owner); + + assert_eq!(factory.overwrites.len(), 1); + assert!(factory + .overwrites + .values() + .any(|&v| v == balance)); + } + + #[test] + fn test_set_allowance() { + let mut factory = setup_factory(); + let owner = Address::random(); + let spender = Address::random(); + let allowance = U256::from(500); + + factory.set_allowance(allowance, spender, owner); + + assert_eq!(factory.overwrites.len(), 1); + assert!(factory + .overwrites + .values() + .any(|&v| v == allowance)); + } + + #[test] + fn test_set_total_supply() { + let mut factory = setup_factory(); + let supply = U256::from(1_000_000); + + factory.set_total_supply(supply); + + assert_eq!(factory.overwrites.len(), 1); + assert_eq!(factory.overwrites[&factory.total_supply_slot], supply); + } + + #[test] + fn test_get_protosim_overwrites() { + let mut factory = setup_factory(); + let supply = U256::from(1_000_000); + factory.set_total_supply(supply); + + let overwrites = factory.get_protosim_overwrites(); + + assert_eq!(overwrites.len(), 1); + assert!(overwrites.contains_key(&factory.token_address)); + assert_eq!(overwrites[&factory.token_address].len(), 1); + assert_eq!(overwrites[&factory.token_address][&factory.total_supply_slot], supply); + } + + #[test] + fn test_get_geth_overwrites() { + let mut factory = setup_factory(); + + let storage_slot = H256::from_low_u64_be(1); + let val = U256::from(123456); + factory + .overwrites + .insert(storage_slot, val); + + let result = factory.get_geth_overwrites(); + + assert_eq!(result.len(), 1); + + let geth_overwrite = result + .get(&factory.token_address) + .expect("Missing token address"); + assert_eq!(geth_overwrite.state_diff.len(), 1); + + let expected_key = + String::from("0000000000000000000000000000000000000000000000000000000000000001"); + let expected_val = + String::from("0x000000000000000000000000000000000000000000000000000000000001e240"); + assert_eq!( + geth_overwrite + .state_diff + .get(&expected_key), + Some(&expected_val) + ); + assert_eq!(geth_overwrite.code.len(), 8752); + } }