Add icicle mmio support

src/common/utils/object.hpp

@@ -0,0 +1,15 @@
+#pragma once
+
+namespace utils
+{
+    struct object
+    {
+        object() = default;
+        virtual ~object() = default;
+
+        object(object&&) = default;
+        object(const object&) = default;
+        object& operator=(object&&) = default;
+        object& operator=(const object&) = default;
+    };
+}

src/icicle-emulator/icicle_x64_emulator.cpp

@@ -1,13 +1,20 @@
 #define ICICLE_EMULATOR_IMPL
 #include "icicle_x64_emulator.hpp"
 
+#include <utils/object.hpp>
+
 using icicle_emulator = struct icicle_emulator_;
 
 extern "C"
 {
+    using icicle_mmio_read_func = void(void* user, uint64_t address, size_t length, void* data);
+    using icicle_mmio_write_func = void(void* user, uint64_t address, size_t length, const void* data);
+
     icicle_emulator* icicle_create_emulator();
     int32_t icicle_protect_memory(icicle_emulator*, uint64_t address, uint64_t length, uint8_t permissions);
     int32_t icicle_map_memory(icicle_emulator*, uint64_t address, uint64_t length, uint8_t permissions);
+    int32_t icicle_map_mmio(icicle_emulator*, uint64_t address, uint64_t length, icicle_mmio_read_func* read_callback,
+                            void* read_data, icicle_mmio_write_func* write_callback, void* write_data);
     int32_t icicle_unmap_memory(icicle_emulator*, uint64_t address, uint64_t length);
     int32_t icicle_read_memory(icicle_emulator*, uint64_t address, void* data, size_t length);
     int32_t icicle_write_memory(icicle_emulator*, uint64_t address, const void* data, size_t length);
@@ -61,6 +68,9 @@ namespace icicle
                 timeout = {};
             }
 
+            (void)start;
+            (void)end;
+            (void)count;
             icicle_start(this->emu_);
         }
 
@@ -113,8 +123,47 @@ namespace icicle
         void map_mmio(const uint64_t address, const size_t size, mmio_read_callback read_cb,
                       mmio_write_callback write_cb) override
         {
-            this->map_memory(address, size, memory_permission::read_write);
-            // throw std::runtime_error("Not implemented");
+            struct mmio_wrapper : utils::object
+            {
+                uint64_t base{};
+                mmio_read_callback read_cb{};
+                mmio_write_callback write_cb{};
+            };
+
+            auto wrapper = std::make_unique<mmio_wrapper>();
+            wrapper->base = address;
+            wrapper->read_cb = std::move(read_cb);
+            wrapper->write_cb = std::move(write_cb);
+
+            auto* ptr = wrapper.get();
+            this->storage_.push_back(std::move(wrapper));
+
+            auto* read_wrapper = +[](void* user, const uint64_t address, const size_t length, void* data) {
+                constexpr auto limit = sizeof(uint64_t);
+                const auto* w = static_cast<mmio_wrapper*>(user);
+
+                for (size_t offset = 0; offset < length; offset += limit)
+                {
+                    const auto max_read = std::min(limit, length - offset);
+                    const auto value = w->read_cb(address + offset - w->base, max_read);
+                    memcpy(static_cast<uint8_t*>(data) + offset, &value, max_read);
+                }
+            };
+
+            auto* write_wrapper = +[](void* user, const uint64_t address, const size_t length, const void* data) {
+                constexpr auto limit = sizeof(uint64_t);
+                const auto* w = static_cast<mmio_wrapper*>(user);
+
+                for (size_t offset = 0; offset < length; offset += limit)
+                {
+                    uint64_t value{};
+                    const auto max_read = std::min(limit, length - offset);
+                    memcpy(&value, static_cast<const uint8_t*>(data) + offset, max_read);
+                    w->write_cb(address + offset - w->base, max_read, value);
+                }
+            };
+
+            icicle_map_mmio(this->emu_, address, size, read_wrapper, ptr, write_wrapper, ptr);
         }
 
         void map_memory(const uint64_t address, const size_t size, memory_permission permissions) override
@@ -187,6 +236,7 @@ namespace icicle
 
         emulator_hook* hook_interrupt(interrupt_hook_callback callback) override
         {
+            (void)callback;
             return nullptr;
             // throw std::runtime_error("Not implemented");
         }
@@ -194,6 +244,9 @@ namespace icicle
         emulator_hook* hook_memory_violation(uint64_t address, size_t size,
                                              memory_violation_hook_callback callback) override
         {
+            (void)address;
+            (void)size;
+            (void)callback;
             return nullptr;
             // throw std::runtime_error("Not implemented");
         }
@@ -206,6 +259,9 @@ namespace icicle
                 return nullptr;
             }
 
+            (void)address;
+            (void)size;
+            (void)callback;
             return nullptr;
             // throw std::runtime_error("Not implemented");
         }
@@ -225,12 +281,16 @@ namespace icicle
 
         void serialize_state(utils::buffer_serializer& buffer, const bool is_snapshot) const override
         {
-            // throw std::runtime_error("Not implemented");
+            (void)buffer;
+            (void)is_snapshot;
+            throw std::runtime_error("Not implemented");
         }
 
         void deserialize_state(utils::buffer_deserializer& buffer, const bool is_snapshot) override
         {
-            // throw std::runtime_error("Not implemented");
+            (void)buffer;
+            (void)is_snapshot;
+            throw std::runtime_error("Not implemented");
         }
 
         std::vector<std::byte> save_registers() override
@@ -241,6 +301,7 @@ namespace icicle
 
         void restore_registers(const std::vector<std::byte>& register_data) override
         {
+            (void)register_data;
             // throw std::runtime_error("Not implemented");
         }
 
@@ -250,6 +311,7 @@ namespace icicle
         }
 
       private:
+        std::list<std::unique_ptr<utils::object>> storage_{};
         using syscall_hook_storage = std::unique_ptr<std::function<void()>>;
         std::unordered_map<uint32_t, syscall_hook_storage> syscall_hooks_{};
         icicle_emulator* emu_{};

src/icicle/src/icicle.rs

@@ -10,7 +10,7 @@ fn create_x64_vm() -> icicle_vm::Vm {
     cpu_config.track_uninitialized = false;
     cpu_config.optimize_instructions = true;
     cpu_config.optimize_block = false;
-    
+
     return icicle_vm::build(&cpu_config).unwrap();
 }
 
@@ -36,22 +36,107 @@ fn map_permissions(foreign_permissions: u8) -> u8 {
     return permissions;
 }
 
+#[repr(u8)]
+#[allow(dead_code)]
+#[derive(PartialEq)]
+enum HookType {
+    Syscall = 1,
+    Read,
+    Write,
+    Execute,
+    Unknown,
+}
+
+fn u8_to_hook_type_unsafe(value: u8) -> HookType {
+    // This is unsafe because it assumes the value is valid
+    unsafe { std::mem::transmute(value) }
+}
+
+fn split_hook_id(id: u32) -> (u32, HookType) {
+    let hook_id = id & 0xFFFFFF;
+    let hook_type = u8_to_hook_type_unsafe((id >> 24) as u8);
+
+    return (hook_id, hook_type);
+}
+
+fn qualify_hook_id(hook_id: u32, hook_type: HookType) -> u32 {
+    let hook_type: u32 = (hook_type as u8).into();
+    let hook_type_mask: u32 = hook_type << 24;
+    return (hook_id | hook_type_mask).into();
+}
+
+pub struct HookContainer<Func: ?Sized> {
+    hook_id: u32,
+    hooks: HashMap<u32, Box<Func>>,
+}
+
+impl<Func: ?Sized> HookContainer<Func> {
+    pub fn new() -> Self {
+        Self {
+            hook_id: 0,
+            hooks: HashMap::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);
+
+        return id;
+    }
+
+    pub fn get_hooks(&self) -> &HashMap<u32, Box<Func>> {
+        return &self.hooks;
+    }
+
+    pub fn remove_hook(&mut self, id: u32) {
+        self.hooks.remove(&id);
+    }
+}
+
 pub struct IcicleEmulator {
     vm: icicle_vm::Vm,
     reg: X64RegisterNodes,
-    
-    syscall_hook_id: u32,
-    syscall_hooks: HashMap<u32, Box<dyn Fn()>>,
+    syscall_hooks: HookContainer<dyn Fn()>,
+}
+
+pub struct MmioHandler {
+    read_handler: Box<dyn Fn(u64, &mut [u8])>,
+    write_handler: Box<dyn Fn(u64, &[u8])>,
+}
+
+impl MmioHandler {
+    pub fn new(
+        read_function: Box<dyn Fn(u64, &mut [u8])>,
+        write_function: Box<dyn Fn(u64, &[u8])>,
+    ) -> Self {
+        Self {
+            read_handler: read_function,
+            write_handler: write_function,
+        }
+    }
+}
+
+impl icicle_cpu::mem::IoMemory for MmioHandler {
+    fn read(&mut self, addr: u64, buf: &mut [u8]) -> icicle_cpu::mem::MemResult<()> {
+        (self.read_handler)(addr, buf);
+        return Ok(());
+    }
+
+    fn write(&mut self, addr: u64, value: &[u8]) -> icicle_cpu::mem::MemResult<()> {
+        (self.write_handler)(addr, value);
+        return Ok(());
+    }
 }
 
 impl IcicleEmulator {
     pub fn new() -> Self {
-        let vm_i = create_x64_vm();
+        let virtual_machine = create_x64_vm();
         Self {
-            reg: X64RegisterNodes::new(&vm_i.cpu.arch),
-            vm: vm_i,
-            syscall_hook_id: 0,
-            syscall_hooks: HashMap::new(),
+            reg: X64RegisterNodes::new(&virtual_machine.cpu.arch),
+            vm: virtual_machine,
+            syscall_hooks: HookContainer::new(),
         }
     }
 
@@ -64,7 +149,9 @@ impl IcicleEmulator {
             let reason = self.vm.run();
 
             let invoke_syscall = match reason {
-                icicle_vm::VmExit::UnhandledException((code, _)) => code == icicle_cpu::ExceptionCode::Syscall,
+                icicle_vm::VmExit::UnhandledException((code, _)) => {
+                    code == icicle_cpu::ExceptionCode::Syscall
+                }
                 _ => false,
             };
 
@@ -72,7 +159,7 @@ impl IcicleEmulator {
                 break;
             }
 
-            for (_key, func) in &self.syscall_hooks {
+            for (_key, func) in self.syscall_hooks.get_hooks() {
                 func();
             }
 
@@ -81,14 +168,17 @@ impl IcicleEmulator {
     }
 
     pub fn add_syscall_hook(&mut self, callback: Box<dyn Fn()>) -> u32 {
-        self.syscall_hook_id += 1;
-        let id = self.syscall_hook_id;
-        self.syscall_hooks.insert(id, callback);
-        return id;
+        let hook_id = self.syscall_hooks.add_hook(callback);
+        return qualify_hook_id(hook_id, HookType::Syscall);
     }
 
-    pub fn remove_syscall_hook(&mut self, id: u32) {
-        self.syscall_hooks.remove(&id);
+    pub fn remove_hook(&mut self, id: u32) {
+        let (hook_id, hook_type) = split_hook_id(id);
+
+        match hook_type {
+            HookType::Syscall => self.syscall_hooks.remove_hook(hook_id),
+            _ => {}
+        }
     }
 
     pub fn map_memory(&mut self, address: u64, length: u64, permissions: u8) -> bool {
@@ -113,6 +203,28 @@ impl IcicleEmulator {
         return res.is_ok();
     }
 
+    pub fn map_mmio(
+        &mut self,
+        address: u64,
+        length: u64,
+        read_function: Box<dyn Fn(u64, &mut [u8])>,
+        write_function: Box<dyn Fn(u64, &[u8])>,
+    ) -> bool {
+        let mem = self.get_mem();
+
+        let handler = MmioHandler::new(read_function, write_function);
+        let handler_id = mem.register_io_handler(handler);
+
+        let layout = icicle_vm::cpu::mem::AllocLayout {
+            addr: Some(address),
+            size: length,
+            align: 0x1000,
+        };
+
+        let res = mem.alloc_memory(layout, handler_id);
+        return res.is_ok();
+    }
+
     pub fn unmap_memory(&mut self, address: u64, length: u64) -> bool {
         return self.get_mem().unmap_memory_len(address, length);
     }
@@ -162,22 +274,70 @@ impl IcicleEmulator {
         //self.vm.cpu.write_trunc(reg_node, value);
 
         match reg_node.size {
-            1 => self.vm.cpu.write_var::<[u8; 1]>(reg_node, buffer[..1].try_into().expect("")),
-            2 => self.vm.cpu.write_var::<[u8; 2]>(reg_node, buffer[..2].try_into().expect("")),
-            3 => self.vm.cpu.write_var::<[u8; 3]>(reg_node, buffer[..3].try_into().expect("")),
-            4 => self.vm.cpu.write_var::<[u8; 4]>(reg_node, buffer[..4].try_into().expect("")),
-            5 => self.vm.cpu.write_var::<[u8; 5]>(reg_node, buffer[..5].try_into().expect("")),
-            6 => self.vm.cpu.write_var::<[u8; 6]>(reg_node, buffer[..6].try_into().expect("")),
-            7 => self.vm.cpu.write_var::<[u8; 7]>(reg_node, buffer[..7].try_into().expect("")),
-            8 => self.vm.cpu.write_var::<[u8; 8]>(reg_node, buffer[..8].try_into().expect("")),
-            9 => self.vm.cpu.write_var::<[u8; 9]>(reg_node, buffer[..9].try_into().expect("")),
-            10 => self.vm.cpu.write_var::<[u8; 10]>(reg_node, buffer[..10].try_into().expect("")),
-            11 => self.vm.cpu.write_var::<[u8; 11]>(reg_node, buffer[..11].try_into().expect("")),
-            12 => self.vm.cpu.write_var::<[u8; 12]>(reg_node, buffer[..12].try_into().expect("")),
-            13 => self.vm.cpu.write_var::<[u8; 13]>(reg_node, buffer[..13].try_into().expect("")),
-            14 => self.vm.cpu.write_var::<[u8; 14]>(reg_node, buffer[..14].try_into().expect("")),
-            15 => self.vm.cpu.write_var::<[u8; 15]>(reg_node, buffer[..15].try_into().expect("")),
-            16 => self.vm.cpu.write_var::<[u8; 16]>(reg_node, buffer[..16].try_into().expect("")),
+            1 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 1]>(reg_node, buffer[..1].try_into().expect("")),
+            2 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 2]>(reg_node, buffer[..2].try_into().expect("")),
+            3 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 3]>(reg_node, buffer[..3].try_into().expect("")),
+            4 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 4]>(reg_node, buffer[..4].try_into().expect("")),
+            5 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 5]>(reg_node, buffer[..5].try_into().expect("")),
+            6 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 6]>(reg_node, buffer[..6].try_into().expect("")),
+            7 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 7]>(reg_node, buffer[..7].try_into().expect("")),
+            8 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 8]>(reg_node, buffer[..8].try_into().expect("")),
+            9 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 9]>(reg_node, buffer[..9].try_into().expect("")),
+            10 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 10]>(reg_node, buffer[..10].try_into().expect("")),
+            11 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 11]>(reg_node, buffer[..11].try_into().expect("")),
+            12 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 12]>(reg_node, buffer[..12].try_into().expect("")),
+            13 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 13]>(reg_node, buffer[..13].try_into().expect("")),
+            14 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 14]>(reg_node, buffer[..14].try_into().expect("")),
+            15 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 15]>(reg_node, buffer[..15].try_into().expect("")),
+            16 => self
+                .vm
+                .cpu
+                .write_var::<[u8; 16]>(reg_node, buffer[..16].try_into().expect("")),
             _ => panic!("invalid dynamic value size"),
         }
 

src/icicle/src/lib.rs

@@ -25,6 +25,37 @@ pub fn icicle_start(ptr: *mut c_void) {
     }
 }
 
+type MmioReadFunction = extern "C" fn(*mut c_void, u64, usize, *mut c_void);
+type MmioWriteFunction = extern "C" fn(*mut c_void, u64, usize, *const c_void);
+
+#[unsafe(no_mangle)]
+pub fn icicle_map_mmio(
+    ptr: *mut c_void,
+    address: u64,
+    length: u64,
+    read_cb: MmioReadFunction,
+    read_data: *mut c_void,
+    write_cb: MmioWriteFunction,
+    write_data: *mut c_void,
+) -> i32 {
+    unsafe {
+        let emulator = &mut *(ptr as *mut IcicleEmulator);
+
+        let read_wrapper = Box::new(move |addr: u64, data: &mut [u8]| {
+            let raw_pointer: *mut u8 = data.as_mut_ptr();
+            read_cb(read_data, addr, data.len(), raw_pointer as *mut c_void);
+        });
+
+        let write_wrapper = Box::new(move |addr: u64, data: &[u8]| {
+            let raw_pointer: *const u8 = data.as_ptr();
+            write_cb(write_data, addr, data.len(), raw_pointer as *const c_void);
+        });
+
+        let res = emulator.map_mmio(address, length, read_wrapper, write_wrapper);
+        return to_cbool(res);
+    }
+}
+
 #[unsafe(no_mangle)]
 pub fn icicle_map_memory(ptr: *mut c_void, address: u64, length: u64, permissions: u8) -> i32 {
     unsafe {
@@ -83,7 +114,7 @@ type CFunctionPointer = extern "C" fn(*mut c_void);
 pub fn icicle_add_syscall_hook(ptr: *mut c_void, callback: CFunctionPointer, data: *mut c_void) {
     unsafe {
         let emulator = &mut *(ptr as *mut IcicleEmulator);
-       emulator.add_syscall_hook(Box::new(move || callback(data)));
+        emulator.add_syscall_hook(Box::new(move || callback(data)));
     }
 }
 
@@ -91,7 +122,7 @@ pub fn icicle_add_syscall_hook(ptr: *mut c_void, callback: CFunctionPointer, dat
 pub fn icicle_remove_syscall_hook(ptr: *mut c_void, id: u32) {
     unsafe {
         let emulator = &mut *(ptr as *mut IcicleEmulator);
-        emulator.remove_syscall_hook(id);
+        emulator.remove_hook(id);
     }
 }
 

src/unicorn-emulator/function_wrapper.hpp

@@ -3,10 +3,10 @@
 #include <memory>
 #include <functional>
 
-#include "object.hpp"
+#include <utils/object.hpp>
 
 template <typename ReturnType, typename... Args>
-class function_wrapper : public object
+class function_wrapper : public utils::object
 {
   public:
     using user_data_pointer = void*;

src/unicorn-emulator/object.hpp

@@ -1,12 +0,0 @@
-#pragma once
-
-struct object
-{
-    object() = default;
-    virtual ~object() = default;
-
-    object(object&&) = default;
-    object(const object&) = default;
-    object& operator=(object&&) = default;
-    object& operator=(const object&) = default;
-};

src/unicorn-emulator/unicorn_x64_emulator.cpp

@@ -75,7 +75,7 @@ namespace unicorn
             }
         }
 
-        struct hook_object : object
+        struct hook_object : utils::object
         {
             emulator_hook* as_opaque_hook()
             {
@@ -87,7 +87,7 @@ namespace unicorn
         {
           public:
             template <typename T>
-                requires(std::is_base_of_v<object, T> && std::is_move_constructible_v<T>)
+                requires(std::is_base_of_v<utils::object, T> && std::is_move_constructible_v<T>)
             void add(T data, unicorn_hook hook)
             {
                 hook_entry entry{};
@@ -101,7 +101,7 @@ namespace unicorn
           private:
             struct hook_entry
             {
-                std::unique_ptr<object> data{};
+                std::unique_ptr<utils::object> data{};
                 unicorn_hook hook{};
             };