From a9f0dabc64bd0c0d9289538765b16b572b2519f1 Mon Sep 17 00:00:00 2001 From: Philipp Takacs Date: Wed, 5 Apr 2023 11:26:02 +0200 Subject: [PATCH] rust add tlb callback --- bindings/rust/src/ffi.rs | 24 +++++++++++++++- bindings/rust/src/lib.rs | 27 ++++++++++++++++++ bindings/rust/src/unicorn_const.rs | 9 ++++++ tests/rust-tests/main.rs | 45 +++++++++++++++++++++++++++++- 4 files changed, 103 insertions(+), 2 deletions(-) diff --git a/bindings/rust/src/ffi.rs b/bindings/rust/src/ffi.rs index 37602392..fb87c211 100644 --- a/bindings/rust/src/ffi.rs +++ b/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(uc: uc_handle, vaddr: u64, mem_type: MemType, result: *mut TlbEntry, user_data: *mut UcHook) -> bool +where + F: FnMut(&mut crate::Unicorn, u64, MemType) -> Option, +{ + 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(); +} diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs index d72f82ad..5b4e7d33 100644 --- a/bindings/rust/src/lib.rs +++ b/bindings/rust/src/lib.rs @@ -885,6 +885,33 @@ impl<'a, D> Unicorn<'a, D> { } } + pub fn add_tlb_hook(&mut self, begin: u64, end: u64, callback: F) -> Result + where + F: FnMut(&mut crate::Unicorn, u64, MemType) -> Option + '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:: 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. diff --git a/bindings/rust/src/unicorn_const.rs b/bindings/rust/src/unicorn_const.rs index d5c60c1c..2ef7622e 100644 --- a/bindings/rust/src/unicorn_const.rs +++ b/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, +} diff --git a/tests/rust-tests/main.rs b/tests/rust-tests/main.rs index b64b6aa6..1267b86f 100644 --- a/tests/rust-tests/main.rs +++ b/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 { + 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 = 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(())); +}