diff --git a/libafl_frida/src/executor.rs b/libafl_frida/src/executor.rs index 8a2eb15b20..ab6b7448cb 100644 --- a/libafl_frida/src/executor.rs +++ b/libafl_frida/src/executor.rs @@ -76,7 +76,7 @@ where libc::raise(libc::SIGABRT); } } - self.helper.post_exec(input); + self.helper.post_exec(input)?; res } } diff --git a/libafl_frida/src/helper.rs b/libafl_frida/src/helper.rs index 480ab63064..83c222695d 100644 --- a/libafl_frida/src/helper.rs +++ b/libafl_frida/src/helper.rs @@ -2,7 +2,7 @@ use ahash::AHasher; use std::hash::Hasher; use libafl::inputs::{HasTargetBytes, Input}; - +use libafl::Error; use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter}; #[cfg(target_arch = "aarch64")] @@ -69,7 +69,7 @@ pub trait FridaHelper<'a> { fn pre_exec(&mut self, input: &I); /// Called after execution of an input - fn post_exec(&mut self, input: &I); + fn post_exec(&mut self, input: &I) -> Result<(), Error>; /// Returns `true` if stalker is enabled fn stalker_enabled(&self) -> bool; @@ -124,13 +124,14 @@ impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> { } } - fn post_exec(&mut self, input: &I) { + fn post_exec(&mut self, input: &I) -> Result<(), Error> { if self.options.drcov_enabled() { let mut hasher = AHasher::new_with_keys(0, 0); hasher.write(input.target_bytes().as_slice()); let filename = format!("./coverage/{:016x}.drcov", hasher.finish(),); - DrCovWriter::new(&filename, &self.ranges, &mut self.drcov_basic_blocks).write(); + DrCovWriter::new(&self.ranges).write(&filename, &self.drcov_basic_blocks)?; + self.drcov_basic_blocks.clear(); } #[cfg(unix)] @@ -145,6 +146,7 @@ impl<'a> FridaHelper<'a> for FridaInstrumentationHelper<'a> { .poison(slice.as_ptr() as usize, slice.len()); self.asan_runtime.reset_allocations(); } + Ok(()) } fn stalker_enabled(&self) -> bool { diff --git a/libafl_targets/src/drcov.rs b/libafl_targets/src/drcov.rs index 4703e38dc3..9176a1bc2b 100644 --- a/libafl_targets/src/drcov.rs +++ b/libafl_targets/src/drcov.rs @@ -2,26 +2,22 @@ //! writing basic-block trace files to be read by coverage analysis tools, such as [Lighthouse](https://github.com/gaasedelen/lighthouse), //! [bncov](https://github.com/ForAllSecure/bncov), [dragondance](https://github.com/0ffffffffh/dragondance), etc. +use libafl::Error; use rangemap::RangeMap; use std::{ fs::File, io::{BufWriter, Write}, + path::Path, }; /// A basic block struct -#[derive(Clone, Copy)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct DrCovBasicBlock { - start: usize, - end: usize, -} - -/// A writer for `DrCov` files -pub struct DrCovWriter<'a> { - writer: BufWriter, - module_mapping: &'a RangeMap, - basic_blocks: &'a mut Vec, + pub start: usize, + pub end: usize, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] #[repr(C)] struct DrCovBasicBlockEntry { start: u32, @@ -29,46 +25,54 @@ struct DrCovBasicBlockEntry { mod_id: u16, } +/// A writer for `DrCov` files +pub struct DrCovWriter<'a> { + module_mapping: &'a RangeMap, +} + impl DrCovBasicBlock { /// Create a new [`DrCovBasicBlock`] with the given `start` and `end` addresses. #[must_use] pub fn new(start: usize, end: usize) -> Self { Self { start, end } } + + /// Create a new [`DrCovBasicBlock`] with a given `start` address and a block size. + #[must_use] + pub fn new_with_size(start: usize, size: usize) -> Self { + Self::new(start, start + size) + } } + impl<'a> DrCovWriter<'a> { /// Create a new [`DrCovWriter`] - pub fn new( - path: &str, - module_mapping: &'a RangeMap, - basic_blocks: &'a mut Vec, - ) -> Self { - Self { - writer: BufWriter::new( - File::create(path).expect("Unable to create file for coverage data"), - ), - module_mapping, - basic_blocks, - } + #[must_use] + pub fn new(module_mapping: &'a RangeMap) -> Self { + Self { module_mapping } } - /// Write the `DrCov` file. - pub fn write(&mut self) { - self.writer + /// Write the list of basic blocks to a `DrCov` file. + pub fn write

(&mut self, path: P, basic_blocks: &[DrCovBasicBlock]) -> Result<(), Error> + where + P: AsRef, + { + let mut writer = BufWriter::new(File::create(path)?); + + writer .write_all(b"DRCOV VERSION: 2\nDRCOV FLAVOR: libafl\n") .unwrap(); let modules: Vec<(&std::ops::Range, &(u16, String))> = self.module_mapping.iter().collect(); - self.writer + writer .write_all(format!("Module Table: version 2, count {}\n", modules.len()).as_bytes()) .unwrap(); - self.writer + writer .write_all(b"Columns: id, base, end, entry, checksum, timestamp, path\n") .unwrap(); for module in modules { let (range, (id, path)) = module; - self.writer + writer .write_all( format!( "{:03}, 0x{:x}, 0x{:x}, 0x00000000, 0x00000000, 0x00000000, {}\n", @@ -78,23 +82,24 @@ impl<'a> DrCovWriter<'a> { ) .unwrap(); } - self.writer - .write_all(format!("BB Table: {} bbs\n", self.basic_blocks.len()).as_bytes()) + writer + .write_all(format!("BB Table: {} bbs\n", basic_blocks.len()).as_bytes()) .unwrap(); - for block in self.basic_blocks.drain(0..) { + for block in basic_blocks { let (range, (id, _)) = self.module_mapping.get_key_value(&block.start).unwrap(); let basic_block = DrCovBasicBlockEntry { start: (block.start - range.start) as u32, size: (block.end - block.start) as u16, mod_id: *id, }; - self.writer + writer .write_all(unsafe { std::slice::from_raw_parts(&basic_block as *const _ as *const u8, 8) }) .unwrap(); } - self.writer.flush().unwrap(); + writer.flush()?; + Ok(()) } }