Make drcov_dump_address.rs accept list of directories (#2904)

* accept folders

* lol

* use walkdir instead of my impl

utils/drcov_utils/Cargo.toml

@@ -12,6 +12,7 @@ keywords = ["fuzzing", "libafl", "drcov"]
 env_logger = "0.11.6"
 libafl_targets = { workspace = true, default-features = true }
 clap = { workspace = true, features = ["derive", "wrap_help"] }
+walkdir = "2.5"
 
 [lints]
 workspace = true

utils/drcov_utils/src/bin/drcov_dump_addrs.rs

@@ -1,11 +1,12 @@
 use std::{
     fs::{create_dir_all, File},
-    io::{Error, Write},
+    io::Write,
     path::PathBuf,
 };
 
 use clap::Parser;
 use libafl_targets::drcov::DrCovReader;
+use walkdir::WalkDir;
 
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
@@ -15,7 +16,12 @@ use libafl_targets::drcov::DrCovReader;
     long_about = "Writes a list of all addresses from a DrCovFile"
 )]
 pub struct Opt {
-    #[arg(short, long, help = "DrCov traces to read", required = true)]
+    #[arg(
+        short,
+        long,
+        help = "DrCov traces or directories to read",
+        required = true
+    )]
     pub inputs: Vec<PathBuf>,
 
     #[arg(
@@ -35,24 +41,11 @@ pub struct Opt {
     pub sort: bool,
 }
 
-fn main() -> Result<(), Error> {
+fn process(opts: &Opt, input: &PathBuf) -> Result<(), std::io::Error> {
-    let opts = Opt::parse();
-
-    if let Some(out_dir) = &opts.out_dir {
-        if !out_dir.exists() {
-            if let Err(err) = create_dir_all(out_dir) {
-                eprint!("Failed to create dir {out_dir:?}: {err:?}");
-            }
-        }
-
-        assert!(out_dir.is_dir(), "Out_dir {out_dir:?} not a directory!");
-    }
-
-    for input in opts.inputs {
     let Ok(drcov) = DrCovReader::read(&input)
-            .map_err(|err| eprint!("Ignored coverage file {input:?}, reason: {err:?}"))
+        .map_err(|err| eprintln!("Ignored coverage file {input:?}, reason: {err:?}"))
     else {
-            continue;
+        return Ok(());
     };
 
     let mut blocks = drcov.basic_block_addresses_u64();
@@ -72,7 +65,7 @@ fn main() -> Result<(), Error> {
         let Ok(file) = File::create_new(&out_file).map_err(|err| {
             eprintln!("Could not create file {out_file:?} - continuing: {err:?}");
         }) else {
-                continue;
+            return Ok(());
         };
 
         println!("Dumping traces from drcov file {input:?} to {out_file:?}",);
@@ -109,7 +102,53 @@ fn main() -> Result<(), Error> {
             writeln!(writer, "{line:#x}")?;
         }
     }
-    }
 
     Ok(())
 }
+
+#[must_use]
+pub fn find_drcov_files(dir: &PathBuf) -> Vec<PathBuf> {
+    let mut drcov_files = Vec::new();
+
+    for entry in WalkDir::new(dir) {
+        let entry = entry.unwrap().into_path();
+        if let Some(ext) = entry.extension() {
+            if ext == "drcov" {
+                drcov_files.push(entry);
+            }
+        }
+    }
+
+    drcov_files
+}
+
+fn main() {
+    let opts = Opt::parse();
+
+    if let Some(out_dir) = &opts.out_dir {
+        if !out_dir.exists() {
+            if let Err(err) = create_dir_all(out_dir) {
+                eprintln!("Failed to create dir {out_dir:?}: {err:?}");
+            }
+        }
+
+        assert!(out_dir.is_dir(), "Out_dir {out_dir:?} not a directory!");
+    }
+
+    for input in &opts.inputs {
+        let drcovs = if input.is_dir() {
+            find_drcov_files(input)
+        } else {
+            let mut files = vec![];
+            if let Some(ext) = input.extension() {
+                if ext == "drcov" {
+                    files.push(input.clone());
+                }
+            }
+            files
+        };
+        for drcov_file in drcovs {
+            let _ = process(&opts, &drcov_file);
+        }
+    }
+}