From 851914c8d013d9e7e84424b54386dce8eceb74ee Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sun, 6 Oct 2024 23:30:32 +0800
Subject: [PATCH] Fix segfault if tlb is flushed in the hooks

---
 qemu/accel/tcg/cputlb.c | 15 ++++++++++++++-
 tests/unit/test_x86.c   | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/qemu/accel/tcg/cputlb.c b/qemu/accel/tcg/cputlb.c
index 16ba4ed2..f3f89aef 100644
--- a/qemu/accel/tcg/cputlb.c
+++ b/qemu/accel/tcg/cputlb.c
@@ -1600,12 +1600,25 @@ load_helper(CPUArchState *env, target_ulong addr, TCGMemOpIdx oi,
             if (!HOOK_BOUND_CHECK(hook, paddr))
                 continue;
             JIT_CALLBACK_GUARD(((uc_cb_hookmem_t)hook->callback)(env->uc, UC_MEM_READ, paddr, size, 0, hook->user_data));
-
             // the last callback may already asked to stop emulation
             if (uc->stop_request)
                 break;
         }
 
+        /* Unicorn: Previous callbacks may invalidate TLB, reload everything.
+                    This may have impact on performance but generally fine.
+                    A better approach is not always invalidating tlb but this
+                    might cause more chaos regarding re-entry (nested uc_emu_start).
+        */
+        if (tlb_entry_is_empty(entry)) {
+            tlb_fill(env_cpu(env), addr, size,
+                        access_type, mmu_idx, retaddr);
+            index = tlb_index(env, mmu_idx, addr);
+            entry = tlb_entry(env, mmu_idx, addr);
+            tlb_addr = code_read ? entry->addr_code : entry->addr_read;
+            tlb_addr &= ~TLB_INVALID_MASK;
+        }
+
         // callback on non-readable memory
         if (mr != NULL && !(mr->perms & UC_PROT_READ)) {  //non-readable
             handled = false;
diff --git a/tests/unit/test_x86.c b/tests/unit/test_x86.c
index a95259e6..a7d3909a 100644
--- a/tests/unit/test_x86.c
+++ b/tests/unit/test_x86.c
@@ -1798,6 +1798,39 @@ static void test_rex_x64(void)
     }
 }
 
+static bool test_x86_ro_segfault_cb(uc_engine *uc, uc_mem_type type,
+                                              uint64_t address, int size,
+                                              uint64_t value, void *user_data)
+{
+    const char code[] =
+        "\xA1\x00\x10\x00\x00\xA1\x00\x10\x00\x00";
+    OK(uc_mem_write(uc, address, code, sizeof(code) - 1));
+    return true;
+}
+
+static void test_x86_ro_segfault(void)
+{
+    uc_engine* uc;
+    // mov eax, [0x1000]
+    // mov eax, [0x1000]
+    const char code[] =
+        "\xA1\x00\x10\x00\x00\xA1\x00\x10\x00\x00";
+    uint32_t out;
+    uc_hook hh;
+
+    OK(uc_open(UC_ARCH_X86, UC_MODE_32, &uc));
+    OK(uc_mem_map(uc, 0, 0x1000, UC_PROT_ALL));
+    OK(uc_mem_write(uc, 0, code, sizeof(code) - 1));
+    OK(uc_mem_map(uc, 0x1000, 0x1000, UC_PROT_READ));
+    
+    OK(uc_hook_add(uc, &hh, UC_HOOK_MEM_READ, test_x86_ro_segfault_cb, NULL, 1, 0));
+    OK(uc_emu_start(uc, 0, sizeof(code) - 1, 0, 0));
+
+    OK(uc_reg_read(uc, UC_X86_REG_EAX, (void*)&out));
+    TEST_CHECK(out == 0x001000a1);
+    OK(uc_close(uc));
+}
+
 TEST_LIST = {
     {"test_x86_in", test_x86_in},
     {"test_x86_out", test_x86_out},
@@ -1851,4 +1884,5 @@ TEST_LIST = {
     {"test_fxsave_fpip_x64", test_fxsave_fpip_x64},
     {"test_bswap_x64", test_bswap_ax},
     {"test_rex_x64", test_rex_x64},
+    {"test_x86_ro_segfault", test_x86_ro_segfault},
     {NULL, NULL}};