Add map_mmio to rust bindings

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<D, F> (
+    uc: uc_handle,
+    offset: u64,
+    size: usize,
+    user_data: *mut UcHook<D, F>,
+) -> u64 where
+    F: FnMut(&mut crate::Unicorn<D>, 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<D, F> (
+    uc: uc_handle,
+    offset: u64,
+    size: usize,
+    value: u64,
+    user_data: *mut UcHook<D, F>,
+) where
+    F: FnMut(&mut crate::Unicorn<D>, 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<D, F>(
     uc: uc_handle,
     address: u64,

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<dyn ffi::IsUcHook<'a> + 'a>)>,
+    /// To keep ownership over the mmio callbacks for this uc instance's lifetime
+    pub mmio_callbacks: Vec<(Option<Box<dyn ffi::IsUcHook<'a> + 'a>>, Option<Box<dyn ffi::IsUcHook<'a> + '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<R: 'a, W: 'a>(
+        &mut self,
+        address: u64, 
+        size: libc::size_t,
+        read_callback: Option<R>,
+        write_callback: Option<W>,
+    ) -> Result<(), uc_error>
+    where
+        R: FnMut(&mut Unicorn<D>, u64, usize) -> u64,
+        W: FnMut(&mut Unicorn<D>, 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::<D, R> as _,
+                match read_data {
+                    Some(ref mut d) => d.as_mut() as *mut _ as _,
+                    None => ptr::null_mut(),
+                },
+                ffi::mmio_write_callback_proxy::<D, W> 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<dyn ffi::IsUcHook> );
+            let wd = write_data.map( |c| c as Box<dyn ffi::IsUcHook> );
+            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<F: 'a>(
+        &mut self,
+        address: u64,
+        size: libc::size_t,
+        callback: F,
+    ) -> Result<(), uc_error>
+    where
+        F: FnMut(&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.
+    ///
+    /// `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<F: 'a>(
+        &mut self,
+        address: u64,
+        size: libc::size_t,
+        callback: F,
+    ) -> Result<(), uc_error>
+    where
+        F: FnMut(&mut Unicorn<D>, u64, usize, u64),
+    {
+        self.mmio_map(address, size, None::<fn(&mut Unicorn<D>, u64, usize) -> u64>, Some(callback))
+    }
+
     /// Unmap a memory region.
     ///
     /// `address` must be aligned to 4kb or this will return `Error::ARG`.

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<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();
+        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<u8> = 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<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 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<u8> = vec![0x83, 0xb0]; // sub    sp, #0xc