From 8b2c477578c15fbaff5f08536f42e4d1ceab60a5 Mon Sep 17 00:00:00 2001
From: "Takacs, Philipp" <philipp.takacs@iosb.fraunhofer.de>
Date: Fri, 17 Feb 2023 14:22:25 +0100
Subject: [PATCH] clear the TLB cache in uc_ctl_flush_tlb

uc_ctl_flush_tlb implies that the tlb is flushed. This change adds
UC_CTL_TLB_FLUSH which clears the TLB and set the uc_ctl_flush_tlb
alias to UC_CTL_TLB_FLUSH. Also adds a uc_ctl_flush_tb alias for
UC_CTL_TB_FLUSH.
---
 include/unicorn/unicorn.h |  6 +++++-
 tests/unit/test_ctl.c     | 37 +++++++++++++++++++++++++++++++++++++
 uc.c                      | 11 +++++++++++
 3 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/include/unicorn/unicorn.h b/include/unicorn/unicorn.h
index 3e1f9248..2ddb8b9e 100644
--- a/include/unicorn/unicorn.h
+++ b/include/unicorn/unicorn.h
@@ -567,6 +567,9 @@ typedef enum uc_control_type {
     // Invalidate all translation blocks.
     // No arguments.
     UC_CTL_TB_FLUSH,
+    // Invalidate all TLB cache entries and translation blocks.
+    // No arguments
+    UC_CTL_TLB_FLUSH,
     // Change the tlb implementation
     // see uc_tlb_type for current implemented types
     // Write: @args = (int)
@@ -645,7 +648,8 @@ See sample_ctl.c for a detailed example.
     uc_ctl(uc, UC_CTL_WRITE(UC_CTL_TB_REMOVE_CACHE, 2), (address), (end))
 #define uc_ctl_request_cache(uc, address, tb)                                  \
     uc_ctl(uc, UC_CTL_READ_WRITE(UC_CTL_TB_REQUEST_CACHE, 2), (address), (tb))
-#define uc_ctl_flush_tlb(uc) uc_ctl(uc, UC_CTL_WRITE(UC_CTL_TB_FLUSH, 0))
+#define uc_ctl_flush_tb(uc) uc_ctl(uc, UC_CTL_WRITE(UC_CTL_TB_FLUSH, 0))
+#define uc_ctl_flush_tlb(uc) uc_ctl(uc, UC_CTL_WRITE(UC_CTL_TLB_FLUSH, 0))
 #define uc_ctl_tlb_mode(uc, mode) uc_ctl(uc, UC_CTL_WRITE(UC_CTL_TLB_TYPE, 1), (mode))
 // Opaque storage for CPU context, used with uc_context_*()
 struct uc_context;
diff --git a/tests/unit/test_ctl.c b/tests/unit/test_ctl.c
index e1e467d3..9cfcd7dd 100644
--- a/tests/unit/test_ctl.c
+++ b/tests/unit/test_ctl.c
@@ -338,6 +338,42 @@ static void test_uc_emu_stop_set_ip(void)
     OK(uc_close(uc));
 }
 
+static bool test_tlb_clear_tlb(uc_engine *uc, uint64_t addr, uc_mem_type type, uc_tlb_entry *result, void *user_data)
+{
+    size_t *tlbcount = (size_t*)user_data;
+    *tlbcount += 1;
+    result->paddr = addr;
+    result->perms = UC_PROT_ALL;
+    return true;
+}
+
+static void test_tlb_clear_syscall(uc_engine *uc, void *user_data)
+{
+    OK(uc_ctl_flush_tlb(uc));
+}
+
+static void test_tlb_clear(void)
+{
+    uc_engine *uc;
+    uc_hook hook1, hook2;
+    size_t tlbcount = 0;
+    char code[] = "\xa3\x00\x00\x20\x00\x00\x00\x00\x00\x0f\x05\xa3\x00\x00\x20\x00\x00\x00\x00\x00"; //movabs  dword ptr [0x200000], eax; syscall; movabs  dword ptr [0x200000], eax
+
+    uc_common_setup(&uc, UC_ARCH_X86, UC_MODE_64, code, sizeof(code) - 1);
+    OK(uc_mem_map(uc, 0x200000, 0x1000, UC_PROT_ALL));
+
+    OK(uc_ctl_tlb_mode(uc, UC_TLB_VIRTUAL));
+    OK(uc_hook_add(uc, &hook1, UC_HOOK_TLB_FILL, test_tlb_clear_tlb, &tlbcount, 1, 0));
+    OK(uc_hook_add(uc, &hook2, UC_HOOK_INSN, test_tlb_clear_syscall, NULL, 1, 0, UC_X86_INS_SYSCALL));
+
+    OK(uc_emu_start(uc, code_start, code_start + sizeof(code) - 1, 0, 0));
+
+    TEST_CHECK(tlbcount == 4);
+
+    OK(uc_close(uc));
+}
+
+
 TEST_LIST = {{"test_uc_ctl_mode", test_uc_ctl_mode},
              {"test_uc_ctl_page_size", test_uc_ctl_page_size},
              {"test_uc_ctl_arch", test_uc_ctl_arch},
@@ -350,4 +386,5 @@ TEST_LIST = {{"test_uc_ctl_mode", test_uc_ctl_mode},
 #endif
              {"test_uc_hook_cached_uaf", test_uc_hook_cached_uaf},
              {"test_uc_emu_stop_set_ip", test_uc_emu_stop_set_ip},
+             {"test_tlb_clear", test_tlb_clear},
              {NULL, NULL}};
diff --git a/uc.c b/uc.c
index b0326a01..0a8626cf 100644
--- a/uc.c
+++ b/uc.c
@@ -2350,6 +2350,17 @@ uc_err uc_ctl(uc_engine *uc, uc_control_type control, ...)
         }
         break;
 
+    case UC_CTL_TLB_FLUSH:
+
+        UC_INIT(uc);
+
+        if (rw == UC_CTL_IO_WRITE) {
+            uc->tcg_flush_tlb(uc);
+        } else {
+            err = UC_ERR_ARG;
+        }
+        break;
+
     case UC_CTL_TLB_TYPE: {
 
         UC_INIT(uc);