foundry_test_generator -> solution

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;
@@ -17,13 +17,13 @@ use serde::{Deserialize, Deserializer, Serialize};
 
 use crate::generic_vm::vm_executor::ExecutionResult;
 use crate::generic_vm::vm_state::VMStateT;
-use crate::test_generator::TestTx;
 use bytes::Bytes;
 use std::cell::RefCell;
 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 {
@@ -311,17 +311,12 @@ impl ConciseEVMInput {
     }
 }
 
-impl TestTx for ConciseEVMInput {
+impl SolutionTx for ConciseEVMInput {
     #[cfg(feature = "flashloan_v2")]
     fn is_borrow(&self) -> bool {
         self.input_type == EVMInputTy::Borrow
     }
 
-    #[cfg(not(feature = "flashloan_v2"))]
-    fn is_borrow(&self) -> bool {
-        false
-    }
-
     fn caller(&self) -> String {
         format!("0x{}", hex::encode(self.caller))
     }
@@ -368,11 +363,6 @@ impl TestTx for ConciseEVMInput {
     fn liq_percent(&self) -> u8 {
         self.liquidation_percent
     }
-
-    #[cfg(not(feature = "flashloan_v2"))]
-    fn liq_percent(&self) -> u8 {
-        0
-    }
 }
 
 impl HasLen for EVMInput {

src/evm/mod.rs

@@ -21,7 +21,7 @@ pub mod srcmap;
 pub mod types;
 pub mod uniswap;
 pub mod vm;
-pub mod foundry_test_generator;
+pub mod solution;
 
 use crate::fuzzers::evm_fuzzer::evm_fuzzer;
 use crate::oracle::{Oracle, Producer};
@@ -49,8 +49,6 @@ use std::str::FromStr;
 use types::{EVMAddress, EVMFuzzState, EVMU256};
 use vm::EVMState;
 
-use self::foundry_test_generator::FoundryTestGenerator;
-
 pub fn parse_constructor_args_string(input: String) -> HashMap<String, Vec<String>> {
     let mut map = HashMap::new();
 
@@ -277,7 +275,7 @@ enum EVMTargetType {
 }
 
 pub fn evm_main(args: EvmArgs) {
-    let test_generator = FoundryTestGenerator::new(&args);
+    solution::init_cli_args(&args);
 
     let mut target_type: EVMTargetType = match args.target_type {
         Some(v) => match v.as_str() {
@@ -607,7 +605,7 @@ pub fn evm_main(args: EvmArgs) {
     };
 
     match config.fuzzer_type {
-        FuzzerTypes::CMP => evm_fuzzer(config, &mut state, test_generator),
+        FuzzerTypes::CMP => evm_fuzzer(config, &mut state),
         // FuzzerTypes::BASIC => basic_fuzzer(config)
         _ => {}
     }

src/evm/solution.rs

@@ -0,0 +1,124 @@
+use std::{fs::{File, self}, time::SystemTime, sync::OnceLock};
+
+use handlebars::Handlebars;
+use serde::Serialize;
+
+use crate::input::SolutionTx;
+use super::EvmArgs;
+
+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(args: &EvmArgs) {
+    let cli_args = CliArgs {
+        chain: args.chain_type.clone().unwrap_or_default(),
+        target: args.target.clone(),
+        block_number: args.onchain_block_number.unwrap_or_default(),
+        output_dir: format!("{}/vulnerabilities", args.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 mut trace: Vec<Tx> = inputs.iter().map(|x| Tx::from(x)).collect();
+    if trace.is_empty() {
+        return;
+    }
+    let args = TemplateArgs::new(solution, trace);
+    if args.is_none() {
+        return;
+    }
+    let args = args.unwrap();
+    if fs::create_dir_all(&args.output_dir).is_err() {
+        return;
+    }
+    let mut handlebars = Handlebars::new();
+    if handlebars.register_template_file("foundry_test", TEMPLATE_PATH).is_err() {
+        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();
+    let _ = handlebars.render_to_write("foundry_test", &args, &mut output);
+}
+
+
+#[derive(Debug, Clone)]
+struct CliArgs {
+    chain: String,
+    target: String,
+    block_number: u64,
+    output_dir: String,
+}
+
+#[derive(Debug, Serialize, Default)]
+pub struct Tx {
+    is_borrow: bool,
+    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(),
+        }
+    }
+}
+
+#[derive(Debug, Serialize, Default)]
+pub struct TemplateArgs {
+    target: String,
+    chain: String,
+    block_number: u64,
+    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>) -> Option<Self> {
+        let cli_args = CLI_ARGS.get();
+        if cli_args.is_none() {
+            return None;
+        }
+        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;
+        }
+
+        Some(Self {
+            target: cli_args.target.clone(),
+            chain: cli_args.chain.clone(),
+            block_number: cli_args.block_number,
+            etherscan_keyname: format!("{}_ETHERSCAN_API_KEY", cli_args.chain.to_uppercase()),
+            solution,
+            trace,
+            stepping_with_return,
+            output_dir: cli_args.output_dir.clone(),
+        })
+    }
+}