Fix Foundry fork test not picking up targetContracts (#469)

* remove dbg

* chore: forge init

* forge install: forge-std

v1.8.1

* fix foundry onchain not taking targetContracts

* remove tests

* bump

src/evm/concolic/concolic_host.rs

@@ -1,4 +1,5 @@
 use std::{
+    any,
     borrow::Borrow,
     collections::{HashMap, HashSet},
     fmt::{Debug, Display},
@@ -1398,4 +1399,8 @@ where
     fn get_type(&self) -> MiddlewareType {
         Concolic
     }
+
+    fn as_any(&self) -> &dyn any::Any {
+        self
+    }
 }

src/evm/contract_utils.rs

@@ -4,21 +4,24 @@ use std::{
     collections::{HashMap, HashSet},
     fs::File,
     io::Read,
+    ops::Deref,
     path::{Path, PathBuf},
     rc::Rc,
     str::FromStr,
+    sync::RwLockReadGuard,
 };
 
 use bytes::Bytes;
 /// Load contract from file system or remote
 use glob::glob;
 use itertools::Itertools;
 use libafl::{schedulers::StdScheduler, state::HasMetadata};
-use revm_primitives::{Bytecode, Env};
+use revm_primitives::{bitvec::vec, Bytecode, Env};
 use serde_json::Value;
 
 use crate::{
     evm::{
+        middlewares::middleware::MiddlewareType,
         tokens::constant_pair::ConstantPairMetadata,
         types::{fixed_address, generate_random_address, EVMAddress, EVMFuzzState},
         vm::{IN_DEPLOY, SETCODE_ONLY},
@@ -53,13 +56,13 @@ use crate::evm::{
         offchain_config::OffchainConfig,
     },
     bytecode_iterator::all_bytecode,
-    onchain::endpoints::OnChainConfig,
+    onchain::{endpoints::OnChainConfig, OnChain},
 };
 
 // to use this address, call rand_utils::fixed_address(FIX_DEPLOYER)
 pub static FIX_DEPLOYER: &str = "8b21e662154b4bbc1ec0754d0238875fe3d22fa6";
 pub static FOUNDRY_DEPLOYER: &str = "1804c8AB1F12E6bbf3894d4083f33e07309d1f38";
-pub static FOUNDRY_SETUP_ADDR: &str = "7FA9385bE102ac3EAc297483Dd6233D62b3e1496";
+pub static FOUNDRY_SETUP_ADDR: &str = "e1A425f1AC34A8a441566f93c82dD730639c8510";
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ABIConfig {
@@ -120,7 +123,7 @@ impl ContractLoader {
     }
 }
 
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Default)]
 pub struct SetupData {
     pub evmstate: EVMState,
     pub env: Env,
@@ -136,6 +139,26 @@ pub struct SetupData {
     // Flashloan specific
     pub v2_pairs: Vec<EVMAddress>,
     pub constant_pairs: Vec<ConstantPairMetadata>,
+
+    pub onchain_middleware: Option<OnChain>,
+}
+
+impl Clone for SetupData {
+    fn clone(&self) -> Self {
+        Self {
+            evmstate: self.evmstate.clone(),
+            env: self.env.clone(),
+            code: self.code.clone(),
+            excluded_contracts: self.excluded_contracts.clone(),
+            excluded_senders: self.excluded_senders.clone(),
+            target_contracts: self.target_contracts.clone(),
+            target_senders: self.target_senders.clone(),
+            target_selectors: self.target_selectors.clone(),
+            v2_pairs: self.v2_pairs.clone(),
+            constant_pairs: self.constant_pairs.clone(),
+            onchain_middleware: self.onchain_middleware.clone(),
+        }
+    }
 }
 
 pub fn set_hash(name: &str, out: &mut [u8]) {
@@ -802,7 +825,49 @@ impl ContractLoader {
             }
         }
 
+        let mut known_addr = HashSet::new();
+
+        if let Some(onchain) = setup_data.clone().unwrap().onchain_middleware {
+            debug!("Adding onchain middleware");
+            let mut onchain_config = onchain.endpoint.clone();
+
+            for addr in setup_data.clone().unwrap().target_contracts {
+                if known_addr.contains(&addr) || addr == fixed_address(FOUNDRY_SETUP_ADDR) {
+                    continue;
+                }
+                known_addr.insert(addr);
+                let code = onchain_config.get_contract_code(addr, false);
+                if code.is_empty() {
+                    error!("Failed to get code for contract at address {:?}", addr);
+                    continue;
+                }
+                let abi = Self::parse_abi_str(&onchain_config.fetch_abi(addr).unwrap());
+
+                contracts.push(ContractInfo {
+                    name: format!("{}", addr),
+                    code: hex::decode(&code).expect("code is not hex"),
+                    abi: abi.clone(),
+                    is_code_deployed: true,
+                    constructor_args: vec![],
+                    deployed_address: addr,
+                    build_artifact: None,
+                    files: vec![],
+                    source_map_replacements: None,
+                    raw_source_map: None,
+                });
+
+                abis.push(ABIInfo {
+                    source: format!("{}", addr),
+                    abi,
+                });
+            }
+        }
+
         for (addr, code) in setup_data.clone().unwrap().code {
+            if known_addr.contains(&addr) {
+                continue;
+            }
+            known_addr.insert(addr);
             let (artifact_idx, slug) = Self::find_contract_artifact(code.to_vec(), offchain_artifacts);
 
             debug!(
@@ -835,6 +900,7 @@ impl ContractLoader {
                 raw_source_map: Some(more_info.source_map.clone()),
             });
         }
+
         Self {
             contracts,
             abis,
@@ -1078,6 +1144,19 @@ impl ContractLoader {
             .map(|(k, v)| (k, Bytes::from_iter(v.bytecode().iter().cloned())))
             .collect();
 
+        let mut onchain_middleware = None;
+
+        {
+            for middleware in evm_executor.host.middlewares.read().unwrap().clone().into_iter() {
+                if middleware.borrow().get_type() == MiddlewareType::OnChain {
+                    unsafe {
+                        let onchain = middleware.borrow().as_any().downcast_ref_unchecked::<OnChain>().clone();
+                        onchain_middleware = Some(onchain);
+                    }
+                }
+            }
+        }
+
         SetupData {
             evmstate: new_vm_state,
             env: evm_executor.host.env.clone(),
@@ -1089,6 +1168,7 @@ impl ContractLoader {
             target_selectors,
             v2_pairs,
             constant_pairs,
+            onchain_middleware,
         }
     }
 }

src/evm/host.rs

@@ -828,7 +828,7 @@ where
         res: (InstructionResult, Gas, Bytes),
     ) -> (InstructionResult, Gas, Bytes) {
         // Check assert result
-        if let Some(res) = self.check_assert_result(&res) {
+        if let Some(res) = self.check_assert_result() {
             return res;
         }
 
@@ -846,12 +846,9 @@ where
         self.check_expected_calls(res)
     }
 
-    fn check_assert_result(
-        &mut self,
-        res: &(InstructionResult, Gas, Bytes),
-    ) -> Option<(InstructionResult, Gas, Bytes)> {
+    pub fn check_assert_result(&mut self) -> Option<(InstructionResult, Gas, Bytes)> {
         if let Some(ref msg) = self.assert_msg {
-            return Some((InstructionResult::Revert, res.1, msg.abi_encode().into()));
+            return Some((InstructionResult::Revert, Gas::new(0), msg.abi_encode().into()));
         }
 
         None
@@ -1242,6 +1239,7 @@ where
     fn sload(&mut self, address: EVMAddress, index: EVMU256) -> Option<(EVMU256, bool)> {
         if let Some(account) = self.evmstate.get_mut(&address) {
             if let Some(slot) = account.get(&index) {
+                // println!("sload: {:?} -> {:?} = {:?}", address, index, slot);
                 return Some((*slot, true));
             } else {
                 account.insert(index, self.next_slot);
@@ -1251,7 +1249,7 @@ where
             account.insert(index, self.next_slot);
             self.evmstate.insert(address, account);
         }
-
+        // println!("sload(c): {:?} -> {:?} = {:?}", address, index, self.next_slot);
         Some((self.next_slot, true))
     }
 

src/evm/middlewares/call_printer.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, fmt::Debug};
+use std::{any, collections::HashMap, fmt::Debug};
 
 use bytes::Bytes;
 use itertools::Itertools;
@@ -259,4 +259,8 @@ where
     fn get_type(&self) -> MiddlewareType {
         MiddlewareType::CallPrinter
     }
+
+    fn as_any(&self) -> &dyn any::Any {
+        self
+    }
 }

src/evm/middlewares/cheatcode/mod.rs

@@ -95,7 +95,7 @@ macro_rules! cheat_call_error {
 
 impl<SC> Middleware<SC> for Cheatcode<SC>
 where
-    SC: Scheduler<State = EVMFuzzState> + Clone + Debug,
+    SC: Scheduler<State = EVMFuzzState> + Clone + Debug + 'static,
 {
     unsafe fn on_step(&mut self, interp: &mut Interpreter, host: &mut FuzzHost<SC>, _state: &mut EVMFuzzState) {
         let op = interp.current_opcode();
@@ -111,6 +111,10 @@ where
     fn get_type(&self) -> MiddlewareType {
         MiddlewareType::Cheatcode
     }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
 }
 
 impl<SC> Cheatcode<SC>

src/evm/middlewares/coverage.rs

@@ -1,4 +1,5 @@
 use std::{
+    any,
     collections::{HashMap, HashSet},
     fmt::{Debug, Display, Formatter},
     fs,
@@ -360,6 +361,10 @@ where
     fn get_type(&self) -> MiddlewareType {
         MiddlewareType::InstructionCoverage
     }
+
+    fn as_any(&self) -> &dyn any::Any {
+        self
+    }
 }
 
 #[cfg(test)]

src/evm/middlewares/middleware.rs

@@ -1,4 +1,4 @@
-use std::{clone::Clone, fmt::Debug, time::Duration};
+use std::{any, clone::Clone, fmt::Debug, time::Duration};
 
 use bytes::Bytes;
 use libafl::{
@@ -113,4 +113,6 @@ where
     ) {
     }
     fn get_type(&self) -> MiddlewareType;
+
+    fn as_any(&self) -> &dyn any::Any;
 }

src/evm/middlewares/reentrancy.rs

@@ -1,4 +1,5 @@
 use std::{
+    any,
     collections::{HashMap, HashSet},
     fmt::Debug,
 };
@@ -182,6 +183,9 @@ where
                 depths.retain(|&x| x <= depth);
             });
     }
+    fn as_any(&self) -> &dyn any::Any {
+        self
+    }
 }
 
 #[cfg(test)]

src/evm/middlewares/sha3_bypass.rs

@@ -1,4 +1,5 @@
 use std::{
+    any,
     cell::RefCell,
     collections::{HashMap, HashSet},
     fmt::Debug,
@@ -102,6 +103,10 @@ impl Sha3TaintAnalysis {
         self.dirty_storage = ctx.dirty_storage;
         self.dirty_stack = ctx.dirty_stack;
     }
+
+    fn as_any(&self) -> &dyn any::Any {
+        self
+    }
 }
 
 impl<SC> Middleware<SC> for Sha3TaintAnalysis
@@ -386,6 +391,9 @@ where
     fn get_type(&self) -> MiddlewareType {
         MiddlewareType::Sha3TaintAnalysis
     }
+    fn as_any(&self) -> &dyn any::Any {
+        self
+    }
 }
 
 #[derive(Debug)]
@@ -421,6 +429,9 @@ where
     fn get_type(&self) -> MiddlewareType {
         MiddlewareType::Sha3Bypass
     }
+    fn as_any(&self) -> &dyn any::Any {
+        self
+    }
 }
 
 #[cfg(test)]

src/evm/onchain/flashloan.rs

@@ -2,8 +2,8 @@
 // when approval, balanceof, give 2000e18 token
 // when transfer, transferFrom, and src is our, return success, add owed
 // when transfer, transferFrom, and src is not our, return success, reduce owed
-
 use std::{
+    any,
     cell::RefCell,
     collections::{HashMap, HashSet},
     fmt::Debug,
@@ -297,6 +297,10 @@ where
     fn get_type(&self) -> MiddlewareType {
         MiddlewareType::Flashloan
     }
+
+    fn as_any(&self) -> &dyn any::Any {
+        self
+    }
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, Default)]

src/evm/onchain/mod.rs

@@ -62,6 +62,7 @@ pub trait ChainConfig {
     fn get_pegged_token(&self) -> HashMap<String, String>;
 }
 
+#[derive(Clone)]
 pub struct OnChain {
     pub loaded_data: HashSet<(EVMAddress, EVMU256)>,
     pub loaded_code: HashSet<EVMAddress>,
@@ -219,6 +220,7 @@ where
                             .get_contract_slot(address, slot_idx, force_cache!(self.locs, slot_idx))
                     }
                 };
+                // println!("SLOAD {:?} {:?}", slot_idx, host.next_slot);
             }
             #[cfg(feature = "real_balance")]
             // BALANCE
@@ -301,6 +303,10 @@ where
     fn get_type(&self) -> MiddlewareType {
         MiddlewareType::OnChain
     }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
 }
 
 impl OnChain {