libafl_qemu: add jmp instrumentation

libafl_qemu/src/emu.rs

@@ -399,6 +399,15 @@ extern "C" {
         data: *const (),
     );
     fn libafl_qemu_gdb_reply(buf: *const u8, len: usize);
+
+// void libafl_add_jmp_hook(uint64_t (*gen)(target_ulong src, target_ulong dst, uint64_t data),
+//                           void (*exec)(target_ulong src, target_ulong dst, uint64_t id, uint64_t data),
+//                           uint64_t data);
+    fn libafl_add_jmp_hook(
+        gen: Option<extern "C" fn(GuestAddr, GuestAddr, u64) -> u64>,
+        exec: Option<extern "C" fn(GuestAddr, GuestAddr, u64, u64)>,
+        data: u64,
+    );
 }
 
 #[cfg(emulation_mode = "usermode")]
@@ -1054,6 +1063,15 @@ impl Emulator {
         }
     }
 
+    pub fn add_jmp_hooks(
+        &self,
+        gen: Option<extern "C" fn(GuestAddr, GuestAddr, u64) -> u64>,
+        exec: Option<extern "C" fn(GuestAddr, GuestAddr, u64, u64)>,
+        data: u64,
+    ) {
+        unsafe { libafl_add_jmp_hook(gen, exec, data) }
+    }
+
     #[cfg(emulation_mode = "systemmode")]
     pub fn save_snapshot(&self, name: &str, sync: bool) {
         let s = CString::new(name).expect("Invalid snapshot name");

libafl_qemu/src/hooks.rs

@@ -655,6 +655,78 @@ where
     }
 }
 
+static mut JMP_HOOKS: Vec<(Hook, Hook)> = vec![];
+
+extern "C" fn gen_jmp_hook_wrapper<QT, S>(src: GuestAddr, dst: GuestAddr, index: u64) -> u64
+where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    unsafe {
+        let hooks = get_qemu_hooks::<QT, S>();
+        let (gen, _) = &mut JMP_HOOKS[index as usize];
+        match gen {
+            Hook::Function(ptr) => {
+                let func: fn(
+                    &mut QemuHooks<'_, QT, S>,
+                    Option<&mut S>,
+                    GuestAddr,
+                    GuestAddr,
+                ) -> Option<u64> = transmute(*ptr);
+                (func)(hooks, inprocess_get_state::<S>(), src, dst).map_or(SKIP_EXEC_HOOK, |id| id)
+            }
+            Hook::Closure(ptr) => {
+                let func: &mut Box<
+                    dyn FnMut(
+                        &mut QemuHooks<'_, QT, S>,
+                        Option<&mut S>,
+                        GuestAddr,
+                        GuestAddr,
+                    ) -> Option<u64>,
+                > = transmute(ptr);
+                (func)(hooks, inprocess_get_state::<S>(), src, dst).map_or(SKIP_EXEC_HOOK, |id| id)
+            }
+            _ => 0,
+        }
+    }
+}
+
+extern "C" fn exec_jmp_hook_wrapper<QT, S>(src: GuestAddr, dst: GuestAddr, id: u64, index: u64)
+where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    unsafe {
+        let hooks = get_qemu_hooks::<QT, S>();
+        let (_, exec) = &mut JMP_HOOKS[index as usize];
+        match exec {
+            Hook::Function(ptr) => {
+                let func: fn(
+                    &mut QemuHooks<'_, QT, S>,
+                    Option<&mut S>,
+                    GuestAddr,
+                    GuestAddr,
+                    u64,
+                ) = transmute(*ptr);
+                (func)(hooks, inprocess_get_state::<S>(), src, dst, id);
+            }
+            Hook::Closure(ptr) => {
+                let func: &mut Box<
+                    dyn FnMut(
+                        &mut QemuHooks<'_, QT, S>,
+                        Option<&mut S>,
+                        GuestAddr,
+                        GuestAddr,
+                        u64,
+                    ),
+                > = transmute(ptr);
+                (func)(hooks, inprocess_get_state::<S>(), src, dst, id);
+            }
+            _ => (),
+        }
+    }
+}
+
 static mut HOOKS_IS_INITIALIZED: bool = false;
 
 pub struct QemuHooks<'a, QT, S>
@@ -1556,4 +1628,37 @@ where
         self.emulator
             .set_post_syscall_hook(syscall_after_hooks_wrapper::<QT, S>);
     }
+
+    pub fn jmps(
+        &self,
+        generation_hook: Option<
+            fn(&mut Self, Option<&mut S>, src: GuestAddr, dest: GuestAddr) -> Option<u64>,
+        >,
+        execution_hook: Option<fn(&mut Self, Option<&mut S>, src: GuestAddr, dest: GuestAddr, id: u64)>,
+    ) {
+        unsafe {
+            let index = JMP_HOOKS.len();
+            self.emulator.add_jmp_hooks(
+                if generation_hook.is_none() {
+                    None
+                } else {
+                    Some(gen_jmp_hook_wrapper::<QT, S>)
+                },
+                if execution_hook.is_none() {
+                    None
+                } else {
+                    Some(exec_jmp_hook_wrapper::<QT, S>)
+                },
+                index as u64,
+            );
+            JMP_HOOKS.push((
+                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)
+                }),
+            ));
+        }
+    }
 }