From 5559c097d58dff824f221e8063c51dd21bb7c391 Mon Sep 17 00:00:00 2001 From: Bet4 <0xbet4@gmail.com> Date: Mon, 17 Jan 2022 21:36:45 +0800 Subject: [PATCH] rust: Allow to remove self inside a hook --- bindings/rust/src/ffi.rs | 7 ++-- bindings/rust/src/lib.rs | 68 +++++++++++++++++++++------------- bindings/rust/tests/unicorn.rs | 18 ++++++--- 3 files changed, 59 insertions(+), 34 deletions(-) diff --git a/bindings/rust/src/ffi.rs b/bindings/rust/src/ffi.rs index e2b7c078..47240b43 100644 --- a/bindings/rust/src/ffi.rs +++ b/bindings/rust/src/ffi.rs @@ -96,12 +96,13 @@ pub trait IsUcHook<'a> {} impl<'a, D, F> IsUcHook<'a> for UcHook<'a, D, F> {} -pub extern "C" fn mmio_read_callback_proxy ( +pub extern "C" fn mmio_read_callback_proxy( uc: uc_handle, offset: u64, size: usize, user_data: *mut UcHook, -) -> u64 where +) -> u64 +where F: FnMut(&mut crate::Unicorn, u64, usize) -> u64, { let user_data = unsafe { &mut *user_data }; @@ -109,7 +110,7 @@ pub extern "C" fn mmio_read_callback_proxy ( (user_data.callback)(&mut user_data.uc, offset, size) } -pub extern "C" fn mmio_write_callback_proxy ( +pub extern "C" fn mmio_write_callback_proxy( uc: uc_handle, offset: u64, size: usize, diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs index 9ec402e0..203b6e0d 100644 --- a/bindings/rust/src/lib.rs +++ b/bindings/rust/src/lib.rs @@ -84,29 +84,32 @@ pub struct MmioCallbackScope<'a> { impl<'a> MmioCallbackScope<'a> { fn has_regions(&self) -> bool { - self.regions.len() > 0 + !self.regions.is_empty() } fn unmap(&mut self, begin: u64, size: usize) { let end: u64 = begin + size as u64; - self.regions = self.regions.iter().flat_map( |(b, s)| { - let e: u64 = b + *s as u64; - if begin > *b { - if begin >= e { - // The unmapped region is completely after this region - vec![(*b, *s)] - } else { - if end >= e { + self.regions = self + .regions + .iter() + .flat_map(|(b, s)| { + let e: u64 = b + *s as u64; + if begin > *b { + if begin >= e { + // The unmapped region is completely after this region + vec![(*b, *s)] + } else if end >= e { // The unmapped region overlaps with the end of this region vec![(*b, (begin - *b) as usize)] } else { // The unmapped region is in the middle of this region let second_b = end + 1; - vec![(*b, (begin - *b) as usize), (second_b, (e - second_b) as usize)] + vec![ + (*b, (begin - *b) as usize), + (second_b, (e - second_b) as usize), + ] } - } - } else { - if end > *b { + } else if end > *b { if end >= e { // The unmapped region completely contains this region vec![] @@ -118,8 +121,8 @@ impl<'a> MmioCallbackScope<'a> { // The unmapped region is completely before this region vec![(*b, *s)] } - } - }).collect(); + }) + .collect(); } } @@ -318,7 +321,7 @@ impl<'a, D> Unicorn<'a, D> { /// `size` must be a multiple of 4kb or this will return `Error::ARG`. pub fn mmio_map( &mut self, - address: u64, + address: u64, size: libc::size_t, read_callback: Option, write_callback: Option, @@ -327,7 +330,7 @@ impl<'a, D> Unicorn<'a, D> { R: FnMut(&mut Unicorn, u64, usize) -> u64, W: FnMut(&mut Unicorn, u64, usize, u64), { - let mut read_data = read_callback.map( |c| { + let mut read_data = read_callback.map(|c| { Box::new(ffi::UcHook { callback: c, uc: Unicorn { @@ -335,7 +338,7 @@ impl<'a, D> Unicorn<'a, D> { }, }) }); - let mut write_data = write_callback.map( |c| { + let mut write_data = write_callback.map(|c| { Box::new(ffi::UcHook { callback: c, uc: Unicorn { @@ -363,9 +366,9 @@ impl<'a, D> Unicorn<'a, D> { }; if err == uc_error::OK { - let rd = read_data.map( |c| c as Box ); - let wd = write_data.map( |c| c as Box ); - self.inner_mut().mmio_callbacks.push(MmioCallbackScope{ + let rd = read_data.map(|c| c as Box); + let wd = write_data.map(|c| c as Box); + self.inner_mut().mmio_callbacks.push(MmioCallbackScope { regions: vec![(address, size)], read_callback: rd, write_callback: wd, @@ -390,7 +393,12 @@ impl<'a, D> Unicorn<'a, D> { where F: FnMut(&mut Unicorn, u64, usize) -> u64, { - self.mmio_map(address, size, Some(callback), None::, u64, usize, u64)>) + self.mmio_map( + address, + size, + Some(callback), + None::, u64, usize, u64)>, + ) } /// Map in a write-only MMIO region backed by a callback. @@ -406,7 +414,12 @@ impl<'a, D> Unicorn<'a, D> { where F: FnMut(&mut Unicorn, u64, usize, u64), { - self.mmio_map(address, size, None::, u64, usize) -> u64>, Some(callback)) + self.mmio_map( + address, + size, + None::, u64, usize) -> u64>, + Some(callback), + ) } /// Unmap a memory region. @@ -429,7 +442,9 @@ impl<'a, D> Unicorn<'a, D> { for scope in self.inner_mut().mmio_callbacks.iter_mut() { scope.unmap(address, size); } - self.inner_mut().mmio_callbacks.retain( |scope| scope.has_regions() ); + self.inner_mut() + .mmio_callbacks + .retain(|scope| scope.has_regions()); } /// Set the memory permissions for an existing memory region. @@ -817,11 +832,12 @@ impl<'a, D> Unicorn<'a, D> { let err: uc_error; // drop the hook - self.inner_mut() + let inner = self.inner_mut(); + inner .hooks .retain(|(hook_ptr, _hook_impl)| hook_ptr != &hook); - err = unsafe { ffi::uc_hook_del(self.inner().uc, hook) }; + err = unsafe { ffi::uc_hook_del(inner.uc, hook) }; if err == uc_error::OK { Ok(()) diff --git a/bindings/rust/tests/unicorn.rs b/bindings/rust/tests/unicorn.rs index 79f659aa..dba25c49 100644 --- a/bindings/rust/tests/unicorn.rs +++ b/bindings/rust/tests/unicorn.rs @@ -427,7 +427,10 @@ fn x86_mmio() { { // MOV eax, [0x2004]; MOV [0x2008], ax; - let x86_code: Vec = vec![0x8B, 0x04, 0x25, 0x04, 0x20, 0x00, 0x00, 0x66, 0x89, 0x04, 0x25, 0x08, 0x20, 0x00, 0x00]; + let x86_code: Vec = vec![ + 0x8B, 0x04, 0x25, 0x04, 0x20, 0x00, 0x00, 0x66, 0x89, 0x04, 0x25, 0x08, 0x20, 0x00, + 0x00, + ]; let read_cell = Rc::new(RefCell::new(MmioReadExpectation(0, 0))); let cb_read_cell = read_cell.clone(); @@ -436,7 +439,7 @@ fn x86_mmio() { 42 }; - let write_cell = Rc::new(RefCell::new(MmioWriteExpectation(0,0,0))); + let write_cell = Rc::new(RefCell::new(MmioWriteExpectation(0, 0, 0))); let cb_write_cell = write_cell.clone(); let write_callback = move |_: &mut Unicorn<'_, ()>, offset, size, value| { *cb_write_cell.borrow_mut() = MmioWriteExpectation(offset, size, value); @@ -444,7 +447,10 @@ fn x86_mmio() { assert_eq!(emu.mem_write(0x1000, &x86_code), Ok(())); - assert_eq!(emu.mmio_map(0x2000, 0x1000, Some(read_callback), Some(write_callback)), Ok(())); + assert_eq!( + emu.mmio_map(0x2000, 0x1000, Some(read_callback), Some(write_callback)), + Ok(()) + ); assert_eq!( emu.emu_start( @@ -494,9 +500,11 @@ fn x86_mmio() { { // MOV ax, 42; MOV [0x2008], ax; - let x86_code: Vec = vec![0x66, 0xB8, 0x2A, 0x00, 0x66, 0x89, 0x04, 0x25, 0x08, 0x20, 0x00, 0x00]; + let x86_code: Vec = vec![ + 0x66, 0xB8, 0x2A, 0x00, 0x66, 0x89, 0x04, 0x25, 0x08, 0x20, 0x00, 0x00, + ]; - let write_cell = Rc::new(RefCell::new(MmioWriteExpectation(0,0,0))); + let write_cell = Rc::new(RefCell::new(MmioWriteExpectation(0, 0, 0))); let cb_write_cell = write_cell.clone(); let write_callback = move |_: &mut Unicorn<'_, ()>, offset, size, value| { *cb_write_cell.borrow_mut() = MmioWriteExpectation(offset, size, value);