Use trait for token context

src/evm/onchain/endpoints.rs

@@ -23,7 +23,7 @@ use crate::{
     cache::{Cache, FileSystemCache},
     evm::{
         types::{EVMAddress, EVMU256},
-        uniswap::{get_uniswap_info, PairContext, PathContext, TokenContext, UniswapProvider},
+        uniswap::{get_uniswap_info, PairContext, PathContext, UniswapProvider, UniswapTokenContext},
     },
 };
 
@@ -248,7 +248,7 @@ pub struct OnChainConfig {
     price_cache: HashMap<EVMAddress, Option<(u32, u32)>>,
     abi_cache: HashMap<EVMAddress, Option<String>>,
     storage_dump_cache: HashMap<EVMAddress, Option<Arc<HashMap<EVMU256, EVMU256>>>>,
-    uniswap_path_cache: HashMap<EVMAddress, TokenContext>,
+    uniswap_path_cache: HashMap<EVMAddress, UniswapTokenContext>,
     rpc_cache: FileSystemCache,
 }
 
@@ -708,14 +708,14 @@ impl OnChainConfig {
         slot_value
     }
 
-    pub fn fetch_uniswap_path(&mut self, network: &str, token_address: EVMAddress) -> TokenContext {
+    pub fn fetch_uniswap_path(&mut self, network: &str, token_address: EVMAddress) -> UniswapTokenContext {
         let token = format!("{:?}", token_address);
         let info: Info = self.find_path_subgraph(network, &token);
 
         let basic_info = info.basic_info;
         if basic_info.weth.is_empty() {
             warn!("failed to find weth address");
-            return TokenContext::default();
+            return UniswapTokenContext::default();
         }
         let weth = EVMAddress::from_str(&basic_info.weth).unwrap();
         let is_weth = basic_info.is_weth;
@@ -781,15 +781,15 @@ impl OnChainConfig {
             })
             .collect();
 
-        TokenContext {
+        UniswapTokenContext {
             swaps: paths_parsed,
             is_weth,
             weth_address: weth,
             address: token_address,
         }
     }
 
-    pub fn fetch_uniswap_path_cached(&mut self, token: EVMAddress) -> &TokenContext {
+    pub fn fetch_uniswap_path_cached(&mut self, token: EVMAddress) -> &UniswapTokenContext {
         if self.uniswap_path_cache.contains_key(&token) {
             return self.uniswap_path_cache.get(&token).unwrap();
         }

src/evm/onchain/flashloan.rs

@@ -28,7 +28,7 @@ use revm_interpreter::Interpreter;
 use serde::{Deserialize, Serialize};
 use tracing::debug;
 
-use crate::evm::{types::EVMFuzzState, uniswap::TokenContext};
+use crate::evm::{types::EVMFuzzState, uniswap::UniswapTokenContext};
 // Some components are used when `flashloan_v2` feature is not enabled
 #[allow(unused_imports)]
 use crate::{
@@ -166,7 +166,7 @@ impl Flashloan {
             .map(|price| Self::calculate_usd_value(price, amount))
     }
 
-    fn get_token_context(&mut self, addr: EVMAddress) -> Option<TokenContext> {
+    fn get_token_context(&mut self, addr: EVMAddress) -> Option<UniswapTokenContext> {
         self.endpoint
             .as_mut()
             .map(|endpoint| endpoint.fetch_uniswap_path_cached(addr).clone())
@@ -205,7 +205,7 @@ 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, token_ctx);
+                            oracle.unwrap().register_token(*addr, Rc::new(RefCell::new(token_ctx)));
                             self.erc20_address.insert(*addr);
                             is_erc20 = true;
                         } else {

src/evm/oracles/erc20.rs

@@ -11,7 +11,7 @@ use crate::{
         oracles::ERC20_BUG_IDX,
         producers::erc20::ERC20Producer,
         types::{EVMAddress, EVMFuzzState, EVMOracleCtx, EVMU256, EVMU512},
-        uniswap::{generate_uniswap_router_sell, TokenContext},
+        uniswap::TokenContextT,
         vm::EVMState,
     },
     oracle::Oracle,
@@ -20,7 +20,7 @@ use crate::{
 
 pub struct IERC20OracleFlashloan {
     pub balance_of: Vec<u8>,
-    pub known_tokens: HashMap<EVMAddress, TokenContext>,
+    pub known_tokens: HashMap<EVMAddress, Rc<RefCell<dyn TokenContextT<EVMFuzzState>>>>,
     pub known_pair_reserve_slot: HashMap<EVMAddress, EVMU256>,
     pub erc20_producer: Rc<RefCell<ERC20Producer>>,
 }
@@ -35,7 +35,7 @@ impl IERC20OracleFlashloan {
         }
     }
 
-    pub fn register_token(&mut self, token: EVMAddress, token_ctx: TokenContext) {
+    pub fn register_token(&mut self, token: EVMAddress, token_ctx: Rc<RefCell<dyn TokenContextT<EVMFuzzState>>>) {
         self.known_tokens.insert(token, token_ctx);
     }
 
@@ -79,27 +79,16 @@ impl
             let mut liquidation_txs = vec![];
 
             for (caller, _token_info, _amount) in liquidations_earned {
-                // let txs = _token_info.borrow().sell(
-                //     ctx.fuzz_state,
-                //     _amount,
-                //     ctx.fuzz_state.callers_pool[0],
-                //     ctx.input.get_randomness().as_slice(),
-                // );
-
-                let txs = generate_uniswap_router_sell(_token_info, _path_idx, _amount, ctx.fuzz_state.callers_pool[0]);
-                if txs.is_none() {
-                    continue;
-                }
-
-                // liquidation_txs.extend(
-                //     txs.iter()
-                //         .map(|(addr, abi, _)| (caller, *addr, Bytes::from(abi.get_bytes()))),
-                // );
+                let txs = _token_info.borrow().sell(
+                    ctx.fuzz_state,
+                    _amount,
+                    ctx.fuzz_state.callers_pool[0],
+                    ctx.input.get_randomness().as_slice(),
+                );
 
                 liquidation_txs.extend(
-                    txs.unwrap()
-                        .iter()
-                        .map(|(abi, _, addr)| (caller, *addr, Bytes::from(abi.get_bytes()))),
+                    txs.iter()
+                        .map(|(addr, abi, _)| (caller, *addr, Bytes::from(abi.get_bytes()))),
                 );
             }