From d08330fcee8fc21ae2fa8d597bb9f103f5a04aa2 Mon Sep 17 00:00:00 2001 From: hyckomatej Date: Tue, 21 Jan 2025 12:51:13 +0100 Subject: [PATCH 1/5] Add option to select which program to fuzz --- crates/cli/src/command/fuzz.rs | 17 ++++++++++++--- crates/cli/src/command/init.rs | 9 ++++---- crates/cli/src/lib.rs | 12 +++++++++-- crates/client/src/commander/mod.rs | 12 ++++++++--- crates/client/src/idl_loader.rs | 14 ++++++++++++- crates/client/src/test_generator.rs | 26 +++++++++++------------ crates/client/src/utils.rs | 32 ++++++++++++++++++++--------- 7 files changed, 85 insertions(+), 37 deletions(-) diff --git a/crates/cli/src/command/fuzz.rs b/crates/cli/src/command/fuzz.rs index 0b9251ca..e147782b 100644 --- a/crates/cli/src/command/fuzz.rs +++ b/crates/cli/src/command/fuzz.rs @@ -14,7 +14,16 @@ pub const TRIDENT_TOML: &str = "Trident.toml"; #[allow(non_camel_case_types)] pub enum FuzzCommand { #[command(about = "Generate new Fuzz Test template.")] - Add, + Add { + #[arg( + short, + long, + required = false, + help = "Specify the name of the program for which the fuzz test will be generated.", + value_name = "FILE" + )] + program_name: Option, + }, #[command( about = "Run the AFL on desired fuzz test.", override_usage = "Specify the desired fuzz \x1b[92m\x1b[0m.\ @@ -134,9 +143,11 @@ pub async fn fuzz(subcmd: FuzzCommand) { commander.run_hfuzz_debug(target, crash_file_path).await?; } - FuzzCommand::Add => { + FuzzCommand::Add { + program_name + } => { let mut generator = TestGenerator::new_with_root(&root)?; - generator.add_fuzz_test().await?; + generator.add_fuzz_test(program_name).await?; show_howto(); } }; diff --git a/crates/cli/src/command/init.rs b/crates/cli/src/command/init.rs index 330dbd13..31734270 100644 --- a/crates/cli/src/command/init.rs +++ b/crates/cli/src/command/init.rs @@ -11,7 +11,7 @@ pub const TRIDENT_TOML: &str = "Trident.toml"; pub const SKIP: &str = "\x1b[33mSkip\x1b[0m"; #[throws] -pub async fn init(force: bool) { +pub async fn init(force: bool, program_name: Option) { // look for Anchor.toml let root = if let Some(r) = _discover(ANCHOR_TOML)? { r @@ -22,9 +22,10 @@ pub async fn init(force: bool) { let mut generator: TestGenerator = TestGenerator::new_with_root(&root)?; if force { - generator.initialize().await?; + generator.initialize(program_name).await?; show_howto(); - } else { + } + else { let root_path = Path::new(&root).join(TRIDENT_TOML); if root_path.exists() { println!( @@ -34,7 +35,7 @@ pub async fn init(force: bool) { root ); } else { - generator.initialize().await?; + generator.initialize(program_name).await?; show_howto(); } } diff --git a/crates/cli/src/lib.rs b/crates/cli/src/lib.rs index 7b29ceb8..e97a8377 100644 --- a/crates/cli/src/lib.rs +++ b/crates/cli/src/lib.rs @@ -42,6 +42,14 @@ enum Command { help = "Force Trident initialization. Trident dependencies will be updated based on the version of Trident CLI." )] force: bool, + #[arg( + short, + long, + required = false, + help = "Specify the name of the program for which fuzz test will be generated.", + value_name = "FILE" + )] + program_name: Option, }, #[command( about = "Run fuzz subcommands.", @@ -66,8 +74,8 @@ pub async fn start() { match cli.command { Command::How => command::howto()?, - Command::Fuzz { subcmd } => command::fuzz(subcmd).await?, - Command::Init { force } => command::init(force).await?, + Command::Fuzz { subcmd} => command::fuzz(subcmd).await?, + Command::Init { force, program_name} => command::init(force, program_name).await?, Command::Clean => command::clean().await?, } } diff --git a/crates/client/src/commander/mod.rs b/crates/client/src/commander/mod.rs index 684fecf1..b1ba9bc0 100644 --- a/crates/client/src/commander/mod.rs +++ b/crates/client/src/commander/mod.rs @@ -58,9 +58,15 @@ impl Commander { } #[throws] - pub async fn build_anchor_project() { - let success = Command::new("anchor") - .arg("build") + pub async fn build_anchor_project(program_name: Option) { + let mut cmd = Command::new("anchor"); + cmd.arg("build"); + + if let Some(name) = program_name { + cmd.args(["-p", name.as_str()]); + } + + let success = cmd .spawn()? .wait() .await? diff --git a/crates/client/src/idl_loader.rs b/crates/client/src/idl_loader.rs index 9c5114cd..83260d87 100644 --- a/crates/client/src/idl_loader.rs +++ b/crates/client/src/idl_loader.rs @@ -5,14 +5,26 @@ use std::path::PathBuf; use trident_idl_spec::Idl; -pub fn load_idls(dir_path: PathBuf) -> Result, Box> { +pub fn load_idls(dir_path: PathBuf, program_name: Option) -> Result, Box> { let mut idls = Vec::new(); + + // Read the directory and iterate over each entry for entry in fs::read_dir(dir_path)? { let entry = entry?; let path = entry.path(); + if let Some(ref program_name) = program_name { + if path.is_file() && !path.file_name() + .and_then(|name| name.to_str()) + .map(|name| name.contains(program_name)) + .unwrap_or(false) + { + continue; + } + } + // Only process .json files if path.is_file() && path.extension().and_then(|ext| ext.to_str()) == Some("json") { // Remove the .json extension to get the package name diff --git a/crates/client/src/test_generator.rs b/crates/client/src/test_generator.rs index 6222433a..1d8e34b0 100644 --- a/crates/client/src/test_generator.rs +++ b/crates/client/src/test_generator.rs @@ -61,11 +61,10 @@ impl TestGenerator { } } #[throws] - pub async fn initialize(&mut self) { - Commander::build_anchor_project().await?; - - self.get_program_packages().await?; - self.load_programs_idl()?; + pub async fn initialize(&mut self, program_name: Option) { + Commander::build_anchor_project(program_name.clone()).await?; + self.get_program_packages(program_name.clone()).await?; + self.load_programs_idl(program_name.clone())?; self.generate_source_codes().await?; self.initialize_new_fuzz_test().await?; @@ -74,11 +73,10 @@ impl TestGenerator { } #[throws] - pub async fn add_fuzz_test(&mut self) { - Commander::build_anchor_project().await?; - - self.get_program_packages().await?; - self.load_programs_idl()?; + pub async fn add_fuzz_test(&mut self, program_name: Option) { + Commander::build_anchor_project(program_name.clone()).await?; + self.get_program_packages(program_name.clone()).await?; + self.load_programs_idl(program_name.clone())?; self.generate_source_codes().await?; self.add_new_fuzz_test().await?; @@ -86,9 +84,9 @@ impl TestGenerator { } #[throws] - async fn get_program_packages(&mut self) { + async fn get_program_packages(&mut self, program_name: Option) { // TODO consider optionally excluding packages - self.program_packages = collect_program_packages().await?; + self.program_packages = collect_program_packages(program_name).await?; } #[throws] @@ -119,11 +117,11 @@ impl TestGenerator { } #[throws] - fn load_programs_idl(&mut self) { + fn load_programs_idl(&mut self, program_name: Option) { let target_path = construct_path!(self.root, "target/idl/"); // TODO consider optionally excluding packages - self.anchor_idls = crate::idl_loader::load_idls(target_path).unwrap(); + self.anchor_idls = crate::idl_loader::load_idls(target_path, program_name).unwrap(); } #[throws] diff --git a/crates/client/src/utils.rs b/crates/client/src/utils.rs index 123eb6f0..ed6eafab 100644 --- a/crates/client/src/utils.rs +++ b/crates/client/src/utils.rs @@ -82,27 +82,39 @@ pub fn get_fuzz_id(fuzz_dir_path: &Path) -> i32 { } } #[throws] -pub async fn collect_program_packages() -> Vec { - let packages: Vec = program_packages().collect(); +pub async fn collect_program_packages( + program_name: Option, +) -> Vec { + let packages: Vec = program_packages(program_name).collect(); if packages.is_empty() { throw!(Error::NoProgramsFound) } else { packages } } -pub fn program_packages() -> impl Iterator { +pub fn program_packages( + program_name: Option, +) -> Box> { let cargo_toml_data = cargo_metadata::MetadataCommand::new() .no_deps() .exec() .expect("Cargo.toml reading failed"); - cargo_toml_data.packages.into_iter().filter(|package| { - // TODO less error-prone test if the package is a _program_? - if let Some("programs") = package.manifest_path.iter().nth_back(2) { - return true; - } - false - }) + match program_name { + Some(name) => Box::new( + cargo_toml_data + .packages + .into_iter() + .filter(move |package| package.name == name), + ), + None => Box::new(cargo_toml_data.packages.into_iter().filter(|package| { + // TODO less error-prone test if the package is a _program_? + if let Some("programs") = package.manifest_path.iter().nth_back(2) { + return true; + } + false + })), + } } #[throws] From ab31e74aabf8095ad0e3f1b6ff9cbb2aaf404fc3 Mon Sep 17 00:00:00 2001 From: hyckomatej Date: Tue, 21 Jan 2025 12:54:21 +0100 Subject: [PATCH 2/5] Fix formatting --- crates/cli/src/command/init.rs | 3 +-- crates/cli/src/lib.rs | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/crates/cli/src/command/init.rs b/crates/cli/src/command/init.rs index 31734270..8baddec2 100644 --- a/crates/cli/src/command/init.rs +++ b/crates/cli/src/command/init.rs @@ -24,8 +24,7 @@ pub async fn init(force: bool, program_name: Option) { if force { generator.initialize(program_name).await?; show_howto(); - } - else { + } else { let root_path = Path::new(&root).join(TRIDENT_TOML); if root_path.exists() { println!( diff --git a/crates/cli/src/lib.rs b/crates/cli/src/lib.rs index e97a8377..c8a0d809 100644 --- a/crates/cli/src/lib.rs +++ b/crates/cli/src/lib.rs @@ -74,7 +74,7 @@ pub async fn start() { match cli.command { Command::How => command::howto()?, - Command::Fuzz { subcmd} => command::fuzz(subcmd).await?, + Command::Fuzz { subcmd } => command::fuzz(subcmd).await?, Command::Init { force, program_name} => command::init(force, program_name).await?, Command::Clean => command::clean().await?, } From fc0df8a7d81f5cc59f7ec69ac18956734dc1ea89 Mon Sep 17 00:00:00 2001 From: hyckomatej Date: Tue, 21 Jan 2025 14:32:21 +0100 Subject: [PATCH 3/5] Fix formatting --- crates/cli/src/command/fuzz.rs | 4 +--- crates/cli/src/lib.rs | 5 ++++- crates/client/src/commander/mod.rs | 6 +----- crates/client/src/idl_loader.rs | 17 ++++++++++------- 4 files changed, 16 insertions(+), 16 deletions(-) diff --git a/crates/cli/src/command/fuzz.rs b/crates/cli/src/command/fuzz.rs index e147782b..74ab68fd 100644 --- a/crates/cli/src/command/fuzz.rs +++ b/crates/cli/src/command/fuzz.rs @@ -143,9 +143,7 @@ pub async fn fuzz(subcmd: FuzzCommand) { commander.run_hfuzz_debug(target, crash_file_path).await?; } - FuzzCommand::Add { - program_name - } => { + FuzzCommand::Add { program_name } => { let mut generator = TestGenerator::new_with_root(&root)?; generator.add_fuzz_test(program_name).await?; show_howto(); diff --git a/crates/cli/src/lib.rs b/crates/cli/src/lib.rs index c8a0d809..141e2ca7 100644 --- a/crates/cli/src/lib.rs +++ b/crates/cli/src/lib.rs @@ -75,7 +75,10 @@ pub async fn start() { match cli.command { Command::How => command::howto()?, Command::Fuzz { subcmd } => command::fuzz(subcmd).await?, - Command::Init { force, program_name} => command::init(force, program_name).await?, + Command::Init { + force, + program_name, + } => command::init(force, program_name).await?, Command::Clean => command::clean().await?, } } diff --git a/crates/client/src/commander/mod.rs b/crates/client/src/commander/mod.rs index b1ba9bc0..ad35a13e 100644 --- a/crates/client/src/commander/mod.rs +++ b/crates/client/src/commander/mod.rs @@ -66,11 +66,7 @@ impl Commander { cmd.args(["-p", name.as_str()]); } - let success = cmd - .spawn()? - .wait() - .await? - .success(); + let success = cmd.spawn()?.wait().await?.success(); if !success { throw!(Error::BuildProgramsFailed); } diff --git a/crates/client/src/idl_loader.rs b/crates/client/src/idl_loader.rs index 83260d87..2f5ba680 100644 --- a/crates/client/src/idl_loader.rs +++ b/crates/client/src/idl_loader.rs @@ -5,21 +5,24 @@ use std::path::PathBuf; use trident_idl_spec::Idl; -pub fn load_idls(dir_path: PathBuf, program_name: Option) -> Result, Box> { +pub fn load_idls( + dir_path: PathBuf, + program_name: Option, +) -> Result, Box> { let mut idls = Vec::new(); - - // Read the directory and iterate over each entry for entry in fs::read_dir(dir_path)? { let entry = entry?; let path = entry.path(); if let Some(ref program_name) = program_name { - if path.is_file() && !path.file_name() - .and_then(|name| name.to_str()) - .map(|name| name.contains(program_name)) - .unwrap_or(false) + if path.is_file() + && !path + .file_name() + .and_then(|name| name.to_str()) + .map(|name| name.contains(program_name)) + .unwrap_or(false) { continue; } From 8a9f5bf065911a627b0c72fa38aa73b16752675d Mon Sep 17 00:00:00 2001 From: hyckomatej Date: Fri, 24 Jan 2025 12:55:31 +0100 Subject: [PATCH 4/5] Fix test name conversion --- Cargo.lock | 1 + crates/client/Cargo.toml | 1 + crates/client/src/idl_loader.rs | 10 ++++++---- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ef82ecc7..8834e7cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3873,6 +3873,7 @@ dependencies = [ "anyhow", "cargo_metadata", "fehler", + "heck 0.4.1", "pathdiff", "pretty_assertions", "rand 0.8.5", diff --git a/crates/client/Cargo.toml b/crates/client/Cargo.toml index ed620b93..5fb5154b 100644 --- a/crates/client/Cargo.toml +++ b/crates/client/Cargo.toml @@ -30,6 +30,7 @@ cargo_metadata = "0.18" toml = { version = "0.8", features = ["preserve_order"] } pathdiff = "0.2" rand = "0.8" +heck = "0.4.0" [dev-dependencies] pretty_assertions = "1.1.0" diff --git a/crates/client/src/idl_loader.rs b/crates/client/src/idl_loader.rs index 2f5ba680..e2dc2424 100644 --- a/crates/client/src/idl_loader.rs +++ b/crates/client/src/idl_loader.rs @@ -2,6 +2,7 @@ use std::error::Error; use std::fs::{self, File}; use std::io::Read; use std::path::PathBuf; +use heck::ToSnakeCase; use trident_idl_spec::Idl; @@ -18,10 +19,11 @@ pub fn load_idls( if let Some(ref program_name) = program_name { if path.is_file() - && !path - .file_name() - .and_then(|name| name.to_str()) - .map(|name| name.contains(program_name)) + && !path + .file_name() + .and_then(|name| name.to_str()) + // convert program_name to match case of IDL names + .map(|name| name.trim_end_matches(".json") == program_name.to_snake_case()) .unwrap_or(false) { continue; From 5cdde1be8c47b382278cf5687bd488eccf644e97 Mon Sep 17 00:00:00 2001 From: hyckomatej Date: Fri, 24 Jan 2025 12:58:15 +0100 Subject: [PATCH 5/5] Fix formatting --- crates/client/src/idl_loader.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/client/src/idl_loader.rs b/crates/client/src/idl_loader.rs index e2dc2424..4e6182ba 100644 --- a/crates/client/src/idl_loader.rs +++ b/crates/client/src/idl_loader.rs @@ -1,8 +1,8 @@ +use heck::ToSnakeCase; use std::error::Error; use std::fs::{self, File}; use std::io::Read; use std::path::PathBuf; -use heck::ToSnakeCase; use trident_idl_spec::Idl; @@ -19,9 +19,9 @@ pub fn load_idls( if let Some(ref program_name) = program_name { if path.is_file() - && !path - .file_name() - .and_then(|name| name.to_str()) + && !path + .file_name() + .and_then(|name| name.to_str()) // convert program_name to match case of IDL names .map(|name| name.trim_end_matches(".json") == program_name.to_snake_case()) .unwrap_or(false)