JIT fast path for edge cov hooks in libafl_qemu (#1696)

* JIT fast path for edge cov hooks in libafl_qemu

* fix

* fmt

* fix

* unify hooks as opt

libafl_qemu/libafl_qemu_build/src/bindings.rs

@@ -78,6 +78,7 @@ const WRAPPER_HEADER: &str = r#"
 
 #include "libafl_extras/exit.h"
 #include "libafl_extras/hook.h"
+#include "libafl_extras/jit.h"
 
 "#;
 

libafl_qemu/libafl_qemu_sys/src/x86_64_stub_bindings.rs

@@ -13442,6 +13442,12 @@ extern "C" {
         invalidate: ::std::os::raw::c_int,
     ) -> ::std::os::raw::c_int;
 }
+extern "C" {
+    pub fn libafl_qemu_edge_hook_set_jit(
+        num: usize,
+        jit: ::std::option::Option<unsafe extern "C" fn(arg1: u64, arg2: u64) -> usize>,
+    );
+}
 extern "C" {
     pub fn libafl_add_block_hook(
         gen: ::std::option::Option<extern "C" fn(data: u64, pc: target_ulong) -> u64>,
@@ -13613,6 +13619,12 @@ extern "C" {
 extern "C" {
     pub fn libafl_qemu_remove_new_thread_hook(num: usize) -> ::std::os::raw::c_int;
 }
+extern "C" {
+    pub fn libafl_jit_trace_edge_hitcount(data: u64, id: u64) -> usize;
+}
+extern "C" {
+    pub fn libafl_jit_trace_edge_single(data: u64, id: u64) -> usize;
+}
 #[repr(C)]
 #[derive(Debug, Default, Copy, Clone)]
 pub struct kvm_dirty_gfn {

libafl_qemu/src/asan.rs

@@ -23,6 +23,7 @@ use crate::{
     emu::{EmuError, Emulator, MemAccessInfo, SyscallHookResult},
     helper::{HasInstrumentationFilter, QemuHelper, QemuHelperTuple, QemuInstrumentationFilter},
     hooks::{Hook, QemuHooks},
+    snapshot::QemuSnapshotHelper,
     GuestAddr, Regs,
 };
 
@@ -961,6 +962,7 @@ where
             Hook::Function(trace_read_n_asan::<QT, S>),
         );
 
+        if hooks.match_helper::<QemuSnapshotHelper>().is_none() {
             hooks.writes(
                 Hook::Function(gen_readwrite_asan::<QT, S>),
                 Hook::Function(trace_write1_asan::<QT, S>),
@@ -969,6 +971,17 @@ where
                 Hook::Function(trace_write8_asan::<QT, S>),
                 Hook::Function(trace_write_n_asan::<QT, S>),
             );
+        } else {
+            // track writes for both helpers as opt
+            hooks.writes(
+                Hook::Function(gen_write_asan_snapshot::<QT, S>),
+                Hook::Function(trace_write1_asan_snapshot::<QT, S>),
+                Hook::Function(trace_write2_asan_snapshot::<QT, S>),
+                Hook::Function(trace_write4_asan_snapshot::<QT, S>),
+                Hook::Function(trace_write8_asan_snapshot::<QT, S>),
+                Hook::Function(trace_write_n_asan_snapshot::<QT, S>),
+            );
+        }
     }
 
     fn pre_exec(&mut self, emulator: &Emulator, _input: &S::Input) {
@@ -1164,6 +1177,115 @@ pub fn trace_write_n_asan<QT, S>(
     h.read_n(&emulator, id as GuestAddr, addr, size);
 }
 
+pub fn gen_write_asan_snapshot<QT, S>(
+    hooks: &mut QemuHooks<QT, S>,
+    _state: Option<&mut S>,
+    pc: GuestAddr,
+    _info: MemAccessInfo,
+) -> Option<u64>
+where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    let h = hooks.match_helper_mut::<QemuAsanHelper>().unwrap();
+    if h.must_instrument(pc) {
+        Some(pc.into())
+    } else {
+        Some(0)
+    }
+}
+
+pub fn trace_write1_asan_snapshot<QT, S>(
+    hooks: &mut QemuHooks<QT, S>,
+    _state: Option<&mut S>,
+    id: u64,
+    addr: GuestAddr,
+) where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    if id != 0 {
+        let emulator = hooks.emulator().clone();
+        let h = hooks.match_helper_mut::<QemuAsanHelper>().unwrap();
+        h.write_1(&emulator, id as GuestAddr, addr);
+    }
+    let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
+    h.access(addr, 1);
+}
+
+pub fn trace_write2_asan_snapshot<QT, S>(
+    hooks: &mut QemuHooks<QT, S>,
+    _state: Option<&mut S>,
+    id: u64,
+    addr: GuestAddr,
+) where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    if id != 0 {
+        let emulator = hooks.emulator().clone();
+        let h = hooks.match_helper_mut::<QemuAsanHelper>().unwrap();
+        h.write_2(&emulator, id as GuestAddr, addr);
+    }
+    let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
+    h.access(addr, 2);
+}
+
+pub fn trace_write4_asan_snapshot<QT, S>(
+    hooks: &mut QemuHooks<QT, S>,
+    _state: Option<&mut S>,
+    id: u64,
+    addr: GuestAddr,
+) where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    if id != 0 {
+        let emulator = hooks.emulator().clone();
+        let h = hooks.match_helper_mut::<QemuAsanHelper>().unwrap();
+        h.write_4(&emulator, id as GuestAddr, addr);
+    }
+    let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
+    h.access(addr, 4);
+}
+
+pub fn trace_write8_asan_snapshot<QT, S>(
+    hooks: &mut QemuHooks<QT, S>,
+    _state: Option<&mut S>,
+    id: u64,
+    addr: GuestAddr,
+) where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    if id != 0 {
+        let emulator = hooks.emulator().clone();
+        let h = hooks.match_helper_mut::<QemuAsanHelper>().unwrap();
+        h.write_8(&emulator, id as GuestAddr, addr);
+    }
+    let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
+    h.access(addr, 8);
+}
+
+pub fn trace_write_n_asan_snapshot<QT, S>(
+    hooks: &mut QemuHooks<QT, S>,
+    _state: Option<&mut S>,
+    id: u64,
+    addr: GuestAddr,
+    size: usize,
+) where
+    S: UsesInput,
+    QT: QemuHelperTuple<S>,
+{
+    if id != 0 {
+        let emulator = hooks.emulator().clone();
+        let h = hooks.match_helper_mut::<QemuAsanHelper>().unwrap();
+        h.read_n(&emulator, id as GuestAddr, addr, size);
+    }
+    let h = hooks.match_helper_mut::<QemuSnapshotHelper>().unwrap();
+    h.access(addr, size);
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn qasan_fake_syscall<QT, S>(
     hooks: &mut QemuHooks<QT, S>,

libafl_qemu/src/edges.rs

@@ -92,17 +92,31 @@ where
         QT: QemuHelperTuple<S>,
     {
         if self.use_hitcounts {
-            hooks.edges(
+            // hooks.edges(
-                Hook::Function(gen_unique_edge_ids::<QT, S>),
+            //     Hook::Function(gen_unique_edge_ids::<QT, S>),
-                Hook::Raw(trace_edge_hitcount),
+            //     Hook::Raw(trace_edge_hitcount),
+            // );
+            let hook_id = hooks.edges(Hook::Function(gen_unique_edge_ids::<QT, S>), Hook::Empty);
+            unsafe {
+                libafl_qemu_sys::libafl_qemu_edge_hook_set_jit(
+                    hook_id.0,
+                    Some(libafl_qemu_sys::libafl_jit_trace_edge_hitcount),
                 );
+            }
         } else {
-            hooks.edges(
+            // hooks.edges(
-                Hook::Function(gen_unique_edge_ids::<QT, S>),
+            //     Hook::Function(gen_unique_edge_ids::<QT, S>),
-                Hook::Raw(trace_edge_single),
+            //     Hook::Raw(trace_edge_single),
+            // );
+            let hook_id = hooks.edges(Hook::Function(gen_unique_edge_ids::<QT, S>), Hook::Empty);
+            unsafe {
+                libafl_qemu_sys::libafl_qemu_edge_hook_set_jit(
+                    hook_id.0,
+                    Some(libafl_qemu_sys::libafl_jit_trace_edge_single),
                 );
             }
         }
+    }
 }
 
 pub type QemuCollidingEdgeCoverageHelper = QemuEdgeCoverageChildHelper;

libafl_qemu/src/snapshot.rs

@@ -22,6 +22,7 @@ use crate::SYS_mmap2;
 )))]
 use crate::SYS_newfstatat;
 use crate::{
+    asan::QemuAsanHelper,
     emu::{Emulator, MmapPerms, SyscallHookResult},
     helper::{QemuHelper, QemuHelperTuple},
     hooks::{Hook, QemuHooks},
@@ -491,6 +492,8 @@ where
     where
         QT: QemuHelperTuple<S>,
     {
+        if hooks.match_helper::<QemuAsanHelper>().is_none() {
+            // The ASan helper, if present, will call the tracer hook for the snpahsot helper as opt
             hooks.writes(
                 Hook::Empty,
                 Hook::Function(trace_write1_snapshot::<QT, S>),
@@ -499,6 +502,7 @@ where
                 Hook::Function(trace_write8_snapshot::<QT, S>),
                 Hook::Function(trace_write_n_snapshot::<QT, S>),
             );
+        }
 
         if !self.accurate_unmap {
             hooks.syscalls(Hook::Function(filter_mmap_snapshot::<QT, S>));