rust add tlb callback

bindings/rust/src/ffi.rs

@@ -3,7 +3,7 @@
 
 use crate::{Unicorn, UnicornInner};
 
-use super::unicorn_const::{uc_error, Arch, HookType, MemRegion, MemType, Mode, Query};
+use super::unicorn_const::{uc_error, Arch, HookType, MemRegion, MemType, Mode, Query, TlbEntry};
 use alloc::rc::Weak;
 use core::{cell::UnsafeCell, ffi::c_void};
 use libc::{c_char, c_int};
@@ -252,3 +252,25 @@ where
     debug_assert_eq!(uc, user_data_uc.get_handle());
     (user_data.callback)(&mut user_data_uc);
 }
+
+pub extern "C" fn tlb_lookup_hook_proxy<D, F>(uc: uc_handle, vaddr: u64, mem_type: MemType, result: *mut TlbEntry, user_data: *mut UcHook<D, F>) -> bool
+where
+    F: FnMut(&mut crate::Unicorn<D>, u64, MemType) -> Option<TlbEntry>,
+{
+    let user_data = unsafe { &mut *user_data };
+    let mut user_data_uc = Unicorn {
+        inner: user_data.uc.upgrade().unwrap(),
+    };
+    debug_assert_eq!(uc, user_data_uc.get_handle());
+    let r = (user_data.callback)(&mut user_data_uc, vaddr, mem_type);
+    match r {
+        Some(ref e) => {
+            unsafe {
+                let ref_result: &mut TlbEntry = &mut *result;
+                *ref_result = *e;
+            }
+        },
+        None => {},
+    };
+    return r.is_some();
+}

bindings/rust/src/lib.rs

@@ -885,6 +885,33 @@ impl<'a, D> Unicorn<'a, D> {
         }
     }
 
+    pub fn add_tlb_hook<F>(&mut self, begin: u64, end: u64, callback: F) -> Result<ffi::uc_hook, uc_error>
+    where
+        F: FnMut(&mut crate::Unicorn<D>, u64, MemType) -> Option<TlbEntry> + 'a,
+    {
+        let mut hook_ptr = core::ptr::null_mut();
+        let mut user_data = Box::new(ffi::UcHook {
+            callback,
+            uc: Rc::downgrade(&self.inner),
+        });
+        let err = unsafe {
+            ffi::uc_hook_add(self.get_handle(),
+                &mut hook_ptr,
+                HookType::TLB,
+                ffi::tlb_lookup_hook_proxy::<D, F> as _,
+                user_data.as_mut() as *mut _ as _,
+                begin,
+                end,
+            )
+        };
+        if err == uc_error::OK {
+            self.inner_mut().hooks.push((hook_ptr, user_data));
+            Ok(hook_ptr)
+        } else {
+            Err(err)
+        }
+    }
+
     /// Remove a hook.
     ///
     /// `hook` is the value returned by `add_*_hook` functions.

bindings/rust/src/unicorn_const.rs

@@ -93,6 +93,8 @@ bitflags! {
         const MEM_INVALID = Self::MEM_READ_INVALID.bits | Self::MEM_WRITE_INVALID.bits | Self::MEM_FETCH_INVALID.bits;
 
         const MEM_ALL = Self::MEM_VALID.bits | Self::MEM_INVALID.bits;
+
+        const TLB = (1 << 17);
     }
 }
 
@@ -239,3 +241,10 @@ pub enum ControlType {
     UC_CTL_IO_READ = 1<<31,
     UC_CTL_IO_WRITE = 1<<30,
 }
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct TlbEntry {
+    pub paddr: u64,
+    pub perms: Permission,
+}

tests/rust-tests/main.rs

@@ -3,7 +3,7 @@ extern crate alloc;
 use alloc::rc::Rc;
 use core::cell::RefCell;
 use unicorn_engine::unicorn_const::{
-    uc_error, Arch, HookType, MemType, Mode, Permission, SECOND_SCALE,
+    uc_error, Arch, HookType, MemType, Mode, Permission, SECOND_SCALE, TlbEntry, TlbType
 };
 use unicorn_engine::{InsnSysX86, RegisterARM, RegisterMIPS, RegisterPPC, RegisterX86, Unicorn};
 
@@ -772,3 +772,46 @@ fn x86_block_callback() {
     assert_eq!(expects, *blocks_cell.borrow());
     assert_eq!(emu.remove_hook(hook), Ok(()));
 }
+
+#[test]
+fn x86_tlb_callback() {
+    #[derive(PartialEq, Debug)]
+    struct BlockExpectation(u64, u32);
+    let expects:u64 = 4;
+    let count: u64 = 0;
+    let count_cell = Rc::new(RefCell::new(count));
+
+    let callback_counter = count_cell.clone();
+    let tlb_callback = move |_: &mut Unicorn<'_, ()>, address: u64, _: MemType| -> Option<TlbEntry> {
+        let mut blocks = callback_counter.borrow_mut();
+        *blocks += 1;
+        return Some(TlbEntry{paddr: address, perms: Permission::ALL});
+    };
+
+    let syscall_callback = move |uc:  &mut Unicorn<'_, ()>| {
+        assert_eq!(uc.ctl_flush_tlb(), Ok(()));
+    };
+
+    let code: Vec<u8> = vec![0xa3,0x00,0x00,0x20,0x00,0x00,0x00,0x00,0x00,0x0f,0x05,0xa3,0x00,0x00,0x20,0x00,0x00,0x00,0x00,0x00]; // movabs  dword ptr [0x200000], eax; syscall; movabs  dword ptr [0x200000], eax
+
+    let mut emu = unicorn_engine::Unicorn::new(Arch::X86, Mode::MODE_64)
+        .expect("failed to initialize unicorn instance");
+    assert_eq!(emu.ctl_tlb_type(TlbType::VIRTUAL), Ok(()));
+    assert_eq!(emu.mem_map(0x1000, 0x1000, Permission::ALL), Ok(()));
+    assert_eq!(emu.mem_map(0x200000, 0x1000, Permission::ALL), Ok(()));
+    assert_eq!(emu.mem_write(0x1000, &code), Ok(()));
+
+    let tlb_hook = emu
+        .add_tlb_hook(0, !0u64, tlb_callback)
+        .expect("failed to add tlb hook");
+    let syscall_hook = emu
+        .add_insn_sys_hook(InsnSysX86::SYSCALL, 0, !0u64, syscall_callback)
+        .expect("failed to add syscall hook");
+    assert_eq!(
+        emu.emu_start(0x1000, (0x1000 + code.len()) as u64, 0, 0),
+        Ok(())
+    );
+    assert_eq!(expects, *count_cell.borrow());
+    assert_eq!(emu.remove_hook(tlb_hook), Ok(()));
+    assert_eq!(emu.remove_hook(syscall_hook), Ok(()));
+}