From 3978eeed2a8d804ef577c700419cc70d979f8406 Mon Sep 17 00:00:00 2001 From: momo5502 Date: Sat, 5 Apr 2025 15:28:14 +0200 Subject: [PATCH] Finish execution hook support --- src/icicle-emulator/icicle_x64_emulator.cpp | 19 +++-- src/icicle/src/icicle.rs | 90 ++++++++++++++++++--- src/icicle/src/lib.rs | 12 ++- 3 files changed, 101 insertions(+), 20 deletions(-) diff --git a/src/icicle-emulator/icicle_x64_emulator.cpp b/src/icicle-emulator/icicle_x64_emulator.cpp index 6ab52138..77101bc5 100644 --- a/src/icicle-emulator/icicle_x64_emulator.cpp +++ b/src/icicle-emulator/icicle_x64_emulator.cpp @@ -27,7 +27,8 @@ extern "C" int32_t icicle_save_registers(icicle_emulator*, data_accessor_func* accessor, void* accessor_data); int32_t icicle_restore_registers(icicle_emulator*, const void* data, size_t length); uint32_t icicle_add_syscall_hook(icicle_emulator*, raw_func* callback, void* data); - uint32_t icicle_add_execution_hook(icicle_emulator*, ptr_func* callback, void* data); + uint32_t icicle_add_execution_hook(icicle_emulator*, uint64_t address, ptr_func* callback, void* data); + uint32_t icicle_add_generic_execution_hook(icicle_emulator*, ptr_func* callback, void* data); uint32_t icicle_add_violation_hook(icicle_emulator*, violation_func* callback, void* data); uint32_t icicle_add_read_hook(icicle_emulator*, uint64_t start, uint64_t end, memory_access_func* cb, void* data); uint32_t icicle_add_write_hook(icicle_emulator*, uint64_t start, uint64_t end, memory_access_func* cb, void* data); @@ -288,9 +289,17 @@ namespace icicle emulator_hook* hook_memory_execution(const uint64_t address, memory_execution_hook_callback callback) override { - // TODO - (void)address; - throw std::runtime_error("Not implemented"); + auto object = make_function_object(std::move(callback)); + auto* ptr = object.get(); + auto* wrapper = +[](void* user, const uint64_t addr) { + auto& func = *static_cast(user); + (func)(addr); + }; + + const auto id = icicle_add_execution_hook(this->emu_, address, wrapper, ptr); + this->hooks_[id] = std::move(object); + + return wrap_hook(id); } emulator_hook* hook_memory_execution(memory_execution_hook_callback callback) override @@ -302,7 +311,7 @@ namespace icicle (func)(addr); }; - const auto id = icicle_add_execution_hook(this->emu_, wrapper, ptr); + const auto id = icicle_add_generic_execution_hook(this->emu_, wrapper, ptr); this->hooks_[id] = std::move(object); return wrap_hook(id); diff --git a/src/icicle/src/icicle.rs b/src/icicle/src/icicle.rs index 3582ca63..05daf8ae 100644 --- a/src/icicle/src/icicle.rs +++ b/src/icicle/src/icicle.rs @@ -46,7 +46,8 @@ enum HookType { Syscall = 1, Read, Write, - Execute, + ExecuteGeneric, + ExecuteSpecific, Violation, Unknown, } @@ -128,12 +129,72 @@ impl icicle_vm::CodeInjector for InstructionHookInjector { } } +struct ExecutionHooks { + generic_hooks: HookContainer, + specific_hooks: HookContainer, + address_mapping: HashMap>, +} + +impl ExecutionHooks { + pub fn new() -> Self { + Self { + generic_hooks: HookContainer::new(), + specific_hooks: HookContainer::new(), + address_mapping: HashMap::new(), + } + } + + pub fn execute(&self, address: u64) { + for (_key, func) in self.generic_hooks.get_hooks() { + func(address); + } + + let mapping = self.address_mapping.get(&address); + if mapping.is_none(){ + return; + } + + for id in mapping.unwrap() { + let func = self.specific_hooks.get_hooks().get(&id); + if func.is_some() { + func.unwrap()(address); + } + } + } + + pub fn add_generic_hook(&mut self, callback: Box) -> u32 { + self.generic_hooks.add_hook(callback) + } + + pub fn add_specific_hook(&mut self, address: u64, callback: Box) -> u32 { + let id = self.specific_hooks.add_hook(callback); + + let mapping = self.address_mapping.entry(address).or_insert_with(Vec::new); + mapping.push(id); + + return id; + } + + pub fn remove_generic_hook(&mut self, id: u32) { + self.generic_hooks.remove_hook(id); + } + + pub fn remove_specific_hook(&mut self, id: u32) { + self.address_mapping.retain(|_, vec| { + vec.retain(|&x| x != id); + !vec.is_empty() + }); + + self.specific_hooks.remove_hook(id); + } +} + pub struct IcicleEmulator { vm: icicle_vm::Vm, reg: registers::X64RegisterNodes, syscall_hooks: HookContainer, violation_hooks: HookContainer bool>, - execution_hooks: Rc>>, + execution_hooks: Rc>, } struct MemoryHook { @@ -184,15 +245,12 @@ impl icicle_cpu::mem::IoMemory for MmioHandler { impl IcicleEmulator { pub fn new() -> Self { let mut virtual_machine = create_x64_vm(); - let exec_hooks: Rc>> = - Rc::new(RefCell::new(HookContainer::new())); + let exec_hooks = Rc::new(RefCell::new(ExecutionHooks::new())); let exec_hooks_clone = Rc::clone(&exec_hooks); let hook = icicle_cpu::InstHook::new(move |_: &mut icicle_cpu::Cpu, addr: u64| { - for (_key, func) in exec_hooks_clone.borrow().get_hooks() { - func(addr); - } + exec_hooks_clone.borrow().execute(addr); }); let hook = virtual_machine.cpu.add_hook(hook); @@ -279,10 +337,15 @@ impl IcicleEmulator { let hook_id = self.violation_hooks.add_hook(callback); return qualify_hook_id(hook_id, HookType::Violation); } + + pub fn add_execution_hook(&mut self, address:u64, callback: Box) -> u32 { + let hook_id = self.execution_hooks.borrow_mut().add_specific_hook(address, callback); + return qualify_hook_id(hook_id, HookType::ExecuteSpecific); + } - pub fn add_execution_hook(&mut self, callback: Box) -> u32 { - let hook_id = self.execution_hooks.borrow_mut().add_hook(callback); - return qualify_hook_id(hook_id, HookType::Execute); + pub fn add_generic_execution_hook(&mut self, callback: Box) -> u32 { + let hook_id = self.execution_hooks.borrow_mut().add_generic_hook(callback); + return qualify_hook_id(hook_id, HookType::ExecuteGeneric); } pub fn add_syscall_hook(&mut self, callback: Box) -> u32 { @@ -301,7 +364,7 @@ impl IcicleEmulator { return 0; } - return qualify_hook_id(id.expect("Hook id needed"), HookType::Read); + return qualify_hook_id(id.unwrap(), HookType::Read); } pub fn add_write_hook( @@ -315,7 +378,7 @@ impl IcicleEmulator { return 0; } - return qualify_hook_id(id.expect("Hook id needed"), HookType::Write); + return qualify_hook_id(id.unwrap(), HookType::Write); } pub fn remove_hook(&mut self, id: u32) { @@ -324,7 +387,8 @@ impl IcicleEmulator { match hook_type { HookType::Syscall => self.syscall_hooks.remove_hook(hook_id), HookType::Violation => self.violation_hooks.remove_hook(hook_id), - HookType::Execute => self.execution_hooks.borrow_mut().remove_hook(hook_id), + HookType::ExecuteGeneric => self.execution_hooks.borrow_mut().remove_generic_hook(hook_id), + HookType::ExecuteSpecific => self.execution_hooks.borrow_mut().remove_specific_hook(hook_id), HookType::Read => {self.get_mem().remove_read_after_hook(hook_id);()}, HookType::Write => {self.get_mem().remove_write_hook(hook_id);()}, _ => {} diff --git a/src/icicle/src/lib.rs b/src/icicle/src/lib.rs index 655f83ca..a494b806 100644 --- a/src/icicle/src/lib.rs +++ b/src/icicle/src/lib.rs @@ -215,10 +215,18 @@ pub fn icicle_add_syscall_hook(ptr: *mut c_void, callback: RawFunction, data: *m } #[unsafe(no_mangle)] -pub fn icicle_add_execution_hook(ptr: *mut c_void, callback: PtrFunction, data: *mut c_void) -> u32 { +pub fn icicle_add_generic_execution_hook(ptr: *mut c_void, callback: PtrFunction, data: *mut c_void) -> u32 { unsafe { let emulator = &mut *(ptr as *mut IcicleEmulator); - return emulator.add_execution_hook(Box::new(move |ptr: u64| callback(data, ptr))); + return emulator.add_generic_execution_hook(Box::new(move |ptr: u64| callback(data, ptr))); + } +} + +#[unsafe(no_mangle)] +pub fn icicle_add_execution_hook(ptr: *mut c_void, address: u64, callback: PtrFunction, data: *mut c_void) -> u32 { + unsafe { + let emulator = &mut *(ptr as *mut IcicleEmulator); + return emulator.add_execution_hook(address, Box::new(move |ptr: u64| callback(data, ptr))); } }