From d177a55033edd417ac362d95ad98ac61752d0e6b Mon Sep 17 00:00:00 2001 From: Romain Malmain Date: Mon, 4 Dec 2023 17:13:46 +0100 Subject: [PATCH] Added jit for block coverage. --- accel/tcg/translate-all.c | 3 ++ libafl_extras/hook.c | 18 ++++++++++-- libafl_extras/hook.h | 4 ++- libafl_extras/jit.c | 58 +++++++++++++++++++++++++++++++++++++++ libafl_extras/jit.h | 2 ++ 5 files changed, 82 insertions(+), 3 deletions(-) diff --git a/accel/tcg/translate-all.c b/accel/tcg/translate-all.c index 31528d4772..b9f81e1606 100644 --- a/accel/tcg/translate-all.c +++ b/accel/tcg/translate-all.c @@ -298,6 +298,9 @@ static int setjmp_gen_code(CPUArchState *env, TranslationBlock *tb, tcg_temp_free_i64(tmp0); tcg_temp_free_i64(tmp1); } + if (cur_id != (uint64_t)-1 && hook->jit) { + hook->jit(hook->data, cur_id); + } hook = hook->next; } diff --git a/libafl_extras/hook.c b/libafl_extras/hook.c index 82be46e352..0d836812c1 100644 --- a/libafl_extras/hook.c +++ b/libafl_extras/hook.c @@ -221,16 +221,17 @@ size_t libafl_add_edge_hook(uint64_t (*gen)(uint64_t data, target_ulong src, tar GEN_REMOVE_HOOK(edge) -void libafl_qemu_edge_hook_set_jit(size_t num, size_t (*jit)(uint64_t data, uint64_t id)) { +bool libafl_qemu_edge_hook_set_jit(size_t num, size_t (*jit)(uint64_t data, uint64_t id)) { struct libafl_edge_hook* hk = libafl_edge_hooks; while (hk) { if (hk->num == num) { hk->jit = jit; - return; + return true; } else { hk = hk->next; } } + return false; } static TCGHelperInfo libafl_exec_block_hook_info = { @@ -270,6 +271,19 @@ size_t libafl_add_block_hook(uint64_t (*gen)(uint64_t data, target_ulong pc), GEN_REMOVE_HOOK(block) +bool libafl_qemu_block_hook_set_jit(size_t num, size_t (*jit)(uint64_t data, uint64_t id)) { + struct libafl_block_hook* hk = libafl_block_hooks; + while (hk) { + if (hk->num == num) { + hk->jit = jit; + return true; + } else { + hk = hk->next; + } + } + return false; +} + static TCGHelperInfo libafl_exec_read_hook1_info = { .func = NULL, .name = "libafl_exec_read_hook1", \ .flags = dh_callflag(void), \ diff --git a/libafl_extras/hook.h b/libafl_extras/hook.h index f16030ad88..e6d56dc819 100644 --- a/libafl_extras/hook.h +++ b/libafl_extras/hook.h @@ -71,12 +71,13 @@ size_t libafl_add_edge_hook(uint64_t (*gen)(uint64_t data, target_ulong src, tar void (*exec)(uint64_t data, uint64_t id), uint64_t data); int libafl_qemu_remove_edge_hook(size_t num, int invalidate); -void libafl_qemu_edge_hook_set_jit(size_t num, size_t (*jit)(uint64_t, uint64_t)); // no param names to avoid to be marked as safe +bool libafl_qemu_edge_hook_set_jit(size_t num, size_t (*jit)(uint64_t, uint64_t)); // no param names to avoid to be marked as safe struct libafl_block_hook { uint64_t (*gen)(uint64_t data, target_ulong pc); void (*post_gen)(uint64_t data, target_ulong pc, target_ulong block_length); // void (*exec)(uint64_t data, uint64_t id); + size_t (*jit)(uint64_t data, uint64_t id); // optional opt uint64_t data; size_t num; TCGHelperInfo helper_info; @@ -89,6 +90,7 @@ size_t libafl_add_block_hook(uint64_t (*gen)(uint64_t data, target_ulong pc), void (*post_gen)(uint64_t data, target_ulong pc, target_ulong block_length), void (*exec)(uint64_t data, uint64_t id), uint64_t data); int libafl_qemu_remove_block_hook(size_t num, int invalidate); +bool libafl_qemu_block_hook_set_jit(size_t num, size_t (*jit)(uint64_t, uint64_t)); // no param names to avoid to be marked as safe struct libafl_rw_hook { uint64_t (*gen)(uint64_t data, target_ulong pc, MemOpIdx oi); diff --git a/libafl_extras/jit.c b/libafl_extras/jit.c index 15277613da..0fbf047e02 100644 --- a/libafl_extras/jit.c +++ b/libafl_extras/jit.c @@ -26,6 +26,7 @@ pub extern "C" fn trace_edge_single(id: u64, _data: u64) { // from libafl_targets coverage.rs // correct size doesn't matter here uint8_t __afl_area_ptr_local[65536] __attribute__((weak)); +size_t __afl_map_size __attribute__((weak)); size_t libafl_jit_trace_edge_hitcount(uint64_t data, uint64_t id) { TCGv_ptr map_ptr = tcg_constant_ptr(__afl_area_ptr_local); @@ -43,3 +44,60 @@ size_t libafl_jit_trace_edge_single(uint64_t data, uint64_t id) { tcg_gen_st8_i32(counter, map_ptr, (tcg_target_long)id); return 2; // # instructions } + +uint64_t __prev_loc = 0; + +size_t libafl_jit_trace_block_hitcount(uint64_t data, uint64_t id) { + TCGv_ptr map_ptr = tcg_constant_ptr(__afl_area_ptr_local); + TCGv_ptr prev_loc_ptr = tcg_constant_ptr(&__prev_loc); + + TCGv_i32 counter = tcg_temp_new_i32(); + TCGv_i64 id_r = tcg_temp_new_i64(); + TCGv_i64 prev_loc = tcg_temp_new_i64(); + TCGv_ptr prev_loc2 = tcg_temp_new_ptr(); + + // Compute location => 5 insn + tcg_gen_ld_i64(prev_loc, prev_loc_ptr, 0); + tcg_gen_xori_i64(prev_loc, prev_loc, (int64_t) id); + tcg_gen_andi_i64(prev_loc, prev_loc, (int64_t) (__afl_map_size - 1)); + tcg_gen_trunc_i64_ptr(prev_loc2, prev_loc); + tcg_gen_add_ptr(prev_loc2, map_ptr, prev_loc2); + + // Update map => 3 insn + tcg_gen_ld8u_i32(counter, prev_loc2, 0); + tcg_gen_addi_i32(counter, counter, 1); + tcg_gen_st8_i32(counter, prev_loc2, 0); + + // Update prev_loc => 3 insn + tcg_gen_movi_i64(id_r, (int64_t) id); + tcg_gen_shri_i64(id_r, id_r, 1); + tcg_gen_st_i64(id_r, prev_loc_ptr, 0); + return 11; // # instructions +} + +size_t libafl_jit_trace_block_single(uint64_t data, uint64_t id) { + TCGv_ptr map_ptr = tcg_constant_ptr(__afl_area_ptr_local); + TCGv_ptr prev_loc_ptr = tcg_constant_ptr(&__prev_loc); + + TCGv_i32 counter = tcg_temp_new_i32(); + TCGv_i64 id_r = tcg_temp_new_i64(); + TCGv_i64 prev_loc = tcg_temp_new_i64(); + TCGv_ptr prev_loc2 = tcg_temp_new_ptr(); + + // Compute location => 5 insn + tcg_gen_ld_i64(prev_loc, prev_loc_ptr, 0); + tcg_gen_xori_i64(prev_loc, prev_loc, (int64_t) id); + tcg_gen_andi_i64(prev_loc, prev_loc, (int64_t) (__afl_map_size - 1)); + tcg_gen_trunc_i64_ptr(prev_loc2, prev_loc); + tcg_gen_add_ptr(map_ptr, map_ptr, prev_loc2); + + // Update map => 2 insn + tcg_gen_movi_i32(counter, 1); + tcg_gen_st8_i32(counter, map_ptr, 0); + + // Update prev_loc => 3 insn + tcg_gen_movi_i64(id_r, (int64_t) id); + tcg_gen_shri_i64(id_r, id_r, 1); + tcg_gen_st_i64(id_r, prev_loc_ptr, 0); + return 10; // # instructions +} diff --git a/libafl_extras/jit.h b/libafl_extras/jit.h index 24a8374397..a205992643 100644 --- a/libafl_extras/jit.h +++ b/libafl_extras/jit.h @@ -10,3 +10,5 @@ size_t libafl_jit_trace_edge_hitcount(uint64_t data, uint64_t id); size_t libafl_jit_trace_edge_single(uint64_t data, uint64_t id); +size_t libafl_jit_trace_block_hitcount(uint64_t data, uint64_t id); +size_t libafl_jit_trace_block_single(uint64_t data, uint64_t id);