From 345b63ee96e82b7c98828b6efbcf87b1eff33e13 Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sat, 7 May 2022 00:16:31 +0200
Subject: [PATCH] Only exit TB if pc is within the memory range

---
 include/uc_priv.h                 |  3 +++
 qemu/target/arm/unicorn_aarch64.c |  6 ++++++
 qemu/target/arm/unicorn_arm.c     |  9 ++++++++-
 qemu/target/i386/unicorn.c        | 10 ++++++++++
 qemu/target/m68k/unicorn.c        |  6 ++++++
 qemu/target/mips/unicorn.c        |  6 ++++++
 qemu/target/ppc/unicorn.c         |  6 ++++++
 qemu/target/riscv/unicorn.c       |  5 +++++
 qemu/target/s390x/unicorn.c       |  8 +++++++-
 qemu/target/sparc/unicorn.c       |  6 ++++++
 qemu/target/sparc/unicorn64.c     |  6 ++++++
 qemu/target/tricore/unicorn.c     |  9 ++++++++-
 uc.c                              |  8 ++++++--
 13 files changed, 83 insertions(+), 5 deletions(-)

diff --git a/include/uc_priv.h b/include/uc_priv.h
index 95996080..7a90faf9 100644
--- a/include/uc_priv.h
+++ b/include/uc_priv.h
@@ -93,6 +93,8 @@ typedef void (*uc_args_uc_long_t)(struct uc_struct *, unsigned long);
 
 typedef void (*uc_args_uc_u64_t)(struct uc_struct *, uint64_t addr);
 
+typedef uint64_t (*uc_get_pc_t)(struct uc_struct*);
+
 typedef MemoryRegion *(*uc_args_uc_ram_size_t)(struct uc_struct *, hwaddr begin,
                                                size_t size, uint32_t perms);
 
@@ -258,6 +260,7 @@ struct uc_struct {
     uc_read_mem_t read_mem;
     uc_args_void_t release;  // release resource when uc_close()
     uc_args_uc_u64_t set_pc; // set PC for tracecode
+    uc_get_pc_t get_pc;
     uc_args_int_t
         stop_interrupt; // check if the interrupt should stop emulation
     uc_memory_map_io_t memory_map_io;
diff --git a/qemu/target/arm/unicorn_aarch64.c b/qemu/target/arm/unicorn_aarch64.c
index 94534de0..f24c9506 100644
--- a/qemu/target/arm/unicorn_aarch64.c
+++ b/qemu/target/arm/unicorn_aarch64.c
@@ -18,6 +18,11 @@ static void arm64_set_pc(struct uc_struct *uc, uint64_t address)
     ((CPUARMState *)uc->cpu->env_ptr)->pc = address;
 }
 
+static uint64_t arm64_get_pc(struct uc_struct *uc)
+{
+    return ((CPUARMState *)uc->cpu->env_ptr)->pc;
+}
+
 static void arm64_release(void *ctx)
 {
     int i;
@@ -431,6 +436,7 @@ void arm64_uc_init(struct uc_struct *uc)
     uc->reg_write = arm64_reg_write;
     uc->reg_reset = arm64_reg_reset;
     uc->set_pc = arm64_set_pc;
+    uc->get_pc = arm64_get_pc;
     uc->release = arm64_release;
     uc->cpus_init = arm64_cpus_init;
     uc->cpu_context_size = offsetof(CPUARMState, cpu_watchpoint);
diff --git a/qemu/target/arm/unicorn_arm.c b/qemu/target/arm/unicorn_arm.c
index 87007f42..89233b27 100644
--- a/qemu/target/arm/unicorn_arm.c
+++ b/qemu/target/arm/unicorn_arm.c
@@ -16,7 +16,13 @@ ARMCPU *cpu_arm_init(struct uc_struct *uc);
 static void arm_set_pc(struct uc_struct *uc, uint64_t address)
 {
     ((CPUARMState *)uc->cpu->env_ptr)->pc = address;
-    ((CPUARMState *)uc->cpu->env_ptr)->regs[15] = address;
+    ((CPUARMState *)uc->cpu->env_ptr)->regs[15] = address & ~1;
+    ((CPUARMState *)uc->cpu->env_ptr)->thumb = address & 1;
+}
+
+static uint64_t arm_get_pc(struct uc_struct *uc)
+{
+    return ((CPUARMState *)uc->cpu->env_ptr)->regs[15] | ((CPUARMState *)uc->cpu->env_ptr)->thumb;
 }
 
 static void arm_release(void *ctx)
@@ -591,6 +597,7 @@ void arm_uc_init(struct uc_struct *uc)
     uc->reg_write = arm_reg_write;
     uc->reg_reset = arm_reg_reset;
     uc->set_pc = arm_set_pc;
+    uc->get_pc = arm_get_pc;
     uc->stop_interrupt = arm_stop_interrupt;
     uc->release = arm_release;
     uc->query = arm_query;
diff --git a/qemu/target/i386/unicorn.c b/qemu/target/i386/unicorn.c
index b846cd69..b416488f 100644
--- a/qemu/target/i386/unicorn.c
+++ b/qemu/target/i386/unicorn.c
@@ -33,6 +33,15 @@ static void x86_set_pc(struct uc_struct *uc, uint64_t address)
         ((CPUX86State *)uc->cpu->env_ptr)->eip = address;
 }
 
+static uint64_t x86_get_pc(struct uc_struct *uc) 
+{
+    if (uc->mode == UC_MODE_16) {
+        return X86_CPU(uc->cpu)->env.segs[R_CS].selector * 16 + ((CPUX86State *)uc->cpu->env_ptr)->eip;
+    } else {
+        return ((CPUX86State *)uc->cpu->env_ptr)->eip;
+    }
+}
+
 static void x86_release(void *ctx)
 {
     int i;
@@ -1644,6 +1653,7 @@ void x86_uc_init(struct uc_struct *uc)
     uc->reg_reset = x86_reg_reset;
     uc->release = x86_release;
     uc->set_pc = x86_set_pc;
+    uc->get_pc = x86_get_pc;
     uc->stop_interrupt = x86_stop_interrupt;
     uc->insn_hook_validate = x86_insn_hook_validate;
     uc->opcode_hook_invalidate = x86_opcode_hook_invalidate;
diff --git a/qemu/target/m68k/unicorn.c b/qemu/target/m68k/unicorn.c
index 0ff4213a..d748ace7 100644
--- a/qemu/target/m68k/unicorn.c
+++ b/qemu/target/m68k/unicorn.c
@@ -15,6 +15,11 @@ static void m68k_set_pc(struct uc_struct *uc, uint64_t address)
     ((CPUM68KState *)uc->cpu->env_ptr)->pc = address;
 }
 
+static uint64_t m68k_get_pc(struct uc_struct *uc)
+{
+    return ((CPUM68KState *)uc->cpu->env_ptr)->pc;
+}
+
 static void m68k_release(void *ctx)
 {
     int i;
@@ -170,6 +175,7 @@ void m68k_uc_init(struct uc_struct *uc)
     uc->reg_write = m68k_reg_write;
     uc->reg_reset = m68k_reg_reset;
     uc->set_pc = m68k_set_pc;
+    uc->get_pc = m68k_get_pc;
     uc->cpus_init = m68k_cpus_init;
     uc->cpu_context_size = offsetof(CPUM68KState, end_reset_fields);
     uc_common_init(uc);
diff --git a/qemu/target/mips/unicorn.c b/qemu/target/mips/unicorn.c
index 2f8ed5f9..bd4d5595 100644
--- a/qemu/target/mips/unicorn.c
+++ b/qemu/target/mips/unicorn.c
@@ -37,6 +37,11 @@ static void mips_set_pc(struct uc_struct *uc, uint64_t address)
     ((CPUMIPSState *)uc->cpu->env_ptr)->active_tc.PC = address;
 }
 
+static uint64_t mips_get_pc(struct uc_struct *uc)
+{
+    return ((CPUMIPSState *)uc->cpu->env_ptr)->active_tc.PC;
+}
+
 static void mips_release(void *ctx)
 {
     int i;
@@ -266,6 +271,7 @@ void mipsel_uc_init(struct uc_struct *uc)
     uc->reg_reset = mips_reg_reset;
     uc->release = mips_release;
     uc->set_pc = mips_set_pc;
+    uc->get_pc = mips_get_pc;
     uc->mem_redirect = mips_mem_redirect;
     uc->cpus_init = mips_cpus_init;
     uc->cpu_context_size = offsetof(CPUMIPSState, end_reset_fields);
diff --git a/qemu/target/ppc/unicorn.c b/qemu/target/ppc/unicorn.c
index 5907923e..ffdde3fe 100644
--- a/qemu/target/ppc/unicorn.c
+++ b/qemu/target/ppc/unicorn.c
@@ -97,6 +97,11 @@ static void ppc_set_pc(struct uc_struct *uc, uint64_t address)
     ((CPUPPCState *)uc->cpu->env_ptr)->nip = address;
 }
 
+static uint64_t ppc_get_pc(struct uc_struct *uc)
+{
+    return ((CPUPPCState *)uc->cpu->env_ptr)->nip;
+}
+
 void ppc_cpu_instance_finalize(CPUState *obj);
 void ppc_cpu_unrealize(CPUState *dev);
 static void ppc_release(void *ctx)
@@ -414,6 +419,7 @@ void ppc_uc_init(struct uc_struct *uc)
     uc->reg_reset = ppc_reg_reset;
     uc->release = ppc_release;
     uc->set_pc = ppc_set_pc;
+    uc->get_pc = ppc_get_pc;
     uc->mem_redirect = ppc_mem_redirect;
     uc->cpus_init = ppc_cpus_init;
     uc->cpu_context_size = offsetof(CPUPPCState, uc);
diff --git a/qemu/target/riscv/unicorn.c b/qemu/target/riscv/unicorn.c
index 82c4826d..684cfb46 100644
--- a/qemu/target/riscv/unicorn.c
+++ b/qemu/target/riscv/unicorn.c
@@ -52,6 +52,11 @@ static void riscv_set_pc(struct uc_struct *uc, uint64_t address)
     RISCV_CPU(uc->cpu)->env.pc = address;
 }
 
+static uint64_t riscv_get_pc(struct uc_struct *uc)
+{
+    return RISCV_CPU(uc->cpu)->env.pc;
+}
+
 static void riscv_release(void *ctx)
 {
     int i;
diff --git a/qemu/target/s390x/unicorn.c b/qemu/target/s390x/unicorn.c
index 81d64e0c..469cda7c 100644
--- a/qemu/target/s390x/unicorn.c
+++ b/qemu/target/s390x/unicorn.c
@@ -12,7 +12,12 @@ S390CPU *cpu_s390_init(struct uc_struct *uc, const char *cpu_model);
 
 static void s390_set_pc(struct uc_struct *uc, uint64_t address)
 {
-    // ((CPUS390XState *)uc->cpu->env_ptr)->pc = address;
+    ((CPUS390XState *)uc->cpu->env_ptr)->psw.addr = address;
+}
+
+static uint64_t s390_get_pc(struct uc_struct *uc)
+{
+    return ((CPUS390XState *)uc->cpu->env_ptr)->psw.addr;
 }
 
 static void s390_release(void *ctx)
@@ -183,6 +188,7 @@ void s390_uc_init(struct uc_struct *uc)
     uc->reg_write = s390_reg_write;
     uc->reg_reset = s390_reg_reset;
     uc->set_pc = s390_set_pc;
+    uc->get_pc = s390_get_pc;
     uc->cpus_init = s390_cpus_init;
     uc->cpu_context_size = offsetof(CPUS390XState, end_reset_fields);
     uc_common_init(uc);
diff --git a/qemu/target/sparc/unicorn.c b/qemu/target/sparc/unicorn.c
index 49253c86..2df5b9b3 100644
--- a/qemu/target/sparc/unicorn.c
+++ b/qemu/target/sparc/unicorn.c
@@ -24,6 +24,11 @@ static void sparc_set_pc(struct uc_struct *uc, uint64_t address)
     ((CPUSPARCState *)uc->cpu->env_ptr)->npc = address + 4;
 }
 
+static uint64_t sparc_get_pc(struct uc_struct *uc)
+{
+    return ((CPUSPARCState *)uc->cpu->env_ptr)->pc;
+}
+
 static void sparc_release(void *ctx)
 {
     int i;
@@ -190,6 +195,7 @@ void sparc_uc_init(struct uc_struct *uc)
     uc->reg_write = sparc_reg_write;
     uc->reg_reset = sparc_reg_reset;
     uc->set_pc = sparc_set_pc;
+    uc->get_pc = sparc_get_pc;
     uc->stop_interrupt = sparc_stop_interrupt;
     uc->cpus_init = sparc_cpus_init;
     uc->cpu_context_size = offsetof(CPUSPARCState, irq_manager);
diff --git a/qemu/target/sparc/unicorn64.c b/qemu/target/sparc/unicorn64.c
index 03eced6d..5b53943f 100644
--- a/qemu/target/sparc/unicorn64.c
+++ b/qemu/target/sparc/unicorn64.c
@@ -26,6 +26,11 @@ static void sparc_set_pc(struct uc_struct *uc, uint64_t address)
     ((CPUSPARCState *)uc->cpu->env_ptr)->npc = address + 4;
 }
 
+static uint64_t sparc_get_pc(struct uc_struct *uc)
+{
+    return ((CPUSPARCState *)uc->cpu->env_ptr)->pc;
+}
+
 static void sparc_release(void *ctx)
 {
     release_common(ctx);
@@ -196,6 +201,7 @@ void sparc64_uc_init(struct uc_struct* uc)
     uc->reg_write = sparc_reg_write;
     uc->reg_reset = sparc_reg_reset;
     uc->set_pc = sparc_set_pc;
+    uc->get_pc = sparc_get_pc;
     uc->stop_interrupt = sparc_stop_interrupt;
     uc->cpus_init = sparc_cpus_init;
     uc_common_init(uc);
diff --git a/qemu/target/tricore/unicorn.c b/qemu/target/tricore/unicorn.c
index f217335f..a49f0084 100644
--- a/qemu/target/tricore/unicorn.c
+++ b/qemu/target/tricore/unicorn.c
@@ -17,11 +17,17 @@
 
 TriCoreCPU *cpu_tricore_init(struct uc_struct *uc);
 
-void tricore_set_pc(struct uc_struct *uc, uint64_t address)
+static void tricore_set_pc(struct uc_struct *uc, uint64_t address)
 {
     ((CPUTriCoreState *)uc->cpu->env_ptr)->PC = address;
 }
 
+static uint64_t tricore_get_pc(struct uc_struct *uc)
+{
+    return ((CPUTriCoreState *)uc->cpu->env_ptr)->PC;
+}
+
+
 void tricore_reg_reset(struct uc_struct *uc)
 {
     CPUTriCoreState *env;
@@ -264,6 +270,7 @@ void tricore_uc_init(struct uc_struct *uc)
     uc->reg_write = tricore_reg_write;
     uc->reg_reset = tricore_reg_reset;
     uc->set_pc = tricore_set_pc;
+    uc->get_pc = tricore_get_pc;
     uc->cpus_init = tricore_cpus_init;
     uc->cpu_context_size = offsetof(CPUTriCoreState, end_reset_fields);
     uc_common_init(uc);
diff --git a/uc.c b/uc.c
index 1abe4400..e84bb95e 100644
--- a/uc.c
+++ b/uc.c
@@ -1358,6 +1358,7 @@ uc_err uc_mem_protect(struct uc_struct *uc, uint64_t address, size_t size,
 {
     MemoryRegion *mr;
     uint64_t addr = address;
+    uint64_t pc;
     size_t count, len;
     bool remove_exec = false;
 
@@ -1419,8 +1420,11 @@ uc_err uc_mem_protect(struct uc_struct *uc, uint64_t address, size_t size,
     // if EXEC permission is removed, then quit TB and continue at the same
     // place
     if (remove_exec) {
-        uc->quit_request = true;
-        uc_emu_stop(uc);
+        pc = uc->get_pc(uc);
+        if (pc < address + size && pc >= address) {
+            uc->quit_request = true;
+            uc_emu_stop(uc);
+        }
     }
 
     return UC_ERR_OK;