From b3bdfc9d6b2af91de002971b0ef6a9660c9705f2 Mon Sep 17 00:00:00 2001 From: momo5502 Date: Sat, 9 Aug 2025 20:44:09 +0200 Subject: [PATCH] Support hooking/unhooking within hooks for icicle --- .../icicle-bridge/src/icicle.rs | 125 ++++++++-- .../icicle-emulator/icicle-bridge/src/lib.rs | 8 + .../icicle_x86_64_emulator.cpp | 232 ++++++++++++++---- 3 files changed, 297 insertions(+), 68 deletions(-) diff --git a/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs b/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs index 660ade77..3eec79dd 100644 --- a/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs +++ b/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs @@ -1,5 +1,6 @@ use icicle_cpu::ExceptionCode; use icicle_cpu::ValueSource; +use std::collections::HashSet; use std::{cell::RefCell, collections::HashMap, rc::Rc}; use crate::registers; @@ -73,31 +74,96 @@ fn qualify_hook_id(hook_id: u32, hook_type: HookType) -> u32 { pub struct HookContainer { hook_id: u32, + is_iterating: bool, hooks: HashMap>, + hooks_to_add: HashMap>, + hooks_to_remove: HashSet, } impl HookContainer { pub fn new() -> Self { Self { hook_id: 0, + is_iterating: false, hooks: HashMap::new(), + hooks_to_add: HashMap::new(), + hooks_to_remove: HashSet::new(), } } pub fn add_hook(&mut self, callback: Box) -> u32 { self.hook_id += 1; let id = self.hook_id; - self.hooks.insert(id, callback); + + if self.is_iterating { + self.hooks_to_add.insert(id, callback); + } else { + self.hooks.insert(id, callback); + } return id; } - pub fn get_hooks(&self) -> &HashMap> { - return &self.hooks; + pub fn for_each_hook(&mut self, mut callback: F) + where + F: FnMut(&Func), + { + let was_iterating = self.do_pre_access_work(); + + for (_, func) in &self.hooks { + callback(func.as_ref()); + } + + self.do_post_access_work(was_iterating); + } + + pub fn access_hook(&mut self, id: u32, mut callback: F) + where + F: FnMut(&Func), + { + let was_iterating = self.do_pre_access_work(); + + let hook = self.hooks.get(&id); + if hook.is_some() { + callback(hook.unwrap().as_ref()); + } + + self.do_post_access_work(was_iterating); + } + + pub fn is_empty(&self) -> bool { + return self.hooks.is_empty(); } pub fn remove_hook(&mut self, id: u32) { - self.hooks.remove(&id); + if self.is_iterating { + self.hooks_to_remove.insert(id); + } else { + self.hooks.remove(&id); + } + } + + fn do_pre_access_work(&mut self) -> bool { + let was_iterating = self.is_iterating; + self.is_iterating = true; + return was_iterating; + } + + fn do_post_access_work(&mut self, was_iterating: bool) { + self.is_iterating = was_iterating; + if self.is_iterating { + return; + } + + let to_remove = std::mem::take(&mut self.hooks_to_remove); + for id in &to_remove { + self.hooks.remove(&id); + } + + let to_add = std::mem::take(&mut self.hooks_to_add); + for (id, func) in to_add { + self.hooks.insert(id, func); + } } } @@ -157,6 +223,7 @@ struct ExecutionHooks { specific_hooks: HookContainer, block_hooks: HookContainer, address_mapping: HashMap>, + one_time_callbacks: Vec>, } impl ExecutionHooks { @@ -167,31 +234,38 @@ impl ExecutionHooks { specific_hooks: HookContainer::new(), block_hooks: HookContainer::new(), address_mapping: HashMap::new(), + one_time_callbacks: Vec::new(), } } - fn run_hooks(&self, address: u64) { - for (_key, func) in self.generic_hooks.get_hooks() { - func(address); + fn run_hooks(&mut self, address: u64) { + if !self.one_time_callbacks.is_empty() { + let callbacks = std::mem::take(&mut self.one_time_callbacks); + for cb in callbacks { + cb.as_ref()(); + } } + self.generic_hooks.for_each_hook(|func| { + 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); - } + self.specific_hooks.access_hook(*id, |func| { + func(address); + }); } } pub fn on_block(&mut self, address: u64, instructions: u64) { - for (_key, func) in self.block_hooks.get_hooks() { + self.block_hooks.for_each_hook(|func| { func(address, instructions); - } + }); } pub fn execute(&mut self, cpu: &mut icicle_cpu::Cpu, address: u64) { @@ -224,6 +298,10 @@ impl ExecutionHooks { return id; } + pub fn schedule(&mut self, callback: Box) { + self.one_time_callbacks.push(callback); + } + pub fn remove_generic_hook(&mut self, id: u32) { self.generic_hooks.remove_hook(id); } @@ -368,10 +446,10 @@ impl IcicleEmulator { } } - fn handle_interrupt(&self, code: i32) -> bool { - for (_key, func) in self.interrupt_hooks.get_hooks() { + fn handle_interrupt(&mut self, code: i32) -> bool { + self.interrupt_hooks.for_each_hook(|func| { func(code); - } + }); return true; } @@ -393,16 +471,15 @@ impl IcicleEmulator { } fn handle_violation(&mut self, address: u64, permission: u8, unmapped: bool) -> bool { - let hooks = &self.violation_hooks.get_hooks(); - if hooks.is_empty() { + if self.violation_hooks.is_empty() { return false; } let mut continue_execution = true; - for (_key, func) in self.violation_hooks.get_hooks() { + self.violation_hooks.for_each_hook(|func| { continue_execution &= func(address, permission, unmapped); - } + }); return continue_execution; } @@ -412,9 +489,9 @@ impl IcicleEmulator { return self.handle_interrupt(value as i32); } - for (_key, func) in self.syscall_hooks.get_hooks() { + self.syscall_hooks.for_each_hook(|func| { func(); - } + }); self.vm.cpu.write_pc(self.vm.cpu.read_pc() + 2); return true; @@ -521,6 +598,10 @@ impl IcicleEmulator { } } + pub fn run_on_next_instruction(&mut self, callback: Box) { + self.execution_hooks.borrow_mut().schedule(callback); + } + pub fn map_memory(&mut self, address: u64, length: u64, permissions: u8) -> bool { const MAPPING_PERMISSIONS: u8 = icicle_vm::cpu::mem::perm::MAP | icicle_vm::cpu::mem::perm::INIT diff --git a/src/backends/icicle-emulator/icicle-bridge/src/lib.rs b/src/backends/icicle-emulator/icicle-bridge/src/lib.rs index 74e07b54..c8efe281 100644 --- a/src/backends/icicle-emulator/icicle-bridge/src/lib.rs +++ b/src/backends/icicle-emulator/icicle-bridge/src/lib.rs @@ -307,6 +307,14 @@ pub fn icicle_remove_hook(ptr: *mut c_void, id: u32) { } } +#[unsafe(no_mangle)] +pub fn icicle_run_on_next_instruction(ptr: *mut c_void, callback: RawFunction, data: *mut c_void) { + unsafe { + let emulator = &mut *(ptr as *mut IcicleEmulator); + emulator.run_on_next_instruction(Box::new(move || callback(data))); + } +} + #[unsafe(no_mangle)] pub fn icicle_read_register( ptr: *mut c_void, diff --git a/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp b/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp index b01ed360..20ed80bb 100644 --- a/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp +++ b/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp @@ -1,7 +1,9 @@ #define ICICLE_EMULATOR_IMPL #include "icicle_x86_64_emulator.hpp" +#include #include +#include using icicle_emulator = struct icicle_emulator_; @@ -44,6 +46,7 @@ extern "C" void icicle_start(icicle_emulator*, size_t count); void icicle_stop(icicle_emulator*); void icicle_destroy_emulator(icicle_emulator*); + void icicle_run_on_next_instruction(icicle_emulator*, raw_func* callback, void* data); } namespace icicle @@ -58,24 +61,35 @@ namespace icicle } } - emulator_hook* wrap_hook(const uint32_t id) - { - return reinterpret_cast(static_cast(id)); - } - template struct function_object : utils::object { + bool* hook_state{}; std::function func{}; - function_object(std::function f = {}) - : func(std::move(f)) + function_object(std::function f = {}, bool* state = nullptr) + : hook_state(state), + func(std::move(f)) { } template auto operator()(Args&&... args) const { + bool old_state{}; + if (this->hook_state) + { + old_state = *this->hook_state; + *this->hook_state = true; + } + + const auto _ = utils::finally([&] { + if (this->hook_state) + { + *this->hook_state = old_state; + } + }); + return this->func.operator()(std::forward(args)...); } @@ -83,10 +97,18 @@ namespace icicle }; template - std::unique_ptr> make_function_object(std::function func) + std::unique_ptr> make_function_object(std::function func, bool& hook_state) { - return std::make_unique>(std::move(func)); + return std::make_unique>(std::move(func), &hook_state); } + + struct memory_access_hook + { + uint64_t address{}; + uint64_t size{}; + memory_access_hook_callback callback{}; + bool is_read{}; + }; } class icicle_x86_64_emulator : public x86_64_emulator @@ -105,6 +127,7 @@ namespace icicle { reset_object_with_delayed_destruction(this->hooks_); reset_object_with_delayed_destruction(this->storage_); + utils::reset_object_with_delayed_destruction(this->hooks_to_install_); if (this->emu_) { @@ -239,7 +262,7 @@ namespace icicle return nullptr; } - auto obj = make_function_object(std::move(callback)); + auto obj = make_function_object(std::move(callback), this->is_in_hook_); auto* ptr = obj.get(); const auto invoker = +[](void* cb) { @@ -255,7 +278,7 @@ namespace icicle emulator_hook* hook_basic_block(basic_block_hook_callback callback) override { - auto object = make_function_object(std::move(callback)); + auto object = make_function_object(std::move(callback), this->is_in_hook_); auto* ptr = object.get(); auto* wrapper = +[](void* user, const uint64_t addr, const uint64_t instructions) { basic_block block{}; @@ -274,7 +297,7 @@ namespace icicle emulator_hook* hook_interrupt(interrupt_hook_callback callback) override { - auto obj = make_function_object(std::move(callback)); + auto obj = make_function_object(std::move(callback), this->is_in_hook_); auto* ptr = obj.get(); auto* wrapper = +[](void* user, const int32_t code) { const auto& func = *static_cast(user); @@ -289,7 +312,7 @@ namespace icicle emulator_hook* hook_memory_violation(memory_violation_hook_callback callback) override { - auto obj = make_function_object(std::move(callback)); + auto obj = make_function_object(std::move(callback), this->is_in_hook_); auto* ptr = obj.get(); auto* wrapper = +[](void* user, const uint64_t address, const uint8_t operation, const int32_t unmapped) -> int32_t { @@ -310,7 +333,7 @@ namespace icicle emulator_hook* hook_memory_execution(const uint64_t address, memory_execution_hook_callback callback) override { - auto object = make_function_object(std::move(callback)); + auto object = make_function_object(std::move(callback), this->is_in_hook_); auto* ptr = object.get(); auto* wrapper = +[](void* user, const uint64_t addr) { const auto& func = *static_cast(user); @@ -325,7 +348,7 @@ namespace icicle emulator_hook* hook_memory_execution(memory_execution_hook_callback callback) override { - auto object = make_function_object(std::move(callback)); + auto object = make_function_object(std::move(callback), this->is_in_hook_); auto* ptr = object.get(); auto* wrapper = +[](void* user, const uint64_t addr) { const auto& func = *static_cast(user); @@ -341,48 +364,36 @@ namespace icicle emulator_hook* hook_memory_read(const uint64_t address, const uint64_t size, memory_access_hook_callback callback) override { - auto obj = make_function_object(std::move(callback)); - auto* ptr = obj.get(); - auto* wrapper = +[](void* user, const uint64_t address, const void* data, size_t length) { - const auto& func = *static_cast(user); - func(address, data, length); - }; - - const auto id = icicle_add_read_hook(this->emu_, address, address + size, wrapper, ptr); - this->hooks_[id] = std::move(obj); - - return wrap_hook(id); + return this->try_install_memory_access_hook(memory_access_hook{ + .address = address, + .size = size, + .callback = std::move(callback), + .is_read = true, + }); } emulator_hook* hook_memory_write(const uint64_t address, const uint64_t size, memory_access_hook_callback callback) override { - auto obj = make_function_object(std::move(callback)); - auto* ptr = obj.get(); - auto* wrapper = +[](void* user, const uint64_t address, const void* data, size_t length) { - const auto& func = *static_cast(user); - func(address, data, length); - }; - - const auto id = icicle_add_write_hook(this->emu_, address, address + size, wrapper, ptr); - this->hooks_[id] = std::move(obj); - - return wrap_hook(id); + return this->try_install_memory_access_hook(memory_access_hook{ + .address = address, + .size = size, + .callback = std::move(callback), + .is_read = false, + }); } void delete_hook(emulator_hook* hook) override { - const auto id = static_cast(reinterpret_cast(hook)); - const auto entry = this->hooks_.find(id); - if (entry == this->hooks_.end()) + if (this->is_in_hook_) { - return; + this->hooks_to_delete_.insert(hook); + this->schedule_action_execution(); + } + else + { + this->delete_hook_internal(hook); } - - icicle_remove_hook(this->emu_, id); - const auto obj = std::move(entry->second); - this->hooks_.erase(entry); - (void)obj; } void serialize_state(utils::buffer_serializer& buffer, const bool is_snapshot) const override @@ -442,9 +453,138 @@ namespace icicle } private: + bool is_in_hook_{false}; std::list> storage_{}; std::unordered_map> hooks_{}; + std::unordered_map> id_mapping_{}; icicle_emulator* emu_{}; + uint32_t index_{0}; + + std::unordered_set hooks_to_delete_{}; + std::unordered_map hooks_to_install_{}; + + emulator_hook* wrap_hook(const std::optional icicle_id) + { + const auto id = ++this->index_; + auto* hook = reinterpret_cast(static_cast(id)); + + this->id_mapping_[hook] = icicle_id; + + return hook; + } + + emulator_hook* hook_memory_access(memory_access_hook hook, emulator_hook* hook_id) + { + auto obj = make_function_object(std::move(hook.callback), this->is_in_hook_); + auto* ptr = obj.get(); + auto* wrapper = +[](void* user, const uint64_t address, const void* data, size_t length) { + const auto& func = *static_cast(user); + func(address, data, length); + }; + + auto* installer = hook.is_read ? &icicle_add_read_hook : &icicle_add_write_hook; + const auto id = installer(this->emu_, hook.address, hook.address + hook.size, wrapper, ptr); + this->hooks_[id] = std::move(obj); + + if (hook_id) + { + this->id_mapping_[hook_id] = id; + return hook_id; + } + + return wrap_hook(id); + } + + void delete_hook_internal(emulator_hook* hook) + { + auto hook_id = this->id_mapping_.find(hook); + if (hook_id == this->id_mapping_.end()) + { + return; + } + + if (!hook_id->second.has_value()) + { + this->hooks_to_delete_.insert(hook); + return; + } + + const auto id = *hook_id->second; + this->id_mapping_.erase(hook_id); + + const auto entry = this->hooks_.find(id); + if (entry == this->hooks_.end()) + { + return; + } + + icicle_remove_hook(this->emu_, id); + const auto obj = std::move(entry->second); + this->hooks_.erase(entry); + (void)obj; + } + + void perform_pending_actions() + { + const auto hooks_to_delete = std::move(this->hooks_to_delete_); + const auto hooks_to_install = std::move(this->hooks_to_install_); + + this->hooks_to_delete_ = {}; + this->hooks_to_install_ = {}; + + for (auto& hook : hooks_to_install) + { + this->hook_memory_access(std::move(hook.second), hook.first); + } + + for (auto* hook : hooks_to_delete) + { + this->delete_hook_internal(hook); + } + } + + emulator_hook* try_install_memory_access_hook(memory_access_hook hook) + { + if (!this->is_in_hook_) + { + return this->hook_memory_access(std::move(hook), nullptr); + } + + auto* hook_id = wrap_hook(std::nullopt); + this->hooks_to_install_[hook_id] = std::move(hook); + + this->schedule_action_execution(); + + return hook_id; + } + + void schedule_action_execution() + { + this->run_on_next_instruction([this] { + this->perform_pending_actions(); // + }); + } + + void run_on_next_instruction(std::function func) const + { + auto* heap_func = new std::function(std::move(func)); + auto* callback = +[](void* data) { + auto* cb = static_cast*>(data); + + try + { + (*cb)(); + } + catch (...) + { + // Ignore + } + + delete cb; + }; + + icicle_run_on_next_instruction(this->emu_, callback, heap_func); + } }; std::unique_ptr create_x86_64_emulator()