fix: defer removed hook callback drops

crates/unicorn/src/lib.rs

@@ -143,6 +143,9 @@ pub struct UnicornInner<'a, D> {
     pub arch: Arch,
     /// to keep ownership over the hook for this uc instance's lifetime
     pub hooks: Vec<(UcHookId, Box<dyn hook::IsUcHook<'a> + 'a>)>,
+    // Hooks removed while emulation is active, retained until the outermost emulation returns.
+    pub deleted_hooks: Vec<Box<dyn hook::IsUcHook<'a> + 'a>>,
+    pub emu_depth: usize,
     /// To keep ownership over the mmio callbacks for this uc instance's lifetime
     pub mmio_callbacks: Vec<MmioCallbackScope<'a>>,
     pub data: D,
@@ -198,6 +201,8 @@ where
                     arch,
                     data,
                     hooks: vec![],
+                    deleted_hooks: vec![],
+                    emu_depth: 0,
                     mmio_callbacks: vec![],
                 })),
             })
@@ -229,6 +234,8 @@ where
                 arch: arch.try_into()?,
                 data,
                 hooks: vec![],
+                deleted_hooks: vec![],
+                emu_depth: 0,
                 mmio_callbacks: vec![],
             })),
         })
@@ -1182,12 +1189,19 @@ impl<'a, D> Unicorn<'a, D> {
     /// Remove a hook.
     ///
     /// `hook_id` is the value returned by `add_*_hook` functions.
+    ///
+    /// Unicorn core marks a deleted hook and releases it after the outermost
+    /// emulation call returns.
     pub fn remove_hook(&mut self, hook_id: UcHookId) -> Result<(), uc_error> {
-        // drop the hook
         let inner = self.inner_mut();
-        inner.hooks.retain(|(id, _)| id != &hook_id);
+        unsafe { uc_hook_del(inner.handle, hook_id.0) }.and_then(|| {
+            if let Some(index) = inner.hooks.iter().position(|(id, _)| id == &hook_id) {
+                let (_, callback) = inner.hooks.remove(index);
+                inner.deleted_hooks.push(callback);
+            }
 
-        unsafe { uc_hook_del(inner.handle, hook_id.0) }.into()
+            Ok(())
+        })
     }
 
     /// Allocate and return an empty Unicorn context.
@@ -1247,7 +1261,15 @@ impl<'a, D> Unicorn<'a, D> {
         timeout: u64,
         count: usize,
     ) -> Result<(), uc_error> {
-        unsafe { uc_emu_start(self.get_handle(), begin, until, timeout, count as _) }.into()
+        let inner = self.inner_mut();
+        inner.emu_depth += 1;
+        let result =
+            unsafe { uc_emu_start(inner.handle, begin, until, timeout, count as _) }.into();
+        inner.emu_depth -= 1;
+        if inner.emu_depth == 0 {
+            inner.deleted_hooks.clear();
+        }
+        result
     }
 
     /// Stop the emulation.

crates/unicorn/src/tests/ctl.rs

@@ -1,4 +1,8 @@
-use std::time::{Duration, Instant};
+use std::{
+    cell::Cell,
+    rc::Rc,
+    time::{Duration, Instant},
+};
 
 use unicorn_engine_sys::{RegisterX86, X86Insn};
 
@@ -183,6 +187,53 @@ fn test_uc_hook_cached_uaf() {
     assert_eq!(*uc.get_data(), 4);
 }
 
+#[test]
+fn test_self_removed_hook_defers_drop() {
+    struct DropFlag(Rc<Cell<bool>>);
+
+    impl Drop for DropFlag {
+        fn drop(&mut self) {
+            self.0.set(true);
+        }
+    }
+
+    let mut uc = Unicorn::new(Arch::X86, Mode::MODE_32).unwrap();
+    uc.mem_map(CODE_START, CODE_LEN, Prot::ALL).unwrap();
+    uc.mem_write(CODE_START, b"\x90\x90").unwrap();
+
+    let hook_id = Rc::new(Cell::new(None));
+    let dropped = Rc::new(Cell::new(false));
+    let dropped_during_nested_emu = Rc::new(Cell::new(None));
+    let nested_emu_dropped = Rc::clone(&dropped);
+    let nested_emu_result = Rc::clone(&dropped_during_nested_emu);
+
+    // observes whether the removed callback remains alive in nested emulation.
+    uc.add_code_hook(
+        CODE_START + 1,
+        CODE_START + 1,
+        move |_, _, _| nested_emu_result.set(Some(nested_emu_dropped.get())),
+    )
+    .unwrap();
+
+    let callback_hook_id = Rc::clone(&hook_id);
+    let drop_flag = DropFlag(Rc::clone(&dropped));
+
+    // registers a callback that removes itself, then starts nested emulation.
+    let hook = uc
+        .add_code_hook(CODE_START, CODE_START, move |uc, _, _| {
+            let _ = &drop_flag;
+            uc.remove_hook(callback_hook_id.get().unwrap()).unwrap();
+            uc.emu_start(CODE_START + 1, CODE_START + 2, 0, 0).unwrap();
+        })
+        .unwrap();
+    hook_id.set(Some(hook));
+
+    uc.emu_start(CODE_START, CODE_START + 1, 0, 0).unwrap();
+
+    assert_eq!(dropped_during_nested_emu.get(), Some(false));
+    assert!(dropped.get());
+}
+
 #[test]
 fn test_uc_emu_stop_set_ip() {
     #[rustfmt::skip]