From bbbc7856accbae7697db84d10a28c5a6aec81856 Mon Sep 17 00:00:00 2001
From: Mio <mio@lazym.io>
Date: Wed, 12 Apr 2023 20:56:54 +0800
Subject: [PATCH] Invalidate tb cache once mapping is removed

---
 qemu/accel/tcg/translate-all.c |  8 ++++--
 qemu/softmmu/memory.c          | 11 +++++---
 tests/unit/test_arm64.c        | 51 ++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/qemu/accel/tcg/translate-all.c b/qemu/accel/tcg/translate-all.c
index 56ddb35e..ac71e391 100644
--- a/qemu/accel/tcg/translate-all.c
+++ b/qemu/accel/tcg/translate-all.c
@@ -1018,11 +1018,13 @@ static void uc_invalidate_tb(struct uc_struct *uc, uint64_t start_addr, size_t l
         return;
     }
 
-    // GPA to GVA
+    // GPA to ram addr
+    // https://raw.githubusercontent.com/android/platform_external_qemu/master/docs/QEMU-MEMORY-MANAGEMENT.TXT
     // start_addr : GPA
-    // addr: GVA
+    // start (returned): ram addr
     // (GPA -> HVA via memory_region_get_ram_addr(mr) + GPA + block->host,
-    // HVA->HPA via host mmu)
+    // GVA -> GPA via tlb & softmmu
+    // HVA -> HPA via host mmu)
     start = get_page_addr_code(uc->cpu->env_ptr, start_addr) & (target_ulong)(-1);
 
     uc->nested_level--;
diff --git a/qemu/softmmu/memory.c b/qemu/softmmu/memory.c
index c58bd79f..9d856472 100644
--- a/qemu/softmmu/memory.c
+++ b/qemu/softmmu/memory.c
@@ -153,12 +153,15 @@ void memory_unmap(struct uc_struct *uc, MemoryRegion *mr)
     int i;
     hwaddr addr;
 
-    // Make sure all pages associated with the MemoryRegion are flushed
-    // Only need to do this if we are in a running state
     if (uc->cpu) {
-        for (addr = mr->addr; addr < mr->end; addr += uc->target_page_size) {
+        // We also need to remove all tb cache
+        uc->uc_invalidate_tb(uc, mr->addr, mr->size);
+
+        // Make sure all pages associated with the MemoryRegion are flushed
+        // Only need to do this if we are in a running state
+        for (addr = mr->addr; (int64_t)(mr->end - addr) > 0; addr += uc->target_page_size) {
            tlb_flush_page(uc->cpu, addr);
-        }
+        }        
     }
     memory_region_del_subregion(uc->system_memory, mr);
 
diff --git a/tests/unit/test_arm64.c b/tests/unit/test_arm64.c
index 8a03945c..44f0ad9c 100644
--- a/tests/unit/test_arm64.c
+++ b/tests/unit/test_arm64.c
@@ -2,6 +2,7 @@
 #include "unicorn/unicorn.h"
 #include "unicorn_test.h"
 #include <stdbool.h>
+#include <stdint.h>
 #include <stdio.h>
 
 const uint64_t code_start = 0x1000;
@@ -473,6 +474,55 @@ static void test_arm64_mmu(void)
     free(data);
 }
 
+static void test_arm64_pc_wrap(void)
+{
+    uc_engine *uc;
+    // add x1 x2
+    char add_x1_x2[] = "\x20\x00\x02\x8b";
+    // add x1 x3
+    char add_x1_x3[] = "\x20\x00\x03\x8b";
+    uint64_t x0, x1, x2, x3;
+    uint64_t pc = 0xFFFFFFFFFFFFFFFCULL;
+    uint64_t page = 0xFFFFFFFFFFFFF000ULL;
+    
+    OK(uc_open(UC_ARCH_ARM64, UC_MODE_ARM, &uc));
+    OK(uc_mem_map(uc, page, 4096, UC_PROT_READ | UC_PROT_EXEC));
+    OK(uc_mem_write(uc, pc, add_x1_x2, sizeof(add_x1_x2) - 1));
+
+    x1 = 1;
+    x2 = 2;
+    OK(uc_reg_write(uc, UC_ARM64_REG_X1, &x1));
+    OK(uc_reg_write(uc, UC_ARM64_REG_X2, &x2));
+
+    OK(uc_emu_start(uc, pc, pc + 4, 0, 1));
+
+    OK(uc_mem_unmap(uc, page, 4096));
+
+    OK(uc_reg_read(uc, UC_ARM64_REG_X0, &x0));
+
+    TEST_CHECK( (x0 == 1 + 2) );
+
+    OK(uc_mem_map(uc, page, 4096, UC_PROT_READ | UC_PROT_EXEC));
+    OK(uc_mem_write(uc, pc, add_x1_x3, sizeof(add_x1_x3) - 1));
+
+    x1 = 5;
+    x2 = 0;
+    x3 = 5;
+    OK(uc_reg_write(uc, UC_ARM64_REG_X1, &x1));
+    OK(uc_reg_write(uc, UC_ARM64_REG_X2, &x2));
+    OK(uc_reg_write(uc, UC_ARM64_REG_X3, &x3));
+
+    OK(uc_emu_start(uc, pc, pc + 4, 0, 1));
+
+    OK(uc_mem_unmap(uc, page, 4096));
+
+    OK(uc_reg_read(uc, UC_ARM64_REG_X0, &x0));
+
+    TEST_CHECK( (x0 == 5 + 5) );
+
+    OK(uc_close(uc));
+}
+
 TEST_LIST = {{"test_arm64_until", test_arm64_until},
              {"test_arm64_code_patching", test_arm64_code_patching},
              {"test_arm64_code_patching_count", test_arm64_code_patching_count},
@@ -487,4 +537,5 @@ TEST_LIST = {{"test_arm64_until", test_arm64_until},
              {"test_arm64_block_invalid_mem_read_write_sync",
               test_arm64_block_invalid_mem_read_write_sync},
             {"test_arm64_mmu", test_arm64_mmu},
+            {"test_arm64_pc_wrap", test_arm64_pc_wrap},
              {NULL, NULL}};