Post gen (#1282)

* Add post_gen

* Adopt post_gen hooks in DrCovHelper

* Bump qemu-libafl-bridge revision

---------

Co-authored-by: Your Name <you@example.com>
Co-authored-by: Andrea Fioraldi <andreafioraldi@gmail.com>

libafl_qemu/libafl_qemu_build/src/build.rs

@@ -8,7 +8,7 @@ use which::which;
 
 const QEMU_URL: &str = "https://github.com/AFLplusplus/qemu-libafl-bridge";
 const QEMU_DIRNAME: &str = "qemu-libafl-bridge";
-const QEMU_REVISION: &str = "0dc52ed6f3915f727aaec8648706760f278f0571";
+const QEMU_REVISION: &str = "6ae8b5bfb0bc4ac7a0b2ec463cf1000363836c0c";
 
 fn build_dep_check(tools: &[&str]) {
     for tool in tools {

libafl_qemu/src/blocks.rs

@@ -1,100 +0,0 @@
-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,
-    mode: Option<capstone::Mode>,
-) -> 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();
-
-    let mut cs = crate::capstone().detail(true).build().unwrap();
-    if let Some(m) = mode {
-        cs.set_mode(m).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/calls.rs

@@ -414,7 +414,7 @@ where
     where
         QT: QemuHelperTuple<S>,
     {
-        hooks.blocks(Some(Self::gen_blocks_calls::<QT, S>), None);
+        hooks.blocks(Some(Self::gen_blocks_calls::<QT, S>), None, None);
     }
 
     fn pre_exec(&mut self, emulator: &Emulator, input: &S::Input) {

libafl_qemu/src/drcov.rs

@@ -9,8 +9,7 @@ use rangemap::RangeMap;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    blocks::pc2basicblock,
+    emu::{GuestAddr, GuestUsize},
-    emu::GuestAddr,
     helper::{QemuHelper, QemuHelperTuple, QemuInstrumentationFilter},
     hooks::QemuHooks,
     Emulator,
@@ -18,6 +17,7 @@ use crate::{
 
 static DRCOV_IDS: Mutex<Option<Vec<u64>>> = Mutex::new(None);
 static DRCOV_MAP: Mutex<Option<HashMap<GuestAddr, u64>>> = Mutex::new(None);
+static DRCOV_LENGTHS: Mutex<Option<HashMap<GuestAddr, GuestUsize>>> = Mutex::new(None);
 
 #[derive(Debug, Default, Serialize, Deserialize)]
 pub struct QemuDrCovMetadata {
@@ -55,6 +55,7 @@ impl QemuDrCovHelper {
             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,
@@ -80,6 +81,7 @@ where
     {
         hooks.blocks(
             Some(gen_unique_block_ids::<QT, S>),
+            Some(gen_block_lengths::<QT, S>),
             Some(exec_trace_block::<QT, S>),
         );
     }
@@ -88,13 +90,15 @@ where
 
     fn post_exec<OT>(
         &mut self,
-        emulator: &Emulator,
+        _emulator: &Emulator,
         _input: &S::Input,
         _observers: &mut OT,
         _exit_kind: &mut ExitKind,
     ) where
         OT: ObserversTuple<S>,
     {
+        let lengths_opt = DRCOV_LENGTHS.lock().unwrap();
+        let lengths = lengths_opt.as_ref().unwrap();
         if self.full_trace {
             if DRCOV_IDS.lock().unwrap().as_ref().unwrap().len() > self.drcov_len {
                 let mut drcov_vec = Vec::<DrCovBasicBlock>::new();
@@ -114,27 +118,16 @@ where
                             continue 'pcs_full;
                         }
                         if *idm == *id {
-                            #[cfg(cpu_target = "arm")]
+                            match lengths.get(pc) {
-                            let mode = if pc & 1 == 1 {
+                                Some(block_length) => {
-                                Some(capstone::arch::arm::ArchMode::Thumb.into())
-                            } else {
-                                Some(capstone::arch::arm::ArchMode::Arm.into())
-                            };
-                            #[cfg(not(cpu_target = "arm"))]
-                            let mode = None;
-
-                            match pc2basicblock(*pc, emulator, mode) {
-                                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,
+                                        *pc as usize + *block_length as usize,
                                     ));
                                 }
-                                Err(r) => log::info!("{r:#?}"),
+                                None => {
+                                    log::info!("Failed to find block length for: {pc:}");
+                                }
                             }
                         }
                     }
@@ -162,26 +155,16 @@ where
                     if !module_found {
                         continue 'pcs;
                     }
-
+                    match lengths.get(pc) {
-                    #[cfg(cpu_target = "arm")]
+                        Some(block_length) => {
-                    let mode = if pc & 1 == 1 {
+                            drcov_vec.push(DrCovBasicBlock::new(
-                        Some(capstone::arch::arm::ArchMode::Thumb.into())
+                                *pc as usize,
-                    } else {
+                                *pc as usize + *block_length as usize,
-                        Some(capstone::arch::arm::ArchMode::Arm.into())
+                            ));
-                    };
-                    #[cfg(not(cpu_target = "arm"))]
-                    let mode = None;
-
-                    match pc2basicblock(*pc, emulator, mode) {
-                        Ok(block) => {
-                            let mut block_len = 0;
-                            for instr in &block {
-                                block_len += instr.insn_len;
                         }
-                            drcov_vec
+                        None => {
-                                .push(DrCovBasicBlock::new(*pc as usize, *pc as usize + block_len));
+                            log::info!("Failed to find block length for: {pc:}");
                         }
-                        Err(r) => log::info!("{r:#?}"),
                     }
                 }
 
@@ -249,6 +232,31 @@ where
     }
 }
 
+pub fn gen_block_lengths<QT, S>(
+    hooks: &mut QemuHooks<'_, QT, S>,
+    _state: Option<&mut S>,
+    pc: GuestAddr,
+    block_length: GuestUsize,
+) where
+    S: HasMetadata,
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    let drcov_helper = hooks
+        .helpers()
+        .match_first_type::<QemuDrCovHelper>()
+        .unwrap();
+    if !drcov_helper.must_instrument(pc) {
+        return;
+    }
+    DRCOV_LENGTHS
+        .lock()
+        .unwrap()
+        .as_mut()
+        .unwrap()
+        .insert(pc, block_length);
+}
+
 pub fn exec_trace_block<QT, S>(hooks: &mut QemuHooks<'_, QT, S>, _state: Option<&mut S>, id: u64)
 where
     S: HasMetadata,

libafl_qemu/src/emu.rs

@@ -336,6 +336,7 @@ extern "C" {
     // void libafl_add_block_hook(uint64_t (*gen)(target_ulong pc), void (*exec)(uint64_t id));
     fn libafl_add_block_hook(
         gen: Option<extern "C" fn(GuestAddr, u64) -> u64>,
+        post_gen: Option<extern "C" fn(GuestAddr, GuestUsize, u64)>,
         exec: Option<extern "C" fn(u64, u64)>,
         data: u64,
     );
@@ -1037,10 +1038,11 @@ impl Emulator {
     pub fn add_block_hooks(
         &self,
         gen: Option<extern "C" fn(GuestAddr, u64) -> u64>,
+        post_gen: Option<extern "C" fn(GuestAddr, GuestUsize, u64)>,
         exec: Option<extern "C" fn(u64, u64)>,
         data: u64,
     ) {
-        unsafe { libafl_add_block_hook(gen, exec, data) }
+        unsafe { libafl_add_block_hook(gen, post_gen, exec, data) }
     }
 
     pub fn add_read_hooks(

libafl_qemu/src/hooks.rs

@@ -15,7 +15,7 @@ pub use crate::emu::SyscallHookResult;
 use crate::{
     emu::{Emulator, FatPtr, MemAccessInfo, SKIP_EXEC_HOOK},
     helper::QemuHelperTuple,
-    GuestAddr,
+    GuestAddr, GuestUsize,
 };
 
 // all kinds of hooks
@@ -136,7 +136,7 @@ where
     }
 }
 
-static mut BLOCK_HOOKS: Vec<(Hook, Hook)> = vec![];
+static mut BLOCK_HOOKS: Vec<(Hook, Hook, Hook)> = vec![];
 
 extern "C" fn gen_block_hook_wrapper<QT, S>(pc: GuestAddr, index: u64) -> u64
 where
@@ -145,7 +145,7 @@ where
 {
     unsafe {
         let hooks = get_qemu_hooks::<QT, S>();
-        let (gen, _) = &mut BLOCK_HOOKS[index as usize];
+        let (gen, _, _) = &mut BLOCK_HOOKS[index as usize];
         match gen {
             Hook::Function(ptr) => {
                 let func: fn(&mut QemuHooks<'_, QT, S>, Option<&mut S>, GuestAddr) -> Option<u64> =
@@ -163,6 +163,34 @@ where
     }
 }
 
+extern "C" fn gen_post_block_hook_wrapper<QT, S>(
+    pc: GuestAddr,
+    block_length: GuestUsize,
+    index: u64,
+) where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    unsafe {
+        let hooks = get_qemu_hooks::<QT, S>();
+        let (_, post_gen, _) = &mut BLOCK_HOOKS[index as usize];
+        match post_gen {
+            Hook::Function(ptr) => {
+                let func: fn(&mut QemuHooks<'_, QT, S>, Option<&mut S>, GuestAddr, GuestUsize) =
+                    transmute(*ptr);
+                (func)(hooks, inprocess_get_state::<S>(), pc, block_length);
+            }
+            Hook::Closure(ptr) => {
+                let func: &mut Box<
+                    dyn FnMut(&mut QemuHooks<'_, QT, S>, Option<&mut S>, GuestAddr, GuestUsize),
+                > = transmute(ptr);
+                (func)(hooks, inprocess_get_state::<S>(), pc, block_length);
+            }
+            _ => (),
+        }
+    }
+}
+
 extern "C" fn exec_block_hook_wrapper<QT, S>(id: u64, index: u64)
 where
     S: UsesInput,
@@ -170,7 +198,7 @@ where
 {
     unsafe {
         let hooks = get_qemu_hooks::<QT, S>();
-        let (_, exec) = &mut BLOCK_HOOKS[index as usize];
+        let (_, _, exec) = &mut BLOCK_HOOKS[index as usize];
         match exec {
             Hook::Function(ptr) => {
                 let func: fn(&mut QemuHooks<'_, QT, S>, Option<&mut S>, u64) = transmute(*ptr);
@@ -859,6 +887,9 @@ where
     pub fn blocks(
         &self,
         generation_hook: Option<fn(&mut Self, Option<&mut S>, pc: GuestAddr) -> Option<u64>>,
+        post_generation_hook: Option<
+            fn(&mut Self, Option<&mut S>, pc: GuestAddr, block_length: GuestUsize),
+        >,
         execution_hook: Option<fn(&mut Self, Option<&mut S>, id: u64)>,
     ) {
         unsafe {
@@ -869,6 +900,11 @@ where
                 } else {
                     Some(gen_block_hook_wrapper::<QT, S>)
                 },
+                if post_generation_hook.is_none() {
+                    None
+                } else {
+                    Some(gen_post_block_hook_wrapper::<QT, S>)
+                },
                 if execution_hook.is_none() {
                     None
                 } else {
@@ -880,6 +916,9 @@ where
                 generation_hook.map_or(Hook::Empty, |hook| {
                     Hook::Function(hook as *const libc::c_void)
                 }),
+                post_generation_hook.map_or(Hook::Empty, |hook| {
+                    Hook::Function(hook as *const libc::c_void)
+                }),
                 execution_hook.map_or(Hook::Empty, |hook| {
                     Hook::Function(hook as *const libc::c_void)
                 }),
@@ -892,6 +931,9 @@ where
         generation_hook: Option<
             Box<dyn FnMut(&'a mut Self, Option<&'a mut S>, GuestAddr) -> Option<u64>>,
         >,
+        post_generation_hook: Option<
+            Box<dyn FnMut(&'a mut Self, Option<&'a mut S>, GuestAddr, GuestUsize)>,
+        >,
         execution_hook: Option<Box<dyn FnMut(&'a mut Self, Option<&'a mut S>, u64)>>,
     ) {
         let index = BLOCK_HOOKS.len();
@@ -901,6 +943,11 @@ where
             } else {
                 Some(gen_block_hook_wrapper::<QT, S>)
             },
+            if post_generation_hook.is_none() {
+                None
+            } else {
+                Some(gen_post_block_hook_wrapper::<QT, S>)
+            },
             if execution_hook.is_none() {
                 None
             } else {
@@ -910,6 +957,7 @@ where
         );
         BLOCK_HOOKS.push((
             generation_hook.map_or(Hook::Empty, |hook| Hook::Closure(transmute(hook))),
+            post_generation_hook.map_or(Hook::Empty, |hook| Hook::Closure(transmute(hook))),
             execution_hook.map_or(Hook::Empty, |hook| Hook::Closure(transmute(hook))),
         ));
     }
@@ -917,6 +965,9 @@ where
     pub fn blocks_raw(
         &self,
         generation_hook: Option<fn(&mut Self, Option<&mut S>, pc: GuestAddr) -> Option<u64>>,
+        post_generation_hook: Option<
+            fn(&mut Self, Option<&mut S>, pc: GuestAddr, block_length: GuestUsize),
+        >,
         execution_hook: Option<extern "C" fn(id: u64, data: u64)>,
     ) {
         unsafe {
@@ -927,6 +978,11 @@ where
                 } else {
                     Some(gen_block_hook_wrapper::<QT, S>)
                 },
+                if post_generation_hook.is_none() {
+                    None
+                } else {
+                    Some(gen_post_block_hook_wrapper::<QT, S>)
+                },
                 execution_hook,
                 index as u64,
             );
@@ -934,6 +990,9 @@ where
                 generation_hook.map_or(Hook::Empty, |hook| {
                     Hook::Function(hook as *const libc::c_void)
                 }),
+                post_generation_hook.map_or(Hook::Empty, |hook| {
+                    Hook::Function(hook as *const libc::c_void)
+                }),
                 Hook::Empty,
             ));
         }

libafl_qemu/src/lib.rs

@@ -74,8 +74,6 @@ pub mod asan;
 #[cfg(emulation_mode = "usermode")]
 pub use asan::{init_with_asan, QemuAsanHelper};
 
-pub mod blocks;
-
 pub mod calls;
 pub mod drcov;