Merge pull request #3 from Yashiru/feat/fork

Fork support

Cargo.toml

@@ -6,8 +6,9 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+tokio = { version = "1", features = ["full"] }
 colored = "2.0.4"
 ethers = "2.0.7"
 mem_storage = "0.1.1"
 primitive-types = "0.12.1"
-hex = "0.4"
+hex = "0.4"

bytecode.bin

@@ -1 +1 @@
-746b60ff6001525a336000505050600052600c6014f36000526015600b6000f060006000600060008461fffffa746b60ff6001525a336000505050600052600c6014f36000526015600b6000f060006000600060008461fffffa746b60ff6001525a336000505050600052600c6014f36000526015600b6000f060006000600060008461fffffa746b60ff6001525a336000505050600052600c6014f36000526015600b6000f060006000600060008461fffffa746b60ff6001525a336000505050600052600c6014f36000526015600b6000f060006000600060008461fffffa746b60ff6001525a336000505050600052600c6014f36000526015600b6000f060006000600060008461fffffa746b60ff6001525a336000505050600052600c6014f36000526015600b6000f060006000600060008461fffffa
+7f70a08231000000000000000000000000f04a5cc80b1e94c69b48f5ee68a08cd26000527ff09a7c3e00000000000000000000000000000000000000000000000000000000602052602060006024600073c02aaa39b223fe8d0a0e5c4f27ead9083c756cc260fffa

src/core_module/op_codes/system.rs

@@ -381,8 +381,10 @@ pub fn return_(runner: &mut Runner) -> Result<(), ExecutionError> {
         runner.print_debug(&format!("{}", "RETURN".red()));
     }
 
-    // Increment PC
-    runner.increment_pc(1)
+    // Set the program counter to the end of the bytecode
+    runner.set_pc(runner.bytecode.len());
+
+    Ok(())
 }
 
 #[cfg(test)]
@@ -585,9 +587,8 @@ mod tests {
         assert_eq!(result, pad_left(&[0x01]));
 
         let stored_code = runner.state.get_code_at(bytes32_to_address(&result));
-
-        assert!(stored_code.is_err());
-        assert_eq!(stored_code.unwrap_err(), ExecutionError::AccountNotFound);
+        assert!(!stored_code.is_err());
+        assert_eq!(stored_code.unwrap().len(), 0);
 
         let balance_result = get_balance(bytes32_to_address(&result), &mut runner);
         assert!(balance_result.is_err());

src/core_module/runner.rs

@@ -52,7 +52,7 @@ impl Runner {
             state: if state.is_some() {
                 state.unwrap()
             } else {
-                EvmState::new()
+                EvmState::new(Some("https://mainnet.infura.io/v3/48de6103d9864b9bb17cf47d6cacf6ed".to_owned()))
             },
             // Create an empty memory
             memory: Memory::new(None),
@@ -450,19 +450,14 @@ impl Runner {
         };
 
         // Interpret the bytecode
-        let interpret_result =
-            self.interpret(self.state.get_code_at(to)?.to_owned(), self.debug_level, false);
+        let code = self.state.get_code_at(to)?.to_owned();
+        let interpret_result = self.interpret(code, self.debug_level, false);
 
         // Check if the interpretation was successful
         if interpret_result.is_err() {
             error = Some(interpret_result.unwrap_err());
         }
 
-        // Check if the call was successful
-        if !self.stack.stack.is_empty() {
-            error = Some(ExecutionError::NotEmptyStack);
-        }
-
         // Get the return data
         let return_data = self.returndata.heap.clone();
 
@@ -613,7 +608,7 @@ impl Runner {
         println!();
     }
 
-    fn debug_storage(&self) {
+    fn debug_storage(&mut self) {
         self.state.debug_state();
     }
 }

src/core_module/state.rs

@@ -1,5 +1,6 @@
 use std::{collections::HashMap, fmt};
 
+use ethers::prelude::*;
 use ethers::{types::U256, utils::keccak256};
 
 use crate::core_module::utils;
@@ -9,7 +10,6 @@ use super::utils::errors::ExecutionError;
 // Colored output
 use colored::*;
 
-
 /* -------------------------------------------------------------------------- */
 /*                             AccountState struct                            */
 /* -------------------------------------------------------------------------- */
@@ -68,7 +68,12 @@ pub struct Log {
 
 impl fmt::Debug for Log {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        writeln!(f, "{}: {}", "Address".magenta(), utils::debug::to_hex_address(self.address))?;
+        writeln!(
+            f,
+            "{}: {}",
+            "Address".magenta(),
+            utils::debug::to_hex_address(self.address)
+        )?;
 
         write!(f, "{}: ", "Topics".magenta())?;
         if !self.topics.is_empty() {
@@ -82,7 +87,12 @@ impl fmt::Debug for Log {
             writeln!(f, "{}", "No topics".red())?;
         }
 
-        writeln!(f, "{}: {}", "Data".magenta(), utils::debug::vec_to_hex_string(self.data.clone()))?;
+        writeln!(
+            f,
+            "{}: {}",
+            "Data".magenta(),
+            utils::debug::vec_to_hex_string(self.data.clone())
+        )?;
 
         Ok(())
     }
@@ -98,15 +108,21 @@ pub struct EvmState {
     pub codes: HashMap<[u8; 32], Vec<u8>>,
     pub logs: Vec<Log>,
     pub static_mode: bool,
+    pub provider: Option<Provider<Http>>,
 }
 
 impl EvmState {
-    pub fn new() -> Self {
+    pub fn new(fork_url: Option<String>) -> Self {
         Self {
             accounts: HashMap::new(),
             codes: HashMap::new(),
             logs: Vec::new(),
-            static_mode: false
+            static_mode: false,
+            provider: if fork_url.is_some() {
+                Some(Provider::<Http>::try_from(fork_url.unwrap()).unwrap())
+            } else {
+                None
+            },
         }
     }
 
@@ -164,14 +180,43 @@ impl EvmState {
         Ok(())
     }
 
-    // Load a 32 bytes word from storage of a specific account
-    pub fn sload(&self, account: [u8; 20], slot: [u8; 32]) -> Result<[u8; 32], ExecutionError> {
+    pub fn sload(&mut self, account: [u8; 20], slot: [u8; 32]) -> Result<[u8; 32], ExecutionError> {
         match self.accounts.get(&account) {
-            Some(account_state) => match account_state.storage.get(&slot) {
-                Some(value) => Ok(*value),
-                None => Ok([0u8; 32]),
-            },
-            None => Ok([0u8; 32]),
+            Some(account_state) => {
+                match account_state.storage.get(&slot) {
+                    Some(value) => Ok(*value),
+                    None => Ok([0u8; 32])
+                }
+            }
+            None => {
+                let provider = self.provider.as_ref();
+
+                if provider.is_none() {
+                    return Ok([0u8; 32]);
+                }
+
+                let contract_address = Address::from(account);
+                let future = provider.unwrap().get_storage_at(contract_address, H256::from(&slot), None);
+
+                // Block on the future and get the result
+                let storage_result = tokio::runtime::Runtime::new()
+                    .expect("Could not create a Runtime")
+                    .block_on(future);
+
+                match storage_result {
+                    Ok(storage) => {
+                        let storage_bytes = storage.to_fixed_bytes();
+
+                        // Save the fetched storage data locally
+                        if let Some(account_state) = self.accounts.get_mut(&account) {
+                            account_state.storage.insert(slot, storage_bytes);
+                        }
+
+                        Ok(storage_bytes)
+                    }
+                    Err(_) => Ok([0u8; 32]),
+                }
+            }
         }
     }
 
@@ -197,13 +242,41 @@ impl EvmState {
     }
 
     // Get the code of an account
-    pub fn get_code_at(&self, address: [u8; 20]) -> Result<&Vec<u8>, ExecutionError> {
+    pub fn get_code_at(&mut self, address: [u8; 20]) -> Result<&Vec<u8>, ExecutionError> {
         match self.accounts.get(&address) {
             Some(account_state) => {
                 let code_hash = account_state.code_hash;
                 self.get_code(code_hash)
             }
-            None => Err(ExecutionError::AccountNotFound),
+            None => {
+                let provider = self.provider.as_ref();
+
+                if provider.is_none() {
+                    return Err(ExecutionError::CodeNotFound);
+                }
+
+                // Asynchronously fetch the code from the blockchain
+                let contract_address = Address::from(address);
+
+                let future = provider.unwrap().get_code(contract_address, None);
+
+                // Block on the future and get the result
+                let code_result = tokio::runtime::Runtime::new()
+                    .expect("Could not create a Runtime")
+                    .block_on(future);
+
+                match code_result {
+                    Ok(code) => {
+                        let code_hash = keccak256(&code.0);
+                        if let Some(account) = self.accounts.get_mut(&address) {
+                            account.code_hash = code_hash;
+                        }
+                        self.codes.insert(code_hash, code.to_vec());
+                        Ok(&self.codes[&code_hash])
+                    }
+                    Err(_) => Err(ExecutionError::CodeNotFound),
+                }
+            }
         }
     }
 
@@ -239,7 +312,7 @@ impl EvmState {
         Ok(code_hash)
     }
 
-    pub fn debug_state(&self) {
+    pub fn debug_state(&mut self) {
         let border_line =
             "\n╔═══════════════════════════════════════════════════════════════════════════════════════════════════════╗";
         let footer_line =
@@ -255,21 +328,23 @@ impl EvmState {
         );
         println!("{}", footer_line.clone().red());
 
-        for (address, account_state) in &self.accounts {
+        // ... other code ...
+
+        // Create a vector of all addresses
+        let addresses: Vec<_> = self.accounts.keys().cloned().collect();
+
+        // Iterate over the vector of addresses
+        for address in addresses {
+            let account_state = &self.accounts[&address];
             let hex: String = utils::debug::to_hex_address(address.to_owned());
-            println!(
-                "{}",
-                hex.blue()
-            );
+            println!("{}", hex.blue());
             println!("{:?}", account_state);
-            // Print the code of the contract
             let code_hash = account_state.code_hash;
             if code_hash != [0u8; 32] {
                 let code = self.get_code_at(address.to_owned()).unwrap();
                 let code_hex: String = utils::debug::vec_to_hex_string(code.to_owned());
                 println!("  {}: {}", "Code".magenta(), code_hex);
             }
-
             println!("\n");
         }
 

src/core_module/utils/bytes.rs

@@ -45,7 +45,6 @@ pub fn _hex_string_to_bytes(hex: &str) -> Vec<u8> {
             panic!("Error: {}", e.to_string().red());
         },
     }
-
 }
 
 /* -------------------------------------------------------------------------- */

src/core_module/utils/errors.rs

@@ -12,7 +12,6 @@ pub enum ExecutionError {
     InsufficientBalance,
 
     // Flow errors
-    NotEmptyStack,
     StaticCallStateChanged,
     InvalidOpcode(u8),
 
@@ -45,7 +44,6 @@ impl fmt::Display for ExecutionError {
             ExecutionError::CodeNotFound => write!(f, "Trying to access non-existent account code"),
             ExecutionError::RevertWithoutData => write!(f, "Execution revert without data"),
             ExecutionError::InsufficientBalance => write!(f, "Insufficient balance to transfer"),
-            ExecutionError::NotEmptyStack => write!(f, "Stack is not empty after the call"),
             ExecutionError::InvalidOpcode(op_code) => {
                 write!(f, "Invalid op code 0x{:X}", op_code)
             }
@@ -74,7 +72,6 @@ impl PartialEq for ExecutionError {
             | (CodeNotFound, CodeNotFound)
             | (EmptyByteCode, EmptyByteCode)
             | (InsufficientBalance, InsufficientBalance)
-            | (NotEmptyStack, NotEmptyStack)
             | (StaticCallStateChanged, StaticCallStateChanged)
             | (StackTooSmall, StackTooSmall)
             | (StackTooDeep, StackTooDeep)