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

* accept folders

* lol

* use walkdir instead of my impl

Author: tokatoka

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,81 +41,114 @@ pub struct Opt {
     pub sort: bool,
 }
 
-fn main() -> Result<(), Error> {
-    let opts = Opt::parse();
+fn process(opts: &Opt, input: &PathBuf) -> Result<(), std::io::Error> {
+    let Ok(drcov) = DrCovReader::read(&input)
+        .map_err(|err| eprintln!("Ignored coverage file {input:?}, reason: {err:?}"))
+    else {
+        return Ok(());
+    };
 
-    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:?}");
-            }
-        }
+    let mut blocks = drcov.basic_block_addresses_u64();
 
-        assert!(out_dir.is_dir(), "Out_dir {out_dir:?} not a directory!");
+    if opts.sort {
+        blocks.sort_unstable();
     }
 
-    for input in opts.inputs {
-        let Ok(drcov) = DrCovReader::read(&input)
-            .map_err(|err| eprint!("Ignored coverage file {input:?}, reason: {err:?}"))
-        else {
-            continue;
+    let mut writer: Box<dyn Write> = if let Some(out_dir) = &opts.out_dir {
+        // Write files to a directory
+        let out_file = out_dir.join(
+            input
+                .file_name()
+                .expect("File without filename shouldn't exist"),
+        );
+
+        let Ok(file) = File::create_new(&out_file).map_err(|err| {
+            eprintln!("Could not create file {out_file:?} - continuing: {err:?}");
+        }) else {
+            return Ok(());
         };
 
-        let mut blocks = drcov.basic_block_addresses_u64();
-
-        if opts.sort {
-            blocks.sort_unstable();
+        println!("Dumping traces from drcov file {input:?} to {out_file:?}",);
+
+        Box::new(file)
+    } else {
+        Box::new(std::io::stdout())
+    };
+
+    // dump to stdout
+    let modules = &drcov.module_entries;
+
+    if opts.metadata {
+        writeln!(writer, "# {} Modules:", modules.len())?;
+        for module in &drcov.module_entries {
+            writeln!(
+                writer,
+                "\t{} - [{:#020x}-{:#020x}] {}",
+                module.id,
+                module.base,
+                module.end,
+                module.path.display()
+            )?;
         }
+        writeln!(writer, "# {} Blocks covered in {input:?}.", blocks.len())?;
 
-        let mut writer: Box<dyn Write> = if let Some(out_dir) = &opts.out_dir {
-            // Write files to a directory
-            let out_file = out_dir.join(
-                input
-                    .file_name()
-                    .expect("File without filename shouldn't exist"),
-            );
+        if opts.addrs {
+            writeln!(writer)?;
+        }
+    }
 
-            let Ok(file) = File::create_new(&out_file).map_err(|err| {
-                eprintln!("Could not create file {out_file:?} - continuing: {err:?}");
-            }) else {
-                continue;
-            };
+    if opts.addrs {
+        for line in blocks {
+            writeln!(writer, "{line:#x}")?;
+        }
+    }
 
-            println!("Dumping traces from drcov file {input:?} to {out_file:?}",);
+    Ok(())
+}
 
-            Box::new(file)
-        } else {
-            Box::new(std::io::stdout())
-        };
+#[must_use]
+pub fn find_drcov_files(dir: &PathBuf) -> Vec<PathBuf> {
+    let mut drcov_files = Vec::new();
 
-        // dump to stdout
-        let modules = &drcov.module_entries;
-
-        if opts.metadata {
-            writeln!(writer, "# {} Modules:", modules.len())?;
-            for module in &drcov.module_entries {
-                writeln!(
-                    writer,
-                    "\t{} - [{:#020x}-{:#020x}] {}",
-                    module.id,
-                    module.base,
-                    module.end,
-                    module.path.display()
-                )?;
+    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);
             }
-            writeln!(writer, "# {} Blocks covered in {input:?}.", blocks.len())?;
+        }
+    }
+
+    drcov_files
+}
+
+fn main() {
+    let opts = Opt::parse();
 
-            if opts.addrs {
-                writeln!(writer)?;
+    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:?}");
             }
         }
 
-        if opts.addrs {
-            for line in blocks {
-                writeln!(writer, "{line:#x}")?;
+        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);
         }
     }
-
-    Ok(())
 }