From fdd129fd30bab69aa8a45c5b3b75374fd3d7c1ed Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Thu, 2 Jun 2022 14:46:02 +0200
Subject: [PATCH] Remember the regions a hook has intrumented and clear cache
 on deletion

---
 include/uc_priv.h           | 49 +++++++++++++++++++++++++++++++++++++
 qemu/accel/tcg/translator.c |  2 ++
 uc.c                        | 12 +++++++++
 3 files changed, 63 insertions(+)

diff --git a/include/uc_priv.h b/include/uc_priv.h
index 978e71fe..a56ea944 100644
--- a/include/uc_priv.h
+++ b/include/uc_priv.h
@@ -9,6 +9,7 @@
 #include <stdio.h>
 
 #include "qemu.h"
+#include "qemu/xxhash.h"
 #include "unicorn/unicorn.h"
 #include "list.h"
 
@@ -157,6 +158,7 @@ struct hook {
                          // address (depends on hook type)
     void *callback;      // a uc_cb_* type
     void *user_data;
+    GHashTable *hooked_regions; // The regions this hook instrumented on
 };
 
 // Add an inline hook to helper_table
@@ -414,6 +416,53 @@ static inline int uc_addr_is_exit(uc_engine *uc, uint64_t addr)
     }
 }
 
+typedef struct HookedRegion {
+    uint64_t start;
+    uint64_t length;
+} HookedRegion;
+
+// hooked_regions related functions
+static inline guint hooked_regions_hash(const void* p) {
+    HookedRegion *region = (HookedRegion*)p;
+
+    return qemu_xxhash4(region->start, region->length);
+}
+
+static inline gboolean hooked_regions_equal(const void* lhs, const void* rhs) {
+    HookedRegion *l = (HookedRegion*)lhs;
+    HookedRegion *r = (HookedRegion*)rhs;
+
+    return l->start == r->start && l->length == r->length;
+}
+
+static inline void hooked_regions_add(struct hook* h, uint64_t start, uint64_t length) {
+    HookedRegion tmp;
+    tmp.start = start;
+    tmp.length = length;
+
+    if (!g_hash_table_lookup(h->hooked_regions, (void*)&tmp)) {
+        HookedRegion* r = malloc(sizeof(HookedRegion));
+        r->start = start;
+        r->length = length;
+        g_hash_table_insert(h->hooked_regions, (void*)r, (void*)1);
+    }
+}
+
+static inline void hooked_regions_check_single(struct list_item *cur, uint64_t start, uint64_t length) {
+    while (cur != NULL) {
+        if (HOOK_BOUND_CHECK((struct hook *)cur->data, start)) {
+            hooked_regions_add((struct hook *)cur->data, start, length);
+        }
+        cur = cur->next;
+    }
+}
+
+static inline void hooked_regions_check(uc_engine *uc, uint64_t start, uint64_t length) {
+    // Only UC_HOOK_BLOCK and UC_HOOK_CODE might be wrongle cached!
+    hooked_regions_check_single(uc->hook[UC_HOOK_CODE_IDX].head, start, length);
+    hooked_regions_check_single(uc->hook[UC_HOOK_BLOCK_IDX].head, start, length);
+}
+
 #ifdef UNICORN_TRACER
 #define UC_TRACE_START(loc) trace_start(get_tracer(), loc)
 #define UC_TRACE_END(loc, fmt, ...)                                            \
diff --git a/qemu/accel/tcg/translator.c b/qemu/accel/tcg/translator.c
index 9ed272fd..eccf06a7 100644
--- a/qemu/accel/tcg/translator.c
+++ b/qemu/accel/tcg/translator.c
@@ -147,6 +147,8 @@ _end_loop:
     db->tb->size = db->pc_next - db->pc_first;
     db->tb->icount = db->num_insns;
 
+    hooked_regions_check(uc, db->tb->pc, db->tb->size);
+
     if (block_hook) {
         TCGOp *tcg_op;
 
diff --git a/uc.c b/uc.c
index cb4c0d2a..e8eee462 100644
--- a/uc.c
+++ b/uc.c
@@ -49,6 +49,14 @@ static void *hook_append(struct list *l, struct hook *h)
     return item;
 }
 
+static void hook_invalidate_region(void* key, void* data, void* opaq)
+{
+    uc_engine* uc = (uc_engine*)opaq;
+    HookedRegion* region = (HookedRegion*)key;
+
+    uc->uc_invalidate_tb(uc, region->start, region->length);
+}
+
 static void hook_delete(void *data)
 {
     struct hook *h = (struct hook *)data;
@@ -56,6 +64,7 @@ static void hook_delete(void *data)
     h->refs--;
 
     if (h->refs == 0) {
+        g_hash_table_destroy(h->hooked_regions);
         free(h);
     }
 }
@@ -1561,6 +1570,7 @@ uc_err uc_hook_add(uc_engine *uc, uc_hook *hh, int type, void *callback,
     hook->user_data = user_data;
     hook->refs = 0;
     hook->to_delete = false;
+    hook->hooked_regions = g_hash_table_new_full(hooked_regions_hash, hooked_regions_equal, g_free, NULL);
     *hh = (uc_hook)hook;
 
     // UC_HOOK_INSN has an extra argument for instruction ID
@@ -1670,6 +1680,8 @@ uc_err uc_hook_del(uc_engine *uc, uc_hook hh)
     // and store the type mask in the hook pointer.
     for (i = 0; i < UC_HOOK_MAX; i++) {
         if (list_exists(&uc->hook[i], (void *)hook)) {
+            g_hash_table_foreach(hook->hooked_regions, hook_invalidate_region, uc);
+            g_hash_table_remove_all(hook->hooked_regions);
             hook->to_delete = true;
             uc->hooks_count[i]--;
             hook_append(&uc->hooks_to_del, hook);