From 09b15e907185863efb7468f71bcb1928daf003ea Mon Sep 17 00:00:00 2001
From: lazymio <mio@lazym.io>
Date: Sun, 6 Mar 2022 23:40:34 +0100
Subject: [PATCH] Fix exits wrongly cleared in nested uc_emu_start

---
 include/uc_priv.h   | 11 ++++++++---
 qemu/softmmu/cpus.c |  7 +++++--
 uc.c                | 15 +++++++--------
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/include/uc_priv.h b/include/uc_priv.h
index bd9865de..81692e4c 100644
--- a/include/uc_priv.h
+++ b/include/uc_priv.h
@@ -341,7 +341,8 @@ struct uc_struct {
     int invalid_error;     // invalid memory code: 1 = READ, 2 = WRITE, 3 = CODE
 
     int use_exits;
-    GTree *exits; // addresses where emulation stops (@until param of
+    uint64_t exits[UC_MAX_NESTED_LEVEL]; // When multiple exits is not enabled.
+    GTree *ctl_exits; // addresses where emulation stops (@until param of
                   // uc_emu_start()) Also see UC_CTL_USE_EXITS for more details.
 
     int thumb; // thumb mode for ARM
@@ -390,14 +391,18 @@ static inline void uc_add_exit(uc_engine *uc, uint64_t addr)
 {
     uint64_t *new_exit = g_malloc(sizeof(uint64_t));
     *new_exit = addr;
-    g_tree_insert(uc->exits, (gpointer)new_exit, (gpointer)1);
+    g_tree_insert(uc->ctl_exits, (gpointer)new_exit, (gpointer)1);
 }
 
 // This function has to exist since we would like to accept uint32_t or
 // it's complex to achieve so.
 static inline int uc_addr_is_exit(uc_engine *uc, uint64_t addr)
 {
-    return g_tree_lookup(uc->exits, (gpointer)(&addr)) == (gpointer)1;
+    if (uc->use_exits) {
+        return g_tree_lookup(uc->ctl_exits, (gpointer)(&addr)) == (gpointer)1;
+    } else {
+        return uc->exits[uc->nested_level - 1] == addr;
+    }
 }
 
 #ifdef UNICORN_TRACER
diff --git a/qemu/softmmu/cpus.c b/qemu/softmmu/cpus.c
index e8de7941..571a6a35 100644
--- a/qemu/softmmu/cpus.c
+++ b/qemu/softmmu/cpus.c
@@ -214,8 +214,11 @@ void resume_all_vcpus(struct uc_struct* uc)
     // clear the cache of the exits address, since the generated code
     // at that address is to exit emulation, but not for the instruction there.
     // if we dont do this, next time we cannot emulate at that address
-
-    g_tree_foreach(uc->exits, uc_exit_invalidate_iter, (void*)uc);
+    if (uc->use_exits) {
+        g_tree_foreach(uc->ctl_exits, uc_exit_invalidate_iter, (void*)uc);
+    } else {
+        uc_exit_invalidate_iter((gpointer)&uc->exits[uc->nested_level - 1], NULL, (gpointer)uc);
+    }
 
     cpu->created = false;
 }
diff --git a/uc.c b/uc.c
index 99bfe35b..db10af89 100644
--- a/uc.c
+++ b/uc.c
@@ -209,7 +209,7 @@ static uc_err uc_init(uc_engine *uc)
         uc->hook[i].delete_fn = hook_delete;
     }
 
-    uc->exits = g_tree_new_full(uc_exits_cmp, NULL, g_free, NULL);
+    uc->ctl_exits = g_tree_new_full(uc_exits_cmp, NULL, g_free, NULL);
 
     if (machine_initialize(uc)) {
         return UC_ERR_RESOURCE;
@@ -465,7 +465,7 @@ uc_err uc_close(uc_engine *uc)
 
     free(uc->mapped_blocks);
 
-    g_tree_destroy(uc->exits);
+    g_tree_destroy(uc->ctl_exits);
 
     // finally, free uc itself.
     memset(uc, 0, sizeof(*uc));
@@ -828,8 +828,7 @@ uc_err uc_emu_start(uc_engine *uc, uint64_t begin, uint64_t until,
     // If UC_CTL_UC_USE_EXITS is set, then the @until param won't have any
     // effect. This is designed for the backward compatibility.
     if (!uc->use_exits) {
-        g_tree_remove_all(uc->exits);
-        uc_add_exit(uc, until);
+        uc->exits[uc->nested_level - 1] = until;
     }
 
     if (timeout) {
@@ -2127,7 +2126,7 @@ uc_err uc_ctl(uc_engine *uc, uc_control_type control, ...)
             err = UC_ERR_ARG;
         } else if (rw == UC_CTL_IO_READ) {
             size_t *exits_cnt = va_arg(args, size_t *);
-            *exits_cnt = g_tree_nnodes(uc->exits);
+            *exits_cnt = g_tree_nnodes(uc->ctl_exits);
         } else {
             err = UC_ERR_ARG;
         }
@@ -2143,20 +2142,20 @@ uc_err uc_ctl(uc_engine *uc, uc_control_type control, ...)
         } else if (rw == UC_CTL_IO_READ) {
             uint64_t *exits = va_arg(args, uint64_t *);
             size_t cnt = va_arg(args, size_t);
-            if (cnt < g_tree_nnodes(uc->exits)) {
+            if (cnt < g_tree_nnodes(uc->ctl_exits)) {
                 err = UC_ERR_ARG;
             } else {
                 uc_ctl_exit_request req;
                 req.array = exits;
                 req.len = 0;
 
-                g_tree_foreach(uc->exits, uc_read_exit_iter, (void *)&req);
+                g_tree_foreach(uc->ctl_exits, uc_read_exit_iter, (void *)&req);
             }
         } else if (rw == UC_CTL_IO_WRITE) {
             uint64_t *exits = va_arg(args, uint64_t *);
             size_t cnt = va_arg(args, size_t);
 
-            g_tree_remove_all(uc->exits);
+            g_tree_remove_all(uc->ctl_exits);
 
             for (size_t i = 0; i < cnt; i++) {
                 uc_add_exit(uc, exits[i]);