rust: Allow to remove self inside a hook

This commit is contained in:
Bet4 2022-01-17 21:36:45 +08:00
parent ea9c7425b0
commit 5559c097d5
3 changed files with 59 additions and 34 deletions

View File

@ -96,12 +96,13 @@ pub trait IsUcHook<'a> {}
impl<'a, D, F> IsUcHook<'a> for UcHook<'a, D, F> {} impl<'a, D, F> IsUcHook<'a> for UcHook<'a, D, F> {}
pub extern "C" fn mmio_read_callback_proxy<D, F> ( pub extern "C" fn mmio_read_callback_proxy<D, F>(
uc: uc_handle, uc: uc_handle,
offset: u64, offset: u64,
size: usize, size: usize,
user_data: *mut UcHook<D, F>, user_data: *mut UcHook<D, F>,
) -> u64 where ) -> u64
where
F: FnMut(&mut crate::Unicorn<D>, u64, usize) -> u64, F: FnMut(&mut crate::Unicorn<D>, u64, usize) -> u64,
{ {
let user_data = unsafe { &mut *user_data }; let user_data = unsafe { &mut *user_data };
@ -109,7 +110,7 @@ pub extern "C" fn mmio_read_callback_proxy<D, F> (
(user_data.callback)(&mut user_data.uc, offset, size) (user_data.callback)(&mut user_data.uc, offset, size)
} }
pub extern "C" fn mmio_write_callback_proxy<D, F> ( pub extern "C" fn mmio_write_callback_proxy<D, F>(
uc: uc_handle, uc: uc_handle,
offset: u64, offset: u64,
size: usize, size: usize,

View File

@ -84,29 +84,32 @@ pub struct MmioCallbackScope<'a> {
impl<'a> MmioCallbackScope<'a> { impl<'a> MmioCallbackScope<'a> {
fn has_regions(&self) -> bool { fn has_regions(&self) -> bool {
self.regions.len() > 0 !self.regions.is_empty()
} }
fn unmap(&mut self, begin: u64, size: usize) { fn unmap(&mut self, begin: u64, size: usize) {
let end: u64 = begin + size as u64; let end: u64 = begin + size as u64;
self.regions = self.regions.iter().flat_map( |(b, s)| { self.regions = self
let e: u64 = b + *s as u64; .regions
if begin > *b { .iter()
if begin >= e { .flat_map(|(b, s)| {
// The unmapped region is completely after this region let e: u64 = b + *s as u64;
vec![(*b, *s)] if begin > *b {
} else { if begin >= e {
if end >= 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 // The unmapped region overlaps with the end of this region
vec![(*b, (begin - *b) as usize)] vec![(*b, (begin - *b) as usize)]
} else { } else {
// The unmapped region is in the middle of this region // The unmapped region is in the middle of this region
let second_b = end + 1; 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 { if end >= e {
// The unmapped region completely contains this region // The unmapped region completely contains this region
vec![] vec![]
@ -118,8 +121,8 @@ impl<'a> MmioCallbackScope<'a> {
// The unmapped region is completely before this region // The unmapped region is completely before this region
vec![(*b, *s)] vec![(*b, *s)]
} }
} })
}).collect(); .collect();
} }
} }
@ -327,7 +330,7 @@ impl<'a, D> Unicorn<'a, D> {
R: FnMut(&mut Unicorn<D>, u64, usize) -> u64, R: FnMut(&mut Unicorn<D>, u64, usize) -> u64,
W: FnMut(&mut Unicorn<D>, u64, usize, u64), W: FnMut(&mut Unicorn<D>, u64, usize, u64),
{ {
let mut read_data = read_callback.map( |c| { let mut read_data = read_callback.map(|c| {
Box::new(ffi::UcHook { Box::new(ffi::UcHook {
callback: c, callback: c,
uc: Unicorn { 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 { Box::new(ffi::UcHook {
callback: c, callback: c,
uc: Unicorn { uc: Unicorn {
@ -363,9 +366,9 @@ impl<'a, D> Unicorn<'a, D> {
}; };
if err == uc_error::OK { if err == uc_error::OK {
let rd = read_data.map( |c| c as Box<dyn ffi::IsUcHook> ); let rd = read_data.map(|c| c as Box<dyn ffi::IsUcHook>);
let wd = write_data.map( |c| c as Box<dyn ffi::IsUcHook> ); let wd = write_data.map(|c| c as Box<dyn ffi::IsUcHook>);
self.inner_mut().mmio_callbacks.push(MmioCallbackScope{ self.inner_mut().mmio_callbacks.push(MmioCallbackScope {
regions: vec![(address, size)], regions: vec![(address, size)],
read_callback: rd, read_callback: rd,
write_callback: wd, write_callback: wd,
@ -390,7 +393,12 @@ impl<'a, D> Unicorn<'a, D> {
where where
F: FnMut(&mut Unicorn<D>, u64, usize) -> u64, F: FnMut(&mut Unicorn<D>, u64, usize) -> u64,
{ {
self.mmio_map(address, size, Some(callback), None::<fn(&mut Unicorn<D>, u64, usize, u64)>) self.mmio_map(
address,
size,
Some(callback),
None::<fn(&mut Unicorn<D>, u64, usize, u64)>,
)
} }
/// Map in a write-only MMIO region backed by a callback. /// Map in a write-only MMIO region backed by a callback.
@ -406,7 +414,12 @@ impl<'a, D> Unicorn<'a, D> {
where where
F: FnMut(&mut Unicorn<D>, u64, usize, u64), F: FnMut(&mut Unicorn<D>, u64, usize, u64),
{ {
self.mmio_map(address, size, None::<fn(&mut Unicorn<D>, u64, usize) -> u64>, Some(callback)) self.mmio_map(
address,
size,
None::<fn(&mut Unicorn<D>, u64, usize) -> u64>,
Some(callback),
)
} }
/// Unmap a memory region. /// Unmap a memory region.
@ -429,7 +442,9 @@ impl<'a, D> Unicorn<'a, D> {
for scope in self.inner_mut().mmio_callbacks.iter_mut() { for scope in self.inner_mut().mmio_callbacks.iter_mut() {
scope.unmap(address, size); 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. /// Set the memory permissions for an existing memory region.
@ -817,11 +832,12 @@ impl<'a, D> Unicorn<'a, D> {
let err: uc_error; let err: uc_error;
// drop the hook // drop the hook
self.inner_mut() let inner = self.inner_mut();
inner
.hooks .hooks
.retain(|(hook_ptr, _hook_impl)| hook_ptr != &hook); .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 { if err == uc_error::OK {
Ok(()) Ok(())

View File

@ -427,7 +427,10 @@ fn x86_mmio() {
{ {
// MOV eax, [0x2004]; MOV [0x2008], ax; // MOV eax, [0x2004]; MOV [0x2008], ax;
let x86_code: Vec<u8> = vec![0x8B, 0x04, 0x25, 0x04, 0x20, 0x00, 0x00, 0x66, 0x89, 0x04, 0x25, 0x08, 0x20, 0x00, 0x00]; let x86_code: Vec<u8> = 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 read_cell = Rc::new(RefCell::new(MmioReadExpectation(0, 0)));
let cb_read_cell = read_cell.clone(); let cb_read_cell = read_cell.clone();
@ -436,7 +439,7 @@ fn x86_mmio() {
42 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 cb_write_cell = write_cell.clone();
let write_callback = move |_: &mut Unicorn<'_, ()>, offset, size, value| { let write_callback = move |_: &mut Unicorn<'_, ()>, offset, size, value| {
*cb_write_cell.borrow_mut() = MmioWriteExpectation(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.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!( assert_eq!(
emu.emu_start( emu.emu_start(
@ -494,9 +500,11 @@ fn x86_mmio() {
{ {
// MOV ax, 42; MOV [0x2008], ax; // MOV ax, 42; MOV [0x2008], ax;
let x86_code: Vec<u8> = vec![0x66, 0xB8, 0x2A, 0x00, 0x66, 0x89, 0x04, 0x25, 0x08, 0x20, 0x00, 0x00]; let x86_code: Vec<u8> = 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 cb_write_cell = write_cell.clone();
let write_callback = move |_: &mut Unicorn<'_, ()>, offset, size, value| { let write_callback = move |_: &mut Unicorn<'_, ()>, offset, size, value| {
*cb_write_cell.borrow_mut() = MmioWriteExpectation(offset, size, value); *cb_write_cell.borrow_mut() = MmioWriteExpectation(offset, size, value);