From 5ff654c77b17141b5fb8f377a18e26368737a62d Mon Sep 17 00:00:00 2001 From: Kevin Schneider Date: Thu, 30 Mar 2023 13:44:59 +0200 Subject: [PATCH 1/3] add rust bindings for uc_ctl --- bindings/rust/src/ffi.rs | 1 + bindings/rust/src/lib.rs | 146 +++++++++++++++++++++++++++++ bindings/rust/src/unicorn_const.rs | 43 +++++++++ 3 files changed, 190 insertions(+) diff --git a/bindings/rust/src/ffi.rs b/bindings/rust/src/ffi.rs index b1c04f57..37602392 100644 --- a/bindings/rust/src/ffi.rs +++ b/bindings/rust/src/ffi.rs @@ -86,6 +86,7 @@ extern "C" { pub fn uc_context_alloc(engine: uc_handle, context: *mut uc_context) -> uc_error; pub fn uc_context_save(engine: uc_handle, context: uc_context) -> uc_error; pub fn uc_context_restore(engine: uc_handle, context: uc_context) -> uc_error; + pub fn uc_ctl(engine: uc_handle, control: u32, ...) -> uc_error; } pub struct UcHook<'a, D: 'a, F: 'a> { diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs index cb8497c7..0f1c7049 100644 --- a/bindings/rust/src/lib.rs +++ b/bindings/rust/src/lib.rs @@ -30,7 +30,9 @@ #[macro_use] extern crate alloc; +extern crate std; +#[macro_use] pub mod unicorn_const; mod arm; @@ -1051,4 +1053,148 @@ impl<'a, D> Unicorn<'a, D> { }; self.reg_write(reg, value) } + + pub fn ctl_get_mode(&self) -> Result { + let mut result: i32 = Default::default(); + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_READ!(ControlType::UC_CTL_UC_MODE), &mut result) }; + if err == uc_error::OK { + Ok(Mode::from_bits_truncate(result)) + } else { + Err(err) + } + } + + pub fn ctl_get_page_size(&self) -> Result { + let mut result: u32 = Default::default(); + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_READ!(ControlType::UC_CTL_UC_PAGE_SIZE), &mut result) }; + if err == uc_error::OK { + Ok(result) + } else { + Err(err) + } + } + + pub fn ctl_set_page_size(&self, page_size: u32) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_UC_PAGE_SIZE), page_size) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } + + pub fn ctl_get_arch(&self) -> Result { + let mut result: i32 = Default::default(); + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_READ!(ControlType::UC_CTL_UC_ARCH), &mut result) }; + if err == uc_error::OK { + Arch::try_from(result as usize) + } else { + Err(err) + } + } + + pub fn ctl_get_timeout(&self) -> Result { + let mut result: u64 = Default::default(); + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_READ!(ControlType::UC_CTL_UC_TIMEOUT), &mut result) }; + if err == uc_error::OK { + Ok(result) + } else { + Err(err) + } + } + + pub fn ctl_exits_enable(&self) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_UC_USE_EXITS), 1) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } + + pub fn ctl_exits_disable(&self) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_UC_USE_EXITS), 0) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } + + pub fn ctl_get_exits_count(&self) -> Result { + let mut result: libc::size_t = Default::default(); + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_READ!(ControlType::UC_CTL_UC_EXITS_CNT), &mut result) }; + if err == uc_error::OK { + Ok(result) + } else { + Err(err) + } + } + + pub fn ctl_get_exits(&self) -> Result, uc_error> { + let exits_count: libc::size_t = self.ctl_get_exits_count()?; + let mut exits: Vec = Vec::with_capacity(exits_count); + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_READ!(ControlType::UC_CTL_UC_EXITS), exits.as_mut_ptr(), exits_count) }; + if err == uc_error::OK { + unsafe { exits.set_len(exits_count); } + Ok(exits) + } else { + Err(err) + } + } + + pub fn ctl_set_exits(&self, exits: &[u64]) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_UC_EXITS), exits.as_ptr(), exits.len() as libc::size_t) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } + + pub fn ctl_get_cpu_model(&self) -> Result { + let mut result: i32 = Default::default(); + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_READ!(ControlType::UC_CTL_CPU_MODEL), &mut result) }; + if err == uc_error::OK { + Ok(result) + } else { + Err(err) + } + } + + pub fn ctl_set_cpu_model(&self, cpu_model: i32) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_CPU_MODEL), cpu_model) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } + + pub fn ctl_remove_cache(&self, address: u64, end: u64) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_TB_REMOVE_CACHE), address, end) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } + + pub fn ctl_request_cache(&self, address: u64, tb: &mut TranslationBlock) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_READ_WRITE!(ControlType::UC_CTL_TB_REQUEST_CACHE), address, tb) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } + + pub fn ctl_flush_tlb(&self) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_TB_FLUSH)) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } } diff --git a/bindings/rust/src/unicorn_const.rs b/bindings/rust/src/unicorn_const.rs index 7e8185d2..b0b440c4 100644 --- a/bindings/rust/src/unicorn_const.rs +++ b/bindings/rust/src/unicorn_const.rs @@ -187,3 +187,46 @@ bitflags! { const RISCV64 = Self::MIPS64.bits; } } + +// Represent a TranslationBlock. +#[repr(C)] +pub struct TranslationBlock { + pub pc: u64, + pub icount: u16, + pub size: u16 +} + +macro_rules! UC_CTL_READ { + ($expr:expr) => { + $expr as u32 | ControlType::UC_CTL_IO_READ as u32 + }; +} + +macro_rules! UC_CTL_WRITE { + ($expr:expr) => { + $expr as u32 | ControlType::UC_CTL_IO_WRITE as u32 + }; +} + +macro_rules! UC_CTL_READ_WRITE { + ($expr:expr) => { + $expr as u32 | ControlType::UC_CTL_IO_WRITE as u32 | ControlType::UC_CTL_IO_READ as u32 + }; +} + +#[allow(clippy::upper_case_acronyms)] +pub enum ControlType { + UC_CTL_UC_MODE = 0, + UC_CTL_UC_PAGE_SIZE = 1, + UC_CTL_UC_ARCH = 2, + UC_CTL_UC_TIMEOUT = 3, + UC_CTL_UC_USE_EXITS = 4, + UC_CTL_UC_EXITS_CNT = 5, + UC_CTL_UC_EXITS = 6, + UC_CTL_CPU_MODEL = 7, + UC_CTL_TB_REQUEST_CACHE = 8, + UC_CTL_TB_REMOVE_CACHE = 9, + UC_CTL_TB_FLUSH = 10, + UC_CTL_IO_READ = 1<<31, + UC_CTL_IO_WRITE = 1<<30, +} From 0729dc031247a58c1a37380cb6f124b0d353e990 Mon Sep 17 00:00:00 2001 From: Philipp Takacs Date: Wed, 5 Apr 2023 12:35:05 +0200 Subject: [PATCH 2/3] rust update uc_ctl_flush_tlb and add uc_ctl_tlb_mode --- bindings/rust/src/lib.rs | 20 +++++++++++++++++++- bindings/rust/src/unicorn_const.rs | 9 +++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs index 0f1c7049..d72f82ad 100644 --- a/bindings/rust/src/lib.rs +++ b/bindings/rust/src/lib.rs @@ -1189,7 +1189,7 @@ impl<'a, D> Unicorn<'a, D> { } } - pub fn ctl_flush_tlb(&self) -> Result<(), uc_error> { + pub fn ctl_flush_tb(&self) -> Result<(), uc_error> { let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_TB_FLUSH)) }; if err == uc_error::OK { Ok(()) @@ -1197,4 +1197,22 @@ impl<'a, D> Unicorn<'a, D> { Err(err) } } + + pub fn ctl_flush_tlb(&self) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_TLB_FLUSH)) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } + + pub fn ctl_tlb_type(&self, t: TlbType) -> Result<(), uc_error> { + let err = unsafe { ffi::uc_ctl(self.get_handle(), UC_CTL_WRITE!(ControlType::UC_CTL_TLB_TYPE), t as i32) }; + if err == uc_error::OK { + Ok(()) + } else { + Err(err) + } + } } diff --git a/bindings/rust/src/unicorn_const.rs b/bindings/rust/src/unicorn_const.rs index b0b440c4..d5c60c1c 100644 --- a/bindings/rust/src/unicorn_const.rs +++ b/bindings/rust/src/unicorn_const.rs @@ -53,6 +53,13 @@ pub enum MemType { READ_AFTER = 25, } +#[repr(C)] +#[derive(PartialEq, Debug, Clone, Copy)] +pub enum TlbType { + CPU = 0, + VIRTUAL = 1, +} + bitflags! { #[repr(C)] pub struct HookType: i32 { @@ -227,6 +234,8 @@ pub enum ControlType { UC_CTL_TB_REQUEST_CACHE = 8, UC_CTL_TB_REMOVE_CACHE = 9, UC_CTL_TB_FLUSH = 10, + UC_CTL_TLB_FLUSH = 11, + UC_CTL_TLB_TYPE = 12, UC_CTL_IO_READ = 1<<31, UC_CTL_IO_WRITE = 1<<30, } From a9f0dabc64bd0c0d9289538765b16b572b2519f1 Mon Sep 17 00:00:00 2001 From: Philipp Takacs Date: Wed, 5 Apr 2023 11:26:02 +0200 Subject: [PATCH 3/3] 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(())); +}