From 20d97a0c00e69f20bc51f30c2048ee79290a4eff Mon Sep 17 00:00:00 2001 From: Sven Bartscher Date: Wed, 17 Nov 2021 18:35:30 +0100 Subject: [PATCH 1/2] Add map_mmio to rust bindings --- bindings/rust/src/ffi.rs | 36 +++++++++++ bindings/rust/src/lib.rs | 96 ++++++++++++++++++++++++++++ bindings/rust/tests/unicorn.rs | 110 +++++++++++++++++++++++++++++++++ 3 files changed, 242 insertions(+) diff --git a/bindings/rust/src/ffi.rs b/bindings/rust/src/ffi.rs index 300e1220..e2b7c078 100644 --- a/bindings/rust/src/ffi.rs +++ b/bindings/rust/src/ffi.rs @@ -41,6 +41,15 @@ extern "C" { perms: u32, ptr: *mut c_void, ) -> uc_error; + pub fn uc_mmio_map( + engine: uc_handle, + address: u64, + size: libc::size_t, + read_cb: *mut c_void, + user_data_read: *mut c_void, + write_cb: *mut c_void, + user_data_write: *mut c_void, + ) -> uc_error; pub fn uc_mem_unmap(engine: uc_handle, address: u64, size: libc::size_t) -> uc_error; pub fn uc_mem_protect( engine: uc_handle, @@ -87,6 +96,33 @@ pub trait IsUcHook<'a> {} impl<'a, D, F> IsUcHook<'a> for UcHook<'a, D, F> {} +pub extern "C" fn mmio_read_callback_proxy ( + uc: uc_handle, + offset: u64, + size: usize, + user_data: *mut UcHook, +) -> u64 where + F: FnMut(&mut crate::Unicorn, u64, usize) -> u64, +{ + let user_data = unsafe { &mut *user_data }; + debug_assert_eq!(uc, user_data.uc.inner().uc); + (user_data.callback)(&mut user_data.uc, offset, size) +} + +pub extern "C" fn mmio_write_callback_proxy ( + uc: uc_handle, + offset: u64, + size: usize, + value: u64, + user_data: *mut UcHook, +) where + F: FnMut(&mut crate::Unicorn, u64, usize, u64), +{ + let user_data = unsafe { &mut *user_data }; + debug_assert_eq!(uc, user_data.uc.inner().uc); + (user_data.callback)(&mut user_data.uc, offset, size, value); +} + pub extern "C" fn code_hook_proxy( uc: uc_handle, address: u64, diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs index b2a82252..068695a8 100644 --- a/bindings/rust/src/lib.rs +++ b/bindings/rust/src/lib.rs @@ -81,6 +81,8 @@ pub struct UnicornInner<'a, D> { pub arch: Arch, /// to keep ownership over the hook for this uc instance's lifetime pub hooks: Vec<(ffi::uc_hook, Box + 'a>)>, + /// To keep ownership over the mmio callbacks for this uc instance's lifetime + pub mmio_callbacks: Vec<(Option + 'a>>, Option + 'a>>)>, pub data: D, } @@ -123,6 +125,7 @@ where arch, data, hooks: vec![], + mmio_callbacks: vec![], })), }) } else { @@ -262,6 +265,99 @@ impl<'a, D> Unicorn<'a, D> { } } + /// Map in am MMIO region backed by callbacks. + /// + /// `address` must be aligned to 4kb or this will return `Error::ARG`. + /// `size` must be a multiple of 4kb or this will return `Error::ARG`. + pub fn mmio_map( + &mut self, + address: u64, + size: libc::size_t, + read_callback: Option, + write_callback: Option, + ) -> Result<(), uc_error> + where + R: FnMut(&mut Unicorn, u64, usize) -> u64, + W: FnMut(&mut Unicorn, u64, usize, u64), + { + let mut read_data = read_callback.map( |c| { + Box::new(ffi::UcHook { + callback: c, + uc: Unicorn { + inner: self.inner.clone(), + }, + }) + }); + let mut write_data = write_callback.map( |c| { + Box::new(ffi::UcHook { + callback: c, + uc: Unicorn { + inner: self.inner.clone(), + }, + }) + }); + + let err = unsafe { + ffi::uc_mmio_map( + self.inner().uc, + address, + size, + ffi::mmio_read_callback_proxy:: as _, + match read_data { + Some(ref mut d) => d.as_mut() as *mut _ as _, + None => ptr::null_mut(), + }, + ffi::mmio_write_callback_proxy:: as _, + match write_data { + Some(ref mut d) => d.as_mut() as *mut _ as _, + None => ptr::null_mut(), + }, + ) + }; + + 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((rd, wd)); + + Ok(()) + } else { + Err(err) + } + } + + /// Map in a read-only MMIO region backed by a callback. + /// + /// `address` must be aligned to 4kb or this will return `Error::ARG`. + /// `size` must be a multiple of 4kb or this will return `Error::ARG`. + pub fn mmio_map_ro( + &mut self, + address: u64, + size: libc::size_t, + callback: F, + ) -> Result<(), uc_error> + where + F: FnMut(&mut Unicorn, 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. + /// + /// `address` must be aligned to 4kb or this will return `Error::ARG`. + /// `size` must be a multiple of 4kb or this will return `Error::ARG`. + pub fn mmio_map_wo( + &mut self, + address: u64, + size: libc::size_t, + callback: F, + ) -> Result<(), uc_error> + where + F: FnMut(&mut Unicorn, u64, usize, u64), + { + self.mmio_map(address, size, None::, u64, usize) -> u64>, Some(callback)) + } + /// Unmap a memory region. /// /// `address` must be aligned to 4kb or this will return `Error::ARG`. diff --git a/bindings/rust/tests/unicorn.rs b/bindings/rust/tests/unicorn.rs index 98ca78e1..79f659aa 100644 --- a/bindings/rust/tests/unicorn.rs +++ b/bindings/rust/tests/unicorn.rs @@ -412,6 +412,116 @@ fn x86_insn_sys_callback() { assert_eq!(emu.remove_hook(hook), Ok(())); } +#[test] +fn x86_mmio() { + #[derive(PartialEq, Debug)] + struct MmioReadExpectation(u64, usize); + #[derive(PartialEq, Debug)] + struct MmioWriteExpectation(u64, usize, u64); + let read_expect = MmioReadExpectation(4, 4); + let write_expect = MmioWriteExpectation(8, 2, 42); + + let mut emu = unicorn_engine::Unicorn::new(Arch::X86, Mode::MODE_64) + .expect("failed to initialize unicorn instance"); + assert_eq!(emu.mem_map(0x1000, 0x1000, Permission::ALL), Ok(())); + + { + // 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 read_cell = Rc::new(RefCell::new(MmioReadExpectation(0, 0))); + let cb_read_cell = read_cell.clone(); + let read_callback = move |_: &mut Unicorn<'_, ()>, offset, size| { + *cb_read_cell.borrow_mut() = MmioReadExpectation(offset, size); + 42 + }; + + 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); + }; + + 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.emu_start( + 0x1000, + 0x1000 + x86_code.len() as u64, + 10 * SECOND_SCALE, + 1000 + ), + Ok(()) + ); + + assert_eq!(read_expect, *read_cell.borrow()); + assert_eq!(write_expect, *write_cell.borrow()); + + assert_eq!(emu.mem_unmap(0x2000, 0x1000), Ok(())); + } + + { + // MOV eax, [0x2004]; + let x86_code: Vec = vec![0x8B, 0x04, 0x25, 0x04, 0x20, 0x00, 0x00]; + + let read_cell = Rc::new(RefCell::new(MmioReadExpectation(0, 0))); + let cb_read_cell = read_cell.clone(); + let read_callback = move |_: &mut Unicorn<'_, ()>, offset, size| { + *cb_read_cell.borrow_mut() = MmioReadExpectation(offset, size); + 42 + }; + + assert_eq!(emu.mem_write(0x1000, &x86_code), Ok(())); + + assert_eq!(emu.mmio_map_ro(0x2000, 0x1000, read_callback), Ok(())); + + assert_eq!( + emu.emu_start( + 0x1000, + 0x1000 + x86_code.len() as u64, + 10 * SECOND_SCALE, + 1000 + ), + Ok(()) + ); + + assert_eq!(read_expect, *read_cell.borrow()); + + assert_eq!(emu.mem_unmap(0x2000, 0x1000), Ok(())); + } + + { + // MOV ax, 42; MOV [0x2008], ax; + 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 cb_write_cell = write_cell.clone(); + let write_callback = move |_: &mut Unicorn<'_, ()>, offset, size, value| { + *cb_write_cell.borrow_mut() = MmioWriteExpectation(offset, size, value); + }; + + assert_eq!(emu.mem_write(0x1000, &x86_code), Ok(())); + + assert_eq!(emu.mmio_map_wo(0x2000, 0x1000, write_callback), Ok(())); + + assert_eq!( + emu.emu_start( + 0x1000, + 0x1000 + x86_code.len() as u64, + 10 * SECOND_SCALE, + 1000 + ), + Ok(()) + ); + + assert_eq!(write_expect, *write_cell.borrow()); + + assert_eq!(emu.mem_unmap(0x2000, 0x1000), Ok(())); + } +} + #[test] fn emulate_arm() { let arm_code32: Vec = vec![0x83, 0xb0]; // sub sp, #0xc From a237505adb5e7da4c889544eaccea4a8b4a6a720 Mon Sep 17 00:00:00 2001 From: Sven Bartscher Date: Wed, 17 Nov 2021 23:11:23 +0100 Subject: [PATCH 2/2] rust: Implement deallocation of MMIO callbacks Previously the user data of MMIO callbacks would live until the end of the containing Unicorn engine. Now they are deallocated once all memory referencing those callbacks has been unmapped. --- bindings/rust/src/lib.rs | 65 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 63 insertions(+), 2 deletions(-) diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs index 068695a8..9ec402e0 100644 --- a/bindings/rust/src/lib.rs +++ b/bindings/rust/src/lib.rs @@ -76,13 +76,60 @@ impl Drop for Context { } } +pub struct MmioCallbackScope<'a> { + pub regions: Vec<(u64, usize)>, + pub read_callback: Option + 'a>>, + pub write_callback: Option + 'a>>, +} + +impl<'a> MmioCallbackScope<'a> { + fn has_regions(&self) -> bool { + self.regions.len() > 0 + } + + 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 { + // 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)] + } + } + } else { + if end > *b { + if end >= e { + // The unmapped region completely contains this region + vec![] + } else { + // The unmapped region overlaps with the start of this region + vec![(end, (e - end) as usize)] + } + } else { + // The unmapped region is completely before this region + vec![(*b, *s)] + } + } + }).collect(); + } +} + pub struct UnicornInner<'a, D> { pub uc: uc_handle, pub arch: Arch, /// to keep ownership over the hook for this uc instance's lifetime pub hooks: Vec<(ffi::uc_hook, Box + 'a>)>, /// To keep ownership over the mmio callbacks for this uc instance's lifetime - pub mmio_callbacks: Vec<(Option + 'a>>, Option + 'a>>)>, + pub mmio_callbacks: Vec>, pub data: D, } @@ -318,7 +365,11 @@ 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((rd, wd)); + self.inner_mut().mmio_callbacks.push(MmioCallbackScope{ + regions: vec![(address, size)], + read_callback: rd, + write_callback: wd, + }); Ok(()) } else { @@ -364,6 +415,9 @@ impl<'a, D> Unicorn<'a, D> { /// `size` must be a multiple of 4kb or this will return `Error::ARG`. pub fn mem_unmap(&mut self, address: u64, size: libc::size_t) -> Result<(), uc_error> { let err = unsafe { ffi::uc_mem_unmap(self.inner().uc, address, size) }; + + self.mmio_unmap(address, size); + if err == uc_error::OK { Ok(()) } else { @@ -371,6 +425,13 @@ impl<'a, D> Unicorn<'a, D> { } } + fn mmio_unmap(&mut self, address: u64, size: libc::size_t) { + for scope in self.inner_mut().mmio_callbacks.iter_mut() { + scope.unmap(address, size); + } + self.inner_mut().mmio_callbacks.retain( |scope| scope.has_regions() ); + } + /// Set the memory permissions for an existing memory region. /// /// `address` must be aligned to 4kb or this will return `Error::ARG`.