generate foundry using real swap path

src/evm/abi.rs

@@ -44,6 +44,10 @@ static mut CONCOLIC_COUNTER: u64 = 0;
 
 /// Convert a vector of bytes to hex string
 fn vec_to_hex(v: &Vec<u8>) -> String {
+    if v.is_empty() {
+        return String::from("\"\"");
+    }
+
     let mut s = String::new();
     s.push_str("0x");
     for i in v {

src/evm/corpus_initializer.rs

@@ -542,6 +542,7 @@ where
             direct_data: Default::default(),
             randomness: vec![0],
             repeat: 1,
+            swap_data: HashMap::new(),
         };
         add_input_to_corpus!(self.state, &mut self.scheduler, input.clone(), artifacts);
         #[cfg(feature = "print_txn_corpus")]

src/evm/host.rs

@@ -1361,6 +1361,7 @@ where
                             direct_data: Default::default(),
                             randomness: vec![0],
                             repeat: 1,
+                            swap_data: HashMap::new(),
                         };
                         add_corpus(self, state, &input);
                     });

src/evm/input.rs

@@ -1,4 +1,4 @@
-use std::{cell::RefCell, fmt::Debug, ops::Deref, rc::Rc};
+use std::{cell::RefCell, collections::HashMap, fmt::Debug, ops::Deref, rc::Rc};
 
 use bytes::Bytes;
 use colored::{ColoredString, Colorize};
@@ -22,7 +22,10 @@ use crate::{
         types::{checksum, EVMAddress, EVMStagedVMState, EVMU256, EVMU512},
         vm::EVMState,
     },
-    generic_vm::{vm_executor::ExecutionResult, vm_state::VMStateT},
+    generic_vm::{
+        vm_executor::ExecutionResult,
+        vm_state::{SwapInfo, VMStateT},
+    },
     input::{ConciseSerde, SolutionTx, VMInputT},
     mutation_utils::byte_mutator,
     state::{HasCaller, HasItyState},
@@ -88,6 +91,8 @@ pub trait EVMInputT {
     fn set_liquidation_percent(&mut self, v: u8);
 
     fn get_repeat(&self) -> usize;
+
+    fn get_swap_data(&self) -> HashMap<String, SwapInfo>;
 }
 
 /// EVM Input
@@ -137,6 +142,9 @@ pub struct EVMInput {
 
     /// Execute the transaction multiple times
     pub repeat: usize,
+
+    /// Swap data
+    pub swap_data: HashMap<String, SwapInfo>,
 }
 
 /// EVM Input Minimum for Deserializing
@@ -183,6 +191,9 @@ pub struct ConciseEVMInput {
 
     /// return data
     pub return_data: Option<Vec<u8>>,
+
+    /// Swap data
+    pub swap_data: HashMap<String, SwapInfo>,
 }
 
 /// EVM Input Minimum for Deserializing with human readable ABI
@@ -248,6 +259,8 @@ impl ConciseEVMInput {
             v => Some(v),
         };
 
+        let swap_data = execution_result.new_state.state.get_swap_data();
+
         Self {
             input_type: input.get_input_type(),
             caller: input.get_caller(),
@@ -271,6 +284,7 @@ impl ConciseEVMInput {
                 None => u32::MAX,
             },
             return_data,
+            swap_data,
         }
     }
 
@@ -298,6 +312,7 @@ impl ConciseEVMInput {
             layer: input.get_state().get_post_execution_len(),
             call_leak,
             return_data: None,
+            swap_data: input.get_swap_data(),
         }
     }
 
@@ -324,6 +339,7 @@ impl ConciseEVMInput {
                 direct_data: Bytes::from(hex::decode(&self.direct_data).unwrap_or_default()),
                 randomness: self.randomness.clone(),
                 repeat: self.repeat,
+                swap_data: self.swap_data.clone(),
             },
             self.call_leak,
         )
@@ -360,7 +376,13 @@ impl ConciseEVMInput {
             Some(ref d) => self.as_abi_call(d.to_colored_string()),
             None => match self.input_type {
                 EVMInputTy::ABI | EVMInputTy::ArbitraryCallBoundedAddr => self.as_transfer(),
-                EVMInputTy::Borrow => self.as_borrow(),
+                EVMInputTy::Borrow => {
+                    if self.swap_data.contains_key("deposit") {
+                        self.as_deposit()
+                    } else {
+                        self.as_borrow()
+                    }
+                }
                 EVMInputTy::Liquidate => None,
             },
         }
@@ -439,17 +461,32 @@ impl ConciseEVMInput {
         ))
     }
 
+    #[allow(dead_code)]
+    #[inline]
+    fn as_deposit(&self) -> Option<String> {
+        Some(format!(
+            "WETH.{}{}();",
+            self.colored_fn_name("deposit"),
+            self.colored_value(),
+        ))
+    }
+
     #[inline]
     fn append_liquidation(&self, indent: String, call: String) -> String {
         if self.liquidation_percent == 0 || unsafe { !CAN_LIQUIDATE } {
             return call;
         }
-
-        let liq_call = format!(
-            "{}.{}(100% Balance, 0, path:(* → WETH), address(this), block.timestamp);",
-            colored_address("Router"),
-            self.colored_fn_name("swapExactTokensForETH"),
-        );
+        let liq_call = if self.swap_data.contains_key("withdraw") {
+            format!("WETH.{}(100% Balance);", self.colored_fn_name("withdraw"))
+        } else if self.swap_data.contains_key("sell") {
+            format!(
+                "{}.{}(100% Balance, 0, path:(* → WETH), address(this), block.timestamp);",
+                colored_address("Router"),
+                self.colored_fn_name("swapExactTokensForETH"),
+            )
+        } else {
+            return call;
+        };
 
         let mut liq = indent.clone();
         liq.push_str(format!("├─[{}] {}", self.layer + 1, liq_call).as_str());
@@ -541,6 +578,10 @@ impl SolutionTx for ConciseEVMInput {
     fn liq_percent(&self) -> u8 {
         self.liquidation_percent
     }
+
+    fn swap_data(&self) -> HashMap<String, SwapInfo> {
+        self.swap_data.clone()
+    }
 }
 
 impl HasLen for EVMInput {
@@ -628,6 +669,10 @@ impl EVMInputT for EVMInput {
     fn get_repeat(&self) -> usize {
         self.repeat
     }
+
+    fn get_swap_data(&self) -> HashMap<String, SwapInfo> {
+        self.swap_data.clone()
+    }
 }
 
 ///

src/evm/middlewares/cheatcode.rs

@@ -1261,6 +1261,7 @@ mod tests {
                     input_type: EVMInputTy::ABI,
                     randomness: vec![],
                     repeat: 1,
+                    swap_data: HashMap::new(),
                 };
                 let mut state = FuzzState::new(0);
                 // deposit some ETH to the test contract

src/evm/middlewares/sha3_bypass.rs

@@ -485,6 +485,7 @@ mod tests {
             input_type: EVMInputTy::ABI,
             randomness: vec![],
             repeat: 1,
+            swap_data: HashMap::new(),
         };
 
         let res = evm_executor.execute(&input, &mut state);

src/evm/mod.rs

@@ -747,8 +747,6 @@ pub fn evm_main(args: EvmArgs) {
 
     let json_str = serde_json::to_string(&abis_map).expect("Failed to serialize ABI map to JSON");
 
-    let work_dir = args.work_dir.clone();
-
     let abis_json = format!("{}/abis.json", args.work_dir.clone().as_str());
 
     let mut file = OpenOptions::new()

src/evm/onchain/endpoints.rs

@@ -1121,7 +1121,7 @@ impl OnChainConfig {
         for (kth, route) in (&mut routes).iter_mut().enumerate() {
             let mut low_liquidity = false;
             for hop in route {
-                low_liquidity |= (!self.add_reserve_info(hop));
+                low_liquidity |= !self.add_reserve_info(hop);
             }
             if !low_liquidity {
                 routes_without_low_liquidity_idx.push(kth);

src/evm/onchain/flashloan.rs

@@ -100,6 +100,7 @@ where
                 direct_data: Default::default(),
                 randomness: vec![0],
                 repeat: 1,
+                swap_data: HashMap::new(),
             }
         }
         .as_any()

src/evm/onchain/mod.rs

@@ -473,6 +473,7 @@ impl OnChain {
                     direct_data: Default::default(),
                     randomness: vec![0],
                     repeat: 1,
+                    swap_data: HashMap::new(),
                 };
                 add_corpus(host, state, &input);
             });

src/evm/oracles/erc20.rs

@@ -86,6 +86,7 @@ impl
             let _path_idx = ctx.input.get_randomness()[0] as usize;
 
             let mut liquidation_txs = vec![];
+            let mut swap_infos = vec![];
 
             for (caller, _token_info, _amount) in liquidations_earned {
                 let txs = _token_info.borrow().sell(
@@ -99,9 +100,20 @@ impl
                     txs.iter()
                         .map(|(addr, abi, _)| (caller, *addr, Bytes::from(abi.get_bytes()))),
                 );
+
+                if let Some(swap_info) = txs.last() {
+                    swap_infos.push(swap_info.clone());
+                }
+            }
+
+            let (_out, mut state) = ctx.call_post_batch_dyn(&liquidation_txs);
+
+            // Record the swap info for generating foundry in the future.
+            state.swap_data = ctx.fuzz_state.get_execution_result().new_state.state.swap_data.clone();
+            for (target, mut abi, _) in swap_infos {
+                state.swap_data.push(&target, &mut abi);
             }
 
-            let (_out, state) = ctx.call_post_batch_dyn(&liquidation_txs);
             ctx.fuzz_state.get_execution_result_mut().new_state.state = state;
         }
 

src/evm/solution/foundry_test.hbs

@@ -24,38 +24,72 @@ contract {{contract_name}} is Test {
     }
 
     function test() public {
-    {{#if include_interface}}
+    {{#if router}}
         address router = {{router}};
-        address weth = {{weth}};
-        {{/if}}{{#each trace}}
+
+    {{/if}}
+{{#each trace}}
         vm.prank({{caller}});
-        {{#with this}}{{#if raw_code}}{{raw_code}}{{! Raw Code }}
-        {{else}}{{#if is_deposit}}vm.deal({{caller}}, {{value}});
-        weth.call{value: {{value}}}(abi.encodeWithSignature("deposit()", {{value}}));{{! Deposit weth }}
-        {{else}}{{#if is_borrow}}address[] memory path{{borrow_idx}} = new address[](2);{{! Borrow token }}
-        path{{borrow_idx}}[0] = weth;
-        path{{borrow_idx}}[1] = {{contract}};
+{{#with this}}
+    {{#if raw_code}}
+        {{raw_code}}
+    {{else}}
+    {{#if (is_deposit buy_type)}}
+        vm.deal({{caller}}, {{value}});
+        {{#with (lookup swap_data "deposit")~}}{{target}}.call{value: {{../value}}}(abi.encodeWithSignature("deposit()"));{{/with}}
+    {{else}}
+    {{#if (is_buy buy_type)}}
+        address[] memory path{{borrow_idx}} = new address[]();
+    {{#with (lookup swap_data "buy")~}}
+        {{#each path}}
+        path{{../../borrow_idx}}[{{@index}}] = {{{this}}};
+        {{/each}}
+    {{/with}}
         vm.deal({{caller}}, {{value}});
         IUniswapV2Router(router).swapExactETHForTokensSupportingFeeOnTransferTokens{
             value: {{value}}
         }(0, path{{borrow_idx}}, address(this), block.timestamp);
-        {{else}}{{#if value}}vm.deal({{caller}}, {{value}});{{/if}}
-        {{#if fn_signature}}{{contract}}.call{{#if value}}{value: {{value}}}{{/if}}(abi.encodeWithSignature({{! Call with signature }}
-            "{{fn_signature}}"{{#if fn_args}},{{fn_args}}{{/if}}
-        )); {{else}}{{contract}}.call{{#if value}}{value: {{value}}}{{/if}}(abi.encodeWithSelector({{! Call with selector }}
-            {{fn_selector}}{{#if fn_args}},{{fn_args}}{{/if}}
-        ));{{/if}}{{/if}}{{/if}}{{/if}}{{#if liq_percent}}{{! Liquidation }}
+    {{else}}
+    {{#if value}}
+        vm.deal({{caller}}, {{value}});
+    {{/if}}
+    {{#if fn_signature}}
+        {{contract}}.call{{#if value}}{value: {{value}}}{{/if}}(abi.encodeWithSignature(
+            "{{fn_signature}}"{{#if fn_args}}, {{{fn_args}}}{{/if}}
+        ));
+    {{else}}
+        {{contract}}.call{{#if value}}{value: {{value}}}{{/if}}(abi.encodeWithSelector(
+            {{fn_selector}}{{#if fn_args}},{{{fn_args}}}{{/if}}
+        ));
+    {{/if}}
+    {{/if}}
+    {{/if}}
+    {{/if}}
+    {{#if (is_withdraw sell_type)}}
+        vm.startPrank({{caller}});
+        uint256 amount{{balance_idx}} = IERC20({{contract}}).balanceOf(address(this));
+        {{#with (lookup swap_data "withdraw")~}}{{target}}.call(abi.encodeWithSignature("withdraw(uint256)", amount{{../balance_idx}}));{{/with}}
+        vm.stopPrank();
+    {{else}}
+    {{#if (is_sell sell_type)}}
         vm.startPrank({{caller}});
-        uint256 amount{{liq_idx}} = IERC20({{contract}}).balanceOf(address(this));
-        IERC20({{contract}}).approve(router, amount{{liq_idx}});
-        address[] memory liq_path{{liq_idx}} = new address[](2);
-        liq_path{{liq_idx}}[0] = {{contract}};
-        liq_path{{liq_idx}}[1] = address(weth);
-        vm.deal({{caller}}, amount{{liq_idx}});
+        uint256 amount{{balance_idx}} = IERC20({{contract}}).balanceOf(address(this));
+        IERC20({{contract}}).approve(router, amount{{balance_idx}});
+        address[] memory liq_path{{balance_idx}} = new address[]();
+    {{#with (lookup swap_data "sell")~}}
+        {{#each path}}
+        liq_path{{../../balance_idx}}[{{@index}}] = {{{this}}};
+        {{/each}}
+    {{/with}}
+        vm.deal({{caller}}, amount{{balance_idx}});
         IUniswapV2Router(router).swapExactTokensForETHSupportingFeeOnTransferTokens(
-            amount{{liq_idx}}, 0, liq_path{{liq_idx}}, address(this), block.timestamp
+            amount{{balance_idx}}, 0, liq_path{{balance_idx}}, address(this), block.timestamp
         );
-        vm.stopPrank();{{/if}}{{/with}}{{/each}}
+        vm.stopPrank();
+    {{/if}}
+    {{/if}}
+{{/with}}
+{{/each}}
     }
 
 {{#if stepping_with_return}}