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> {}
pub extern "C" fn mmio_read_callback_proxy<D, F> (
pub extern "C" fn mmio_read_callback_proxy<D, F>(
uc: uc_handle,
offset: u64,
size: usize,
user_data: *mut UcHook<D, F>,
) -> u64 where
) -> u64
where
F: FnMut(&mut crate::Unicorn<D>, u64, usize) -> u64,
{
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)
}
pub extern "C" fn mmio_write_callback_proxy<D, F> (
pub extern "C" fn mmio_write_callback_proxy<D, F>(
uc: uc_handle,
offset: u64,
size: usize,

View File

@ -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<R: 'a, W: 'a>(
&mut self,
address: u64,
address: u64,
size: libc::size_t,
read_callback: Option<R>,
write_callback: Option<W>,
@ -327,7 +330,7 @@ impl<'a, D> Unicorn<'a, D> {
R: 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 {
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<dyn ffi::IsUcHook> );
let wd = write_data.map( |c| c as Box<dyn ffi::IsUcHook> );
self.inner_mut().mmio_callbacks.push(MmioCallbackScope{
let rd = read_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 {
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<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.
@ -406,7 +414,12 @@ impl<'a, D> Unicorn<'a, D> {
where
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.
@ -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(())

View File

@ -427,7 +427,10 @@ fn x86_mmio() {
{
// 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 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<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 write_callback = move |_: &mut Unicorn<'_, ()>, offset, size, value| {
*cb_write_cell.borrow_mut() = MmioWriteExpectation(offset, size, value);