diff --git a/fuzzers/binary_only/qemu_coverage/src/fuzzer.rs b/fuzzers/binary_only/qemu_coverage/src/fuzzer.rs index 8852109250..a7bc1506ab 100644 --- a/fuzzers/binary_only/qemu_coverage/src/fuzzer.rs +++ b/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) diff --git a/libafl_qemu/src/modules/usermode/drcov.rs b/libafl_qemu/src/modules/drcov.rs similarity index 59% rename from libafl_qemu/src/modules/usermode/drcov.rs rename to libafl_qemu/src/modules/drcov.rs index 93d0036046..22ced9a12d 100644 --- a/libafl_qemu/src/modules/usermode/drcov.rs +++ b/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,6 +9,8 @@ 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}, @@ -35,10 +39,71 @@ impl DrCovMetadata { libafl_bolts::impl_serdeany!(DrCovMetadata); +#[derive(Debug)] +pub struct DrCovModuleBuilder { + filter: Option, + module_mapping: Option>, + filename: Option, + full_trace: Option, +} + +impl DrCovModuleBuilder +where + F: AddressFilter, +{ + pub fn build(self) -> DrCovModule { + DrCovModule::new( + self.filter.unwrap(), + self.filename.unwrap(), + self.module_mapping, + self.full_trace.unwrap(), + ) + } + + pub fn filter(self, filter: F2) -> DrCovModuleBuilder { + 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) -> 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 { filter: F, - module_mapping: RangeMap, + module_mapping: Option>, filename: PathBuf, full_trace: bool, drcov_len: usize, @@ -48,9 +113,24 @@ impl DrCovModule where F: AddressFilter, { + #[must_use] + pub fn builder() -> DrCovModuleBuilder { + DrCovModuleBuilder { + filter: None, + module_mapping: None, + full_trace: None, + filename: None, + } + } + #[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>, + full_trace: bool, + ) -> Self { if full_trace { let _ = DRCOV_IDS.lock().unwrap().insert(vec![]); } @@ -58,7 +138,7 @@ where let _ = DRCOV_LENGTHS.lock().unwrap().insert(HashMap::new()); Self { filter, - module_mapping: RangeMap::new(), + module_mapping, filename, full_trace, drcov_len: 0, @@ -77,6 +157,8 @@ where S: Unpin + UsesInput + HasMetadata, { type ModuleAddressFilter = F; + #[cfg(emulation_mode = "systemmode")] + type ModulePageFilter = NopPageFilter; fn init_module(&self, emulator_modules: &mut EmulatorModules) where @@ -89,25 +171,47 @@ where ); } + #[cfg(emulation_mode = "usermode")] fn first_exec(&mut self, emulator_modules: &mut EmulatorModules, _state: &mut S) where ET: EmulatorModuleTuple, { - 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 = 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(&mut self, _emulator_modules: &mut EmulatorModules, _state: &mut S) + where + ET: EmulatorModuleTuple, + { + assert!( + self.module_mapping.is_some(), + "DrCov should have a module mapping already set." + ); + } + fn post_exec( &mut self, _emulator_modules: &mut EmulatorModules, @@ -127,13 +231,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,9 +264,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_IDS.lock().unwrap().as_ref().unwrap().len(); } else { @@ -165,13 +279,18 @@ where let mut drcov_vec = Vec::::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 +309,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 +329,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( diff --git a/libafl_qemu/src/modules/mod.rs b/libafl_qemu/src/modules/mod.rs index 572de025df..7b5f68a2c9 100644 --- a/libafl_qemu/src/modules/mod.rs +++ b/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`. diff --git a/libafl_qemu/src/modules/usermode/mod.rs b/libafl_qemu/src/modules/usermode/mod.rs index 1b53eb77b5..1d9cc503d2 100644 --- a/libafl_qemu/src/modules/usermode/mod.rs +++ b/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;