Full Foundry Support

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "tests/evm_manual/foundry1/lib/forge-std"]
+	path = tests/evm_manual/foundry1/lib/forge-std
+	url = https://github.com/foundry-rs/forge-std

src/evm/contract_utils.rs

@@ -120,6 +120,13 @@ pub struct SetupData {
     pub evmstate: EVMState,
     pub env: Env,
     pub code: HashMap<EVMAddress, Bytes>,
+
+    // Foundry specific
+    pub excluded_contracts: Vec<EVMAddress>,
+    pub excluded_senders: Vec<EVMAddress>,
+    pub target_contracts: Vec<EVMAddress>,
+    pub target_senders: Vec<EVMAddress>,
+    pub target_selectors: HashMap<EVMAddress, Vec<Vec<u8>>>,
 }
 
 pub fn set_hash(name: &str, out: &mut [u8]) {
@@ -827,6 +834,128 @@ impl ContractLoader {
         }
         assert!(res[0].1, "setUp() failed");
 
+        // now get Foundry invariant test config by calling
+        // * excludeContracts() => array of addresses
+        // * excludeSenders() => array of sender
+        // * targetContracts() => array of addresses
+        // * targetSenders() => array of sender
+        // * targetSelectors() => array of selectors
+        macro_rules! abi_sig {
+            ($name: expr) => {
+                Bytes::from_iter(ethers::abi::short_signature($name, &[]).iter().cloned())
+            };
+        }
+
+        let calls = vec![
+            (deployer, deployed_addr, abi_sig!("excludeContracts")),
+            (deployer, deployed_addr, abi_sig!("excludeSenders")),
+            (deployer, deployed_addr, abi_sig!("targetContracts")),
+            (deployer, deployed_addr, abi_sig!("targetSenders")),
+            (deployer, deployed_addr, abi_sig!("targetSelectors")),
+        ];
+
+        let (res, _) = evm_executor.fast_call(&calls, &new_vm_state, &mut state);
+
+        macro_rules! parse_nth_result_addr {
+            ($nth: expr) => {{
+                let (res_bys, succ) = res[$nth].clone();
+                let mut addrs = vec![];
+                if succ {
+                    let contracts = ethers::abi::decode(
+                        &[ethers::abi::ParamType::Array(Box::new(
+                            ethers::abi::ParamType::Address,
+                        ))],
+                        &res_bys,
+                    )
+                    .unwrap();
+                    addrs = if let ethers::abi::Token::Array(contracts) = contracts[0].clone() {
+                        contracts
+                            .iter()
+                            .map(|x| {
+                                if let ethers::abi::Token::Address(addr) = x {
+                                    EVMAddress::from_slice(addr.as_bytes())
+                                } else {
+                                    panic!("invalid address")
+                                }
+                            })
+                            .collect::<Vec<_>>()
+                    } else {
+                        panic!("invalid array")
+                    };
+                }
+                addrs
+            }};
+        }
+
+        // (address addr, bytes4[] selectors)[]
+        macro_rules! parse_nth_result_selector {
+            ($nth: expr) => {{
+                let (res_bys, succ) = res[$nth].clone();
+                let mut sigs = vec![];
+                if succ {
+                    let sigs_parsed = ethers::abi::decode(
+                        &[ethers::abi::ParamType::Array(Box::new(
+                            ethers::abi::ParamType::Tuple(vec![
+                                ethers::abi::ParamType::Address,
+                                ethers::abi::ParamType::Array(Box::new(ethers::abi::ParamType::FixedBytes(4))),
+                            ]),
+                        ))],
+                        &res_bys,
+                    )
+                    .unwrap();
+
+                    let sigs_parsed = if let ethers::abi::Token::Array(sigs_parsed) = sigs_parsed[0].clone() {
+                        sigs_parsed
+                    } else {
+                        panic!("invalid array")
+                    };
+
+                    sigs_parsed.iter().for_each(|x| {
+                        if let ethers::abi::Token::Tuple(tuple) = x {
+                            let addr = tuple[0].clone();
+                            let selectors = tuple[1].clone();
+                            if let (ethers::abi::Token::Address(addr), ethers::abi::Token::Array(selectors)) =
+                                (addr, selectors)
+                            {
+                                let addr = EVMAddress::from_slice(addr.as_bytes());
+                                let selectors = selectors
+                                    .iter()
+                                    .map(|x| {
+                                        if let ethers::abi::Token::FixedBytes(bytes) = x {
+                                            bytes.clone()
+                                        } else {
+                                            panic!("invalid bytes")
+                                        }
+                                    })
+                                    .collect::<Vec<_>>();
+                                sigs.push((addr, selectors));
+                            } else {
+                                panic!("invalid tuple: {:?}", tuple)
+                            }
+                        } else {
+                            panic!("invalid tuple")
+                        }
+                    });
+                }
+                sigs
+            }};
+        }
+
+        let excluded_contracts = parse_nth_result_addr!(0);
+        let excluded_senders = parse_nth_result_addr!(1);
+        let target_contracts = parse_nth_result_addr!(2);
+        let target_senders = parse_nth_result_addr!(3);
+        let target_selectors_vec = parse_nth_result_selector!(4);
+        let target_selectors = {
+            let mut map = HashMap::new();
+            target_selectors_vec.iter().for_each(|(addr, selectors)| {
+                map.entry(*addr)
+                    .or_insert_with(std::vec::Vec::new)
+                    .extend(selectors.clone());
+            });
+            map
+        };
+
         // get the newly deployed code by setUp()
         let code: HashMap<EVMAddress, Bytes> = evm_executor
             .host
@@ -839,6 +968,11 @@ impl ContractLoader {
             evmstate: new_vm_state,
             env: evm_executor.host.env.clone(),
             code,
+            excluded_contracts,
+            excluded_senders,
+            target_contracts,
+            target_senders,
+            target_selectors,
         }
     }
 }

src/evm/corpus_initializer.rs

@@ -179,8 +179,8 @@ where
 
     pub fn initialize(&mut self, loader: &mut ContractLoader) -> EVMInitializationArtifacts {
         self.state.metadata_map_mut().insert(ABIMap::new());
-        self.setup_default_callers();
-        self.setup_contract_callers();
+        self.setup_default_callers(loader);
+        self.setup_contract_callers(loader);
         self.init_cheatcode_contract();
         self.initialize_contract(loader);
         self.initialize_source_map(loader);
@@ -352,6 +352,36 @@ where
                 continue;
             }
 
+            let mut target_sig = None; // none means all selectors are targeted
+
+            if let Some(setup_data) = &loader.setup_data {
+                // Check if this contract is included by Foundry targetContracts
+                if !setup_data.target_contracts.is_empty() &&
+                    !setup_data.target_contracts.contains(&contract.deployed_address)
+                {
+                    continue;
+                }
+                // Check if this contract is excluded by Foundry excludeContracts
+                if !setup_data.excluded_contracts.is_empty() &&
+                    setup_data.excluded_contracts.contains(&contract.deployed_address)
+                {
+                    continue;
+                }
+
+                // Check if this contract and sig is included by Foundry targetSelectors
+                if !setup_data.target_selectors.is_empty() {
+                    target_sig = Some(
+                        setup_data
+                            .target_selectors
+                            .get(&contract.deployed_address)
+                            .cloned()
+                            .unwrap_or(
+                                vec![], // empty vec means none of the selectors are targeted
+                            ),
+                    );
+                }
+            }
+
             for abi in contract.abi.clone() {
                 let name = &abi.function_name;
 
@@ -361,6 +391,13 @@ where
                     continue;
                 }
 
+                if let Some(target_sig) = target_sig.clone() {
+                    if !target_sig.contains(&Vec::from_iter(abi.function.iter().cloned())) {
+                        debug!("Skipping function because of targetSelectors: {}", name);
+                        continue;
+                    }
+                }
+
                 self.add_abi(&abi, contract.deployed_address, &mut artifacts);
             }
         }
@@ -379,7 +416,21 @@ where
         artifacts
     }
 
-    pub fn setup_default_callers(&mut self) {
+    pub fn setup_default_callers(&mut self, loader: &mut ContractLoader) {
+        // We override default callers when target senders are specified
+        if let Some(setup_data) = &loader.setup_data {
+            if !setup_data.target_senders.is_empty() {
+                for caller in setup_data.target_senders.iter() {
+                    self.state.add_caller(caller);
+                    self.executor
+                        .host
+                        .evmstate
+                        .set_balance(*caller, EVMU256::from(INITIAL_BALANCE));
+                }
+                return;
+            }
+        }
+
         let default_callers = HashSet::from([
             fixed_address("8EF508Aca04B32Ff3ba5003177cb18BfA6Cd79dd"),
             fixed_address("35c9dfd76bf02107ff4f7128Bd69716612d31dDb"),
@@ -395,7 +446,16 @@ where
         }
     }
 
-    pub fn setup_contract_callers(&mut self) {
+    pub fn setup_contract_callers(&mut self, loader: &mut ContractLoader) {
+        // We no longer need to setup contract callers when target senders are specified
+        // The target senders are already setup in setup_default_callers. Existence
+        // of the code in those addresses depends on setUp function of the contract.
+        if let Some(setup_data) = &loader.setup_data {
+            if !setup_data.target_senders.is_empty() {
+                return;
+            }
+        }
+
         let contract_callers = HashSet::from([
             fixed_address("e1A425f1AC34A8a441566f93c82dD730639c8510"),
             fixed_address("68Dd4F5AC792eAaa5e36f4f4e0474E0625dc9024"),

src/evm/input.rs

@@ -11,7 +11,10 @@ use libafl_bolts::{prelude::Rand, HasLen};
 use revm_primitives::Env;
 use serde::{Deserialize, Deserializer, Serialize};
 
-use super::utils::{colored_address, colored_sender, prettify_value};
+use super::{
+    onchain::flashloan::CAN_LIQUIDATE,
+    utils::{colored_address, colored_sender, prettify_value},
+};
 use crate::{
     evm::{
         abi::{AEmpty, AUnknown, BoxedABI},
@@ -366,15 +369,14 @@ impl ConciseEVMInput {
 
     #[inline]
     fn append_liquidation(&self, indent: String, call: String) -> String {
-        if self.liquidation_percent == 0 {
+        if self.liquidation_percent == 0 || unsafe { !CAN_LIQUIDATE } {
             return call;
         }
 
         let liq_call = format!(
-            "{}.{}(100% Balance, 0, path:({} → WETH), address(this), block.timestamp);",
+            "{}.{}(100% Balance, 0, path:(* → WETH), address(this), block.timestamp);",
             colored_address("Router"),
             self.colored_fn_name("swapExactTokensForETH"),
-            colored_address(&self.contract())
         );
 
         let mut liq = indent.clone();

src/evm/mutator.rs

@@ -12,6 +12,7 @@ use libafl_bolts::{prelude::Rand, Named};
 use revm_interpreter::Interpreter;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
+use super::onchain::flashloan::CAN_LIQUIDATE;
 /// Mutator for EVM inputs
 use crate::evm::input::EVMInputT;
 use crate::{
@@ -293,10 +294,15 @@ where
             if input.is_step() {
                 let res = match state.rand_mut().below(100) {
                     0..=5 => {
-                        let prev_percent = input.get_liquidation_percent();
-                        input.set_liquidation_percent(if state.rand_mut().below(100) < 80 { 10 } else { 0 } as u8);
-                        if prev_percent != input.get_liquidation_percent() {
-                            MutationResult::Mutated
+                        // only when there are more than one liquidation path, we attempt to liquidate
+                        if unsafe { CAN_LIQUIDATE } {
+                            let prev_percent = input.get_liquidation_percent();
+                            input.set_liquidation_percent(if state.rand_mut().below(100) < 80 { 10 } else { 0 } as u8);
+                            if prev_percent != input.get_liquidation_percent() {
+                                MutationResult::Mutated
+                            } else {
+                                MutationResult::Skipped
+                            }
                         } else {
                             MutationResult::Skipped
                         }

src/evm/onchain/flashloan.rs

@@ -51,6 +51,8 @@ use crate::{
     state::{HasCaller, HasItyState},
 };
 
+pub static mut CAN_LIQUIDATE: bool = false;
+
 macro_rules! scale {
     () => {
         EVMU512::from(1_000_000)
@@ -205,7 +207,10 @@ impl Flashloan {
                         let oracle = self.flashloan_oracle.deref().try_borrow_mut();
                         // avoid delegate call on token -> make oracle borrow multiple times
                         if oracle.is_ok() {
-                            oracle.unwrap().register_token(*addr, Rc::new(RefCell::new(token_ctx)));
+                            let can_liquidate = !token_ctx.swaps.is_empty(); // if there is more than one liquidation path, we can liquidate
+                            oracle
+                                .unwrap()
+                                .register_token(*addr, Rc::new(RefCell::new(token_ctx)), can_liquidate);
                             self.erc20_address.insert(*addr);
                             is_erc20 = true;
                         } else {

src/evm/oracles/arb_call.rs

@@ -86,7 +86,7 @@ impl
                 caller.hash(&mut hasher);
                 target.hash(&mut hasher);
                 pc.hash(&mut hasher);
-                let real_bug_idx = hasher.finish() << (8 + ARB_CALL_BUG_IDX);
+                let real_bug_idx = (hasher.finish() << 8) + ARB_CALL_BUG_IDX;
 
                 let name = self
                     .address_to_name

src/evm/oracles/arb_transfer.rs

@@ -71,7 +71,7 @@ Oracle<
                 let mut hasher = DefaultHasher::new();
                 caller.hash(&mut hasher);
                 pc.hash(&mut hasher);
-                let real_bug_idx = (hasher.finish() as u64) << 8 + ARB_TRANSFER_BUG_IDX;
+                let real_bug_idx = (hasher.finish() << 8 as u64) + ARB_TRANSFER_BUG_IDX;
 
                 let mut name = self.address_to_name
                     .get(caller)

src/evm/oracles/erc20.rs

@@ -7,6 +7,7 @@ use tracing::debug;
 use crate::{
     evm::{
         input::{ConciseEVMInput, EVMInput},
+        onchain::flashloan::CAN_LIQUIDATE,
         oracle::EVMBugResult,
         oracles::{u512_div_float, ERC20_BUG_IDX},
         producers::erc20::ERC20Producer,
@@ -36,7 +37,16 @@ impl IERC20OracleFlashloan {
         }
     }
 
-    pub fn register_token(&mut self, token: EVMAddress, token_ctx: Rc<RefCell<dyn TokenContextT<EVMFuzzState>>>) {
+    pub fn register_token(
+        &mut self,
+        token: EVMAddress,
+        token_ctx: Rc<RefCell<dyn TokenContextT<EVMFuzzState>>>,
+        can_liquidate: bool,
+    ) {
+        // setting can_liquidate to true to turn on liquidation
+        unsafe {
+            CAN_LIQUIDATE |= can_liquidate;
+        }
         self.known_tokens.insert(token, token_ctx);
     }