Added jit for block coverage.

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;
     }
     

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), \

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);

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
+}

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);