From 473405797d1ade8d01f79503641ff3031a01787c Mon Sep 17 00:00:00 2001 From: Nikolas Eller Date: Mon, 21 Sep 2020 04:36:58 +0200 Subject: [PATCH] Improved Rust bindings: (#1309) * Added basic block hooking * Changed confusing struct naming. Before: Protection::All -> R,W,X, Now: Permission::All -> R,W,X * Fixed issue with remove_hook(..). Implementation tried to remove hook from incorrect hashmap. * Made unused private vmmap(..) public. --- bindings/rust/README.md | 4 +- bindings/rust/src/ffi.rs | 12 ++++ bindings/rust/src/lib.rs | 89 ++++++++++++++++++++++++++---- bindings/rust/src/unicorn_const.rs | 4 +- bindings/rust/src/utils.rs | 8 +-- bindings/rust/tests/unicorn.rs | 60 +++++++++++++++----- 6 files changed, 142 insertions(+), 35 deletions(-) diff --git a/bindings/rust/README.md b/bindings/rust/README.md index 9e7b1bda..0da26496 100644 --- a/bindings/rust/README.md +++ b/bindings/rust/README.md @@ -6,14 +6,14 @@ An extended version for fuzzing with AFL++ support can be found in https://githu ```rust use unicorn::RegisterARM; -use unicorn::unicorn_const::{Arch, Mode, Protection, SECOND_SCALE}; +use unicorn::unicorn_const::{Arch, Mode, Permission, SECOND_SCALE}; fn main() { let arm_code32: Vec = vec![0x17, 0x00, 0x40, 0xe2]; // sub r0, #23 let mut unicorn = unicorn::Unicorn::new(Arch::ARM, Mode::LITTLE_ENDIAN, 0).expect("failed to initialize Unicorn instance"); let mut emu = unicorn.borrow(); - emu.mem_map(0x1000, 0x4000, Protection::ALL).expect("failed to map code page"); + emu.mem_map(0x1000, 0x4000, Permission::ALL).expect("failed to map code page"); emu.mem_write(0x1000, &arm_code32).expect("failed to write instructions"); emu.reg_write(RegisterARM::R0 as i32, 123).expect("failed write R0"); diff --git a/bindings/rust/src/ffi.rs b/bindings/rust/src/ffi.rs index cee5bd51..f848c1a7 100644 --- a/bindings/rust/src/ffi.rs +++ b/bindings/rust/src/ffi.rs @@ -80,6 +80,11 @@ pub struct CodeHook { pub callback: Box, u64, u32)> } +pub struct BlockHook { + pub unicorn: *mut crate::UnicornInner, + pub callback: Box, u64, u32)> +} + pub struct MemHook { pub unicorn: *mut crate::UnicornInner, pub callback: Box, MemType, u64, usize, i64)> @@ -112,6 +117,13 @@ pub extern "C" fn code_hook_proxy(uc: uc_handle, address: u64, size: u32, use callback(crate::UnicornHandle { inner: unsafe { Pin::new_unchecked(unicorn) } }, address, size); } +pub extern "C" fn block_hook_proxy(uc: uc_handle, address: u64, size: u32, user_data: *mut BlockHook) { + let unicorn = unsafe { &mut *(*user_data).unicorn }; + let callback = &mut unsafe { &mut *(*user_data).callback }; + assert_eq!(uc, unicorn.uc); + callback(crate::UnicornHandle { inner: unsafe { Pin::new_unchecked(unicorn) } }, address, size); +} + pub extern "C" fn mem_hook_proxy(uc: uc_handle, mem_type: MemType, address: u64, diff --git a/bindings/rust/src/lib.rs b/bindings/rust/src/lib.rs index 9543e157..5d7a82f3 100644 --- a/bindings/rust/src/lib.rs +++ b/bindings/rust/src/lib.rs @@ -8,14 +8,14 @@ //! ```rust //! //! use unicorn::RegisterARM; -//! use unicorn::unicorn_const::{Arch, Mode, Protection, SECOND_SCALE}; +//! use unicorn::unicorn_const::{Arch, Mode, Permission, SECOND_SCALE}; //! //! fn main() { //! let arm_code32: Vec = vec![0x17, 0x00, 0x40, 0xe2]; // sub r0, #23 //! //! let mut unicorn = unicorn::Unicorn::new(Arch::ARM, Mode::LITTLE_ENDIAN, 0).expect("failed to initialize Unicorn instance"); //! let mut emu = unicorn.borrow(); -//! emu.mem_map(0x1000, 0x4000, Protection::ALL).expect("failed to map code page"); +//! emu.mem_map(0x1000, 0x4000, Permission::ALL).expect("failed to map code page"); //! emu.mem_write(0x1000, &arm_code32).expect("failed to write instructions"); //! //! emu.reg_write(RegisterARM::R0 as i32, 123).expect("failed write R0"); @@ -93,6 +93,7 @@ pub struct UnicornInner { pub uc: uc_handle, pub arch: Arch, pub code_hooks: HashMap<*mut libc::c_void, Box>>, + pub block_hooks: HashMap<*mut libc::c_void, Box>>, pub mem_hooks: HashMap<*mut libc::c_void, Box>>, pub intr_hooks: HashMap<*mut libc::c_void, Box>>, pub insn_in_hooks: HashMap<*mut libc::c_void, Box>>, @@ -116,6 +117,7 @@ impl Unicorn { uc: handle, arch: arch, code_hooks: HashMap::new(), + block_hooks: HashMap::new(), mem_hooks: HashMap::new(), intr_hooks: HashMap::new(), insn_in_hooks: HashMap::new(), @@ -231,7 +233,7 @@ impl<'a, D> UnicornHandle<'a, D> { pub fn mem_map_ptr(&mut self, address: u64, size: usize, - perms: Protection, + perms: Permission, ptr: *mut c_void ) -> Result<(), uc_error> { let err = unsafe { ffi::uc_mem_map_ptr(self.inner.uc, address, size, perms.bits(), ptr) }; @@ -249,7 +251,7 @@ impl<'a, D> UnicornHandle<'a, D> { pub fn mem_map(&mut self, address: u64, size: libc::size_t, - perms: Protection + perms: Permission ) -> Result<(), uc_error> { let err = unsafe { ffi::uc_mem_map(self.inner.uc, address, size, perms.bits()) }; if err == uc_error::OK { @@ -282,7 +284,7 @@ impl<'a, D> UnicornHandle<'a, D> { pub fn mem_protect(&mut self, address: u64, size: libc::size_t, - perms: Protection + perms: Permission ) -> Result<(), uc_error> { let err = unsafe { ffi::uc_mem_protect(self.inner.uc, address, size, perms.bits()) }; if err == uc_error::OK { @@ -417,6 +419,38 @@ impl<'a, D> UnicornHandle<'a, D> { } } + /// Add a block hook. + pub fn add_block_hook( + &mut self, + callback: F, + ) -> Result + where F: FnMut(UnicornHandle, u64, u32) + { + let mut hook_ptr = std::ptr::null_mut(); + let mut user_data = Box::new(ffi::BlockHook { + unicorn: unsafe { self.inner.as_mut().get_unchecked_mut() } as _, + callback: Box::new(callback), + }); + + let err = unsafe { + ffi::uc_hook_add( + self.inner.uc, + &mut hook_ptr, + HookType::BLOCK, + ffi::block_hook_proxy:: as _, + user_data.as_mut() as *mut _ as _, + 1, + 0, + ) + }; + if err == uc_error::OK { + unsafe { self.inner.as_mut().get_unchecked_mut() }.block_hooks.insert(hook_ptr, user_data); + Ok(hook_ptr) + } else { + Err(err) + } + } + /// Add a memory hook. pub fn add_mem_hook( &mut self, @@ -596,14 +630,45 @@ impl<'a, D> UnicornHandle<'a, D> { pub fn remove_hook(&mut self, hook: ffi::uc_hook) -> Result<(), uc_error> { let handle = unsafe { self.inner.as_mut().get_unchecked_mut() }; let err: uc_error; - if handle.code_hooks.contains_key(&hook) || - handle.mem_hooks.contains_key(&hook) || - handle.intr_hooks.contains_key(&hook) || - handle.insn_in_hooks.contains_key(&hook) || - handle.insn_out_hooks.contains_key(&hook) || - handle.insn_sys_hooks.contains_key(&hook) { - err = unsafe { ffi::uc_hook_del(handle.uc, hook) }; + let mut in_one_hashmap = false; + + if handle.code_hooks.contains_key(&hook) { + in_one_hashmap = true; + handle.code_hooks.remove(&hook); + } + + if handle.mem_hooks.contains_key(&hook) { + in_one_hashmap = true; handle.mem_hooks.remove(&hook); + } + + if handle.block_hooks.contains_key(&hook) { + in_one_hashmap = true; + handle.block_hooks.remove(&hook); + } + + if handle.intr_hooks.contains_key(&hook) { + in_one_hashmap = true; + handle.intr_hooks.remove(&hook); + } + + if handle.insn_in_hooks.contains_key(&hook) { + in_one_hashmap = true; + handle.insn_in_hooks.remove(&hook); + } + + if handle.insn_out_hooks.contains_key(&hook) { + in_one_hashmap = true; + handle.insn_out_hooks.remove(&hook); + } + + if handle.insn_sys_hooks.contains_key(&hook) { + in_one_hashmap = true; + handle.insn_sys_hooks.remove(&hook); + } + + if in_one_hashmap { + err = unsafe { ffi::uc_hook_del(handle.uc, hook) }; } else { err = uc_error::HOOK; } diff --git a/bindings/rust/src/unicorn_const.rs b/bindings/rust/src/unicorn_const.rs index c1120c4e..5697e1c3 100644 --- a/bindings/rust/src/unicorn_const.rs +++ b/bindings/rust/src/unicorn_const.rs @@ -89,7 +89,7 @@ pub enum Query { bitflags! { #[repr(C)] -pub struct Protection : u32 { +pub struct Permission : u32 { const NONE = 0; const READ = 1; const WRITE = 2; @@ -103,7 +103,7 @@ pub struct Protection : u32 { pub struct MemRegion { pub begin: u64, pub end: u64, - pub perms: Protection, + pub perms: Permission, } #[repr(C)] diff --git a/bindings/rust/src/utils.rs b/bindings/rust/src/utils.rs index 001db0f9..f0208abe 100644 --- a/bindings/rust/src/utils.rs +++ b/bindings/rust/src/utils.rs @@ -8,7 +8,7 @@ use super::x86::RegisterX86; use super::sparc::RegisterSPARC; use super::mips::RegisterMIPS; use super::m68k::RegisterM68K; -use super::{Protection, Mode, Arch, HookType, MemType, uc_error}; +use super::{Permission, Mode, Arch, HookType, MemType, uc_error}; use std::ptr; use std::cell::RefCell; use std::collections::HashMap; @@ -119,7 +119,7 @@ pub fn init_emu_with_heap(arch: Arch, unsafe { // manually mmap space for heap to know location let arena_ptr = mmap(null_ptr, size as usize, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, 0, 0); - uc.mem_map_ptr(base_addr, size as usize, Protection::READ | Protection::WRITE, arena_ptr)?; + uc.mem_map_ptr(base_addr, size as usize, Permission::READ | Permission::WRITE, arena_ptr)?; let h = uc.add_mem_hook(HookType::MEM_VALID, base_addr, base_addr + size as u64, Box::new(heap_unalloc))?; let chunks = HashMap::new(); let heap: &mut Heap = &mut *uc.get_data().borrow_mut(); @@ -163,7 +163,7 @@ pub fn uc_alloc(uc: &mut super::UnicornHandle>, mut size: u64) -> if increase_by % 8 != 0 { increase_by = ((increase_by / 8) + 1) * 8; } - uc.mem_map(uc_base + len as u64, increase_by, Protection::READ | Protection::WRITE)?; + uc.mem_map(uc_base + len as u64, increase_by, Permission::READ | Permission::WRITE)?; uc.get_data().borrow_mut().len += increase_by; len = uc.get_data().borrow_mut().len; } @@ -284,7 +284,7 @@ fn heap_uaf (uc: super::UnicornHandle>, _mem_type: MemType, addr: } -fn vmmap(uc: &mut super::UnicornHandle) { +pub fn vmmap(uc: &mut super::UnicornHandle) { let regions = uc .mem_regions() .expect("failed to retrieve memory mappings"); diff --git a/bindings/rust/tests/unicorn.rs b/bindings/rust/tests/unicorn.rs index 94305863..2fbcb5ad 100644 --- a/bindings/rust/tests/unicorn.rs +++ b/bindings/rust/tests/unicorn.rs @@ -3,7 +3,7 @@ use std::cell::RefCell; use std::rc::Rc; use unicorn::{RegisterARM, RegisterX86, InsnSysX86, RegisterMIPS, RegisterPPC}; -use unicorn::unicorn_const::{Mode, Arch, Protection, MemType, HookType, SECOND_SCALE, uc_error}; +use unicorn::unicorn_const::{Mode, Arch, Permission, MemType, HookType, SECOND_SCALE, uc_error}; pub static X86_REGISTERS: [RegisterX86; 145] = [ RegisterX86::AH, @@ -171,7 +171,7 @@ fn emulate_x86() { ); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); @@ -215,7 +215,7 @@ fn x86_code_callback() { let mut unicorn = unicorn::Unicorn::new(Arch::X86, Mode::MODE_32, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); @@ -248,7 +248,7 @@ fn x86_intr_callback() { let mut unicorn = unicorn::Unicorn::new(Arch::X86, Mode::MODE_32, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); @@ -301,7 +301,7 @@ fn x86_mem_callback() { let mut unicorn = unicorn::Unicorn::new(Arch::X86, Mode::MODE_32, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); @@ -341,7 +341,7 @@ fn x86_insn_in_callback() { let mut unicorn = unicorn::Unicorn::new(Arch::X86, Mode::MODE_32, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); @@ -378,7 +378,7 @@ fn x86_insn_out_callback() { let mut unicorn = unicorn::Unicorn::new(Arch::X86, Mode::MODE_32, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); @@ -420,7 +420,7 @@ fn x86_insn_sys_callback() { let mut unicorn = unicorn::Unicorn::new(Arch::X86, Mode::MODE_64, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code), Ok(())); @@ -456,7 +456,7 @@ fn emulate_arm() { ); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &arm_code32), Ok(())); @@ -490,7 +490,7 @@ fn emulate_mips() { let mut unicorn = unicorn::Unicorn::new(Arch::MIPS, Mode::MODE_32, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &mips_code32), Ok(())); @@ -518,7 +518,7 @@ fn emulate_ppc() { let mut unicorn = unicorn::Unicorn::new(Arch::PPC, Mode::PPC32, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &ppc_code32), Ok(())); @@ -545,7 +545,7 @@ fn mem_unmapping() { let mut unicorn = unicorn::Unicorn::new(Arch::X86, Mode::MODE_32, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_unmap(0x1000, 0x4000), Ok(())); @@ -567,7 +567,7 @@ fn mem_map_ptr() { ); assert_eq!( - emu.mem_map_ptr(0x1000, 0x4000, Protection::ALL, mem.as_mut_ptr() as _), + emu.mem_map_ptr(0x1000, 0x4000, Permission::ALL, mem.as_mut_ptr() as _), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); @@ -603,7 +603,7 @@ fn mem_map_ptr() { ); assert_eq!( - emu.mem_map_ptr(0x1000, 0x4000, Protection::ALL, mem.as_mut_ptr() as _), + emu.mem_map_ptr(0x1000, 0x4000, Permission::ALL, mem.as_mut_ptr() as _), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); @@ -638,7 +638,7 @@ fn x86_context_save_and_restore() { let mut unicorn = unicorn::Unicorn::new(Arch::X86, mode, 0).expect("failed to initialize unicorn instance"); let mut emu = unicorn.borrow(); assert_eq!( - emu.mem_map(0x1000, 0x4000, Protection::ALL), + emu.mem_map(0x1000, 0x4000, Permission::ALL), Ok(()) ); assert_eq!(emu.mem_write(0x1000, &x86_code), Ok(())); @@ -663,3 +663,33 @@ fn x86_context_save_and_restore() { } } } + +#[test] +fn x86_block_callback() { + #[derive(PartialEq, Debug)] + struct BlockExpectation(u64, u32); + let expects = vec![BlockExpectation(0x1000, 2)]; + let blocks: Vec = Vec::new(); + let blocks_cell = Rc::new(RefCell::new(blocks)); + + let callback_blocks = blocks_cell.clone(); + let callback = move |_: Unicorn<'_>, address: u64, size: u32| { + let mut blocks = callback_blocks.borrow_mut(); + blocks.push(BlockExpectation(address, size)); + }; + + let x86_code32: Vec = vec![0x41, 0x4a]; // INC ecx; DEC edx + + let mut unicorn = unicorn::Unicorn::new(Arch::X86, Mode::MODE_32, 0).expect("failed to initialize unicorn instance"); + let mut emu = unicorn.borrow(); + assert_eq!( + emu.mem_map(0x1000, 0x4000, Permission::ALL), + Ok(()) + ); + assert_eq!(emu.mem_write(0x1000, &x86_code32), Ok(())); + + let hook = emu.add_block_hook(callback).expect("failed to add block hook"); + assert_eq!(emu.emu_start(0x1000, 0x1002, 10 * SECOND_SCALE, 1000), Ok(())); + assert_eq!(expects, *blocks_cell.borrow()); + assert_eq!(emu.remove_hook(hook), Ok(())); +}