Support read and write hooks

src/icicle-emulator/icicle_x64_emulator.cpp

@@ -7,13 +7,14 @@ using icicle_emulator = struct icicle_emulator_;
 
 extern "C"
 {
+    using icicle_mmio_read_func = void(void* user, uint64_t address, void* data, size_t length);
+    using icicle_mmio_write_func = void(void* user, uint64_t address, const void* data, size_t length);
+
     using raw_func = void(void*);
     using ptr_func = void(void*, uint64_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);
-
-    using icicle_mmio_read_func = void(void* user, uint64_t address, void* data, size_t length);
-    using icicle_mmio_write_func = void(void* user, uint64_t address, const void* data, size_t length);
+    using memory_access_func = icicle_mmio_write_func;
 
     icicle_emulator* icicle_create_emulator();
     int32_t icicle_protect_memory(icicle_emulator*, uint64_t address, uint64_t length, uint8_t permissions);
@@ -28,7 +29,9 @@ extern "C"
     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_violation_hook(icicle_emulator*, violation_func* callback, void* data);
-    void icicle_remove_syscall_hook(icicle_emulator*, uint32_t id);
+    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);
+    void icicle_remove_hook(icicle_emulator*, uint32_t id);
     size_t icicle_read_register(icicle_emulator*, int reg, void* data, size_t length);
     size_t icicle_write_register(icicle_emulator*, int reg, const void* data, size_t length);
     void icicle_start(icicle_emulator*, size_t count);
@@ -77,21 +80,6 @@ namespace icicle
         {
             return std::make_unique<function_object<T>>(std::move(func));
         }
-
-        template <typename T>
-        std::unique_ptr<utils::object> wrap_shared(std::shared_ptr<T> shared_ptr)
-        {
-            struct shard_wrapper : utils::object
-            {
-                std::shared_ptr<T> ptr{};
-                ~shard_wrapper() override = default;
-            };
-
-            auto wrapper = std::make_unique<shard_wrapper>();
-            wrapper->ptr = std::move(shared_ptr);
-
-            return wrapper;
-        }
     }
 
     class icicle_x64_emulator : public x64_emulator
@@ -307,15 +295,15 @@ namespace icicle
 
         emulator_hook* hook_memory_execution(memory_execution_hook_callback callback) override
         {
-            auto wrapper = make_function_object(std::move(callback));
-            auto* ptr = wrapper.get();
-            auto* func = +[](void* user, const uint64_t addr) {
-                auto& func = *static_cast<memory_execution_hook_callback*>(user);
+            auto object = make_function_object(std::move(callback));
+            auto* ptr = object.get();
+            auto* wrapper = +[](void* user, const uint64_t addr) {
+                auto& func = *static_cast<decltype(ptr)>(user);
                 (func)(addr);
             };
 
-            const auto id = icicle_add_execution_hook(this->emu_, func, ptr);
-            this->hooks_[id] = std::move(wrapper);
+            const auto id = icicle_add_execution_hook(this->emu_, wrapper, ptr);
+            this->hooks_[id] = std::move(object);
 
             return wrap_hook(id);
         }
@@ -323,21 +311,33 @@ namespace icicle
         emulator_hook* hook_memory_read(const uint64_t address, const size_t size,
                                         memory_access_hook_callback callback) override
         {
-            // TODO
-            (void)address;
-            (void)size;
-            (void)callback;
-            throw std::runtime_error("Not implemented");
+            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);
         }
 
         emulator_hook* hook_memory_write(const uint64_t address, const size_t size,
-                                         const memory_access_hook_callback callback) override
+                                         memory_access_hook_callback callback) override
         {
-            // TODO
-            (void)address;
-            (void)size;
-            (void)callback;
-            throw std::runtime_error("Not implemented");
+            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);
         }
 
         void delete_hook(emulator_hook* hook) override
@@ -349,7 +349,7 @@ namespace icicle
                 return;
             }
 
-            icicle_remove_syscall_hook(this->emu_, id);
+            icicle_remove_hook(this->emu_, id);
             this->hooks_.erase(entry);
         }
 

src/icicle/src/icicle.rs

@@ -1,5 +1,5 @@
-use icicle_cpu::ValueSource;
 use icicle_cpu::ExceptionCode;
+use icicle_cpu::ValueSource;
 use std::{cell::RefCell, collections::HashMap, rc::Rc};
 
 use crate::registers;
@@ -123,10 +123,7 @@ impl icicle_vm::CodeInjector for InstructionHookInjector {
                 }
             }
 
-            std::mem::swap(
-                &mut tmp_block.instructions,
-                &mut block.pcode.instructions,
-            );
+            std::mem::swap(&mut tmp_block.instructions, &mut block.pcode.instructions);
         }
     }
 }
@@ -139,6 +136,22 @@ pub struct IcicleEmulator {
     execution_hooks: Rc<RefCell<HookContainer<dyn Fn(u64)>>>,
 }
 
+struct MemoryHook {
+    callback: Box<dyn Fn(u64, &[u8])>,
+}
+
+impl icicle_cpu::mem::WriteHook for MemoryHook {
+    fn write(&mut self, _mem: &mut icicle_cpu::Mmu, addr: u64, value: &[u8]) {
+        (self.callback)(addr, value);
+    }
+}
+
+impl icicle_cpu::mem::ReadAfterHook for MemoryHook {
+    fn read(&mut self, _mem: &mut icicle_cpu::Mmu, addr: u64, value: &[u8]) {
+        (self.callback)(addr, value);
+    }
+}
+
 pub struct MmioHandler {
     read_handler: Box<dyn Fn(u64, &mut [u8])>,
     write_handler: Box<dyn Fn(u64, &[u8])>,
@@ -171,7 +184,8 @@ impl icicle_cpu::mem::IoMemory for MmioHandler {
 impl IcicleEmulator {
     pub fn new() -> Self {
         let mut virtual_machine = create_x64_vm();
-        let exec_hooks: Rc<RefCell<HookContainer<dyn Fn(u64)>>> = Rc::new(RefCell::new(HookContainer::new()));
+        let exec_hooks: Rc<RefCell<HookContainer<dyn Fn(u64)>>> =
+            Rc::new(RefCell::new(HookContainer::new()));
 
         let exec_hooks_clone = Rc::clone(&exec_hooks);
 
@@ -211,9 +225,9 @@ impl IcicleEmulator {
                 icicle_vm::VmExit::UnhandledException((code, value)) => {
                     let continue_execution = self.handle_exception(code, value);
                     if !continue_execution {
-                        break
+                        break;
                     }
-                },
+                }
                 _ => break,
             };
         }
@@ -227,7 +241,7 @@ impl IcicleEmulator {
             ExceptionCode::ReadUnmapped => self.handle_violation(value, FOREIGN_READ, true),
             ExceptionCode::WriteUnmapped => self.handle_violation(value, FOREIGN_WRITE, true),
             ExceptionCode::ExecViolation => self.handle_violation(value, FOREIGN_EXEC, true),
-            _ => false
+            _ => false,
         };
 
         return continue_execution;
@@ -242,13 +256,13 @@ impl IcicleEmulator {
         let mut continue_execution = true;
 
         for (_key, func) in self.violation_hooks.get_hooks() {
-            continue_execution &= func(address, permission, unmapped );
+            continue_execution &= func(address, permission, unmapped);
         }
 
         return continue_execution;
     }
 
-    fn handle_syscall(&mut self) -> bool{
+    fn handle_syscall(&mut self) -> bool {
         for (_key, func) in self.syscall_hooks.get_hooks() {
             func();
         }
@@ -258,7 +272,7 @@ impl IcicleEmulator {
     }
 
     pub fn stop(&mut self) {
-       self.vm.icount_limit = 0;
+        self.vm.icount_limit = 0;
     }
 
     pub fn add_violation_hook(&mut self, callback: Box<dyn Fn(u64, u8, bool) -> bool>) -> u32 {
@@ -276,6 +290,34 @@ impl IcicleEmulator {
         return qualify_hook_id(hook_id, HookType::Syscall);
     }
 
+    pub fn add_read_hook(
+        &mut self,
+        start: u64,
+        end: u64,
+        callback: Box<dyn Fn(u64, &[u8])>,
+    ) -> u32 {
+        let id = self.get_mem().add_read_after_hook(start, end, Box::new(MemoryHook { callback }));
+        if id.is_none() {
+            return 0;
+        }
+
+        return qualify_hook_id(id.expect("Hook id needed"), HookType::Read);
+    }
+
+    pub fn add_write_hook(
+        &mut self,
+        start: u64,
+        end: u64,
+        callback: Box<dyn Fn(u64, &[u8])>,
+    ) -> u32 {
+        let id = self.get_mem().add_write_hook(start, end, Box::new(MemoryHook { callback }));
+        if id.is_none() {
+            return 0;
+        }
+
+        return qualify_hook_id(id.expect("Hook id needed"), HookType::Write);
+    }
+
     pub fn remove_hook(&mut self, id: u32) {
         let (hook_id, hook_type) = split_hook_id(id);
 
@@ -283,6 +325,8 @@ impl IcicleEmulator {
             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::Read => {self.get_mem().remove_read_after_hook(hook_id);()},
+            HookType::Write => {self.get_mem().remove_write_hook(hook_id);()},
             _ => {}
         }
     }

src/icicle/src/lib.rs

@@ -41,6 +41,7 @@ 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);
 type ViolationFunction = extern "C" fn(*mut c_void, u64, u8, i32) -> i32;
+type MemoryAccessFunction = MmioWriteFunction;
 
 #[unsafe(no_mangle)]
 pub fn icicle_map_mmio(
@@ -181,6 +182,30 @@ 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 {
+    unsafe {
+        let emulator = &mut *(ptr as *mut IcicleEmulator);
+        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 {
+    unsafe {
+        let emulator = &mut *(ptr as *mut IcicleEmulator);
+        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());
+            },
+        ));
+    }
+}
+
 #[unsafe(no_mangle)]
 pub fn icicle_add_syscall_hook(ptr: *mut c_void, callback: RawFunction, data: *mut c_void) -> u32 {
     unsafe {
@@ -198,7 +223,7 @@ pub fn icicle_add_execution_hook(ptr: *mut c_void, callback: PtrFunction, data:
 }
 
 #[unsafe(no_mangle)]
-pub fn icicle_remove_syscall_hook(ptr: *mut c_void, id: u32) {
+pub fn icicle_remove_hook(ptr: *mut c_void, id: u32) {
     unsafe {
         let emulator = &mut *(ptr as *mut IcicleEmulator);
         emulator.remove_hook(id);