Feature/select program to fuzz

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<String>,
+    },
     #[command(
         about = "Run the AFL on desired fuzz test.",
         override_usage = "Specify the desired fuzz \x1b[92m<TARGET>\x1b[0m.\
@@ -134,9 +143,9 @@ 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();
         }
     };

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<String>) {
     // look for Anchor.toml
     let root = if let Some(r) = _discover(ANCHOR_TOML)? {
         r
@@ -22,7 +22,7 @@ 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 {
         let root_path = Path::new(&root).join(TRIDENT_TOML);
@@ -34,7 +34,7 @@ pub async fn init(force: bool) {
                 root
             );
         } else {
-            generator.initialize().await?;
+            generator.initialize(program_name).await?;
             show_howto();
         }
     }

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<String>,
     },
     #[command(
         about = "Run fuzz subcommands.",
@@ -67,7 +75,10 @@ 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::Init {
+            force,
+            program_name,
+        } => command::init(force, program_name).await?,
         Command::Clean => command::clean().await?,
     }
 }

crates/client/src/commander/mod.rs

@@ -58,13 +58,15 @@ impl Commander {
     }
 
     #[throws]
-    pub async fn build_anchor_project() {
-        let success = Command::new("anchor")
-            .arg("build")
-            .spawn()?
-            .wait()
-            .await?
-            .success();
+    pub async fn build_anchor_project(program_name: Option<String>) {
+        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?.success();
         if !success {
             throw!(Error::BuildProgramsFailed);
         }

crates/client/src/idl_loader.rs

@@ -5,14 +5,29 @@ use std::path::PathBuf;
 
 use trident_idl_spec::Idl;
 
-pub fn load_idls(dir_path: PathBuf) -> Result<Vec<Idl>, Box<dyn Error>> {
+pub fn load_idls(
+    dir_path: PathBuf,
+    program_name: Option<String>,
+) -> Result<Vec<Idl>, Box<dyn Error>> {
     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

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<String>) {
+        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,21 +73,20 @@ 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<String>) {
+        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?;
 
         // update_package_metadata(&self.program_packages, &self.versions_config).await?;
     }
 
     #[throws]
-    async fn get_program_packages(&mut self) {
+    async fn get_program_packages(&mut self, program_name: Option<String>) {
         // 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<String>) {
         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]

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<cargo_metadata::Package> {
-    let packages: Vec<cargo_metadata::Package> = program_packages().collect();
+pub async fn collect_program_packages(
+    program_name: Option<String>,
+) -> Vec<cargo_metadata::Package> {
+    let packages: Vec<cargo_metadata::Package> = program_packages(program_name).collect();
     if packages.is_empty() {
         throw!(Error::NoProgramsFound)
     } else {
         packages
     }
 }
-pub fn program_packages() -> impl Iterator<Item = cargo_metadata::Package> {
+pub fn program_packages(
+    program_name: Option<String>,
+) -> Box<dyn Iterator<Item = cargo_metadata::Package>> {
     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]