From c0cd48475cf46e78531a7f3763ef12468b55864e Mon Sep 17 00:00:00 2001
From: aarnav <aarnav@srlabs.de>
Date: Tue, 4 Jun 2024 11:31:46 +0200
Subject: [PATCH] feat: add recursive option when running inputs

---
 src/bin/cargo-ziggy/main.rs |  6 +++++-
 src/bin/cargo-ziggy/run.rs  | 41 +++++++++++++++++++++++++++++++++++--
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/bin/cargo-ziggy/main.rs b/src/bin/cargo-ziggy/main.rs
index fce7f4d..da2ca3f 100644
--- a/src/bin/cargo-ziggy/main.rs
+++ b/src/bin/cargo-ziggy/main.rs
@@ -174,6 +174,10 @@ pub struct Run {
     #[clap(short, long, value_name = "DIR", default_value = DEFAULT_CORPUS_DIR)]
     inputs: Vec<PathBuf>,
 
+    /// Recursively run nested directories for all input directories
+    #[clap(short, long)]
+    recursive: bool,
+
     /// Fuzzers output directory
     #[clap(short, long, env="ZIGGY_OUTPUT", value_parser, value_name = "DIR", default_value=DEFAULT_OUTPUT_DIR)]
     ziggy_output: PathBuf,
@@ -298,7 +302,7 @@ fn main() -> Result<(), anyhow::Error> {
     match command {
         Ziggy::Build(args) => args.build().context("Failed to build the fuzzers"),
         Ziggy::Fuzz(mut args) => args.fuzz().context("Failure running fuzzers"),
-        Ziggy::Run(args) => args.run().context("Failure running inputs"),
+        Ziggy::Run(mut args) => args.run().context("Failure running inputs"),
         Ziggy::Minimize(mut args) => args.minimize().context("Failure running minimization"),
         Ziggy::Cover(mut args) => args
             .generate_coverage()
diff --git a/src/bin/cargo-ziggy/run.rs b/src/bin/cargo-ziggy/run.rs
index 6df36ea..19e962f 100644
--- a/src/bin/cargo-ziggy/run.rs
+++ b/src/bin/cargo-ziggy/run.rs
@@ -1,11 +1,16 @@
 use crate::{find_target, Run};
 use anyhow::{anyhow, Context, Result};
 use console::style;
-use std::{env, process};
+use std::{
+    collections::HashSet,
+    env, fs,
+    path::{Path, PathBuf},
+    process,
+};
 
 impl Run {
     // Run inputs
-    pub fn run(&self) -> Result<(), anyhow::Error> {
+    pub fn run(&mut self) -> Result<(), anyhow::Error> {
         let cargo = env::var("CARGO").unwrap_or_else(|_| String::from("cargo"));
         let target = find_target(&self.target)?;
 
@@ -29,6 +34,21 @@ impl Run {
         }
 
         eprintln!("    {} runner", style("Finished").cyan().bold());
+
+        if self.recursive {
+            info!("Finding nested input directories recursively...");
+            let mut all_dirs = HashSet::new();
+            for input in &self.inputs {
+                all_dirs.insert(input.clone());
+                collect_dirs_recursively(input, &mut all_dirs)?;
+            }
+            for dir in all_dirs {
+                if !self.inputs.contains(&dir) {
+                    self.inputs.push(dir);
+                }
+            }
+        }
+
         info!("Running inputs");
         let run_args: Vec<String> = self
             .inputs
@@ -52,3 +72,20 @@ impl Run {
         Ok(())
     }
 }
+
+fn collect_dirs_recursively(
+    dir: &Path,
+    dir_list: &mut HashSet<PathBuf>,
+) -> Result<(), anyhow::Error> {
+    if dir.is_dir() {
+        for entry in fs::read_dir(dir)? {
+            let entry = entry?;
+            let path = entry.path();
+            if path.is_dir() && !dir_list.contains(&path) {
+                dir_list.insert(path.clone());
+                collect_dirs_recursively(&path, dir_list)?;
+            }
+        }
+    }
+    Ok(())
+}