From 1545514ed8bcb62b07bff8fbc1bc5e9c5471e323 Mon Sep 17 00:00:00 2001 From: Andrea Fioraldi Date: Mon, 27 Nov 2023 13:20:20 +0100 Subject: [PATCH] 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 | 1 + .../src/x86_64_stub_bindings.rs | 12 ++ libafl_qemu/src/asan.rs | 138 +++++++++++++++++- libafl_qemu/src/edges.rs | 30 +++- libafl_qemu/src/snapshot.rs | 20 ++- 5 files changed, 177 insertions(+), 24 deletions(-) diff --git a/libafl_qemu/libafl_qemu_build/src/bindings.rs b/libafl_qemu/libafl_qemu_build/src/bindings.rs index fd9128f2a8..26d6eb4871 100644 --- a/libafl_qemu/libafl_qemu_build/src/bindings.rs +++ b/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" "#; diff --git a/libafl_qemu/libafl_qemu_sys/src/x86_64_stub_bindings.rs b/libafl_qemu/libafl_qemu_sys/src/x86_64_stub_bindings.rs index 91c9cef288..bfc3f44b98 100644 --- a/libafl_qemu/libafl_qemu_sys/src/x86_64_stub_bindings.rs +++ b/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 usize>, + ); +} extern "C" { pub fn libafl_add_block_hook( gen: ::std::option::Option 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 { diff --git a/libafl_qemu/src/asan.rs b/libafl_qemu/src/asan.rs index 6c2ca488cc..4a25023b0f 100644 --- a/libafl_qemu/src/asan.rs +++ b/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,14 +962,26 @@ where Hook::Function(trace_read_n_asan::), ); - hooks.writes( - Hook::Function(gen_readwrite_asan::), - Hook::Function(trace_write1_asan::), - Hook::Function(trace_write2_asan::), - Hook::Function(trace_write4_asan::), - Hook::Function(trace_write8_asan::), - Hook::Function(trace_write_n_asan::), - ); + if hooks.match_helper::().is_none() { + hooks.writes( + Hook::Function(gen_readwrite_asan::), + Hook::Function(trace_write1_asan::), + Hook::Function(trace_write2_asan::), + Hook::Function(trace_write4_asan::), + Hook::Function(trace_write8_asan::), + Hook::Function(trace_write_n_asan::), + ); + } else { + // track writes for both helpers as opt + hooks.writes( + Hook::Function(gen_write_asan_snapshot::), + Hook::Function(trace_write1_asan_snapshot::), + Hook::Function(trace_write2_asan_snapshot::), + Hook::Function(trace_write4_asan_snapshot::), + Hook::Function(trace_write8_asan_snapshot::), + Hook::Function(trace_write_n_asan_snapshot::), + ); + } } fn pre_exec(&mut self, emulator: &Emulator, _input: &S::Input) { @@ -1164,6 +1177,115 @@ pub fn trace_write_n_asan( h.read_n(&emulator, id as GuestAddr, addr, size); } +pub fn gen_write_asan_snapshot( + hooks: &mut QemuHooks, + _state: Option<&mut S>, + pc: GuestAddr, + _info: MemAccessInfo, +) -> Option +where + S: UsesInput, + QT: QemuHelperTuple, +{ + let h = hooks.match_helper_mut::().unwrap(); + if h.must_instrument(pc) { + Some(pc.into()) + } else { + Some(0) + } +} + +pub fn trace_write1_asan_snapshot( + hooks: &mut QemuHooks, + _state: Option<&mut S>, + id: u64, + addr: GuestAddr, +) where + S: UsesInput, + QT: QemuHelperTuple, +{ + if id != 0 { + let emulator = hooks.emulator().clone(); + let h = hooks.match_helper_mut::().unwrap(); + h.write_1(&emulator, id as GuestAddr, addr); + } + let h = hooks.match_helper_mut::().unwrap(); + h.access(addr, 1); +} + +pub fn trace_write2_asan_snapshot( + hooks: &mut QemuHooks, + _state: Option<&mut S>, + id: u64, + addr: GuestAddr, +) where + S: UsesInput, + QT: QemuHelperTuple, +{ + if id != 0 { + let emulator = hooks.emulator().clone(); + let h = hooks.match_helper_mut::().unwrap(); + h.write_2(&emulator, id as GuestAddr, addr); + } + let h = hooks.match_helper_mut::().unwrap(); + h.access(addr, 2); +} + +pub fn trace_write4_asan_snapshot( + hooks: &mut QemuHooks, + _state: Option<&mut S>, + id: u64, + addr: GuestAddr, +) where + S: UsesInput, + QT: QemuHelperTuple, +{ + if id != 0 { + let emulator = hooks.emulator().clone(); + let h = hooks.match_helper_mut::().unwrap(); + h.write_4(&emulator, id as GuestAddr, addr); + } + let h = hooks.match_helper_mut::().unwrap(); + h.access(addr, 4); +} + +pub fn trace_write8_asan_snapshot( + hooks: &mut QemuHooks, + _state: Option<&mut S>, + id: u64, + addr: GuestAddr, +) where + S: UsesInput, + QT: QemuHelperTuple, +{ + if id != 0 { + let emulator = hooks.emulator().clone(); + let h = hooks.match_helper_mut::().unwrap(); + h.write_8(&emulator, id as GuestAddr, addr); + } + let h = hooks.match_helper_mut::().unwrap(); + h.access(addr, 8); +} + +pub fn trace_write_n_asan_snapshot( + hooks: &mut QemuHooks, + _state: Option<&mut S>, + id: u64, + addr: GuestAddr, + size: usize, +) where + S: UsesInput, + QT: QemuHelperTuple, +{ + if id != 0 { + let emulator = hooks.emulator().clone(); + let h = hooks.match_helper_mut::().unwrap(); + h.read_n(&emulator, id as GuestAddr, addr, size); + } + let h = hooks.match_helper_mut::().unwrap(); + h.access(addr, size); +} + #[allow(clippy::too_many_arguments)] pub fn qasan_fake_syscall( hooks: &mut QemuHooks, diff --git a/libafl_qemu/src/edges.rs b/libafl_qemu/src/edges.rs index 3eddde39fd..8049f5ec0d 100644 --- a/libafl_qemu/src/edges.rs +++ b/libafl_qemu/src/edges.rs @@ -92,15 +92,29 @@ where QT: QemuHelperTuple, { if self.use_hitcounts { - hooks.edges( - Hook::Function(gen_unique_edge_ids::), - Hook::Raw(trace_edge_hitcount), - ); + // hooks.edges( + // Hook::Function(gen_unique_edge_ids::), + // Hook::Raw(trace_edge_hitcount), + // ); + let hook_id = hooks.edges(Hook::Function(gen_unique_edge_ids::), 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( - Hook::Function(gen_unique_edge_ids::), - Hook::Raw(trace_edge_single), - ); + // hooks.edges( + // Hook::Function(gen_unique_edge_ids::), + // Hook::Raw(trace_edge_single), + // ); + let hook_id = hooks.edges(Hook::Function(gen_unique_edge_ids::), Hook::Empty); + unsafe { + libafl_qemu_sys::libafl_qemu_edge_hook_set_jit( + hook_id.0, + Some(libafl_qemu_sys::libafl_jit_trace_edge_single), + ); + } } } } diff --git a/libafl_qemu/src/snapshot.rs b/libafl_qemu/src/snapshot.rs index 7767d6bccd..69a7d45db2 100644 --- a/libafl_qemu/src/snapshot.rs +++ b/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,14 +492,17 @@ where where QT: QemuHelperTuple, { - hooks.writes( - Hook::Empty, - Hook::Function(trace_write1_snapshot::), - Hook::Function(trace_write2_snapshot::), - Hook::Function(trace_write4_snapshot::), - Hook::Function(trace_write8_snapshot::), - Hook::Function(trace_write_n_snapshot::), - ); + if hooks.match_helper::().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::), + Hook::Function(trace_write2_snapshot::), + Hook::Function(trace_write4_snapshot::), + Hook::Function(trace_write8_snapshot::), + Hook::Function(trace_write_n_snapshot::), + ); + } if !self.accurate_unmap { hooks.syscalls(Hook::Function(filter_mmap_snapshot::));