Feat: offchain get token context (#463)

* WIP: get_token_context add OffChainConfig * WIP: test offchain get_token_context * WIP: offchain get_token_context * WIP: offchain get_token_context * cargo fmt * test_logger ignores error * add v2_pairs test contracts * forge install: openzeppelin-contracts v5.0.2 * add v2_pairs test contracts * deploy preset_weth * feat: offchain get_token_context
fuzzland · Apr 10, 2024 · 6d6629f · 6d6629f
1 parent c878633
commit 6d6629f
Show file tree

Hide file tree

Showing 18 changed files with 1,185 additions and 150 deletions.
diff --git a/.gitmodules b/.gitmodules
@@ -1,3 +1,6 @@
 [submodule "tests/evm_manual/foundry1/lib/forge-std"]
 	path = tests/evm_manual/foundry1/lib/forge-std
 	url = https://github.com/foundry-rs/forge-std
+[submodule "tests/evm_manual/foundry1/lib/openzeppelin-contracts"]
+	path = tests/evm_manual/foundry1/lib/openzeppelin-contracts
+	url = https://github.com/OpenZeppelin/openzeppelin-contracts
diff --git a/src/evm/contract_utils.rs b/src/evm/contract_utils.rs
@@ -22,6 +22,7 @@ use crate::{
         tokens::constant_pair::ConstantPairMetadata,
         types::{fixed_address, generate_random_address, EVMAddress, EVMFuzzState},
         vm::{IN_DEPLOY, SETCODE_ONLY},
+        PRESET_WETH,
     },
     generic_vm::vm_executor::GenericVM,
     state::{FuzzState, HasCaller},
@@ -897,6 +898,13 @@ impl ContractLoader {
         unsafe {
             SETCODE_ONLY = true;
         }
+        // Deploy preset WETH
+        let weth_addr = EVMAddress::from_str(PRESET_WETH).unwrap();
+        let hex_code = include_str!("../../tests/presets/v2_pair/WETH9.bytecode").trim();
+        let weth_code = Bytes::from(hex::decode(hex_code).unwrap());
+        let deployed_weth = evm_executor.deploy(Bytecode::new_raw(weth_code), None, weth_addr, &mut state);
+        assert!(deployed_weth.is_some(), "failed to deploy WETH");
+
         let addr = evm_executor.deploy(
             Bytecode::new_raw(deploy_code),
             None,
@@ -920,7 +928,7 @@ impl ContractLoader {
         );
 
         if !res[0].1 {
-            info!("setUp() failed: {:?}", res[0].0);
+            error!("setUp() failed: {:?}", res[0].0);
         }
 
         // now get Foundry invariant test config by calling

diff --git a/src/evm/mod.rs b/src/evm/mod.rs
@@ -63,6 +63,8 @@ use crate::{
     state::FuzzState,
 };
 
+pub const PRESET_WETH: &str = "0x4200000000000000000000000000000000000006";
+
 pub fn parse_constructor_args_string(input: String) -> HashMap<String, Vec<String>> {
     let mut map = HashMap::new();
 

diff --git a/src/evm/onchain/endpoints.rs b/src/evm/onchain/endpoints.rs
@@ -18,8 +18,9 @@ use revm_interpreter::analysis::to_analysed;
 use revm_primitives::{Bytecode, B160};
 use serde::Deserialize;
 use serde_json::{json, Value};
-use tracing::{debug, error, info};
+use tracing::{debug, error, info, warn};
 
+use super::ChainConfig;
 use crate::{
     cache::{Cache, FileSystemCache},
     evm::{
@@ -321,6 +322,94 @@ impl Debug for OnChainConfig {
     }
 }
 
+impl ChainConfig for OnChainConfig {
+    fn get_pair(&mut self, token: &str, is_pegged: bool) -> Vec<PairData> {
+        let network = self.chain_name.clone();
+        let weth = self.get_weth();
+        self.get_pair(token, &network, is_pegged, weth)
+    }
+
+    fn fetch_reserve(&self, pair: &str) -> Option<(String, String)> {
+        self.fetch_reserve(pair)
+    }
+
+    fn get_contract_code_analyzed(&mut self, address: EVMAddress, force_cache: bool) -> Bytecode {
+        self.get_contract_code_analyzed(address, force_cache)
+    }
+
+    fn get_v3_fee(&mut self, address: EVMAddress) -> u32 {
+        self.get_v3_fee(address)
+    }
+
+    fn get_token_balance(&mut self, token: EVMAddress, address: EVMAddress) -> EVMU256 {
+        self.get_token_balance(token, address)
+    }
+
+    fn get_weth(&self) -> String {
+        let pegged_token = self.get_pegged_token();
+
+        match self.chain_name.as_str() {
+            "eth" => return pegged_token.get("WETH").unwrap().to_string(),
+            "bsc" => return pegged_token.get("WBNB").unwrap().to_string(),
+            "polygon" => return pegged_token.get("WMATIC").unwrap().to_string(),
+            "local" => return pegged_token.get("ZERO").unwrap().to_string(),
+            // "mumbai" => panic!("Not supported"),
+            _ => {
+                warn!("Unknown network");
+                "".to_string()
+            }
+        }
+    }
+
+    fn get_pegged_token(&self) -> HashMap<String, String> {
+        match self.chain_name.as_str() {
+            "eth" => [
+                ("WETH", "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2"),
+                ("USDC", "0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48"),
+                ("USDT", "0xdac17f958d2ee523a2206206994597c13d831ec7"),
+                ("DAI", "0x6b175474e89094c44da98b954eedeac495271d0f"),
+                ("WBTC", "0x2260fac5e5542a773aa44fbcfedf7c193bc2c599"),
+                ("WMATIC", "0x7d1afa7b718fb893db30a3abc0cfc608aacfebb0"),
+            ]
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.to_string()))
+            .collect(),
+            "bsc" => [
+                ("WBNB", "0xbb4cdb9cbd36b01bd1cbaebf2de08d9173bc095c"),
+                ("USDC", "0x8ac76a51cc950d9822d68b83fe1ad97b32cd580d"),
+                ("USDT", "0x55d398326f99059ff775485246999027b3197955"),
+                ("DAI", "0x1af3f329e8be154074d8769d1ffa4ee058b1dbc3"),
+                ("WBTC", "0x7130d2a12b9bcbfae4f2634d864a1ee1ce3ead9c"),
+                ("WETH", "0x2170ed0880ac9a755fd29b2688956bd959f933f8"),
+                ("BUSD", "0xe9e7cea3dedca5984780bafc599bd69add087d56"),
+                ("CAKE", "0x0e09fabb73bd3ade0a17ecc321fd13a19e81ce82"),
+            ]
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.to_string()))
+            .collect(),
+            "polygon" => [
+                ("WMATIC", "0x0d500b1d8e8ef31e21c99d1db9a6444d3adf1270"),
+                ("USDC", "0x2791bca1f2de4661ed88a30c99a7a9449aa84174"),
+                ("USDT", "0xc2132d05d31c914a87c6611c10748aeb04b58e8f"),
+                ("DAI", "0x8f3cf7ad23cd3cadbd9735aff958023239c6a063"),
+                ("WBTC", "0x1bfd67037b42cf73acf2047067bd4f2c47d9bfd6"),
+                ("WETH", "0x7ceb23fd6bc0add59e62ac25578270cff1b9f619"),
+            ]
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.to_string()))
+            .collect(),
+            "local" => [("ZERO", "0x0000000000000000000000000000000000000000")]
+                .iter()
+                .map(|(k, v)| (k.to_string(), v.to_string()))
+                .collect(),
+            _ => {
+                warn!("[Flashloan] Network is not supported");
+                HashMap::new()
+            }
+        }
+    }
+}
+
 impl OnChainConfig {
     pub fn new(chain: Chain, block_number: u64) -> Self {
         Self::new_raw(

diff --git a/src/evm/onchain/flashloan.rs b/src/evm/onchain/flashloan.rs
@@ -25,6 +25,7 @@ use revm_interpreter::Interpreter;
 use serde::{Deserialize, Serialize};
 use tracing::debug;
 
+use super::ChainConfig;
 use crate::{
     evm::{
         contract_utils::ABIConfig,
@@ -33,7 +34,6 @@ use crate::{
         input::{ConciseEVMInput, EVMInput, EVMInputT, EVMInputTy},
         middlewares::middleware::{Middleware, MiddlewareType},
         mutator::AccessPattern,
-        onchain::endpoints::OnChainConfig,
         oracles::erc20::IERC20OracleFlashloan,
         tokens::{uniswap::fetch_uniswap_path, TokenContext},
         types::{convert_u256_to_h160, EVMAddress, EVMFuzzState, EVMU256, EVMU512},
@@ -54,7 +54,7 @@ macro_rules! scale {
 pub struct Flashloan {
     use_contract_value: bool,
     known_addresses: HashSet<EVMAddress>,
-    endpoint: Option<OnChainConfig>,
+    chain_cfg: Option<Box<dyn ChainConfig>>,
     erc20_address: HashSet<EVMAddress>,
     pair_address: HashSet<EVMAddress>,
     pub unbound_tracker: HashMap<usize, HashSet<EVMAddress>>, // pc -> [address called]
@@ -118,13 +118,13 @@ where
 impl Flashloan {
     pub fn new(
         use_contract_value: bool,
-        endpoint: Option<OnChainConfig>,
+        chain_cfg: Option<Box<dyn ChainConfig>>,
         flashloan_oracle: Rc<RefCell<IERC20OracleFlashloan>>,
     ) -> Self {
         Self {
             use_contract_value,
             known_addresses: Default::default(),
-            endpoint,
+            chain_cfg,
             erc20_address: Default::default(),
             pair_address: Default::default(),
             unbound_tracker: Default::default(),
@@ -134,7 +134,7 @@ impl Flashloan {
     }
 
     fn get_token_context(&mut self, addr: EVMAddress) -> Option<TokenContext> {
-        self.endpoint.as_mut().map(|config| fetch_uniswap_path(config, addr))
+        self.chain_cfg.as_mut().map(|config| fetch_uniswap_path(config, addr))
     }
 
     pub fn on_contract_insertion(

diff --git a/src/evm/onchain/mod.rs b/src/evm/onchain/mod.rs
@@ -1,6 +1,7 @@
 pub mod abi_decompiler;
 pub mod endpoints;
 pub mod flashloan;
+pub mod offchain;
 
 use std::{
     cell::RefCell,
@@ -20,6 +21,7 @@ use revm_interpreter::{analysis::to_analysed, Interpreter};
 use revm_primitives::Bytecode;
 use tracing::debug;
 
+use self::endpoints::PairData;
 use super::{corpus_initializer::EnvMetadata, types::EVMFuzzState};
 use crate::{
     evm::{
@@ -50,6 +52,16 @@ pub static mut WHITELIST_ADDR: Option<HashSet<EVMAddress>> = None;
 
 const UNBOUND_THRESHOLD: usize = 30;
 
+pub trait ChainConfig {
+    fn get_pair(&mut self, token: &str, is_pegged: bool) -> Vec<PairData>;
+    fn fetch_reserve(&self, pair: &str) -> Option<(String, String)>;
+    fn get_contract_code_analyzed(&mut self, address: EVMAddress, force_cache: bool) -> Bytecode;
+    fn get_v3_fee(&mut self, address: EVMAddress) -> u32;
+    fn get_token_balance(&mut self, token: EVMAddress, address: EVMAddress) -> EVMU256;
+    fn get_weth(&self) -> String;
+    fn get_pegged_token(&self) -> HashMap<String, String>;
+}
+
 pub struct OnChain {
     pub loaded_data: HashSet<(EVMAddress, EVMU256)>,
     pub loaded_code: HashSet<EVMAddress>,
@@ -135,13 +147,12 @@ impl OnChain {
     }
 }
 
-pub fn keccak_hex(data: EVMU256) -> String {
+pub fn keccak256(input: &[u8]) -> EVMU256 {
     let mut hasher = Sha3::keccak256();
     let mut output = [0u8; 32];
-    let input: [u8; 32] = data.to_be_bytes();
-    hasher.input(input.as_ref());
+    hasher.input(input);
     hasher.result(&mut output);
-    hex::encode(output)
+    EVMU256::from_be_bytes(output)
 }
 
 impl<SC> Middleware<SC> for OnChain