From 7173dfbbd1fe6bb36b74c4ebdf744b38ca79e4e4 Mon Sep 17 00:00:00 2001
From: Alwin Berger <alwin.berger@tu-dortmund.de>
Date: Mon, 19 Dec 2022 13:11:40 +0100
Subject: [PATCH] add jmp instrumentation

---
 accel/tcg/translate-all.c  | 61 ++++++++++++++++++++++++++++++++++++++
 target/arm/tcg/translate.c |  6 ++++
 2 files changed, 67 insertions(+)

diff --git a/accel/tcg/translate-all.c b/accel/tcg/translate-all.c
index 12b9cb3095..45e6969c2d 100644
--- a/accel/tcg/translate-all.c
+++ b/accel/tcg/translate-all.c
@@ -637,6 +637,67 @@ void libafl_add_backdoor_hook(void (*exec)(target_ulong id, uint64_t data),
     libafl_helper_table_add(&hook->helper_info);
 }
 
+static TCGHelperInfo libafl_exec_jmp_hook_info = {
+    .func = NULL, .name = "libafl_exec_jmp_hook", \
+    .flags = dh_callflag(void), \
+    .typemask = dh_typemask(void, 0) | dh_typemask(tl, 2) | dh_typemask(i64, 1)
+};
+
+struct libafl_jmp_hook {
+    uint64_t (*gen)(target_ulong src, target_ulong dst, uint64_t data);
+    void (*exec)(target_ulong src, target_ulong dst, uint64_t id, uint64_t data);
+    uint64_t data;
+    TCGHelperInfo helper_info;
+    struct libafl_jmp_hook* next;
+};
+
+struct libafl_jmp_hook* libafl_jmp_hooks;
+
+void libafl_add_jmp_hook(uint64_t (*gen)(target_ulong src, target_ulong dst, uint64_t data),
+                          void (*exec)(target_ulong src, target_ulong dst, uint64_t id, uint64_t data),
+                          uint64_t data);
+void libafl_add_jmp_hook(uint64_t (*gen)(target_ulong src, target_ulong dst, uint64_t data),
+                          void (*exec)(target_ulong src, target_ulong dst, uint64_t id, uint64_t data),
+                          uint64_t data)
+{
+    struct libafl_jmp_hook* hook = malloc(sizeof(struct libafl_jmp_hook));
+    hook->gen = gen;
+    hook->exec = exec;
+    hook->data = data;
+    hook->next = libafl_jmp_hooks;
+    libafl_jmp_hooks = hook;
+    
+    memcpy(&hook->helper_info, &libafl_exec_jmp_hook_info, sizeof(TCGHelperInfo));
+    hook->helper_info.func = exec;
+    libafl_helper_table_add(&hook->helper_info);
+}
+
+void libafl_gen_jmp(target_ulong src, target_ulong dst);
+void libafl_gen_jmp(target_ulong src, target_ulong dst)
+{
+
+    struct libafl_jmp_hook* hook = libafl_jmp_hooks;
+    while (hook) {
+        uint64_t cur_id = 0;
+        if (hook->gen)
+            cur_id = hook->gen(src, dst, hook->data);
+        void* func = hook->exec;
+        if (cur_id != (uint64_t)-1 && func) {
+            TCGv_i64 tmp0 = tcg_const_i64(src);
+            TCGv_i64 tmp1 = tcg_const_i64(dst);
+            TCGv_i64 tmp2 = tcg_const_i64(cur_id);
+            TCGv_i64 tmp3 = tcg_const_i64(hook->data);
+            TCGTemp *tmp4[4] = { tcgv_i64_temp(tmp0), tcgv_i64_temp(tmp1), tcgv_i64_temp(tmp2), tcgv_i64_temp(tmp3) };
+            tcg_gen_callN(func, NULL,4, tmp4);
+            tcg_temp_free_i64(tmp0);
+            tcg_temp_free_i64(tmp1);
+            tcg_temp_free_i64(tmp2);
+            tcg_temp_free_i64(tmp3);
+        }
+        hook = hook->next;
+    }
+}
+
 //// --- End LibAFL code ---
 
 /* Make sure all possible CPU event bits fit in tb->trace_vcpu_dstate */
diff --git a/target/arm/tcg/translate.c b/target/arm/tcg/translate.c
index 0fab359cae..23b14936ed 100644
--- a/target/arm/tcg/translate.c
+++ b/target/arm/tcg/translate.c
@@ -2679,8 +2679,14 @@ static void gen_jmp_tb(DisasContext *s, target_long diff, int tbno)
     }
 }
 
+//// --- Begin LibAFL code ---
+void libafl_gen_jmp(target_ulong src, target_ulong dst); // see translate-all.c
+//// --- End LibAFL code ---
 static inline void gen_jmp(DisasContext *s, target_long diff)
 {
+//// --- Begin LibAFL code ---
+    libafl_gen_jmp(s->pc_curr,s->pc_curr+diff);
+//// --- End LibAFL code ---
     gen_jmp_tb(s, diff, 0);
 }