Drcov remodelling (#415)

* drcov remodelling

* fmt

* fix

Co-authored-by: tokatoka <tokazerkje@outlook.com>

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
     }
 }

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<I: Input + HasTargetBytes>(&mut self, input: &I);
 
     /// Called after execution of an input
-    fn post_exec<I: Input + HasTargetBytes>(&mut self, input: &I);
+    fn post_exec<I: Input + HasTargetBytes>(&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<I: Input + HasTargetBytes>(&mut self, input: &I) {
+    fn post_exec<I: Input + HasTargetBytes>(&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 {

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<File>,
-    module_mapping: &'a RangeMap<usize, (u16, String)>,
-    basic_blocks: &'a mut Vec<DrCovBasicBlock>,
+    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<usize, (u16, String)>,
+}
+
 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 }
     }
-}
-impl<'a> DrCovWriter<'a> {
-    /// Create a new [`DrCovWriter`]
-    pub fn new(
-        path: &str,
-        module_mapping: &'a RangeMap<usize, (u16, String)>,
-        basic_blocks: &'a mut Vec<DrCovBasicBlock>,
-    ) -> Self {
-        Self {
-            writer: BufWriter::new(
-                File::create(path).expect("Unable to create file for coverage data"),
-            ),
-            module_mapping,
-            basic_blocks,
+
+    /// 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)
     }
 }
 
-    /// Write the `DrCov` file.
-    pub fn write(&mut self) {
-        self.writer
+impl<'a> DrCovWriter<'a> {
+    /// Create a new [`DrCovWriter`]
+    #[must_use]
+    pub fn new(module_mapping: &'a RangeMap<usize, (u16, String)>) -> Self {
+        Self { module_mapping }
+    }
+
+    /// Write the list of basic blocks to a `DrCov` file.
+    pub fn write<P>(&mut self, path: P, basic_blocks: &[DrCovBasicBlock]) -> Result<(), Error>
+    where
+        P: AsRef<Path>,
+    {
+        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<usize>, &(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(())
     }
 }