Support hooking/unhooking within hooks for icicle

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<Func: ?Sized> {
     hook_id: u32,
+    is_iterating: bool,
     hooks: HashMap<u32, Box<Func>>,
+    hooks_to_add: HashMap<u32, Box<Func>>,
+    hooks_to_remove: HashSet<u32>,
 }
 
 impl<Func: ?Sized> HookContainer<Func> {
     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<Func>) -> 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<u32, Box<Func>> {
-        return &self.hooks;
+    pub fn for_each_hook<F>(&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<F>(&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<dyn Fn(u64)>,
     block_hooks: HookContainer<dyn Fn(u64, u64)>,
     address_mapping: HashMap<u64, Vec<u32>>,
+    one_time_callbacks: Vec<Box<dyn Fn()>>,
 }
 
 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<dyn Fn()>) {
+        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<dyn Fn()>) {
+        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

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,

src/backends/icicle-emulator/icicle_x86_64_emulator.cpp

@@ -1,7 +1,9 @@
 #define ICICLE_EMULATOR_IMPL
 #include "icicle_x86_64_emulator.hpp"
 
+#include <unordered_set>
 #include <utils/object.hpp>
+#include <utils/finally.hpp>
 
 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<emulator_hook*>(static_cast<size_t>(id));
-        }
-
         template <typename T>
         struct function_object : utils::object
         {
+            bool* hook_state{};
             std::function<T> func{};
 
-            function_object(std::function<T> f = {})
-                : func(std::move(f))
+            function_object(std::function<T> f = {}, bool* state = nullptr)
+                : hook_state(state),
+                  func(std::move(f))
             {
             }
 
             template <typename... Args>
             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>(args)...);
             }
 
@@ -83,10 +97,18 @@ namespace icicle
         };
 
         template <typename T>
-        std::unique_ptr<function_object<T>> make_function_object(std::function<T> func)
+        std::unique_ptr<function_object<T>> make_function_object(std::function<T> func, bool& hook_state)
         {
-            return std::make_unique<function_object<T>>(std::move(func));
+            return std::make_unique<function_object<T>>(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<decltype(ptr)>(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<decltype(ptr)>(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<decltype(ptr)>(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<decltype(ptr)>(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<decltype(ptr)>(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<uint32_t>(reinterpret_cast<size_t>(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<std::unique_ptr<utils::object>> storage_{};
         std::unordered_map<uint32_t, std::unique_ptr<utils::object>> hooks_{};
+        std::unordered_map<emulator_hook*, std::optional<uint32_t>> id_mapping_{};
         icicle_emulator* emu_{};
+        uint32_t index_{0};
+
+        std::unordered_set<emulator_hook*> hooks_to_delete_{};
+        std::unordered_map<emulator_hook*, memory_access_hook> hooks_to_install_{};
+
+        emulator_hook* wrap_hook(const std::optional<uint32_t> icicle_id)
+        {
+            const auto id = ++this->index_;
+            auto* hook = reinterpret_cast<emulator_hook*>(static_cast<size_t>(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<decltype(ptr)>(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<void()> func) const
+        {
+            auto* heap_func = new std::function(std::move(func));
+            auto* callback = +[](void* data) {
+                auto* cb = static_cast<std::function<void()>*>(data);
+
+                try
+                {
+                    (*cb)();
+                }
+                catch (...)
+                {
+                    // Ignore
+                }
+
+                delete cb;
+            };
+
+            icicle_run_on_next_instruction(this->emu_, callback, heap_func);
+        }
     };
 
     std::unique_ptr<x86_64_emulator> create_x86_64_emulator()