Re-add drcov for both usermode and systemmode. (#2573)

* re-add drcov for both usermode and systemmode.

fuzzers/binary_only/qemu_coverage/src/fuzzer.rs

@@ -233,11 +233,11 @@ pub fn fuzz() {
             let core = core_id.0;
             cov_path.set_file_name(format!("{coverage_name}-{core:03}.{coverage_extension}"));
 
-            let emulator_modules = tuple_list!(DrCovModule::new(
-                StdAddressFilter::default(),
-                cov_path,
-                false,
-            ));
+            let emulator_modules = tuple_list!(DrCovModule::builder()
+                .filter(StdAddressFilter::default())
+                .filename(cov_path)
+                .full_trace(false)
+                .build());
 
             let emulator = Emulator::empty()
                 .qemu(qemu)

libafl_qemu/src/modules/usermode/drcov.rs → libafl_qemu/src/modules/drcov.rs

@@ -1,3 +1,5 @@
+#[cfg(emulation_mode = "systemmode")]
+use std::ptr::addr_of_mut;
 use std::{path::PathBuf, sync::Mutex};
 
 use hashbrown::{hash_map::Entry, HashMap};
@@ -7,9 +9,11 @@ use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter};
 use rangemap::RangeMap;
 use serde::{Deserialize, Serialize};
 
+#[cfg(emulation_mode = "systemmode")]
+use crate::modules::{NopPageFilter, NOP_PAGE_FILTER};
 use crate::{
     emu::EmulatorModules,
-    modules::{AddressFilter, EmulatorModule, EmulatorModuleTuple},
+    modules::{AddressFilter, EmulatorModule, EmulatorModuleTuple, NopAddressFilter},
     qemu::Hook,
 };
 
@@ -35,30 +39,107 @@ impl DrCovMetadata {
 
 libafl_bolts::impl_serdeany!(DrCovMetadata);
 
+#[derive(Debug)]
+pub struct DrCovModuleBuilder<F> {
+    filter: Option<F>,
+    module_mapping: Option<RangeMap<usize, (u16, String)>>,
+    filename: Option<PathBuf>,
+    full_trace: Option<bool>,
+}
+
+impl<F> DrCovModuleBuilder<F>
+where
+    F: AddressFilter,
+{
+    pub fn build(self) -> DrCovModule<F> {
+        DrCovModule::new(
+            self.filter.unwrap(),
+            self.filename.unwrap(),
+            self.module_mapping,
+            self.full_trace.unwrap(),
+        )
+    }
+
+    pub fn filter<F2>(self, filter: F2) -> DrCovModuleBuilder<F2> {
+        DrCovModuleBuilder {
+            filter: Some(filter),
+            module_mapping: self.module_mapping,
+            filename: self.filename,
+            full_trace: self.full_trace,
+        }
+    }
+
+    #[must_use]
+    pub fn module_mapping(self, module_mapping: RangeMap<usize, (u16, String)>) -> Self {
+        Self {
+            filter: self.filter,
+            module_mapping: Some(module_mapping),
+            filename: self.filename,
+            full_trace: self.full_trace,
+        }
+    }
+
+    #[must_use]
+    pub fn filename(self, filename: PathBuf) -> Self {
+        Self {
+            filter: self.filter,
+            module_mapping: self.module_mapping,
+            filename: Some(filename),
+            full_trace: self.full_trace,
+        }
+    }
+
+    #[must_use]
+    pub fn full_trace(self, full_trace: bool) -> Self {
+        Self {
+            filter: self.filter,
+            module_mapping: self.module_mapping,
+            filename: self.filename,
+            full_trace: Some(full_trace),
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct DrCovModule<F> {
     filter: F,
-    module_mapping: RangeMap<usize, (u16, String)>,
+    module_mapping: Option<RangeMap<usize, (u16, String)>>,
     filename: PathBuf,
     full_trace: bool,
     drcov_len: usize,
 }
+impl DrCovModule<NopAddressFilter> {
+    #[must_use]
+    pub fn builder() -> DrCovModuleBuilder<NopAddressFilter> {
+        DrCovModuleBuilder {
+            filter: Some(NopAddressFilter),
+            module_mapping: None,
+            full_trace: None,
+            filename: None,
+        }
+    }
+}
 
 impl<F> DrCovModule<F>
 where
     F: AddressFilter,
 {
     #[must_use]
     #[allow(clippy::let_underscore_untyped)]
-    pub fn new(filter: F, filename: PathBuf, full_trace: bool) -> Self {
+    pub fn new(
+        filter: F,
+        filename: PathBuf,
+        module_mapping: Option<RangeMap<usize, (u16, String)>>,
+        full_trace: bool,
+    ) -> Self {
         if full_trace {
             let _ = DRCOV_IDS.lock().unwrap().insert(vec![]);
         }
         let _ = DRCOV_MAP.lock().unwrap().insert(HashMap::new());
         let _ = DRCOV_LENGTHS.lock().unwrap().insert(HashMap::new());
         Self {
             filter,
-            module_mapping: RangeMap::new(),
+            module_mapping,
             filename,
             full_trace,
             drcov_len: 0,
@@ -77,6 +158,8 @@ where
     S: Unpin + UsesInput + HasMetadata,
 {
     type ModuleAddressFilter = F;
+    #[cfg(emulation_mode = "systemmode")]
+    type ModulePageFilter = NopPageFilter;
 
     fn init_module<ET>(&self, emulator_modules: &mut EmulatorModules<ET, S>)
     where
@@ -89,25 +172,47 @@ where
         );
     }
 
+    #[cfg(emulation_mode = "usermode")]
     fn first_exec<ET>(&mut self, emulator_modules: &mut EmulatorModules<ET, S>, _state: &mut S)
     where
         ET: EmulatorModuleTuple<S>,
     {
-        let qemu = emulator_modules.qemu();
-
-        for (i, (r, p)) in qemu
-            .mappings()
-            .filter_map(|m| {
-                m.path()
-                    .map(|p| ((m.start() as usize)..(m.end() as usize), p.to_string()))
-                    .filter(|(_, p)| !p.is_empty())
-            })
-            .enumerate()
-        {
-            self.module_mapping.insert(r, (i as u16, p));
+        if self.module_mapping.is_none() {
+            log::info!("Auto-filling module mapping for DrCov module from QEMU mapping.");
+
+            let qemu = emulator_modules.qemu();
+
+            let mut module_mapping: RangeMap<usize, (u16, String)> = RangeMap::new();
+
+            for (i, (r, p)) in qemu
+                .mappings()
+                .filter_map(|m| {
+                    m.path()
+                        .map(|p| ((m.start() as usize)..(m.end() as usize), p.to_string()))
+                        .filter(|(_, p)| !p.is_empty())
+                })
+                .enumerate()
+            {
+                module_mapping.insert(r, (i as u16, p));
+            }
+
+            self.module_mapping = Some(module_mapping);
+        } else {
+            log::info!("Using user-provided module mapping for DrCov module.");
         }
     }
 
+    #[cfg(emulation_mode = "systemmode")]
+    fn first_exec<ET>(&mut self, _emulator_modules: &mut EmulatorModules<ET, S>, _state: &mut S)
+    where
+        ET: EmulatorModuleTuple<S>,
+    {
+        assert!(
+            self.module_mapping.is_some(),
+            "DrCov should have a module mapping already set."
+        );
+    }
+
     fn post_exec<OT, ET>(
         &mut self,
         _emulator_modules: &mut EmulatorModules<ET, S>,
@@ -127,13 +232,18 @@ where
                 for id in DRCOV_IDS.lock().unwrap().as_ref().unwrap() {
                     'pcs_full: for (pc, idm) in DRCOV_MAP.lock().unwrap().as_ref().unwrap() {
                         let mut module_found = false;
-                        for module in self.module_mapping.iter() {
-                            let (range, (_, _)) = module;
-                            if *pc >= range.start.try_into().unwrap()
-                                && *pc <= range.end.try_into().unwrap()
-                            {
-                                module_found = true;
-                                break;
+                        // # Safety
+                        //
+                        // Module mapping is already set. It's checked or filled when the module is first run.
+                        unsafe {
+                            for module in self.module_mapping.as_ref().unwrap_unchecked().iter() {
+                                let (range, (_, _)) = module;
+                                if *pc >= range.start.try_into().unwrap()
+                                    && *pc <= range.end.try_into().unwrap()
+                                {
+                                    module_found = true;
+                                    break;
+                                }
                             }
                         }
                         if !module_found {
@@ -155,23 +265,33 @@ where
                     }
                 }
 
-                DrCovWriter::new(&self.module_mapping)
-                    .write(&self.filename, &drcov_vec)
-                    .expect("Failed to write coverage file");
+                // # Safety
+                //
+                // Module mapping is already set. It's checked or filled when the module is first run.
+                unsafe {
+                    DrCovWriter::new(self.module_mapping.as_ref().unwrap_unchecked())
+                        .write(&self.filename, &drcov_vec)
+                        .expect("Failed to write coverage file");
+                }
             }
             self.drcov_len = DRCOV_IDS.lock().unwrap().as_ref().unwrap().len();
         } else {
             if DRCOV_MAP.lock().unwrap().as_ref().unwrap().len() > self.drcov_len {
                 let mut drcov_vec = Vec::<DrCovBasicBlock>::new();
                 'pcs: for (pc, _) in DRCOV_MAP.lock().unwrap().as_ref().unwrap() {
                     let mut module_found = false;
-                    for module in self.module_mapping.iter() {
-                        let (range, (_, _)) = module;
-                        if *pc >= range.start.try_into().unwrap()
-                            && *pc <= range.end.try_into().unwrap()
-                        {
-                            module_found = true;
-                            break;
+                    // # Safety
+                    //
+                    // Module mapping is already set. It's checked or filled when the module is first run.
+                    unsafe {
+                        for module in self.module_mapping.as_ref().unwrap_unchecked().iter() {
+                            let (range, (_, _)) = module;
+                            if *pc >= range.start.try_into().unwrap()
+                                && *pc <= range.end.try_into().unwrap()
+                            {
+                                module_found = true;
+                                break;
+                            }
                         }
                     }
                     if !module_found {
@@ -190,9 +310,14 @@ where
                     }
                 }
 
-                DrCovWriter::new(&self.module_mapping)
-                    .write(&self.filename, &drcov_vec)
-                    .expect("Failed to write coverage file");
+                // # Safety
+                //
+                // Module mapping is already set. It's checked or filled when the module is first run.
+                unsafe {
+                    DrCovWriter::new(self.module_mapping.as_ref().unwrap_unchecked())
+                        .write(&self.filename, &drcov_vec)
+                        .expect("Failed to write coverage file");
+                }
             }
             self.drcov_len = DRCOV_MAP.lock().unwrap().as_ref().unwrap().len();
         }
@@ -205,6 +330,16 @@ where
     fn address_filter_mut(&mut self) -> &mut Self::ModuleAddressFilter {
         &mut self.filter
     }
+
+    #[cfg(emulation_mode = "systemmode")]
+    fn page_filter(&self) -> &Self::ModulePageFilter {
+        &NopPageFilter
+    }
+
+    #[cfg(emulation_mode = "systemmode")]
+    fn page_filter_mut(&mut self) -> &mut Self::ModulePageFilter {
+        unsafe { addr_of_mut!(NOP_PAGE_FILTER).as_mut().unwrap().get_mut() }
+    }
 }
 
 pub fn gen_unique_block_ids<ET, F, S>(

libafl_qemu/src/modules/mod.rs

@@ -18,7 +18,7 @@ pub mod systemmode;
 pub use systemmode::*;
 
 pub mod edges;
-pub use edges::EdgeCoverageModule;
+pub use edges::*;
 
 #[cfg(not(cpu_target = "hexagon"))]
 pub mod calls;
@@ -30,6 +30,11 @@ pub mod cmplog;
 #[cfg(not(any(cpu_target = "mips", cpu_target = "hexagon")))]
 pub use cmplog::CmpLogModule;
 
+#[cfg(not(cpu_target = "hexagon"))]
+pub mod drcov;
+#[cfg(not(cpu_target = "hexagon"))]
+pub use drcov::*;
+
 use crate::{emu::EmulatorModules, Qemu};
 
 /// A module for `libafl_qemu`.

libafl_qemu/src/modules/usermode/mod.rs

@@ -1,9 +1,3 @@
-#[cfg(not(cpu_target = "hexagon"))]
-pub mod drcov;
-
-#[cfg(not(cpu_target = "hexagon"))]
-pub use drcov::DrCovModule;
-
 #[cfg(feature = "injections")]
 pub mod injections;
 #[cfg(feature = "injections")]
@@ -12,9 +6,7 @@ pub use injections::InjectionModule;
 #[cfg(not(cpu_target = "hexagon"))]
 pub mod snapshot;
 #[cfg(not(cpu_target = "hexagon"))]
-pub use snapshot::IntervalSnapshotFilter;
-#[cfg(not(cpu_target = "hexagon"))]
-pub use snapshot::SnapshotModule;
+pub use snapshot::{IntervalSnapshotFilter, SnapshotModule};
 
 #[cfg(not(cpu_target = "hexagon"))]
 pub mod asan;