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
This commit is contained in:
Andrea Fioraldi 2023-11-27 13:20:20 +01:00 committed by GitHub
parent 00740190a7
commit 1545514ed8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 177 additions and 24 deletions

View File

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

View File

@ -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 {

View File

@ -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::<QT, S>),
);
hooks.writes(
Hook::Function(gen_readwrite_asan::<QT, S>),
Hook::Function(trace_write1_asan::<QT, S>),
Hook::Function(trace_write2_asan::<QT, S>),
Hook::Function(trace_write4_asan::<QT, S>),
Hook::Function(trace_write8_asan::<QT, S>),
Hook::Function(trace_write_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>),
Hook::Function(trace_write2_asan::<QT, S>),
Hook::Function(trace_write4_asan::<QT, S>),
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>,

View File

@ -92,15 +92,29 @@ where
QT: QemuHelperTuple<S>,
{
if self.use_hitcounts {
hooks.edges(
Hook::Function(gen_unique_edge_ids::<QT, S>),
Hook::Raw(trace_edge_hitcount),
);
// hooks.edges(
// Hook::Function(gen_unique_edge_ids::<QT, S>),
// 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(
Hook::Function(gen_unique_edge_ids::<QT, S>),
Hook::Raw(trace_edge_single),
);
// hooks.edges(
// Hook::Function(gen_unique_edge_ids::<QT, S>),
// 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),
);
}
}
}
}

View File

@ -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<S>,
{
hooks.writes(
Hook::Empty,
Hook::Function(trace_write1_snapshot::<QT, S>),
Hook::Function(trace_write2_snapshot::<QT, S>),
Hook::Function(trace_write4_snapshot::<QT, S>),
Hook::Function(trace_write8_snapshot::<QT, S>),
Hook::Function(trace_write_n_snapshot::<QT, 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>),
Hook::Function(trace_write2_snapshot::<QT, S>),
Hook::Function(trace_write4_snapshot::<QT, S>),
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>));