From fd018c46b4e20ed21d4720d1d4df7276f20391c5 Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Tue, 1 Apr 2025 18:23:24 +0200
Subject: [PATCH] Support saving and restoring registers

---
 src/icicle-emulator/icicle_x64_emulator.cpp | 19 ++++++++++++----
 src/icicle/src/icicle.rs                    | 20 +++++++++++++++++
 src/icicle/src/lib.rs                       | 24 ++++++++++++++++++---
 3 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/src/icicle-emulator/icicle_x64_emulator.cpp b/src/icicle-emulator/icicle_x64_emulator.cpp
index e2a8f923..53b45139 100644
--- a/src/icicle-emulator/icicle_x64_emulator.cpp
+++ b/src/icicle-emulator/icicle_x64_emulator.cpp
@@ -10,6 +10,8 @@ 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);
 
+    using data_accessor_func = void(void* user, const void* data, size_t length);
+
     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);
@@ -18,6 +20,8 @@ extern "C"
     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);
+    int32_t icicle_save_registers(icicle_emulator*, data_accessor_func* accessor, void* accessor_data);
+    int32_t icicle_restore_registers(icicle_emulator*, const void* data, size_t length);
     uint32_t icicle_add_syscall_hook(icicle_emulator*, void (*callback)(void*), void* data);
     void icicle_remove_syscall_hook(icicle_emulator*, uint32_t id);
     size_t icicle_read_register(icicle_emulator*, int reg, void* data, size_t length);
@@ -295,14 +299,21 @@ namespace icicle
 
         std::vector<std::byte> save_registers() override
         {
-            // throw std::runtime_error("Not implemented");
-            return {};
+            std::vector<std::byte> data{};
+            auto* accessor = +[](void* user, const void* data, const size_t length) {
+                auto& vec = *static_cast<std::vector<std::byte>*>(user);
+                vec.resize(length);
+                memcpy(vec.data(), data, length);
+            };
+
+            icicle_save_registers(this->emu_, accessor, &data);
+
+            return data;
         }
 
         void restore_registers(const std::vector<std::byte>& register_data) override
         {
-            (void)register_data;
-            // throw std::runtime_error("Not implemented");
+            icicle_restore_registers(this->emu_, register_data.data(), register_data.size());
         }
 
         bool has_violation() const override
diff --git a/src/icicle/src/icicle.rs b/src/icicle/src/icicle.rs
index 39f6aa1d..f8ebc1a4 100644
--- a/src/icicle/src/icicle.rs
+++ b/src/icicle/src/icicle.rs
@@ -251,6 +251,26 @@ impl IcicleEmulator {
         return res.is_ok();
     }
 
+    pub fn save_registers(&self) -> Vec<u8> {
+        const REG_SIZE: usize = std::mem::size_of::<icicle_cpu::Regs>();
+        unsafe {
+            let data: [u8; REG_SIZE] = self.vm.cpu.regs.read_at(0);
+            return data.to_vec();
+        }
+    }
+
+    pub fn restore_registers(&mut self, data: &[u8]) {
+        const REG_SIZE: usize = std::mem::size_of::<icicle_cpu::Regs>();
+
+        let mut buffer: [u8; REG_SIZE] = [0; REG_SIZE];
+        let size = std::cmp::min(REG_SIZE, data.len());
+        buffer.copy_from_slice(&data[..size]);
+
+        unsafe {
+            self.vm.cpu.regs.write_at(0, buffer);
+        };
+    }
+
     pub fn read_register(&mut self, reg: X64Register, buffer: &mut [u8]) -> usize {
         let reg_node = self.reg.get_node(reg);
 
diff --git a/src/icicle/src/lib.rs b/src/icicle/src/lib.rs
index adbeee0c..22c40e16 100644
--- a/src/icicle/src/lib.rs
+++ b/src/icicle/src/lib.rs
@@ -25,6 +25,8 @@ pub fn icicle_start(ptr: *mut c_void) {
     }
 }
 
+type RawFunction = extern "C" fn(*mut c_void);
+type DataFunction = extern "C" fn(*mut c_void, *const c_void, usize);
 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);
 
@@ -98,6 +100,24 @@ pub fn icicle_write_memory(
     }
 }
 
+#[unsafe(no_mangle)]
+pub fn icicle_save_registers(ptr: *mut c_void, accessor: DataFunction, accessor_data: *mut c_void) {
+    unsafe {
+        let emulator = &mut *(ptr as *mut IcicleEmulator);
+        let registers = emulator.save_registers();
+        accessor(accessor_data, registers.as_ptr() as *const c_void, registers.len());
+    }
+}
+
+#[unsafe(no_mangle)]
+pub fn icicle_restore_registers(ptr: *mut c_void, data: *const c_void, size: usize) {
+    unsafe {
+        let emulator = &mut *(ptr as *mut IcicleEmulator);
+        let u8_slice = std::slice::from_raw_parts(data as *const u8, size);
+        emulator.restore_registers(u8_slice);
+    }
+}
+
 #[unsafe(no_mangle)]
 pub fn icicle_read_memory(ptr: *mut c_void, address: u64, data: *mut c_void, size: usize) -> i32 {
     unsafe {
@@ -108,10 +128,8 @@ pub fn icicle_read_memory(ptr: *mut c_void, address: u64, data: *mut c_void, siz
     }
 }
 
-type CFunctionPointer = extern "C" fn(*mut c_void);
-
 #[unsafe(no_mangle)]
-pub fn icicle_add_syscall_hook(ptr: *mut c_void, callback: CFunctionPointer, data: *mut c_void) {
+pub fn icicle_add_syscall_hook(ptr: *mut c_void, callback: RawFunction, data: *mut c_void) {
     unsafe {
         let emulator = &mut *(ptr as *mut IcicleEmulator);
         emulator.add_syscall_hook(Box::new(move || callback(data)));