Adding DrCov for qemu (#878)

* Adding DrCov for qemu

* Fixing cargo fmt

* Trying to fix maturin build

* Fixing clippy

* libafl_qemu --no-default-features fix

* Adding make module mapping a user input as suggested from @WorksButNotTested

* Switching from blocks_raw() -> blocks() and full_tracing as an option

* Avoiding get before get_mut

* HashSet to Vec

* Avoiding lazy_static

* Adding DrCov for example fuzzer qemu_arm_launcher

* Removing mut for globals in DrCov

* Using emu.mappings() for drcov module mappings

* Fixing clippy

Co-authored-by: Dominik Maier <domenukk@gmail.com>
Co-authored-by: Andrea Fioraldi <andreafioraldi@gmail.com>
Co-authored-by: Dongjia "toka" Zhang <tokazerkje@outlook.com>

fuzzers/qemu_arm_launcher/.gitignore

@@ -4,3 +4,4 @@ libpng_harness_crashing
 zlib-*
 crashes
 target
+drcov.log

fuzzers/qemu_arm_launcher/Cargo.toml

@@ -17,3 +17,4 @@ debug = true
 [dependencies]
 libafl = { path = "../../libafl/" }
 libafl_qemu = { path = "../../libafl_qemu/", features = ["arm", "usermode"] }
+rangemap = "1.0.3"

fuzzers/qemu_arm_launcher/src/fuzzer.rs

@@ -29,6 +29,7 @@ use libafl::{
     Error,
 };
 use libafl_qemu::{
+    drcov::QemuDrCovHelper,
     //asan::QemuAsanHelper,
     edges,
     edges::QemuEdgeCoverageHelper,
@@ -37,8 +38,10 @@ use libafl_qemu::{
     MmapPerms,
     QemuExecutor,
     QemuHooks,
+    QemuInstrumentationFilter,
     Regs,
 };
+use rangemap::RangeMap;
 
 pub fn fuzz() {
     // Hardcoded parameters
@@ -151,7 +154,30 @@ pub fn fuzz() {
         // A fuzzer with feedbacks and a corpus scheduler
         let mut fuzzer = StdFuzzer::new(scheduler, feedback, objective);
 
-        let mut hooks = QemuHooks::new(&emu, tuple_list!(QemuEdgeCoverageHelper::default()));
+        let mut rangemap = RangeMap::<usize, (u16, String)>::new();
+        let mappings = emu.mappings();
+        let mut idx = 0;
+        for map in mappings {
+            if map.path().unwrap() != "" {
+                rangemap.insert(
+                    (map.start() as usize)..(map.end() as usize),
+                    (idx, map.path().unwrap().to_string()),
+                );
+                idx += 1;
+            }
+        }
+        let mut hooks = QemuHooks::new(
+            &emu,
+            tuple_list!(
+                QemuEdgeCoverageHelper::default(),
+                QemuDrCovHelper::new(
+                    QemuInstrumentationFilter::None,
+                    rangemap,
+                    PathBuf::from("drcov.log"),
+                    false,
+                )
+            ),
+        );
 
         // Create a QEMU in-process executor
         let executor = QemuExecutor::new(

libafl_qemu/Cargo.toml

@@ -46,6 +46,7 @@ thread_local = "1.1.4"
 capstone = "0.11.0"
 #pyo3 = { version = "0.15", features = ["extension-module"], optional = true }
 pyo3 = { version = "0.17", features = ["pyproto"], optional = true }
+rangemap = "1.0"
 
 [build-dependencies]
 cc = { version = "1.0" }

libafl_qemu/src/blocks.rs

@@ -0,0 +1,113 @@
+use capstone::prelude::*;
+
+use crate::{Emulator, GuestAddr};
+
+pub struct Instruction {
+    pub start_addr: GuestAddr,
+    pub mnemonic: String,
+    pub operands: String,
+    pub insn_len: usize,
+}
+
+/*
+ *  Generating the basic block from it's starting address (pc)
+ *  Basic block:
+ *  - Starting at pc
+ *  - Ending at the first branch/jump/interrupt/call instruction
+ *  Output:
+ *  - Vector of instructions
+ *      - Start address
+ *      - mnemonic string
+ *      - operand string
+ *      - instruction length
+ */
+pub fn pc2basicblock(pc: GuestAddr, emu: &Emulator) -> Result<Vec<Instruction>, String> {
+    #[allow(unused_mut)]
+    let mut code = {
+        #[cfg(emulation_mode = "usermode")]
+        unsafe {
+            std::slice::from_raw_parts(emu.g2h(pc), 512)
+        }
+        #[cfg(emulation_mode = "systemmode")]
+        &mut [0; 512]
+    };
+    #[cfg(emulation_mode = "systemmode")]
+    unsafe {
+        emu.read_mem(pc, code)
+    };
+    // TODO better fault handling
+    if code.iter().all(|&x| x == 0) {
+        return Err("Memory region is empty".to_string());
+    }
+
+    let mut iaddr = pc;
+    let mut block = Vec::<Instruction>::new();
+
+    #[cfg(cpu_target = "x86_64")]
+    let cs = Capstone::new()
+        .x86()
+        .mode(capstone::arch::x86::ArchMode::Mode64)
+        .detail(true)
+        .build()
+        .unwrap();
+    #[cfg(cpu_target = "arm")]
+    let cs = Capstone::new()
+        .arm()
+        .mode(capstone::arch::arm::ArchMode::Arm)
+        .detail(true)
+        .build()
+        .unwrap();
+    #[cfg(cpu_target = "aarch64")]
+    let cs = Capstone::new()
+        .arm64()
+        .mode(capstone::arch::arm64::ArchMode::Arm)
+        .detail(true)
+        .build()
+        .unwrap();
+
+    'disasm: while let Ok(insns) = cs.disasm_count(code, iaddr.into(), 1) {
+        if insns.is_empty() {
+            break;
+        }
+        let insn = insns.first().unwrap();
+        let insn_detail: InsnDetail = cs.insn_detail(insn).unwrap();
+        block.push(Instruction {
+            start_addr: insn.address() as GuestAddr,
+            mnemonic: insn.mnemonic().unwrap().to_string(),
+            operands: insn.op_str().unwrap().to_string(),
+            insn_len: insn.len(),
+        });
+        for detail in insn_detail.groups() {
+            match u32::from(detail.0) {
+                capstone::InsnGroupType::CS_GRP_BRANCH_RELATIVE
+                | capstone::InsnGroupType::CS_GRP_CALL
+                | capstone::InsnGroupType::CS_GRP_INT
+                | capstone::InsnGroupType::CS_GRP_INVALID
+                | capstone::InsnGroupType::CS_GRP_IRET
+                | capstone::InsnGroupType::CS_GRP_JUMP
+                | capstone::InsnGroupType::CS_GRP_PRIVILEGE
+                | capstone::InsnGroupType::CS_GRP_RET => {
+                    break 'disasm;
+                }
+                _ => {}
+            }
+        }
+
+        iaddr += insn.bytes().len() as GuestAddr;
+
+        #[cfg(emulation_mode = "usermode")]
+        unsafe {
+            code = std::slice::from_raw_parts(emu.g2h(iaddr), 512);
+        }
+        #[cfg(emulation_mode = "systemmode")]
+        unsafe {
+            emu.read_mem(iaddr, code);
+        }
+        // TODO better fault handling
+        if code.iter().all(|&x| x == 0) {
+            return Err("Memory region is empty".to_string());
+        }
+    }
+
+    Ok(block)
+}

libafl_qemu/src/drcov.rs

@@ -0,0 +1,233 @@
+use std::{path::PathBuf, sync::Mutex};
+
+use hashbrown::{hash_map::Entry, HashMap};
+use libafl::{inputs::UsesInput, state::HasMetadata};
+use libafl_targets::drcov::{DrCovBasicBlock, DrCovWriter};
+use rangemap::RangeMap;
+use serde::{Deserialize, Serialize};
+
+use crate::{
+    blocks::pc2basicblock,
+    emu::GuestAddr,
+    helper::{QemuHelper, QemuHelperTuple, QemuInstrumentationFilter},
+    hooks::QemuHooks,
+    Emulator,
+};
+
+static DRCOV_IDS: Mutex<Option<Vec<u64>>> = Mutex::new(None);
+static DRCOV_MAP: Mutex<Option<HashMap<GuestAddr, u64>>> = Mutex::new(None);
+
+#[derive(Debug, Default, Serialize, Deserialize)]
+pub struct QemuDrCovMetadata {
+    pub current_id: u64,
+}
+
+impl QemuDrCovMetadata {
+    #[must_use]
+    pub fn new() -> Self {
+        Self { current_id: 0 }
+    }
+}
+
+libafl::impl_serdeany!(QemuDrCovMetadata);
+
+#[derive(Debug)]
+pub struct QemuDrCovHelper {
+    filter: QemuInstrumentationFilter,
+    module_mapping: RangeMap<usize, (u16, String)>,
+    filename: PathBuf,
+    full_trace: bool,
+    drcov_len: usize,
+}
+
+impl QemuDrCovHelper {
+    #[must_use]
+    pub fn new(
+        filter: QemuInstrumentationFilter,
+        module_mapping: RangeMap<usize, (u16, String)>,
+        filename: PathBuf,
+        full_trace: bool,
+    ) -> Self {
+        if full_trace {
+            let _ = DRCOV_IDS.lock().unwrap().insert(vec![]);
+        }
+        let _ = DRCOV_MAP.lock().unwrap().insert(HashMap::new());
+        Self {
+            filter,
+            module_mapping,
+            filename,
+            full_trace,
+            drcov_len: 0,
+        }
+    }
+
+    #[must_use]
+    pub fn must_instrument(&self, addr: u64) -> bool {
+        self.filter.allowed(addr)
+    }
+}
+
+impl<S> QemuHelper<S> for QemuDrCovHelper
+where
+    S: UsesInput + HasMetadata,
+{
+    fn init_hooks<QT>(&self, hooks: &QemuHooks<'_, QT, S>)
+    where
+        QT: QemuHelperTuple<S>,
+    {
+        hooks.blocks(
+            Some(gen_unique_block_ids::<QT, S>),
+            Some(exec_trace_block::<QT, S>),
+        );
+    }
+
+    fn pre_exec(&mut self, _emulator: &Emulator, _input: &S::Input) {}
+
+    fn post_exec(&mut self, emulator: &Emulator, _input: &S::Input) {
+        if self.full_trace {
+            if DRCOV_IDS.lock().unwrap().as_ref().unwrap().len() > self.drcov_len {
+                let mut drcov_vec = Vec::<DrCovBasicBlock>::new();
+                for id in DRCOV_IDS.lock().unwrap().as_ref().unwrap().iter() {
+                    'pcs_full: for (pc, idm) in DRCOV_MAP.lock().unwrap().as_ref().unwrap().iter() {
+                        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;
+                            }
+                        }
+                        if !module_found {
+                            continue 'pcs_full;
+                        }
+                        if *idm == *id {
+                            match pc2basicblock(*pc, emulator) {
+                                Ok(block) => {
+                                    let mut block_len = 0;
+                                    for instr in &block {
+                                        block_len += instr.insn_len;
+                                    }
+                                    drcov_vec.push(DrCovBasicBlock::new(
+                                        *pc as usize,
+                                        *pc as usize + block_len,
+                                    ));
+                                }
+                                Err(r) => println!("{r:#?}"),
+                            }
+                        }
+                    }
+                }
+
+                DrCovWriter::new(&self.module_mapping)
+                    .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().iter() {
+                    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;
+                        }
+                    }
+                    if !module_found {
+                        continue 'pcs;
+                    }
+                    match pc2basicblock(*pc, emulator) {
+                        Ok(block) => {
+                            let mut block_len = 0;
+                            for instr in &block {
+                                block_len += instr.insn_len;
+                            }
+                            drcov_vec
+                                .push(DrCovBasicBlock::new(*pc as usize, *pc as usize + block_len));
+                        }
+                        Err(r) => println!("{r:#?}"),
+                    }
+                }
+
+                DrCovWriter::new(&self.module_mapping)
+                    .write(&self.filename, &drcov_vec)
+                    .expect("Failed to write coverage file");
+            }
+            self.drcov_len = DRCOV_MAP.lock().unwrap().as_ref().unwrap().len();
+        }
+    }
+}
+
+pub fn gen_unique_block_ids<QT, S>(
+    hooks: &mut QemuHooks<'_, QT, S>,
+    state: Option<&mut S>,
+    pc: GuestAddr,
+) -> Option<u64>
+where
+    S: HasMetadata,
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    let drcov_helper = hooks
+        .helpers()
+        .match_first_type::<QemuDrCovHelper>()
+        .unwrap();
+    if !drcov_helper.must_instrument(pc.into()) {
+        return None;
+    }
+
+    let state = state.expect("The gen_unique_block_ids hook works only for in-process fuzzing");
+    if state
+        .metadata_mut()
+        .get_mut::<QemuDrCovMetadata>()
+        .is_none()
+    {
+        state.add_metadata(QemuDrCovMetadata::new());
+    }
+    let meta = state.metadata_mut().get_mut::<QemuDrCovMetadata>().unwrap();
+
+    match DRCOV_MAP.lock().unwrap().as_mut().unwrap().entry(pc) {
+        Entry::Occupied(e) => {
+            let id = *e.get();
+            if drcov_helper.full_trace {
+                Some(id)
+            } else {
+                None
+            }
+        }
+        Entry::Vacant(e) => {
+            let id = meta.current_id;
+            e.insert(id);
+            meta.current_id = id + 1;
+            if drcov_helper.full_trace {
+                // GuestAddress is u32 for 32 bit guests
+                #[allow(clippy::unnecessary_cast)]
+                Some(id as u64)
+            } else {
+                None
+            }
+        }
+    }
+}
+
+pub fn exec_trace_block<QT, S>(hooks: &mut QemuHooks<'_, QT, S>, _state: Option<&mut S>, id: u64)
+where
+    S: HasMetadata,
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    if hooks
+        .helpers()
+        .match_first_type::<QemuDrCovHelper>()
+        .unwrap()
+        .full_trace
+    {
+        DRCOV_IDS.lock().unwrap().as_mut().unwrap().push(id);
+    }
+}

libafl_qemu/src/executor.rs

@@ -17,7 +17,6 @@ use libafl::{
     Error,
 };
 
-pub use crate::emu::SyscallHookResult;
 use crate::{emu::Emulator, helper::QemuHelperTuple, hooks::QemuHooks};
 
 pub struct QemuExecutor<'a, H, OT, QT, S>

libafl_qemu/src/lib.rs

@@ -49,6 +49,7 @@ pub use hooks::*;
 pub mod edges;
 pub use edges::QemuEdgeCoverageHelper;
 pub mod cmplog;
+pub mod drcov;
 pub use cmplog::QemuCmpLogHelper;
 #[cfg(emulation_mode = "usermode")]
 pub mod snapshot;
@@ -59,6 +60,7 @@ pub mod asan;
 #[cfg(emulation_mode = "usermode")]
 pub use asan::{init_with_asan, QemuAsanHelper};
 
+pub mod blocks;
 pub mod calls;
 
 pub mod executor;