Merge pull request #461 from fuzzland/fix-460-onchain-fuzzing-with-cu…

…stom-setup

fix #460: onchain-fuzzing-with-custom-setup

src/evm/config.rs

@@ -59,6 +59,7 @@ impl FromStr for FuzzerTypes {
 pub struct Config<VS, Addr, Code, By, Loc, SlotTy, Out, I, S, CI, E> {
     pub onchain: Option<OnChainConfig>,
     pub onchain_storage_fetching: Option<StorageFetchingMode>,
+    pub etherscan_api_key: String,
     pub flashloan: bool,
     pub concolic: bool,
     pub concolic_caller: bool,

src/evm/contract_utils.rs

@@ -13,7 +13,7 @@ use bytes::Bytes;
 /// Load contract from file system or remote
 use glob::glob;
 use itertools::Itertools;
-use libafl::schedulers::StdScheduler;
+use libafl::{schedulers::StdScheduler, state::HasMetadata};
 use revm_primitives::{Bytecode, Env};
 use serde_json::Value;
 
@@ -24,7 +24,7 @@ use crate::{
         vm::{IN_DEPLOY, SETCODE_ONLY},
     },
     generic_vm::vm_executor::GenericVM,
-    state::FuzzState,
+    state::{FuzzState, HasCaller},
 };
 
 extern crate crypto;
@@ -35,7 +35,9 @@ use tracing::{debug, error, info};
 
 use self::crypto::{digest::Digest, sha3::Sha3};
 use super::{
+    abi::ABIAddressToInstanceMap,
     blaz::{is_bytecode_similar_lax, is_bytecode_similar_strict_ranking},
+    corpus_initializer::{EnvMetadata, INITIAL_BALANCE},
     host::FuzzHost,
     input::ConciseEVMInput,
     middlewares::cheatcode::{Cheatcode, CHEATCODE_ADDRESS},
@@ -721,7 +723,12 @@ impl ContractLoader {
         }
     }
 
-    pub fn from_setup(offchain_artifacts: &Vec<OffChainArtifact>, setup_file: String, work_dir: String) -> Self {
+    pub fn from_setup(
+        offchain_artifacts: &Vec<OffChainArtifact>,
+        setup_file: String,
+        work_dir: String,
+        etherscan_api_key: &str,
+    ) -> Self {
         let mut contracts: Vec<ContractInfo> = vec![];
         let mut abis: Vec<ABIInfo> = vec![];
 
@@ -787,6 +794,7 @@ impl ContractLoader {
                     setup_data = Some(Self::call_setup(
                         contract_artifact.deploy_bytecode.clone(),
                         work_dir.clone(),
+                        etherscan_api_key,
                     ));
                     break 'artifacts;
                 }
@@ -836,6 +844,7 @@ impl ContractLoader {
     fn get_vm_with_cheatcode(
         deployer: EVMAddress,
         work_dir: String,
+        etherscan_api_key: &str,
     ) -> (
         EVMExecutor<EVMState, ConciseEVMInput, StdScheduler<EVMFuzzState>>,
         EVMFuzzState,
@@ -850,18 +859,39 @@ impl ContractLoader {
             Bytecode::new_raw(Bytes::from(vec![0xfd, 0x00])),
             &mut state,
         );
-        executor.host.add_middlewares(Rc::new(RefCell::new(Cheatcode::new())));
+        executor
+            .host
+            .add_middlewares(Rc::new(RefCell::new(Cheatcode::new(etherscan_api_key))));
+
+        // Initialize state
+        state
+            .metadata_map_mut()
+            .insert::<ABIAddressToInstanceMap>(ABIAddressToInstanceMap::new());
+        state.metadata_map_mut().insert::<EnvMetadata>(EnvMetadata::default());
+        let default_callers = HashSet::from([
+            fixed_address("8EF508Aca04B32Ff3ba5003177cb18BfA6Cd79dd"),
+            fixed_address("35c9dfd76bf02107ff4f7128Bd69716612d31dDb"),
+            // fixed_address("5E6B78f0748ACd4Fb4868dF6eCcfE41398aE09cb"),
+        ]);
+        for caller in default_callers {
+            state.add_caller(&caller);
+            executor
+                .host
+                .evmstate
+                .set_balance(caller, EVMU256::from(INITIAL_BALANCE));
+        }
+
         (executor, state)
     }
 
     /// Deploy the contract and invoke "setUp()", returns the code, state, and
     /// environment after deployment. Foundry VM Cheatcodes and Hardhat
     /// consoles are enabled here.
-    fn call_setup(deploy_code: Bytes, work_dir: String) -> SetupData {
+    fn call_setup(deploy_code: Bytes, work_dir: String, etherscan_api_key: &str) -> SetupData {
         let deployer = EVMAddress::from_str(FOUNDRY_DEPLOYER).unwrap();
         let deployed_addr = EVMAddress::from_str(FOUNDRY_SETUP_ADDR).unwrap();
 
-        let (mut evm_executor, mut state) = Self::get_vm_with_cheatcode(deployer, work_dir);
+        let (mut evm_executor, mut state) = Self::get_vm_with_cheatcode(deployer, work_dir, etherscan_api_key);
 
         // deploy contract
         unsafe {

src/evm/host.rs

@@ -844,7 +844,6 @@ where
         &mut self,
         res: &(InstructionResult, Gas, Bytes),
     ) -> Option<(InstructionResult, Gas, Bytes)> {
-        debug!("check_assert_result: {:?}, assert_msg: {:?}", res, self.assert_msg);
         if let Some(ref msg) = self.assert_msg {
             return Some((InstructionResult::Revert, res.1, msg.abi_encode().into()));
         }
@@ -854,11 +853,6 @@ where
 
     /// Check expected reverts
     fn check_expected_revert(&mut self, res: (InstructionResult, Gas, Bytes)) -> (InstructionResult, Gas, Bytes) {
-        debug!(
-            "check_expected_revert: {:?}, depth: {}, expected_revert: {:?}",
-            res, self.call_depth, self.expected_revert
-        );
-
         // Check if we should check for reverts
         if self.expected_revert.is_none() {
             return res;
@@ -909,11 +903,6 @@ where
         call: &CallInputs,
         res: (InstructionResult, Gas, Bytes),
     ) -> (InstructionResult, Gas, Bytes) {
-        debug!(
-            "check_expected_emits: {:?}, depth: {}, expected_emits: {:?}",
-            res, self.call_depth, self.expected_emits
-        );
-
         let should_check_emits = self
             .expected_emits
             .iter()
@@ -940,11 +929,6 @@ where
 
     /// Check expected calls
     fn check_expected_calls(&mut self, res: (InstructionResult, Gas, Bytes)) -> (InstructionResult, Gas, Bytes) {
-        debug!(
-            "check_expected_calls: {:?}, depth: {}, expected_calls: {:?}",
-            res, self.call_depth, self.expected_calls
-        );
-
         // Only check expected calls at the root call
         if self.call_depth > 0 {
             return res;

src/evm/middlewares/cheatcode/fork.rs

@@ -55,7 +55,9 @@ where
             Chain::from_str(url_or_alias).ok()?
         };
         let block_number = block.map(|b| b.as_limbs()[0]).unwrap_or_default();
-        let onchain = OnChainConfig::new(chain, block_number);
+        let mut onchain = OnChainConfig::new(chain, block_number);
+        onchain.etherscan_api_key = self.etherscan_api_key.clone();
+
         let storage_fetching = StorageFetchingMode::OneByOne;
         tracing::debug!("createSelectFork(\"{url_or_alias}\", {block_number}), {onchain:?})");
 

src/evm/middlewares/cheatcode/mod.rs

@@ -49,6 +49,8 @@ pub struct Cheatcode<SC> {
     accesses: Option<RecordAccess>,
     /// Recorded logs
     recorded_logs: Option<Vec<Vm::Log>>,
+    /// Etherscan API key
+    etherscan_api_key: Vec<String>,
 
     _phantom: PhantomData<SC>,
 }
@@ -115,10 +117,11 @@ impl<SC> Cheatcode<SC>
 where
     SC: Scheduler<State = EVMFuzzState> + Clone,
 {
-    pub fn new() -> Self {
+    pub fn new(etherscan_api_key: &str) -> Self {
         Self {
             accesses: None,
             recorded_logs: None,
+            etherscan_api_key: etherscan_api_key.split(',').map(|s| s.to_string()).collect(),
             _phantom: PhantomData,
         }
     }
@@ -687,7 +690,7 @@ mod tests {
             std::fs::create_dir(path).unwrap();
         }
         let mut fuzz_host = FuzzHost::new(StdScheduler::new(), "work_dir".to_string());
-        fuzz_host.add_middlewares(Rc::new(RefCell::new(Cheatcode::new())));
+        fuzz_host.add_middlewares(Rc::new(RefCell::new(Cheatcode::new(""))));
         fuzz_host.set_code(
             CHEATCODE_ADDRESS,
             Bytecode::new_raw(Bytes::from(vec![0xfd, 0x00])),

src/evm/mod.rs

@@ -661,6 +661,7 @@ pub fn evm_main(mut args: EvmArgs) {
             &offchain_artifacts.expect("offchain artifacts is required for config target type"),
             args.setup_file,
             args.work_dir.clone(),
+            &etherscan_api_key,
         ),
         EVMTargetType::Address => {
             if onchain.is_none() {
@@ -751,6 +752,7 @@ pub fn evm_main(mut args: EvmArgs) {
         #[cfg(feature = "use_presets")]
         preset_file_path: args.preset_file_path,
         load_corpus: args.load_corpus,
+        etherscan_api_key,
     };
 
     let mut abis_map: HashMap<String, Vec<Vec<serde_json::Value>>> = HashMap::new();

src/evm/onchain/endpoints.rs

@@ -60,25 +60,25 @@ impl FromStr for Chain {
     type Err = ();
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
-            "ETH" | "eth" => Ok(Self::ETH),
-            "GOERLI" | "goerli" => Ok(Self::GOERLI),
-            "SEPOLIA" | "sepolia" => Ok(Self::SEPOLIA),
-            "BSC" | "bsc" => Ok(Self::BSC),
-            "CHAPEL" | "chapel" => Ok(Self::CHAPEL),
-            "POLYGON" | "polygon" => Ok(Self::POLYGON),
-            "MUMBAI" | "mumbai" => Ok(Self::MUMBAI),
-            "FANTOM" | "fantom" => Ok(Self::FANTOM),
-            "AVALANCHE" | "avalanche" => Ok(Self::AVALANCHE),
-            "OPTIMISM" | "optimism" => Ok(Self::OPTIMISM),
-            "ARBITRUM" | "arbitrum" => Ok(Self::ARBITRUM),
-            "GNOSIS" | "gnosis" => Ok(Self::GNOSIS),
-            "BASE" | "base" => Ok(Self::BASE),
-            "CELO" | "celo" => Ok(Self::CELO),
-            "ZKEVM" | "zkevm" => Ok(Self::ZKEVM),
-            "ZKEVM_TESTNET" | "zkevm_testnet" => Ok(Self::ZkevmTestnet),
-            "BLAST" | "blast" => Ok(Self::BLAST),
-            "LOCAL" | "local" => Ok(Self::LOCAL),
+        match s.to_lowercase().as_str() {
+            "eth" | "mainnet" => Ok(Self::ETH),
+            "goerli" => Ok(Self::GOERLI),
+            "sepolia" => Ok(Self::SEPOLIA),
+            "bsc" => Ok(Self::BSC),
+            "chapel" => Ok(Self::CHAPEL),
+            "polygon" => Ok(Self::POLYGON),
+            "mumbai" => Ok(Self::MUMBAI),
+            "fantom" => Ok(Self::FANTOM),
+            "avalanche" => Ok(Self::AVALANCHE),
+            "optimism" => Ok(Self::OPTIMISM),
+            "arbitrum" => Ok(Self::ARBITRUM),
+            "gnosis" => Ok(Self::GNOSIS),
+            "base" => Ok(Self::BASE),
+            "celo" => Ok(Self::CELO),
+            "zkevm" => Ok(Self::ZKEVM),
+            "zkevm_testnet" => Ok(Self::ZkevmTestnet),
+            "blast" => Ok(Self::BLAST),
+            "local" => Ok(Self::LOCAL),
             _ => Err(()),
         }
     }

src/fuzzers/evm_fuzzer.rs

@@ -110,7 +110,7 @@ pub fn evm_fuzzer(
     // **Note**: cheatcode should be the first middleware because it consumes the
     // step if it is a call to cheatcode_address, and this step should not be
     // visible to other middlewares.
-    fuzz_host.add_middlewares(Rc::new(RefCell::new(Cheatcode::new())));
+    fuzz_host.add_middlewares(Rc::new(RefCell::new(Cheatcode::new(&config.etherscan_api_key))));
 
     let onchain_middleware = match config.onchain.clone() {
         Some(onchain) => {

tests/evm_manual/foundry1/foundry.toml

@@ -3,4 +3,6 @@ src = "src"
 out = "out"
 libs = ["lib"]
 
+rpc_endpoints = { mainnet = "https://rpc.ankr.com/eth", eth = "https://rpc.ankr.com/eth", optimism = "https://rpc.ankr.com/optimism", fantom = "https://rpc.ankr.com/fantom", arbitrum = "https://rpc.ankr.com/arbitrum", bsc = "https://rpc.ankr.com/bsc", moonriver = "https://moonriver.public.blastapi.io", gnosis = "https://rpc.ankr.com/gnosis", avax = "https://rpc.ankr.com/avalanche", polygon = "https://rpc.ankr.com/polygon", goerli = "https://rpc.ankr.com/eth_goerli", era = "https://mainnet.era.zksync.io", zksync = "https://mainnet.era.zksync.io" }
+
 # See more config options https://github.com/foundry-rs/foundry/blob/master/crates/config/README.md#all-options

tests/evm_manual/foundry1/test/StaxExploit.t.sol

@@ -0,0 +1,34 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import {Test} from "forge-std/Test.sol";
+
+interface IStaxLP {
+    function balanceOf(address) external returns (uint256);
+    function transfer(address, uint256) external returns (bool);
+    function approve(address, uint256) external returns (bool);
+}
+
+contract StaxExploitTest is Test {
+    uint256 private initialAmount;
+    IStaxLP private StaxLP = IStaxLP(0xBcB8b7FC9197fEDa75C101fA69d3211b5a30dCD9);
+    address private StaxLPStaking = 0xd2869042E12a3506100af1D192b5b04D65137941;
+    address private tokenHolder = address(0xeCb456EA5365865EbAb8a2661B0c503410e9B347);
+
+    function setUp() public {
+        // vm.createSelectFork("http://bsc.internal.fuzz.land", "latest");
+        vm.createSelectFork("mainnet", 15725066);
+
+        targetContract(address(StaxLP));
+        targetContract(address(StaxLPStaking));
+
+        initialAmount = StaxLP.balanceOf(tokenHolder);
+        vm.prank(tokenHolder);
+        StaxLP.transfer(address(this), initialAmount);
+        StaxLP.approve(address(StaxLPStaking), type(uint256).max);
+    }
+
+    function invariant_1() public {
+        assertEq(StaxLP.balanceOf(address(this)), initialAmount + 1);
+    }
+}