From 2ec534a39ce8cf2807b34d88cd2335182f1c5289 Mon Sep 17 00:00:00 2001 From: "Dongjia \"toka\" Zhang" Date: Tue, 28 Jan 2025 19:06:27 +0100 Subject: [PATCH] Make drcov_dump_address.rs accept list of directories (#2904) * accept folders * lol * use walkdir instead of my impl --- utils/drcov_utils/Cargo.toml | 1 + utils/drcov_utils/src/bin/drcov_dump_addrs.rs | 167 +++++++++++------- 2 files changed, 104 insertions(+), 64 deletions(-) diff --git a/utils/drcov_utils/Cargo.toml b/utils/drcov_utils/Cargo.toml index cc72570e3c..ab635e185c 100644 --- a/utils/drcov_utils/Cargo.toml +++ b/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 diff --git a/utils/drcov_utils/src/bin/drcov_dump_addrs.rs b/utils/drcov_utils/src/bin/drcov_dump_addrs.rs index a404eb40fb..121fc20b99 100644 --- a/utils/drcov_utils/src/bin/drcov_dump_addrs.rs +++ b/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, #[arg( @@ -35,81 +41,114 @@ pub struct Opt { pub sort: bool, } -fn main() -> Result<(), Error> { +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(()); + }; + + let mut blocks = drcov.basic_block_addresses_u64(); + + if opts.sort { + blocks.sort_unstable(); + } + + let mut writer: Box = 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(()); + }; + + 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())?; + + if opts.addrs { + writeln!(writer)?; + } + } + + if opts.addrs { + for line in blocks { + writeln!(writer, "{line:#x}")?; + } + } + + Ok(()) +} + +#[must_use] +pub fn find_drcov_files(dir: &PathBuf) -> Vec { + 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) { - eprint!("Failed to create dir {out_dir:?}: {err:?}"); + 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 Ok(drcov) = DrCovReader::read(&input) - .map_err(|err| eprint!("Ignored coverage file {input:?}, reason: {err:?}")) - else { - continue; - }; - - let mut blocks = drcov.basic_block_addresses_u64(); - - if opts.sort { - blocks.sort_unstable(); - } - - let mut writer: Box = 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 { - continue; - }; - - println!("Dumping traces from drcov file {input:?} to {out_file:?}",); - - Box::new(file) + for input in &opts.inputs { + let drcovs = if input.is_dir() { + find_drcov_files(input) } else { - Box::new(std::io::stdout()) + let mut files = vec![]; + if let Some(ext) = input.extension() { + if ext == "drcov" { + files.push(input.clone()); + } + } + files }; - - // 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())?; - - if opts.addrs { - writeln!(writer)?; - } - } - - if opts.addrs { - for line in blocks { - writeln!(writer, "{line:#x}")?; - } + for drcov_file in drcovs { + let _ = process(&opts, &drcov_file); } } - - Ok(()) }