diff --git a/crates/unicorn/src/lib.rs b/crates/unicorn/src/lib.rs index bfef973..494f040 100644 --- a/crates/unicorn/src/lib.rs +++ b/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 + 'a>)>, + // Hooks removed while emulation is active, retained until the outermost emulation returns. + pub deleted_hooks: Vec + 'a>>, + pub emu_depth: usize, /// To keep ownership over the mmio callbacks for this uc instance's lifetime pub mmio_callbacks: Vec>, 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. diff --git a/crates/unicorn/src/tests/ctl.rs b/crates/unicorn/src/tests/ctl.rs index 0acece6..cbbe4e1 100644 --- a/crates/unicorn/src/tests/ctl.rs +++ b/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>); + + 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]