Feature generate foundry test file (#261)

* add foundry-test-generator

* add foundry-test-generator

* ignore errors when generating foundry tests

* foundry_test_generator -> solution

* supporting borrow in foundry test

* test-generator support borrow

Cargo.toml

@@ -109,3 +109,6 @@ move-vm-types = { git = "https://github.com/fuzzland/ityfuzz-sui-fork.git", pack
 sui-move-natives-latest = { git = "https://github.com/fuzzland/ityfuzz-sui-fork.git", package = "sui-move-natives-latest", optional = true }
 sui-protocol-config = { git = "https://github.com/fuzzland/ityfuzz-sui-fork.git", package = "sui-protocol-config", optional = true }
 sui-types = { git = "https://github.com/fuzzland/ityfuzz-sui-fork.git", package = "sui-types", optional = true }
+
+# template engine
+handlebars = "4.4"

foundry_test.hbs

@@ -0,0 +1,84 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity ^0.8.13;
+
+import "forge-std/Test.sol";
+
+{{#if is_onchain}}
+// ityfuzz evm -o -t {{target}} -c {{chain}} --onchain-block-number {{block_number}} -f -i -p --onchain-etherscan-api-key ${{etherscan_keyname}}
+{{/if}}
+{{#unless is_onchain}}
+// ityfuzz evm -t '{{target}}' -f --panic-on-bug
+{{/unless}}
+/*
+
+😊😊 Found violations!
+
+
+{{{solution}}}
+ */
+
+{{#if is_borrow}}
+interface IUniswapV2Router {
+    function swapExactTokensForTokensSupportingFeeOnTransferTokens(
+        uint256 amountIn,
+        uint256 amountOutMin,
+        address[] calldata path,
+        address to,
+        uint256 deadline
+    ) external;
+    function swapExactETHForTokensSupportingFeeOnTransferTokens(
+        uint256 amountOutMin,
+        address[] calldata path,
+        address to,
+        uint256 deadline
+    ) external payable;
+    function swapExactTokensForETHSupportingFeeOnTransferTokens(
+        uint256 amountIn,
+        uint256 amountOutMin,
+        address[] calldata path,
+        address to,
+        uint256 deadline
+    ) external;
+}
+{{/if}}
+
+contract EGD is Test {
+    function setUp() public {
+        {{#if is_onchain}}
+        vm.createSelectFork("{{chain}}", {{block_number}});
+        {{/if}}
+    }
+
+    function test() public {
+        vm.prank(address(this), address(this));
+{{#each trace}}
+
+    {{#with this}}
+        {{#if is_borrow}}
+        address[] memory path{{borrow_idx}} = new address[](2);
+        path{{borrow_idx}}[0] = address({{weth}});
+        path{{borrow_idx}}[1] = address({{contract}});
+        IUniswapV2Router(address({{router}})).swapExactETHForTokensSupportingFeeOnTransferTokens{
+            value: {{value}}
+        }(0, path{{borrow_idx}}, address(this), block.timestamp);
+        {{#if liq_percent}}
+        // swap todo: liq_percent: {{liq_percent}}
+        {{/if}}
+        {{/if}}
+        {{#unless is_borrow}}
+        address({{contract}}).call{{#if value}}{value: {{value}}}{{/if}}(abi.encodeWithSelector(
+            {{fn_selector}}{{#if fn_args}},{{fn_args}}{{/if}}
+        ));
+        {{#if liq_percent}}
+        // swap todo: liq_percent: {{liq_percent}}
+        {{/if}}
+        {{/unless}}
+    {{/with}}
+{{/each}}
+    }
+
+{{#if stepping_with_return}}
+    // Stepping with return
+    receive() external payable {}
+{{/if}}
+}

src/evm/input.rs

@@ -2,7 +2,7 @@ use crate::evm::abi::{AEmpty, AUnknown, BoxedABI};
 use crate::evm::mutator::AccessPattern;
 use crate::evm::types::{EVMAddress, EVMExecutionResult, EVMStagedVMState, EVMU256, EVMU512};
 use crate::evm::vm::EVMState;
-use crate::input::{ConciseSerde, VMInputT};
+use crate::input::{ConciseSerde, VMInputT, SolutionTx};
 use crate::mutation_utils::byte_mutator;
 use crate::state::{HasCaller, HasItyState};
 use crate::state_input::StagedVMState;
@@ -23,6 +23,7 @@ use std::fmt::Debug;
 use std::ops::Deref;
 use std::rc::Rc;
 
+
 /// EVM Input Types
 #[derive(Serialize, Deserialize, Clone, PartialEq, Eq, Debug)]
 pub enum EVMInputTy {
@@ -310,6 +311,60 @@ impl ConciseEVMInput {
     }
 }
 
+impl SolutionTx for ConciseEVMInput {
+    #[cfg(feature = "flashloan_v2")]
+    fn is_borrow(&self) -> bool {
+        self.input_type == EVMInputTy::Borrow
+    }
+
+    fn caller(&self) -> String {
+        format!("0x{}", hex::encode(self.caller))
+    }
+
+    fn contract(&self) -> String {
+        format!("0x{}", hex::encode(self.contract))
+    }
+
+    fn value(&self) -> String {
+        self.txn_value
+            .map(|v| format!("0x{}", hex::encode(v.to_be_bytes_vec())))
+            .unwrap_or(String::new())
+    }
+
+    #[cfg(not(feature = "debug"))]
+    fn fn_selector(&self) -> String {
+        match self.data {
+            Some(ref d) => format!("0x{}", hex::encode(d.function)),
+            None => "TODO".to_string(),
+        }
+    }
+
+    #[cfg(feature = "debug")]
+    fn fn_selector(&self) -> String {
+        "".to_string()
+    }
+
+    #[cfg(not(feature = "debug"))]
+    fn fn_args(&self) -> String {
+        match self.data {
+            Some(ref d) => {
+                d.get().to_string().trim_matches(|c| c == '(' || c == ')').to_string()
+            },
+            None => "TODO".to_string(),
+        }
+    }
+
+    #[cfg(feature = "debug")]
+    fn fn_args(&self) -> String {
+        "".to_string()
+    }
+
+    #[cfg(feature = "flashloan_v2")]
+    fn liq_percent(&self) -> u8 {
+        self.liquidation_percent
+    }
+}
+
 impl HasLen for EVMInput {
     /// Get the length of the ABI encoded input
     fn len(&self) -> usize {

src/evm/mod.rs

@@ -21,6 +21,7 @@ pub mod srcmap;
 pub mod types;
 pub mod uniswap;
 pub mod vm;
+pub mod solution;
 
 use crate::fuzzers::evm_fuzzer::evm_fuzzer;
 use crate::oracle::{Oracle, Producer};
@@ -274,6 +275,9 @@ enum EVMTargetType {
 }
 
 pub fn evm_main(args: EvmArgs) {
+    let target = args.target.clone();
+    let work_dir = args.work_dir.clone();
+
     let mut target_type: EVMTargetType = match args.target_type {
         Some(v) => match v.as_str() {
             "glob" => EVMTargetType::Glob,
@@ -314,6 +318,7 @@ pub fn evm_main(args: EvmArgs) {
         None
     };
 
+    solution::init_cli_args(target, work_dir, &onchain);
     let onchain_clone = onchain.clone();
 
     let etherscan_api_key = match args.onchain_etherscan_api_key {

src/evm/onchain/endpoints.rs

@@ -917,7 +917,7 @@ impl OnChainConfig {
         pairs
     }
 
-    fn get_weth(&self, network: &str) -> String {
+    pub fn get_weth(&self, network: &str) -> String {
         let pegged_token = self.get_pegged_token(network);
 
         match network {

src/evm/oracles/erc20.rs

@@ -117,7 +117,7 @@ impl
             let liquidation_percent = EVMU256::from(liquidation_percent);
             let mut liquidations_earned = Vec::new();
 
-            for ((caller, token), (_prev_balance, new_balance)) in
+            for ((caller, token), (prev_balance, new_balance)) in
                 self.erc20_producer.deref().borrow().balances.iter()
             {
                 let token_info = self.known_tokens.get(token).expect("Token not found");

src/evm/solution.rs

@@ -0,0 +1,159 @@
+use std::{fs::{File, self}, time::SystemTime, sync::OnceLock};
+
+use handlebars::Handlebars;
+use serde::Serialize;
+
+use crate::input::SolutionTx;
+use super::{OnChainConfig, Chain, uniswap::{self, UniswapProvider}};
+
+const TEMPLATE_PATH: &str = "./foundry_test.hbs";
+/// Cli args for generating a test command.
+static CLI_ARGS: OnceLock<CliArgs> = OnceLock::new();
+
+/// Initialize CLI_ARGS.
+pub fn init_cli_args(target: String, work_dir: String, onchain: &Option<OnChainConfig>) {
+    let (chain, weth, block_number) = match onchain {
+        Some(oc) => (oc.chain_name.clone(), oc.get_weth(&oc.chain_name), oc.block_number.clone()),
+        None => (String::from(""), String::from(""), String::from("")),
+    };
+
+    let cli_args = CliArgs {
+        is_onchain: onchain.is_some(),
+        chain,
+        target,
+        block_number,
+        weth,
+        output_dir: format!("{}/vulnerabilities", work_dir),
+    };
+
+    let _ = CLI_ARGS.set(cli_args);
+}
+
+/// Generate a foundry test file.
+pub fn generate_test<T: SolutionTx>(solution: String, inputs: Vec<T>) {
+    let trace: Vec<Tx> = inputs.iter().map(|x| Tx::from(x)).collect();
+    if trace.is_empty() {
+        println!("generate_test error: no trace found.");
+        return;
+    }
+    let args = TemplateArgs::new(solution, trace);
+    if let Err(e) = args {
+        println!("generate_test error: {}", e);
+        return;
+    }
+    let args = args.unwrap();
+    if fs::create_dir_all(&args.output_dir).is_err() {
+        println!("generate_test error: failed to create output dir {:?}.", args.output_dir);
+        return;
+    }
+    let mut handlebars = Handlebars::new();
+    if handlebars.register_template_file("foundry_test", TEMPLATE_PATH).is_err() {
+        println!("generate_test error: failed to register template file.");
+        return;
+    }
+
+    let filename = SystemTime::now()
+        .duration_since(SystemTime::UNIX_EPOCH)
+        .unwrap()
+        .as_secs();
+    let path = format!("{}/{}.t.sol", args.output_dir, filename);
+    let mut output = File::create(&path).unwrap();
+    if let Err(e) = handlebars.render_to_write("foundry_test", &args, &mut output) {
+        println!("generate_test error: failed to render template: {:?}", e);
+    }
+}
+
+
+#[derive(Debug, Clone)]
+struct CliArgs {
+    is_onchain: bool,
+    chain: String,
+    target: String,
+    block_number: String,
+    weth: String,
+    output_dir: String,
+}
+
+#[derive(Debug, Serialize, Default)]
+pub struct Tx {
+    is_borrow: bool,
+    borrow_idx: u32,
+    router: String,
+    weth: String,
+    caller: String,
+    contract: String,
+    value: String,
+    fn_selector: String,
+    fn_args: String,
+    liq_percent: u8,
+}
+
+impl<T: SolutionTx> From<&T> for Tx {
+    fn from(input: &T) -> Self {
+        Self {
+            is_borrow: input.is_borrow(),
+            caller: input.caller(),
+            contract: input.contract(),
+            value: input.value(),
+            fn_selector: input.fn_selector(),
+            fn_args: input.fn_args(),
+            liq_percent: input.liq_percent(),
+            ..Default::default()
+        }
+    }
+}
+
+#[derive(Debug, Serialize, Default)]
+pub struct TemplateArgs {
+    is_onchain: bool,
+    chain: String,
+    target: String,
+    block_number: String,
+    etherscan_keyname: String,
+    solution: String,
+    trace: Vec<Tx>,
+    stepping_with_return: bool,
+    output_dir: String,
+}
+
+impl TemplateArgs {
+    pub fn new(solution: String, mut trace: Vec<Tx>) -> Result<Self, String> {
+        let cli_args = CLI_ARGS.get();
+        if cli_args.is_none() {
+            return Err(String::from("CLI_ARGS is not initialized."));
+        }
+        let cli_args = cli_args.unwrap();
+
+        let mut stepping_with_return = false;
+        if trace.last().unwrap().fn_selector == "0x00000000" {
+            trace.pop();
+            stepping_with_return = true;
+        }
+
+        if let Some(chain) = Chain::from_str(&cli_args.chain) {
+            let router = uniswap::get_uniswap_info(&UniswapProvider::UniswapV2, &chain).router;
+            let router = format!("0x{}", hex::encode(router));
+            let mut borrow_idx = 0;
+            for tx in trace.iter_mut() {
+                if tx.is_borrow {
+                    tx.router = router.clone();
+                    tx.weth = cli_args.weth.clone();
+                    tx.borrow_idx = borrow_idx;
+                    borrow_idx += 1;
+                }
+            }
+        }
+
+        Ok(Self {
+            is_onchain: cli_args.is_onchain,
+            chain: cli_args.chain.clone(),
+            target: cli_args.target.clone(),
+            block_number: cli_args.block_number.clone(),
+            etherscan_keyname: format!("{}_ETHERSCAN_API_KEY", cli_args.chain.to_uppercase()),
+            solution,
+            trace,
+            stepping_with_return,
+            output_dir: cli_args.output_dir.clone(),
+        })
+    }
+}