From ec67e9d43f2db218b6de14bd9fc3db965189a7dc Mon Sep 17 00:00:00 2001
From: momo5502 <mauriceheumann@gmail.com>
Date: Sun, 25 May 2025 16:53:19 +0200
Subject: [PATCH] Support snapshots for icicle

---
 .../icicle-bridge/src/icicle.rs               | 16 ++++++++++++++
 .../icicle-emulator/icicle-bridge/src/lib.rs  | 16 ++++++++++++++
 .../icicle_x86_64_emulator.cpp                | 22 +++++++++++++------
 3 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs b/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs
index 71ac020d..660ade77 100644
--- a/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs
+++ b/src/backends/icicle-emulator/icicle-bridge/src/icicle.rs
@@ -247,6 +247,7 @@ pub struct IcicleEmulator {
     violation_hooks: HookContainer<dyn Fn(u64, u8, bool) -> bool>,
     execution_hooks: Rc<RefCell<ExecutionHooks>>,
     stop: Rc<RefCell<bool>>,
+    snapshots: Vec<Box<icicle_vm::Snapshot>>,
 }
 
 struct MemoryHook {
@@ -329,6 +330,7 @@ impl IcicleEmulator {
             interrupt_hooks: HookContainer::new(),
             violation_hooks: HookContainer::new(),
             execution_hooks: exec_hooks,
+            snapshots: Vec::new(),
         }
     }
 
@@ -643,6 +645,20 @@ impl IcicleEmulator {
         }
     }
 
+    pub fn create_snapshot(&mut self) -> u32 {
+        let snap = self.vm.snapshot();
+
+        let id = self.snapshots.len() as u32;
+        self.snapshots.push(Box::new(snap));
+
+        return id;
+    }
+
+    pub fn restore_snapshot(&mut self, id: u32) {
+        let snap = self.snapshots[id as usize].as_ref();
+        self.vm.restore(&snap);
+    }
+
     fn write_flags<T>(&mut self, data: &[u8]) -> usize {
         const REAL_SIZE: usize = std::mem::size_of::<u64>();
         let limit: usize = std::mem::size_of::<T>();
diff --git a/src/backends/icicle-emulator/icicle-bridge/src/lib.rs b/src/backends/icicle-emulator/icicle-bridge/src/lib.rs
index bd60b9aa..74e07b54 100644
--- a/src/backends/icicle-emulator/icicle-bridge/src/lib.rs
+++ b/src/backends/icicle-emulator/icicle-bridge/src/lib.rs
@@ -149,6 +149,22 @@ pub fn icicle_restore_registers(ptr: *mut c_void, data: *const c_void, size: usi
     }
 }
 
+#[unsafe(no_mangle)]
+pub fn icicle_create_snapshot(ptr: *mut c_void) -> u32 {
+    unsafe {
+        let emulator = &mut *(ptr as *mut IcicleEmulator);
+        return emulator.create_snapshot();
+    }
+}
+
+#[unsafe(no_mangle)]
+pub fn icicle_restore_snapshot(ptr: *mut c_void, id: u32) {
+    unsafe {
+        let emulator = &mut *(ptr as *mut IcicleEmulator);
+        emulator.restore_snapshot(id);
+    }
+}
+
 #[unsafe(no_mangle)]
 pub fn icicle_read_memory(ptr: *mut c_void, address: u64, data: *mut c_void, size: usize) -> i32 {
     if size == 0 {
diff --git a/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp b/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp
index 146b32c0..cb759e3a 100644
--- a/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp
+++ b/src/backends/icicle-emulator/icicle_x86_64_emulator.cpp
@@ -28,6 +28,8 @@ extern "C"
     int32_t icicle_write_memory(icicle_emulator*, uint64_t address, const void* data, size_t length);
     void icicle_save_registers(icicle_emulator*, data_accessor_func* accessor, void* accessor_data);
     void icicle_restore_registers(icicle_emulator*, const void* data, size_t length);
+    uint32_t icicle_create_snapshot(icicle_emulator*);
+    void icicle_restore_snapshot(icicle_emulator*, uint32_t id);
     uint32_t icicle_add_syscall_hook(icicle_emulator*, raw_func* callback, void* data);
     uint32_t icicle_add_interrupt_hook(icicle_emulator*, interrupt_func* callback, void* data);
     uint32_t icicle_add_block_hook(icicle_emulator*, block_func* callback, void* data);
@@ -382,21 +384,27 @@ namespace icicle
         {
             if (is_snapshot)
             {
-                throw std::runtime_error("Not implemented");
+                const auto snapshot = icicle_create_snapshot(this->emu_);
+                buffer.write<uint32_t>(snapshot);
+            }
+            else
+            {
+                buffer.write_vector(this->save_registers());
             }
-
-            buffer.write_vector(this->save_registers());
         }
 
         void deserialize_state(utils::buffer_deserializer& buffer, const bool is_snapshot) override
         {
             if (is_snapshot)
             {
-                throw std::runtime_error("Not implemented");
+                const auto snapshot = buffer.read<uint32_t>();
+                icicle_restore_snapshot(this->emu_, snapshot);
+            }
+            else
+            {
+                const auto data = buffer.read_vector<std::byte>();
+                this->restore_registers(data);
             }
-
-            const auto data = buffer.read_vector<std::byte>();
-            this->restore_registers(data);
         }
 
         std::vector<std::byte> save_registers() const override