diff --git a/src/fuzzer/engine.rs b/src/fuzzer/engine.rs index daf0109..b00ce52 100644 --- a/src/fuzzer/engine.rs +++ b/src/fuzzer/engine.rs @@ -5,7 +5,6 @@ use frame_support::traits::{OnFinalize, OnInitialize}; use prettytable::{row, Table}; use crate::{ - contract::payload::Selector, contract::remote::FullContractResponse, contract::runtime::{ AllPalletsWithSystem, BlockNumber, RuntimeOrigin, Timestamp, SLOT_DURATION, diff --git a/src/fuzzer/fuzz.rs b/src/fuzzer/fuzz.rs index 4e63618..808b59b 100644 --- a/src/fuzzer/fuzz.rs +++ b/src/fuzzer/fuzz.rs @@ -83,7 +83,6 @@ impl Fuzzer { impl FuzzerEngine for Fuzzer { fn fuzz(self) { - let mut transcoder_loader = Mutex::new( ContractMessageTranscoder::load(Path::new(&self.setup.path_to_specs)).unwrap(), ); @@ -94,7 +93,8 @@ impl FuzzerEngine for Fuzzer { let invariants: Vec = PayloadCrafter::extract_invariants(specs) .expect("No invariants found, check your contract"); - let mut selectors_without_invariants: Vec = selectors.clone() + let mut selectors_without_invariants: Vec = selectors + .clone() .into_iter() .filter(|s| !invariants.clone().contains(s)) .collect(); @@ -104,7 +104,6 @@ impl FuzzerEngine for Fuzzer { Self::build_corpus_and_dict(&mut selectors_without_invariants) .expect("๐Ÿ™… Failed to create initial corpus"); - println!( "\n\n๐Ÿš€ Now fuzzing `{}` ({})!\n", self.setup.path_to_specs.as_os_str().to_str().unwrap(), diff --git a/src/fuzzer/instrument.rs b/src/fuzzer/instrument.rs index 28b0979..a8dd570 100644 --- a/src/fuzzer/instrument.rs +++ b/src/fuzzer/instrument.rs @@ -25,7 +25,7 @@ use crate::fuzzer::instrument::instrument::ContractCovUpdater; /// Phink opted for a Rust AST approach. For each code instruction on the smart-contract, Phink will /// automatically add a tracing code, which will then be fetched at the end of the input execution /// in order to get coverage. -#[derive(Default)] +#[derive(Default, Clone)] pub struct InstrumenterEngine { pub contract_dir: PathBuf, } @@ -57,6 +57,60 @@ impl InstrumenterEngine { Self { contract_dir: dir } } + fn get_dirs_to_remove(tmp_dir: &Path, pattern: &str) -> Result, io::Error> { + Ok(fs::read_dir(tmp_dir)? + .filter_map(|entry| { + let entry = entry.ok()?; + let path = entry.path(); + if path.is_dir() && path.file_name()?.to_string_lossy().starts_with(pattern) { + Some(path) + } else { + None + } + }) + .collect::>()) + } + + fn prompt_user_confirmation() -> Result { + print!("๐Ÿ—‘๏ธ Do you really want to remove these directories? (yes/no): "); + io::stdout().flush()?; + let mut input = String::new(); + io::stdin().read_line(&mut input)?; + Ok(input.trim().eq_ignore_ascii_case("yes")) + } + + fn remove_directories(dirs_to_remove: Vec) -> Result<(), io::Error> { + for dir in dirs_to_remove { + fs::remove_dir_all(&dir)?; + println!("โœ… Removed directory: {}", dir.display()); + } + Ok(()) + } + + pub fn clean() -> Result<(), io::Error> { + let tmp_dir = Path::new("/tmp"); + let pattern = "ink_fuzzed_"; + let dirs_to_remove = Self::get_dirs_to_remove(tmp_dir, pattern)?; + + if dirs_to_remove.is_empty() { + println!("โŒ No directories found matching the pattern '{}'", pattern); + return Ok(()); + } + + println!("๐Ÿ” Found the following instrumented ink! contracts:"); + for dir in &dirs_to_remove { + println!("{}", dir.display()); + } + + if Self::prompt_user_confirmation()? { + Self::remove_directories(dirs_to_remove)?; + } else { + println!("โŒ Operation cancelled."); + } + + Ok(()) + } + pub fn find(&self) -> Result { //TODO: Handle this, sometimes the contract is not compiled in `target/ink` let wasm_path = fs::read_dir(self.contract_dir.join("target/ink/")) diff --git a/src/main.rs b/src/main.rs index c78b1d2..89daf82 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ extern crate core; +use env::{set_var, var}; use std::io::BufRead; use std::process::{Command, Stdio}; use std::{env, fs, path::PathBuf}; @@ -23,11 +24,22 @@ mod utils; /// This struct defines the command line arguments expected by Phink. #[derive(Parser, Debug)] -#[clap(author, version, about)] +#[clap( + author, + version, + about = "Phink is a command line tool for fuzzing ink! smart contracts.", + long_about = "๐Ÿ™ Phink, a ink! smart-contract property-based and coverage-guided fuzzer\n\n\ + Phink depends on various environment variables: + + \tPHINK_FROM_ZIGGY : Informs the tooling that the binary is being ran with Ziggy, and not directly from the CLI + \tPHINK_CONTRACT_DIR : Location of the contract code-base. Can be automatically detected. + \tPHINK_START_FUZZING : Tells the harness to start fuzzing. \n" +)] + struct Cli { /// Path where the `lib.rs` is located - #[clap(long, short, value_parser, required = false)] - path: PathBuf, + #[clap(long, short, value_parser)] + path: Option, /// Additional command to specify operation mode #[clap(subcommand)] @@ -39,7 +51,7 @@ struct Cli { enum Commands { /// Starts the fuzzing process Fuzz, - /// Execute one seed + /// Execute one seed, currently in TODO! Execute, /// Instrument the ink! contract, and compile it with Phink features Instrument, @@ -47,17 +59,25 @@ enum Commands { InstrumentAndFuzz, /// Run all seeds Run, - /// Generate a coverage + /// Generate a coverage, currently in TODO! Cover, + /// Remove all the temporary files under `/tmp/ink_fuzzed_XXXX` + Clean, } fn main() { - env::set_var("AFL_FORKSRV_INIT_TMOUT", "10000000"); - if env::var("PHINK_FROM_ZIGGY").is_ok() { - println!("๐Ÿซข Let's use Ziggy"); - let path = - PathBuf::from(env::var("PHINK_CONTRACT_DIR").unwrap_or("sample/dns/".parse().unwrap())); + if var("PHINK_FROM_ZIGGY").is_ok() { + println!("โ„น๏ธ Setting AFL_FORKSRV_INIT_TMOUT to 10000000"); + set_var("AFL_FORKSRV_INIT_TMOUT", "10000000"); + + let path = var("PHINK_CONTRACT_DIR").map(PathBuf::from).expect( + "\n๐Ÿˆฒ๏ธ PHINK_CONTRACT_DIR is not set. \ + You can set it manually, it should contain the source code of your contract, \ + with or without the instrumented binary,\ + depending your options. \n\n", + ); + let mut engine = instrument(path); start_fuzzer(&mut engine); @@ -65,66 +85,63 @@ fn main() { let cli = Cli::parse(); match &cli.command { - //TODO: Handle when CLI is just incorrect command, not just user doing ziggy run Commands::Instrument => { - instrument(cli.path); + set_var("PHINK_CONTRACT_DIR", cli.path.unwrap()); + let contract_dir = PathBuf::from(var("PHINK_CONTRACT_DIR").unwrap()); + instrument(contract_dir); } Commands::Fuzz => { - let mut engine = InstrumenterEngine::new(cli.path.clone()); + set_var("PHINK_CONTRACT_DIR", cli.path.unwrap()); + let contract_dir = PathBuf::from(var("PHINK_CONTRACT_DIR").unwrap()); + let mut engine = InstrumenterEngine::new(contract_dir); - start_cargo_ziggy_fuzz_process(); + start_cargo_ziggy_fuzz_process(engine.clone().contract_dir); - if env::var("PHINK_START_FUZZING").is_ok() { + if var("PHINK_START_FUZZING").is_ok() { start_fuzzer(&mut engine); } } Commands::InstrumentAndFuzz => { - let mut engine = instrument(cli.path); + set_var("PHINK_CONTRACT_DIR", cli.path.unwrap()); + let contract_dir = PathBuf::from(var("PHINK_CONTRACT_DIR").unwrap()); + let mut engine = instrument(contract_dir); - start_cargo_ziggy_fuzz_process(); + start_cargo_ziggy_fuzz_process(engine.clone().contract_dir); - if env::var("PHINK_START_FUZZING").is_ok() { + if var("PHINK_START_FUZZING").is_ok() { start_fuzzer(&mut engine); } } - Commands::Execute => { - todo!(); - } Commands::Run => { - let mut engine = instrument(cli.path); - - start_cargo_ziggy_run_process(); + set_var("PHINK_CONTRACT_DIR", cli.path.unwrap()); + let contract_dir = PathBuf::from(var("PHINK_CONTRACT_DIR").unwrap()); + let engine = instrument(contract_dir); + start_cargo_ziggy_run_process(engine.contract_dir); + } - if env::var("PHINK_START_FUZZING").is_ok() { - start_fuzzer(&mut engine); - } + Commands::Execute => { + todo!(); } + Commands::Cover => { todo!(); } + Commands::Clean => { + InstrumenterEngine::clean().expect("๐Ÿงผ Cannot execute the cleaning properly."); + } }; } - // // Can't use the CLI, it might be a direct Ziggy run - // Err(error) => { - // // return; - // println!("๐Ÿซข You probably used Ziggy directly in CLI, right ?"); - // println!("{:?}", error); - // let mut engine = InstrumenterEngine::new(PathBuf::from( - // env::var("PHINK_CONTRACT_DIR").unwrap_or("sample/dns/".parse().unwrap()), - // )); - // - // start_fuzzer(&mut engine); - // } - // }; } -fn start_cargo_ziggy_fuzz_process() { +fn start_cargo_ziggy_fuzz_process(contract_dir: PathBuf) { let mut child = Command::new("cargo") .arg("ziggy") .arg("fuzz") + .env("PHINK_CONTRACT_DIR", contract_dir) + .env("PHINK_FROM_ZIGGY", "true") .env("PHINK_START_FUZZING", "true") .arg(format!("-g={}", MIN_SEED_LEN)) .arg(format!("-G={}", MAX_SEED_LEN)) @@ -149,10 +166,12 @@ fn start_cargo_ziggy_fuzz_process() { } } -fn start_cargo_ziggy_run_process() { +fn start_cargo_ziggy_run_process(contract_dir: PathBuf) { let mut child = Command::new("cargo") .arg("ziggy") .arg("run") + .env("PHINK_CONTRACT_DIR", contract_dir) + .env("PHINK_FROM_ZIGGY", "true") .env("PHINK_START_FUZZING", "true") .stdout(Stdio::piped()) .spawn()