From 68de01a177e83428b75c6ec41811e6c96ddf572d Mon Sep 17 00:00:00 2001 From: Jvle Date: Thu, 13 Aug 2026 00:30:09 +0800 Subject: [PATCH] refactor: make hook callbacks sound Hook callbacks previously reconstructed Unicorn values from shared Rc> state, which could create multiple mutable aliases to the same emulator state. Introduce HookContext as the callback API, and use Box for the emulator. Signed-off-by: Jvle --- crates/unicorn/src/hook.rs | 136 ++++++++++----------------- crates/unicorn/src/lib.rs | 151 ++++++++++++++++++++---------- crates/unicorn/src/tests/arm.rs | 94 ++++++++++++------- crates/unicorn/src/tests/arm64.rs | 28 +++--- crates/unicorn/src/tests/ctl.rs | 20 ++-- crates/unicorn/src/tests/mem.rs | 30 +++--- 6 files changed, 255 insertions(+), 204 deletions(-) diff --git a/crates/unicorn/src/hook.rs b/crates/unicorn/src/hook.rs index 6d3df3f..6939152 100644 --- a/crates/unicorn/src/hook.rs +++ b/crates/unicorn/src/hook.rs @@ -1,15 +1,17 @@ #![allow(non_camel_case_types)] -use alloc::rc::Weak; -use core::{cell::UnsafeCell, ffi::c_void}; +use core::ffi::c_void; pub use unicorn_engine_sys::{self as sys, uc_context, uc_engine, uc_hook}; -use crate::{Unicorn, UnicornInner}; +use crate::HookContext; pub struct UcHook<'a, D: 'a, F: 'a> { pub callback: F, - pub uc: Weak>>, + // No value of `D` is stored in the hook. This marker keeps the `D` + // type parameter and callback lifetime visible to the type system + // while hook storage is erased behind `IsUcHook`. + pub marker: core::marker::PhantomData<&'a D>, } pub trait IsUcHook<'a> {} @@ -27,14 +29,11 @@ pub extern "C" fn mmio_read_callback_proxy( user_data: *mut c_void, ) -> u64 where - F: FnMut(&mut crate::Unicorn, u64, usize) -> u64, + F: FnMut(&mut HookContext, u64, usize) -> u64, { let user_data = unsafe { &mut *user_data.cast::>() }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, offset, size as usize) + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, offset, size as usize) } /// # Safety @@ -47,14 +46,11 @@ pub unsafe extern "C" fn mmio_write_callback_proxy( value: u64, user_data: *mut c_void, ) where - F: FnMut(&mut crate::Unicorn, u64, usize, u64), + F: FnMut(&mut HookContext, u64, usize, u64), { let user_data = unsafe { &mut *user_data.cast::>() }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, offset, size as usize, value); + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, offset, size as usize, value); } /// # Safety @@ -66,14 +62,11 @@ pub unsafe extern "C" fn code_hook_proxy( size: u32, user_data: *mut UcHook, ) where - F: FnMut(&mut crate::Unicorn, u64, u32), + F: FnMut(&mut HookContext, u64, u32), { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, address, size); + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, address, size); } /// # Safety @@ -85,14 +78,11 @@ pub unsafe extern "C" fn block_hook_proxy( size: u32, user_data: *mut UcHook, ) where - F: FnMut(&mut crate::Unicorn, u64, u32), + F: FnMut(&mut HookContext, u64, u32), { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, address, size); + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, address, size); } /// # Safety @@ -107,14 +97,11 @@ pub unsafe extern "C" fn mem_hook_proxy( user_data: *mut UcHook, ) -> bool where - F: FnMut(&mut crate::Unicorn, sys::MemType, u64, usize, i64) -> bool, + F: FnMut(&mut HookContext, sys::MemType, u64, usize, i64) -> bool, { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, mem_type, address, size as usize, value) + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, mem_type, address, size as usize, value) } /// # Safety @@ -125,14 +112,11 @@ pub unsafe extern "C" fn intr_hook_proxy( value: u32, user_data: *mut UcHook, ) where - F: FnMut(&mut crate::Unicorn, u32), + F: FnMut(&mut HookContext, u32), { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, value); + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, value); } /// # Safety @@ -145,14 +129,11 @@ pub unsafe extern "C" fn insn_in_hook_proxy( user_data: *mut UcHook, ) -> u32 where - F: FnMut(&mut crate::Unicorn, u32, usize) -> u32, + F: FnMut(&mut HookContext, u32, usize) -> u32, { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, port, size) + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, port, size) } /// # Safety @@ -163,14 +144,11 @@ pub unsafe extern "C" fn insn_invalid_hook_proxy( user_data: *mut UcHook, ) -> bool where - F: FnMut(&mut crate::Unicorn) -> bool, + F: FnMut(&mut HookContext) -> bool, { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc) + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context) } /// # Safety @@ -183,14 +161,11 @@ pub unsafe extern "C" fn insn_out_hook_proxy( value: u32, user_data: *mut UcHook, ) where - F: FnMut(&mut crate::Unicorn, u32, usize, u32), + F: FnMut(&mut HookContext, u32, usize, u32), { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, port, size, value); + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, port, size, value); } /// # Safety @@ -198,14 +173,11 @@ pub unsafe extern "C" fn insn_out_hook_proxy( /// This function is unsafe because it dereferences the `user_data` pointer. pub unsafe extern "C" fn insn_sys_hook_proxy(uc: *mut uc_engine, user_data: *mut UcHook) where - F: FnMut(&mut crate::Unicorn), + F: FnMut(&mut HookContext), { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc); + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context); } /// # Safety @@ -219,15 +191,12 @@ pub unsafe extern "C" fn insn_sys_hook_proxy_arm64( user_data: *mut UcHook, ) -> bool where - F: FnMut(&mut crate::Unicorn, sys::RegisterARM64, &sys::RegisterARM64CP) -> bool, + F: FnMut(&mut HookContext, sys::RegisterARM64, &sys::RegisterARM64CP) -> bool, { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); + let mut context = HookContext::from_handle(uc); let cp_reg = unsafe { cp_reg.as_ref() }.unwrap(); - (user_data.callback)(&mut user_data_uc, reg, cp_reg) + (user_data.callback)(&mut context, reg, cp_reg) } /// # Safety @@ -241,14 +210,11 @@ pub unsafe extern "C" fn tlb_lookup_hook_proxy( user_data: *mut UcHook, ) -> bool where - F: FnMut(&mut crate::Unicorn, u64, sys::MemType) -> Option, + F: FnMut(&mut HookContext, u64, sys::MemType) -> Option, { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - let r = (user_data.callback)(&mut user_data_uc, vaddr, mem_type); + let mut context = HookContext::from_handle(uc); + let r = (user_data.callback)(&mut context, vaddr, mem_type); if let Some(ref e) = r { let ref_result: &mut sys::TlbEntry = unsafe { &mut *result }; *ref_result = *e; @@ -267,14 +233,11 @@ pub unsafe extern "C" fn tcg_proxy( size: u32, user_data: *mut UcHook, ) where - F: FnMut(&mut Unicorn, u64, u64, u64, usize), + F: FnMut(&mut HookContext, u64, u64, u64, usize), { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, addr, arg1, arg2, size as usize); + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, addr, arg1, arg2, size as usize); } /// # Safety @@ -286,14 +249,11 @@ pub unsafe extern "C" fn edge_gen_hook_proxy( prev_tb: *mut sys::TranslationBlock, user_data: *mut UcHook, ) where - F: FnMut(&mut Unicorn, &mut sys::TranslationBlock, &mut sys::TranslationBlock), + F: FnMut(&mut HookContext, &mut sys::TranslationBlock, &mut sys::TranslationBlock), { let user_data = unsafe { &mut *user_data }; - let mut user_data_uc = Unicorn { - inner: user_data.uc.upgrade().unwrap(), - }; - debug_assert_eq!(uc, user_data_uc.get_handle()); - (user_data.callback)(&mut user_data_uc, unsafe { &mut *cur_tb }, unsafe { + let mut context = HookContext::from_handle(uc); + (user_data.callback)(&mut context, unsafe { &mut *cur_tb }, unsafe { &mut *prev_tb }); } diff --git a/crates/unicorn/src/lib.rs b/crates/unicorn/src/lib.rs index bfef973..fe09225 100644 --- a/crates/unicorn/src/lib.rs +++ b/crates/unicorn/src/lib.rs @@ -42,8 +42,8 @@ #[macro_use] extern crate alloc; -use alloc::{boxed::Box, rc::Rc, vec::Vec}; -use core::{cell::UnsafeCell, ffi::c_void, ptr}; +use alloc::{boxed::Box, vec::Vec}; +use core::{ffi::c_void, ptr}; #[macro_use] pub mod unicorn_const; @@ -137,6 +137,65 @@ impl MmioCallbackScope<'_> { #[derive(Clone, Copy, PartialEq, Eq, Debug)] pub struct UcHookId(uc_hook); +pub struct HookContext { + handle: *mut uc_engine, +} + +impl HookContext { + pub(crate) fn from_handle(handle: *mut uc_engine) -> Self { + Self { handle } + } + + pub fn get_handle(&self) -> *mut uc_engine { + self.handle + } + + pub fn mem_read(&self, address: u64, buf: &mut [u8]) -> Result<(), uc_error> { + unsafe { + uc_mem_read( + self.handle, + address, + buf.as_mut_ptr().cast(), + buf.len() as u64, + ) + } + .into() + } + + pub fn mem_write(&mut self, address: u64, bytes: &[u8]) -> Result<(), uc_error> { + unsafe { + uc_mem_write( + self.handle, + address, + bytes.as_ptr().cast(), + bytes.len() as u64, + ) + } + .into() + } + + pub fn reg_read>(&self, regid: T) -> Result { + let mut value = 0; + unsafe { uc_reg_read(self.handle, regid.into(), (&raw mut value).cast()) }.and(Ok(value)) + } + + pub fn reg_write>(&mut self, regid: T, value: u64) -> Result<(), uc_error> { + unsafe { uc_reg_write(self.handle, regid.into(), (&raw const value).cast()) }.into() + } + + pub fn emu_stop(&mut self) -> Result<(), uc_error> { + unsafe { uc_emu_stop(self.handle) }.into() + } + + pub fn mem_protect(&mut self, address: u64, size: u64, perms: Prot) -> Result<(), uc_error> { + unsafe { uc_mem_protect(self.handle, address, size, perms.0 as _) }.into() + } + + pub fn ctl_flush_tlb(&mut self) -> Result<(), uc_error> { + unsafe { uc_ctl(self.handle, UC_CTL_WRITE!(ControlType::TLB_FLUSH)) }.into() + } +} + pub struct UnicornInner<'a, D> { pub handle: *mut uc_engine, pub ffi: bool, @@ -158,10 +217,8 @@ impl Drop for UnicornInner<'_, D> { } /// A Unicorn emulator instance. -/// -/// You could clone this instance cheaply, since it has an `Rc` inside. pub struct Unicorn<'a, D: 'a> { - inner: Rc>>, + inner: Box>, } impl<'a> Unicorn<'a, ()> { @@ -192,14 +249,14 @@ where let mut handle = ptr::null_mut(); unsafe { uc_open(arch, mode, &mut handle) }.and_then(|| { Ok(Unicorn { - inner: Rc::new(UnsafeCell::from(UnicornInner { + inner: Box::new(UnicornInner { handle, ffi: false, arch, data, hooks: vec![], mmio_callbacks: vec![], - })), + }), }) }) } @@ -223,14 +280,14 @@ where return Err(err); } Ok(Unicorn { - inner: Rc::new(UnsafeCell::from(UnicornInner { + inner: Box::new(UnicornInner { handle, ffi: true, arch: arch.try_into()?, data, hooks: vec![], mmio_callbacks: vec![], - })), + }), }) } } @@ -241,21 +298,13 @@ impl core::fmt::Debug for Unicorn<'_, D> { } } -impl Clone for Unicorn<'_, D> { - fn clone(&self) -> Self { - Self { - inner: Rc::clone(&self.inner), - } - } -} - impl<'a, D> Unicorn<'a, D> { fn inner(&self) -> &UnicornInner<'a, D> { - unsafe { self.inner.get().as_ref().unwrap() } + &self.inner } fn inner_mut(&mut self) -> &mut UnicornInner<'a, D> { - unsafe { self.inner.get().as_mut().unwrap() } + &mut self.inner } /// Return whatever data was passed during initialization. @@ -442,19 +491,19 @@ impl<'a, D> Unicorn<'a, D> { write_callback: Option, ) -> Result<(), uc_error> where - R: FnMut(&mut Unicorn<'_, D>, u64, usize) -> u64 + 'a, - W: FnMut(&mut Unicorn<'_, D>, u64, usize, u64) + 'a, + R: FnMut(&mut HookContext, u64, usize) -> u64 + 'a, + W: FnMut(&mut HookContext, u64, usize, u64) + 'a, { let mut read_data = read_callback.map(|c| { Box::new(hook::UcHook { callback: c, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }) }); let mut write_data = write_callback.map(|c| { Box::new(hook::UcHook { callback: c, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }) }); @@ -504,13 +553,13 @@ impl<'a, D> Unicorn<'a, D> { /// `size` must be a multiple of 4kb or this will return `Error::ARG`. pub fn mmio_map_ro(&mut self, address: u64, size: u64, callback: F) -> Result<(), uc_error> where - F: FnMut(&mut Unicorn, u64, usize) -> u64 + 'a, + F: FnMut(&mut HookContext, u64, usize) -> u64 + 'a, { self.mmio_map( address, size, Some(callback), - None::, u64, usize, u64)>, + None::, ) } @@ -520,12 +569,12 @@ impl<'a, D> Unicorn<'a, D> { /// `size` must be a multiple of 4kb or this will return `Error::ARG`. pub fn mmio_map_wo(&mut self, address: u64, size: u64, callback: F) -> Result<(), uc_error> where - F: FnMut(&mut Unicorn, u64, usize, u64) + 'a, + F: FnMut(&mut HookContext, u64, usize, u64) + 'a, { self.mmio_map( address, size, - None::, u64, usize) -> u64>, + None:: u64>, Some(callback), ) } @@ -777,12 +826,12 @@ impl<'a, D> Unicorn<'a, D> { callback: F, ) -> Result where - F: FnMut(&mut Unicorn, u64, u32) + 'a, + F: FnMut(&mut HookContext, u64, u32) + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -811,12 +860,12 @@ impl<'a, D> Unicorn<'a, D> { callback: F, ) -> Result where - F: FnMut(&mut Unicorn, u64, u32) + 'a, + F: FnMut(&mut HookContext, u64, u32) + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -846,7 +895,7 @@ impl<'a, D> Unicorn<'a, D> { callback: F, ) -> Result where - F: FnMut(&mut Unicorn, MemType, u64, usize, i64) -> bool + 'a, + F: FnMut(&mut HookContext, MemType, u64, usize, i64) -> bool + 'a, { if hook_type & (HookType::MEM_ALL | HookType::MEM_READ_AFTER) != hook_type { return Err(uc_error::ARG); @@ -855,7 +904,7 @@ impl<'a, D> Unicorn<'a, D> { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -879,12 +928,12 @@ impl<'a, D> Unicorn<'a, D> { /// Add an interrupt hook. pub fn add_intr_hook(&mut self, callback: F) -> Result where - F: FnMut(&mut Unicorn, u32) + 'a, + F: FnMut(&mut HookContext, u32) + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -908,12 +957,12 @@ impl<'a, D> Unicorn<'a, D> { /// Add hook for invalid instructions pub fn add_insn_invalid_hook(&mut self, callback: F) -> Result where - F: FnMut(&mut Unicorn) -> bool + 'a, + F: FnMut(&mut HookContext) -> bool + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -938,12 +987,12 @@ impl<'a, D> Unicorn<'a, D> { #[cfg(feature = "arch_x86")] pub fn add_insn_in_hook(&mut self, callback: F) -> Result where - F: FnMut(&mut Unicorn, u32, usize) -> u32 + 'a, + F: FnMut(&mut HookContext, u32, usize) -> u32 + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -969,12 +1018,12 @@ impl<'a, D> Unicorn<'a, D> { #[cfg(feature = "arch_x86")] pub fn add_insn_out_hook(&mut self, callback: F) -> Result where - F: FnMut(&mut Unicorn, u32, usize, u32) + 'a, + F: FnMut(&mut HookContext, u32, usize, u32) + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -1006,12 +1055,12 @@ impl<'a, D> Unicorn<'a, D> { callback: F, ) -> Result where - F: FnMut(&mut Unicorn) + 'a, + F: FnMut(&mut HookContext) + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -1046,12 +1095,12 @@ impl<'a, D> Unicorn<'a, D> { callback: F, ) -> Result where - F: FnMut(&mut Unicorn, RegisterARM64, &RegisterARM64CP) -> bool + 'a, + F: FnMut(&mut HookContext, RegisterARM64, &RegisterARM64CP) -> bool + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -1080,12 +1129,12 @@ impl<'a, D> Unicorn<'a, D> { callback: F, ) -> Result where - F: FnMut(&mut Unicorn, u64, MemType) -> Option + 'a, + F: FnMut(&mut HookContext, u64, MemType) -> Option + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -1115,12 +1164,12 @@ impl<'a, D> Unicorn<'a, D> { callback: F, ) -> Result where - F: FnMut(&mut Unicorn, u64, u64, u64, usize) + 'a, + F: FnMut(&mut HookContext, u64, u64, u64, usize) + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { @@ -1153,12 +1202,12 @@ impl<'a, D> Unicorn<'a, D> { callback: F, ) -> Result where - F: FnMut(&mut Unicorn, &mut TranslationBlock, &mut TranslationBlock) + 'a, + F: FnMut(&mut HookContext, &mut TranslationBlock, &mut TranslationBlock) + 'a, { let mut hook_id = 0; let mut user_data = Box::new(hook::UcHook { callback, - uc: Rc::downgrade(&self.inner), + marker: core::marker::PhantomData::<&'a D>, }); unsafe { diff --git a/crates/unicorn/src/tests/arm.rs b/crates/unicorn/src/tests/arm.rs index 32dcbd3..aab28d8 100644 --- a/crates/unicorn/src/tests/arm.rs +++ b/crates/unicorn/src/tests/arm.rs @@ -1,5 +1,6 @@ use super::*; use crate::{ArmCpuModel, RegisterARM, RegisterARMCP, TcgOpCode, TcgOpFlag, uc_error}; +use std::{cell::RefCell, rc::Rc}; #[test] fn test_arm_nop() { @@ -146,14 +147,14 @@ fn test_arm_thumb_ite() { let mut r2 = 0u32; let mut r3 = 1u32; let mut pc = CODE_START as u32; - let mut count = 0; + let count = Rc::new(RefCell::new(0u64)); let mut uc = uc_common_setup( Arch::ARM, Mode::THUMB, Some(ArmCpuModel::CORTEX_A15 as i32), code, - count, + (), ); uc.reg_write(RegisterARM::SP, sp).unwrap(); @@ -166,9 +167,14 @@ fn test_arm_thumb_ite() { r2 = 0x4d; uc.mem_write(sp + 4, &r2.to_le_bytes()).unwrap(); - uc.add_code_hook(CODE_START, CODE_START + code.len() as u64, |uc, _, _| { - *uc.get_data_mut() += 1; - }) + let callback_count = Rc::clone(&count); + uc.add_code_hook( + CODE_START, + CODE_START + code.len() as u64, + move |_, _, _| { + *callback_count.borrow_mut() += 1; + }, + ) .unwrap(); // Execute four instructions @@ -177,13 +183,12 @@ fn test_arm_thumb_ite() { r2 = uc.reg_read(RegisterARM::R2).unwrap() as u32; r3 = uc.reg_read(RegisterARM::R3).unwrap() as u32; - count = *uc.get_data(); assert_eq!(r2, 0x68); assert_eq!(r3, 0x78); - assert_eq!(count, 4); + assert_eq!(*count.borrow(), 4); r2 = 0; - *uc.get_data_mut() = 0; + *count.borrow_mut() = 0; uc.reg_write(RegisterARM::R2, r2 as u64).unwrap(); uc.reg_write(RegisterARM::R3, r3 as u64).unwrap(); @@ -198,11 +203,10 @@ fn test_arm_thumb_ite() { r2 = uc.reg_read(RegisterARM::R2).unwrap() as u32; r3 = uc.reg_read(RegisterARM::R3).unwrap() as u32; - count = *uc.get_data(); assert_eq!(r2, 0x68); assert_eq!(r3, 0x78); - assert_eq!(count, 4); + assert_eq!(*count.borrow(), 4); } #[test] @@ -559,19 +563,22 @@ fn test_arm_mem_access_abort() { Mode::ARM, Some(ArmCpuModel::CORTEX_A9 as i32), code, - 0, + (), ); + let callback_pc = Rc::new(RefCell::new(0u64)); uc.reg_write(RegisterARM::R0, r0 as u64).unwrap(); - uc.add_mem_hook(HookType::MEM_UNMAPPED, 1, 0, |uc, _, _, _, _| { - *uc.get_data_mut() = uc.reg_read(RegisterARM::PC).unwrap(); + let mem_callback_pc = Rc::clone(&callback_pc); + uc.add_mem_hook(HookType::MEM_UNMAPPED, 1, 0, move |uc, _, _, _, _| { + *mem_callback_pc.borrow_mut() = uc.reg_read(RegisterARM::PC).unwrap(); false }) .unwrap(); - uc.add_insn_invalid_hook(|uc| { - *uc.get_data_mut() = uc.reg_read(RegisterARM::PC).unwrap(); + let invalid_callback_pc = Rc::clone(&callback_pc); + uc.add_insn_invalid_hook(move |uc| { + *invalid_callback_pc.borrow_mut() = uc.reg_read(RegisterARM::PC).unwrap(); false }) .unwrap(); @@ -580,7 +587,7 @@ fn test_arm_mem_access_abort() { assert_eq!(err, uc_error::READ_UNMAPPED); let pc = uc.reg_read(RegisterARM::PC).unwrap(); - assert_eq!(pc, *uc.get_data()); + assert_eq!(pc, *callback_pc.borrow()); let err = uc .emu_start(CODE_START + 4, CODE_START + 8, 0, 0) @@ -588,13 +595,13 @@ fn test_arm_mem_access_abort() { assert_eq!(err, uc_error::INSN_INVALID); let pc = uc.reg_read(RegisterARM::PC).unwrap(); - assert_eq!(pc, *uc.get_data()); + assert_eq!(pc, *callback_pc.borrow()); let err = uc.emu_start(0x900000, 0x900000 + 8, 0, 0).unwrap_err(); assert_eq!(err, uc_error::FETCH_UNMAPPED); let pc = uc.reg_read(RegisterARM::PC).unwrap(); - assert_eq!(pc, *uc.get_data()); + assert_eq!(pc, *callback_pc.borrow()); } #[test] @@ -805,20 +812,23 @@ fn test_arm_mem_hook_read_write() { Mode::ARM, Some(ArmCpuModel::CORTEX_A15 as i32), code, - [0u64; 2], + (), ); uc.reg_write(RegisterARM::SP, sp).unwrap(); uc.mem_map(0x8000, 1024 * 16, Prot::ALL).unwrap(); - uc.add_mem_hook(HookType::MEM_READ, 1, 0, |uc, _, _, _, _| { - (*uc.get_data_mut())[0] += 1; + let counters = Rc::new(RefCell::new([0u64; 2])); + let read_counters = Rc::clone(&counters); + uc.add_mem_hook(HookType::MEM_READ, 1, 0, move |_, _, _, _, _| { + read_counters.borrow_mut()[0] += 1; false }) .unwrap(); - uc.add_mem_hook(HookType::MEM_WRITE, 1, 0, |uc, _, _, _, _| { - (*uc.get_data_mut())[1] += 1; + let write_counters = Rc::clone(&counters); + uc.add_mem_hook(HookType::MEM_WRITE, 1, 0, move |_, _, _, _, _| { + write_counters.borrow_mut()[1] += 1; false }) .unwrap(); @@ -826,7 +836,7 @@ fn test_arm_mem_hook_read_write() { uc.emu_start(CODE_START, CODE_START + code.len() as u64, 0, 0) .unwrap(); - let [read, write] = *uc.get_data(); + let [read, write] = *counters.borrow(); assert_eq!(read, 2); assert_eq!(write, 2); } @@ -839,8 +849,8 @@ struct CmpInfo { pc: u64, } -fn uc_hook_sub_cmp(uc: &mut Unicorn<'_, CmpInfo>, address: u64, arg1: u64, arg2: u64, size: usize) { - let data = uc.get_data_mut(); +fn uc_hook_sub_cmp(data: &Rc>, address: u64, arg1: u64, arg2: u64, size: usize) { + let mut data = data.borrow_mut(); data.pc = address; data.size = size as u64; data.v0 = arg1; @@ -862,16 +872,24 @@ fn test_arm_tcg_opcode_cmp() { Mode::ARM, Some(ArmCpuModel::CORTEX_A15 as i32), code, - CmpInfo::default(), + (), ); + let cmp_info = Rc::new(RefCell::new(CmpInfo::default())); + let cmp_callback_info = Rc::clone(&cmp_info); - uc.add_tcg_hook(TcgOpCode::SUB, TcgOpFlag::CMP, 1, 0, uc_hook_sub_cmp) - .unwrap(); + uc.add_tcg_hook( + TcgOpCode::SUB, + TcgOpFlag::CMP, + 1, + 0, + move |_, a, b, c, d| uc_hook_sub_cmp(&cmp_callback_info, a, b, c, d), + ) + .unwrap(); uc.emu_start(CODE_START, CODE_START + code.len() as u64, 0, 3) .unwrap(); - let cmp_info = uc.get_data(); + let cmp_info = cmp_info.borrow(); assert_eq!(cmp_info.v0, 5); assert_eq!(cmp_info.v1, 3); assert_eq!(cmp_info.pc, 0x1008); @@ -895,16 +913,24 @@ fn test_arm_thumb_tcg_opcode_cmn() { Mode::THUMB, Some(ArmCpuModel::CORTEX_A15 as i32), code, - CmpInfo::default(), + (), ); + let cmp_info = Rc::new(RefCell::new(CmpInfo::default())); + let cmp_callback_info = Rc::clone(&cmp_info); - uc.add_tcg_hook(TcgOpCode::SUB, TcgOpFlag::CMP, 1, 0, uc_hook_sub_cmp) - .unwrap(); + uc.add_tcg_hook( + TcgOpCode::SUB, + TcgOpFlag::CMP, + 1, + 0, + move |_, a, b, c, d| uc_hook_sub_cmp(&cmp_callback_info, a, b, c, d), + ) + .unwrap(); uc.emu_start(CODE_START | 1, CODE_START + code.len() as u64, 0, 4) .unwrap(); - let cmp_info = uc.get_data(); + let cmp_info = cmp_info.borrow(); assert_eq!(cmp_info.v0, 5); assert_eq!(cmp_info.v1, 3); assert_eq!(cmp_info.pc, 0x1006); diff --git a/crates/unicorn/src/tests/arm64.rs b/crates/unicorn/src/tests/arm64.rs index 7992f6f..ffa4c69 100644 --- a/crates/unicorn/src/tests/arm64.rs +++ b/crates/unicorn/src/tests/arm64.rs @@ -1,3 +1,4 @@ +use std::{cell::RefCell, rc::Rc}; use unicorn_engine_sys::{Arm64CpuModel, Arm64Insn, RegisterARM64, RegisterARM64CP}; use super::*; @@ -251,22 +252,22 @@ fn test_arm64_block_sync_pc() { 0xc1, 0xc5, 0x82, 0xd2, // t: mov x1, #5678 ]; + let first = Rc::new(RefCell::new(true)); + let callback_first = Rc::clone(&first); let mut uc = uc_common_setup( Arch::ARM64, Mode::ARM, Some(Arm64CpuModel::A72 as i32), code, - true, + (), ); - - uc.add_block_hook(CODE_START + 8, CODE_START + 12, |uc, addr, _| { + uc.add_block_hook(CODE_START + 8, CODE_START + 12, move |uc, addr, _| { let pc = uc.reg_read(RegisterARM64::PC).unwrap(); assert_eq!(pc, addr); let val = CODE_START; - let first = *uc.get_data_mut(); - if first { + if *callback_first.borrow() { uc.reg_write(RegisterARM64::PC, val).unwrap(); - *uc.get_data_mut() = false; + *callback_first.borrow_mut() = false; } }) .unwrap(); @@ -510,26 +511,29 @@ fn test_arm64_mem_hook_read_write() { 0xe1, 0x0b, 0x00, 0xa9, // stp x1, x2, [sp] ]; + let counters = Rc::new(RefCell::new([0u64; 2])); + let read_counters = Rc::clone(&counters); + let write_counters = Rc::clone(&counters); let mut uc = uc_common_setup( Arch::ARM64, Mode::ARM, Some(Arm64CpuModel::A72 as i32), code, - [0, 0], + (), ); let sp = 0x16db6a040; uc.reg_write(RegisterARM64::SP, sp).unwrap(); uc.mem_map(0x16db68000, 1024 * 16, Prot::ALL).unwrap(); - uc.add_mem_hook(HookType::MEM_READ, 1, 0, |uc, _, _, _, _| { - (*uc.get_data_mut())[0] += 1; + uc.add_mem_hook(HookType::MEM_READ, 1, 0, move |_, _, _, _, _| { + read_counters.borrow_mut()[0] += 1; false }) .unwrap(); - uc.add_mem_hook(HookType::MEM_WRITE, 1, 0, |uc, _, _, _, _| { - (*uc.get_data_mut())[1] += 1; + uc.add_mem_hook(HookType::MEM_WRITE, 1, 0, move |_, _, _, _, _| { + write_counters.borrow_mut()[1] += 1; false }) .unwrap(); @@ -537,7 +541,7 @@ fn test_arm64_mem_hook_read_write() { uc.emu_start(CODE_START, CODE_START + code.len() as u64, 0, 0) .unwrap(); - let [read, write] = *uc.get_data(); + let [read, write] = *counters.borrow(); assert_eq!(read, 4); assert_eq!(write, 4); } diff --git a/crates/unicorn/src/tests/ctl.rs b/crates/unicorn/src/tests/ctl.rs index 0acece6..995cd35 100644 --- a/crates/unicorn/src/tests/ctl.rs +++ b/crates/unicorn/src/tests/ctl.rs @@ -1,4 +1,5 @@ use std::time::{Duration, Instant}; +use std::{cell::RefCell, rc::Rc}; use unicorn_engine_sys::{RegisterX86, X86Insn}; @@ -152,13 +153,14 @@ fn test_uc_ctl_change_page_size_arm64() { fn test_uc_hook_cached_uaf() { let code = b"\x41\x4a\xeb\x00\x90"; - let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_32, 0u64).unwrap(); + let count = Rc::new(RefCell::new(0u64)); + let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_32, ()).unwrap(); uc.mem_map(CODE_START, CODE_LEN, Prot::ALL).unwrap(); uc.mem_write(CODE_START, code).unwrap(); let hook = uc - .add_code_hook(CODE_START, CODE_START + code.len() as u64, |uc, _, _| { - *uc.get_data_mut() += 1; + .add_code_hook(CODE_START, CODE_START + code.len() as u64, |_, _, _| { + *count.borrow_mut() += 1; }) .unwrap(); @@ -180,7 +182,7 @@ fn test_uc_hook_cached_uaf() { .unwrap(); // Only 4 calls - assert_eq!(*uc.get_data(), 4); + assert_eq!(*count.borrow(), 4); } #[test] @@ -232,7 +234,8 @@ fn test_tlb_clear() { 0xa3, 0x00, 0x00, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, // movabs dword ptr [0x200000], eax ]; - let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_64, 0usize).unwrap(); + let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_64, ()).unwrap(); + let tlbcount = Rc::new(RefCell::new(0usize)); uc.mem_map(CODE_START, CODE_LEN.try_into().unwrap(), Prot::ALL) .unwrap(); uc.mem_write(CODE_START, code).unwrap(); @@ -240,8 +243,9 @@ fn test_tlb_clear() { uc.mem_map(0x200000, 0x1000, Prot::ALL).unwrap(); uc.ctl_set_tlb_type(TlbType::VIRTUAL).unwrap(); - uc.add_tlb_hook(1, 0, |uc, addr, _| { - *uc.get_data_mut() += 1; + let callback_tlbcount = Rc::clone(&tlbcount); + uc.add_tlb_hook(1, 0, move |_, addr, _| { + *callback_tlbcount.borrow_mut() += 1; Some(TlbEntry { paddr: addr, perms: Prot::ALL, @@ -256,7 +260,7 @@ fn test_tlb_clear() { uc.emu_start(CODE_START, CODE_START + code.len() as u64, 0, 0) .unwrap(); - let tlbcount = *uc.get_data(); + let tlbcount = *tlbcount.borrow(); assert_eq!(tlbcount, 4); } diff --git a/crates/unicorn/src/tests/mem.rs b/crates/unicorn/src/tests/mem.rs index 3a34010..c4ab3e3 100644 --- a/crates/unicorn/src/tests/mem.rs +++ b/crates/unicorn/src/tests/mem.rs @@ -1,6 +1,8 @@ +use std::{cell::RefCell, rc::Rc}; use unicorn_engine_sys::{ContextMode, RegisterX86}; use super::*; +use crate::HookContext; #[test] fn test_map_correct() { @@ -175,14 +177,16 @@ fn test_mem_protect_remove_exec() { 0x90, // nop ]; - let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_64, 0).unwrap(); + let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_64, ()).unwrap(); + let count = Rc::new(RefCell::new(0u64)); uc.mem_map(0x1000, 0x1000, Prot::ALL).unwrap(); uc.mem_map(0x2000, 0x1000, Prot::ALL).unwrap(); uc.mem_write(0x1000, &code).unwrap(); - uc.add_block_hook(1, 0, |uc, _, _| { - *uc.get_data_mut() += 1; + let callback_count = Rc::clone(&count); + uc.add_block_hook(1, 0, move |uc, _, _| { + *callback_count.borrow_mut() += 1; uc.mem_protect(0x2000, 0x1000, Prot::READ).unwrap(); }) .unwrap(); @@ -190,7 +194,7 @@ fn test_mem_protect_remove_exec() { uc.emu_start(0x1000, 0x1000 + code.len() as u64, 0, 0) .unwrap(); - assert_eq!(*uc.get_data_mut(), 2); + assert_eq!(*count.borrow(), 2); } #[test] @@ -201,7 +205,8 @@ fn test_mem_protect_mmio() { 0xa3, 0x20, 0x20, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // movabs dword ptr [0x2020], eax ]; - let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_64, 0).unwrap(); + let mut uc = Unicorn::new_with_data(Arch::X86, Mode::MODE_64, ()).unwrap(); + let count = Rc::new(RefCell::new(0i32)); uc.mem_map(0x8000, 0x1000, Prot::ALL).unwrap(); uc.mem_write(0x8000, &code).unwrap(); @@ -209,12 +214,15 @@ fn test_mem_protect_mmio() { uc.mmio_map( 0x1000, 0x3000, - Some(|uc: &mut Unicorn<'_, i32>, addr, _| { - assert_eq!(addr, 0x20); - *uc.get_data_mut() += 1; - 0x114514 + Some({ + let callback_count = Rc::clone(&count); + move |_: &mut HookContext, addr, _| { + assert_eq!(addr, 0x20); + *callback_count.borrow_mut() += 1; + 0x114514 + } }), - Some(|_: &mut Unicorn<'_, i32>, _addr, _size, _val| { + Some(|_: &mut HookContext, _addr, _size, _val| { panic!("Write callback should not be called"); }), ) @@ -227,7 +235,7 @@ fn test_mem_protect_mmio() { ); let eax = uc.reg_read(RegisterX86::RAX).unwrap(); - assert_eq!(*uc.get_data_mut(), 1); + assert_eq!(*count.borrow(), 1); assert_eq!(eax, 0x114514u64); }