add jmp instrumentaion

fuzzers/wcet_qemu_sys/src/sysstate/helpers.rs

@@ -105,17 +105,7 @@ where
     unsafe {
         match SAVED_JUMP.take() {
             Some(s) => {
-                let meta = state
+                sysstate.last_pc = Some(s.0);
-                    .metadata_mut()
-                    .get::<QemuEdgesMapMetadata>()
-                    .unwrap();
-                for (k,v) in meta.map.iter() {
-                    if *v==s {
-                        // println!("Jump Saved {:x}",k.1);
-                        sysstate.last_pc = Some(k.0);   /* Currently save APP entry points */
-                        break;
-                    }
-                }
             },
             None => (),
         }

libafl_qemu/src/edges.rs

@@ -93,9 +93,10 @@ where
         QT: QemuHelperTuple<I, S>,
     {
         executor.hook_edge_generation(gen_unique_edge_ids::<I, QT, S>);
-        match self.app_range {
+        executor.emulator().set_exec_edge_hook(trace_edge_hitcount);
-            None => executor.emulator().set_exec_edge_hook(trace_edge_hitcount),
+        if self.app_range.is_some() {
-            Some(_) => executor.emulator().set_exec_edge_hook(trace_edge_hitcount_and_app_range),
+            executor.hook_jmp_generation(gen_jmp_instrument::<I, QT, S>);
+            executor.emulator().set_exec_jmp_hook(trace_jmp)
         }
     }
     
@@ -113,8 +114,31 @@ fn hash_me(mut x: u64) -> u64 {
     x
 }
 
-pub static mut MUST_SAVE: [bool; EDGES_MAP_SIZE] = [false; EDGES_MAP_SIZE];
+/// Save the source and destination of last interesting jump
-pub static mut SAVED_JUMP: Option<u64> = None;
+pub static mut SAVED_JUMP: Option<(u64, u64)> = None;
+
+pub fn gen_jmp_instrument<I, QT, S>(
+    _emulator: &Emulator,
+    helpers: &mut QT,
+    state: &mut S,
+    src: u64,
+    dest: u64,
+) -> Option<u64>
+where
+    S: HasMetadata,
+    I: Input,
+    QT: QemuHelperTuple<I, S>,
+{
+    if let Some(h) = helpers.match_first_type::<QemuEdgeCoverageHelper>() {
+        if !h.must_instrument(src) && !h.must_instrument(dest) {
+            return None;
+        }
+        if !h.must_save(src, dest) {
+            return None;
+        }
+    }
+    Some(1)
+}
 
 pub fn gen_unique_edge_ids<I, QT, S>(
     _emulator: &Emulator,
@@ -128,12 +152,10 @@ where
     I: Input,
     QT: QemuHelperTuple<I, S>,
 {
-    let mut must_save = false;
     if let Some(h) = helpers.match_first_type::<QemuEdgeCoverageHelper>() {
         if !h.must_instrument(src) && !h.must_instrument(dest) {
             return None;
         }
-        must_save = h.must_save(src, dest);
     }
     if state.metadata().get::<QemuEdgesMapMetadata>().is_none() {
         state.add_metadata(QemuEdgesMapMetadata::new());
@@ -149,9 +171,6 @@ where
             let nxt = (id as usize + 1) & (EDGES_MAP_SIZE - 1);
             unsafe {
                 MAX_EDGES_NUM = max(MAX_EDGES_NUM, nxt);
-                if must_save {
-                    MUST_SAVE[id as usize] = true;
-                }
             }
             Some(id)
         }
@@ -161,9 +180,6 @@ where
             meta.current_id = (id + 1) & (EDGES_MAP_SIZE as u64 - 1);
             unsafe {
                 MAX_EDGES_NUM = meta.current_id as usize;
-                if must_save {
-                    MUST_SAVE[id as usize] = true;
-                }
             }
             Some(id as u64)
         }
@@ -189,11 +205,9 @@ where
     Some(hash_me(src) ^ hash_me(dest))
 }
 
-
+pub extern "C" fn trace_jmp(src: u64, des: u64) {
-pub extern "C" fn trace_edge_hitcount_and_app_range(id: u64) {
     unsafe {
-        if MUST_SAVE[id as usize] { SAVED_JUMP=Some(id); }
+        SAVED_JUMP=Some((src, des));
-        EDGES_MAP[id as usize] = EDGES_MAP[id as usize].wrapping_add(1);
     }
 }
 

libafl_qemu/src/emu.rs

@@ -223,6 +223,10 @@ extern "C" {
 
     static mut libafl_exec_edge_hook: unsafe extern "C" fn(u64);
     static mut libafl_gen_edge_hook: unsafe extern "C" fn(u64, u64) -> u64;
+    #[cfg(feature = "systemmode")]
+    static mut libafl_exec_jmp_hook: unsafe extern "C" fn(u64, u64);
+    #[cfg(feature = "systemmode")]
+    static mut libafl_gen_jmp_hook: unsafe extern "C" fn(u64, u64) -> u64;
     static mut libafl_exec_block_hook: unsafe extern "C" fn(u64);
     static mut libafl_gen_block_hook: unsafe extern "C" fn(u64) -> u64;
 
@@ -532,6 +536,20 @@ impl Emulator {
         }
     }
 
+    #[cfg(feature = "systemmode")]
+    pub fn set_exec_jmp_hook(&self, hook: extern "C" fn(src: u64, dest: u64)) {
+        unsafe {
+            libafl_exec_jmp_hook = hook;
+        }
+    }
+
+    #[cfg(feature = "systemmode")]
+    pub fn set_gen_jmp_hook(&self, hook: extern "C" fn(src: u64, dest: u64) -> u64) {
+        unsafe {
+            libafl_gen_jmp_hook = hook;
+        }
+    }
+
     pub fn set_exec_block_hook(&self, hook: extern "C" fn(pc: u64)) {
         unsafe {
             libafl_exec_block_hook = hook;

libafl_qemu/src/executor.rs

@@ -27,6 +27,37 @@ use crate::{
 
 static mut QEMU_HELPERS_PTR: *const c_void = ptr::null();
 
+static mut GEN_JMP_HOOK_PTR: *const c_void = ptr::null();
+extern "C" fn gen_jmp_hook_wrapper<I, QT, S>(src: u64, dst: u64) -> u64
+where
+    I: Input,
+    QT: QemuHelperTuple<I, S>,
+{
+    unsafe {
+        let helpers = (QEMU_HELPERS_PTR as *mut QT).as_mut().unwrap();
+        let state = inprocess_get_state::<S>().unwrap();
+        let emulator = Emulator::new_empty();
+        let func: fn(&Emulator, &mut QT, &mut S, u64, u64) -> Option<u64> =
+            transmute(GEN_JMP_HOOK_PTR);
+        (func)(&emulator, helpers, state, src, dst).map_or(SKIP_EXEC_HOOK, |id| id)
+    }
+}
+
+static mut JMP_HOOKS: Vec<*const c_void> = vec![];
+extern "C" fn jmp_hooks_wrapper<I, QT, S>(src: u64, dst: u64)
+where
+    I: Input,
+    QT: QemuHelperTuple<I, S>,
+{
+    let helpers = unsafe { (QEMU_HELPERS_PTR as *mut QT).as_mut().unwrap() };
+    let state = inprocess_get_state::<S>().unwrap();
+    let emulator = Emulator::new_empty();
+    for hook in unsafe { &JMP_HOOKS } {
+        let func: fn(&Emulator, &mut QT, &mut S, u64, u64) = unsafe { transmute(*hook) };
+        (func)(&emulator, helpers, state, src, dst);
+    }
+}
+
 static mut GEN_EDGE_HOOK_PTR: *const c_void = ptr::null();
 extern "C" fn gen_edge_hook_wrapper<I, QT, S>(src: u64, dst: u64) -> u64
 where
@@ -530,6 +561,29 @@ where
             .set_exec_edge_hook(edge_hooks_wrapper::<I, QT, S>);
     }
 
+    #[cfg(feature = "systemmode")]
+    #[allow(clippy::unused_self)]
+    pub fn hook_jmp_generation(
+        &self,
+        hook: fn(&Emulator, &mut QT, &mut S, src: u64, dest: u64) -> Option<u64>,
+    ) {
+        unsafe {
+            GEN_JMP_HOOK_PTR = hook as *const _;
+        }
+        self.emulator
+            .set_gen_jmp_hook(gen_jmp_hook_wrapper::<I, QT, S>);
+    }
+
+    #[cfg(feature = "systemmode")]
+    #[allow(clippy::unused_self)]
+    pub fn hook_jmp_execution(&self, hook: fn(&Emulator, &mut QT, &mut S, src: u64, dest: u64)) {
+        unsafe {
+            JMP_HOOKS.push(hook as *const _);
+        }
+        self.emulator
+            .set_exec_jmp_hook(jmp_hooks_wrapper::<I, QT, S>);
+    }
+
     #[allow(clippy::unused_self)]
     pub fn hook_block_generation(
         &self,