From ce431115b5a93cee2374995c53806cb60f0898a5 Mon Sep 17 00:00:00 2001 From: momo5502 Date: Sat, 17 May 2025 16:21:30 +0200 Subject: [PATCH] Add block hooking to icicle --- .../icicle-bridge/src/icicle.rs | 68 +++++++++++++++-- .../icicle-emulator/icicle-bridge/src/lib.rs | 76 ++++++++++++++----- .../icicle_x86_64_emulator.cpp | 20 ++++- 3 files changed, 137 insertions(+), 27 deletions(-) diff --git a/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs b/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs index 19963833..71ac020d 100644 --- a/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs +++ b/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs @@ -50,6 +50,7 @@ enum HookType { ExecuteSpecific, Violation, Interrupt, + Block, Unknown, } @@ -101,7 +102,18 @@ impl HookContainer { } struct InstructionHookInjector { - hook: pcode::HookId, + inst_hook: pcode::HookId, + block_hook: pcode::HookId, +} + +fn count_instructions(block: &icicle_cpu::lifter::Block) -> u64 { + return block.pcode.instructions.iter().fold(0u64, |count, &stmt| { + if let pcode::Op::InstructionMarker = stmt.op { + return count + 1; + } + + return count; + }); } impl icicle_vm::CodeInjector for InstructionHookInjector { @@ -117,10 +129,19 @@ impl icicle_vm::CodeInjector for InstructionHookInjector { let mut tmp_block = pcode::Block::new(); tmp_block.next_tmp = block.pcode.next_tmp; + let mut is_first_inst = true; + let inst_count = count_instructions(&block); + for stmt in block.pcode.instructions.drain(..) { tmp_block.push(stmt); if let pcode::Op::InstructionMarker = stmt.op { - tmp_block.push(pcode::Op::Hook(self.hook)); + if is_first_inst { + is_first_inst = false; + tmp_block.push((pcode::Op::Arg(0), pcode::Inputs::one(inst_count))); + tmp_block.push(pcode::Op::Hook(self.block_hook)); + } + + tmp_block.push(pcode::Op::Hook(self.inst_hook)); code.modified.insert(id); } } @@ -134,6 +155,7 @@ struct ExecutionHooks { stop: Rc>, generic_hooks: HookContainer, specific_hooks: HookContainer, + block_hooks: HookContainer, address_mapping: HashMap>, } @@ -143,6 +165,7 @@ impl ExecutionHooks { stop: stop_value, generic_hooks: HookContainer::new(), specific_hooks: HookContainer::new(), + block_hooks: HookContainer::new(), address_mapping: HashMap::new(), } } @@ -165,6 +188,12 @@ impl ExecutionHooks { } } + pub fn on_block(&mut self, address: u64, instructions: u64) { + for (_key, func) in self.block_hooks.get_hooks() { + func(address, instructions); + } + } + pub fn execute(&mut self, cpu: &mut icicle_cpu::Cpu, address: u64) { self.run_hooks(address); @@ -174,6 +203,14 @@ impl ExecutionHooks { } } + pub fn add_block_hook(&mut self, callback: Box) -> u32 { + self.block_hooks.add_hook(callback) + } + + pub fn remove_block_hook(&mut self, id: u32) { + self.block_hooks.remove_hook(id); + } + pub fn add_generic_hook(&mut self, callback: Box) -> u32 { self.generic_hooks.add_hook(callback) } @@ -263,14 +300,25 @@ impl IcicleEmulator { let stop_value = Rc::new(RefCell::new(false)); let exec_hooks = Rc::new(RefCell::new(ExecutionHooks::new(stop_value.clone()))); - let exec_hooks_clone = Rc::clone(&exec_hooks); + let inst_exec_hooks = Rc::clone(&exec_hooks); - let hook = icicle_cpu::InstHook::new(move |cpu: &mut icicle_cpu::Cpu, addr: u64| { - exec_hooks_clone.borrow_mut().execute(cpu, addr); + let inst_hook = icicle_cpu::InstHook::new(move |cpu: &mut icicle_cpu::Cpu, addr: u64| { + inst_exec_hooks.borrow_mut().execute(cpu, addr); }); - let hook = virtual_machine.cpu.add_hook(hook); - virtual_machine.add_injector(InstructionHookInjector { hook }); + let block_exec_hooks = Rc::clone(&exec_hooks); + + let block_hook = icicle_cpu::InstHook::new(move |cpu: &mut icicle_cpu::Cpu, addr: u64| { + let instructions = cpu.args[0] as u64; + block_exec_hooks.borrow_mut().on_block(addr, instructions); + }); + + let inst_hook_id = virtual_machine.cpu.add_hook(inst_hook); + let block_hook_id = virtual_machine.cpu.add_hook(block_hook); + virtual_machine.add_injector(InstructionHookInjector { + inst_hook: inst_hook_id, + block_hook: block_hook_id, + }); Self { stop: stop_value, @@ -378,6 +426,11 @@ impl IcicleEmulator { } } + pub fn add_block_hook(&mut self, callback: Box) -> u32 { + let hook_id = self.execution_hooks.borrow_mut().add_block_hook(callback); + return qualify_hook_id(hook_id, HookType::Block); + } + pub fn add_violation_hook(&mut self, callback: Box bool>) -> u32 { let hook_id = self.violation_hooks.add_hook(callback); return qualify_hook_id(hook_id, HookType::Violation); @@ -453,6 +506,7 @@ impl IcicleEmulator { .execution_hooks .borrow_mut() .remove_specific_hook(hook_id), + HookType::Block => self.execution_hooks.borrow_mut().remove_block_hook(hook_id), HookType::Read => { self.get_mem().remove_read_after_hook(hook_id); () diff --git a/src/backends/icicle-emulator/icicle-bridge/src/lib.rs b/src/backends/icicle-emulator/icicle-bridge/src/lib.rs index fe5ac7a5..bd60b9aa 100644 --- a/src/backends/icicle-emulator/icicle-bridge/src/lib.rs +++ b/src/backends/icicle-emulator/icicle-bridge/src/lib.rs @@ -37,6 +37,7 @@ pub fn icicle_stop(ptr: *mut c_void) { type RawFunction = extern "C" fn(*mut c_void); type PtrFunction = extern "C" fn(*mut c_void, u64); +type BlockFunction = extern "C" fn(*mut c_void, u64, u64); type DataFunction = extern "C" fn(*mut c_void, *const c_void, usize); type MmioReadFunction = extern "C" fn(*mut c_void, u64, *mut c_void, usize); type MmioWriteFunction = extern "C" fn(*mut c_void, u64, *const c_void, usize); @@ -167,17 +168,23 @@ pub fn icicle_read_memory(ptr: *mut c_void, address: u64, data: *mut c_void, siz } #[unsafe(no_mangle)] -pub fn icicle_add_interrupt_hook(ptr: *mut c_void, callback: InterruptFunction, data: *mut c_void) -> u32 { +pub fn icicle_add_interrupt_hook( + ptr: *mut c_void, + callback: InterruptFunction, + data: *mut c_void, +) -> u32 { unsafe { let emulator = &mut *(ptr as *mut IcicleEmulator); - return emulator.add_interrupt_hook(Box::new( - move |code: i32| callback(data, code), - )); + return emulator.add_interrupt_hook(Box::new(move |code: i32| callback(data, code))); } } #[unsafe(no_mangle)] -pub fn icicle_add_violation_hook(ptr: *mut c_void, callback: ViolationFunction, data: *mut c_void) -> u32 { +pub fn icicle_add_violation_hook( + ptr: *mut c_void, + callback: ViolationFunction, + data: *mut c_void, +) -> u32 { unsafe { let emulator = &mut *(ptr as *mut IcicleEmulator); return emulator.add_violation_hook(Box::new( @@ -194,26 +201,42 @@ pub fn icicle_add_violation_hook(ptr: *mut c_void, callback: ViolationFunction, } #[unsafe(no_mangle)] -pub fn icicle_add_read_hook(ptr: *mut c_void, start: u64, end: u64, callback: MemoryAccessFunction, user: *mut c_void) -> u32 { +pub fn icicle_add_read_hook( + ptr: *mut c_void, + start: u64, + end: u64, + callback: MemoryAccessFunction, + user: *mut c_void, +) -> u32 { unsafe { let emulator = &mut *(ptr as *mut IcicleEmulator); - return emulator.add_read_hook(start, end, Box::new( - move |address: u64, data: &[u8]| { + return emulator.add_read_hook( + start, + end, + Box::new(move |address: u64, data: &[u8]| { callback(user, address, data.as_ptr() as *const c_void, data.len()); - }, - )); + }), + ); } } #[unsafe(no_mangle)] -pub fn icicle_add_write_hook(ptr: *mut c_void, start: u64, end: u64, callback: MemoryAccessFunction, user: *mut c_void) -> u32 { +pub fn icicle_add_write_hook( + ptr: *mut c_void, + start: u64, + end: u64, + callback: MemoryAccessFunction, + user: *mut c_void, +) -> u32 { unsafe { let emulator = &mut *(ptr as *mut IcicleEmulator); - return emulator.add_write_hook(start, end, Box::new( - move |address: u64, data: &[u8]| { + return emulator.add_write_hook( + start, + end, + Box::new(move |address: u64, data: &[u8]| { callback(user, address, data.as_ptr() as *const c_void, data.len()); - }, - )); + }), + ); } } @@ -226,7 +249,21 @@ pub fn icicle_add_syscall_hook(ptr: *mut c_void, callback: RawFunction, data: *m } #[unsafe(no_mangle)] -pub fn icicle_add_generic_execution_hook(ptr: *mut c_void, callback: PtrFunction, data: *mut c_void) -> u32 { +pub fn icicle_add_block_hook(ptr: *mut c_void, callback: BlockFunction, data: *mut c_void) -> u32 { + unsafe { + let emulator = &mut *(ptr as *mut IcicleEmulator); + return emulator.add_block_hook(Box::new(move |address: u64, instructions: u64| { + callback(data, address, instructions) + })); + } +} + +#[unsafe(no_mangle)] +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_generic_execution_hook(Box::new(move |ptr: u64| callback(data, ptr))); @@ -234,7 +271,12 @@ pub fn icicle_add_generic_execution_hook(ptr: *mut c_void, callback: PtrFunction } #[unsafe(no_mangle)] -pub fn icicle_add_execution_hook(ptr: *mut c_void, address: u64, callback: PtrFunction, data: *mut c_void) -> u32 { +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))); diff --git a/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp b/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp index 6f961db7..146b32c0 100644 --- a/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp +++ b/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp @@ -12,6 +12,7 @@ extern "C" using raw_func = void(void*); using ptr_func = void(void*, uint64_t); + using block_func = void(void*, uint64_t, uint64_t); using interrupt_func = void(void*, int32_t); using violation_func = int32_t(void*, uint64_t address, uint8_t operation, int32_t unmapped); using data_accessor_func = void(void* user, const void* data, size_t length); @@ -29,6 +30,7 @@ extern "C" void 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_interrupt_hook(icicle_emulator*, interrupt_func* callback, void* data); + uint32_t icicle_add_block_hook(icicle_emulator*, block_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); @@ -248,9 +250,21 @@ namespace icicle emulator_hook* hook_basic_block(basic_block_hook_callback callback) override { - // TODO - (void)callback; - 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, const uint64_t instructions) { + basic_block block{}; + block.address = addr; + block.instruction_count = static_cast(instructions); + + const auto& func = *static_cast(user); + (func)(block); + }; + + const auto id = icicle_add_block_hook(this->emu_, wrapper, ptr); + this->hooks_[id] = std::move(object); + + return wrap_hook(id); } emulator_hook* hook_interrupt(interrupt_hook_callback callback) override