[hyperlight_host/snapshot] Reorganise structures

This moves `mem::shared_mem_snapshot::SharedMemorySNapshot` into
`sandbox::snapshot::Snapshot`, because that structure has always
contained more than /just/ a snapshot of a shared memory.  It also
removes the Arc reference counting wrapper from the new structure, and
shifts the Arc wrapper into
`sandbox::initialized_multi_use::MultiUseSandbox`, which is where it
was actually needed.

Signed-off-by: Simon Davies <simongdavies@users.noreply.github.com>

Author: syntactically

src/hyperlight_host/benches/benchmarks.rs

@@ -172,7 +172,7 @@ fn bench_guest_call_with_restore(b: &mut criterion::Bencher, size: SandboxSize)
 
     b.iter(|| {
         sbox.call::<String>("Echo", "hello\n".to_string()).unwrap();
-        sbox.restore(&snapshot).unwrap();
+        sbox.restore(snapshot.clone()).unwrap();
     });
 }
 
@@ -340,7 +340,7 @@ fn bench_snapshot_restore(b: &mut criterion::Bencher, size: SandboxSize) {
 
             // Measure only the restore time
             let start = Instant::now();
-            sbox.restore(&snapshot).unwrap();
+            sbox.restore(snapshot.clone()).unwrap();
             total_duration += start.elapsed();
         }
 

src/hyperlight_host/src/mem/mgr.rs

@@ -33,8 +33,8 @@ use super::memory_region::{DEFAULT_GUEST_BLOB_MEM_FLAGS, MemoryRegionType};
 use super::ptr::{GuestPtr, RawPtr};
 use super::ptr_offset::Offset;
 use super::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, HostSharedMemory, SharedMemory};
-use super::shared_mem_snapshot::SharedMemorySnapshot;
 use crate::sandbox::SandboxConfiguration;
+use crate::sandbox::snapshot::Snapshot;
 use crate::sandbox::uninitialized::GuestBlob;
 use crate::{Result, log_then_return, new_error};
 
@@ -285,20 +285,20 @@ where
         &mut self,
         sandbox_id: u64,
         mapped_regions: Vec<MemoryRegion>,
-    ) -> Result<SharedMemorySnapshot> {
-        SharedMemorySnapshot::new(&mut self.shared_mem, sandbox_id, mapped_regions)
+    ) -> Result<Snapshot> {
+        Snapshot::new(&mut self.shared_mem, sandbox_id, mapped_regions)
     }
 
     /// This function restores a memory snapshot from a given snapshot.
-    pub(crate) fn restore_snapshot(&mut self, snapshot: &SharedMemorySnapshot) -> Result<()> {
+    pub(crate) fn restore_snapshot(&mut self, snapshot: &Snapshot) -> Result<()> {
         if self.shared_mem.mem_size() != snapshot.mem_size() {
             return Err(new_error!(
                 "Snapshot size does not match current memory size: {} != {}",
                 self.shared_mem.raw_mem_size(),
                 snapshot.mem_size()
             ));
         }
-        snapshot.restore_from_snapshot(&mut self.shared_mem)?;
+        self.shared_mem.restore_from_snapshot(snapshot)?;
         Ok(())
     }
 }

src/hyperlight_host/src/mem/mod.rs

@@ -35,9 +35,6 @@ pub mod ptr_offset;
 /// A wrapper around unsafe functionality to create and initialize
 /// a memory region for a guest running in a sandbox.
 pub mod shared_mem;
-/// A wrapper around a `SharedMemory` and a snapshot in time
-/// of the memory therein
-pub mod shared_mem_snapshot;
 /// Utilities for writing shared memory tests
 #[cfg(test)]
 pub(crate) mod shared_mem_tests;

src/hyperlight_host/src/mem/shared_mem.rs

@@ -39,6 +39,7 @@ use windows::core::PCSTR;
 use crate::HyperlightError::MemoryAllocationFailed;
 #[cfg(target_os = "windows")]
 use crate::HyperlightError::{MemoryRequestTooBig, WindowsAPIError};
+use crate::sandbox::snapshot::Snapshot;
 use crate::{Result, log_then_return, new_error};
 
 /// Makes sure that the given `offset` and `size` are within the bounds of the memory with size `mem_size`.
@@ -675,6 +676,12 @@ pub trait SharedMemory {
         &mut self,
         f: F,
     ) -> Result<T>;
+
+    /// Restore a SharedMemory from a snapshot with matching size
+    fn restore_from_snapshot(&mut self, snapshot: &Snapshot) -> Result<()> {
+        assert!(self.mem_size() == snapshot.mem_size());
+        self.with_exclusivity(|e| e.copy_from_slice(snapshot.memory(), 0))?
+    }
 }
 
 impl SharedMemory for ExclusiveSharedMemory {

src/hyperlight_host/src/mem/shared_mem_snapshot.rs

@@ -104,7 +104,7 @@ pub struct MultiUseSandbox {
     dbg_mem_access_fn: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
     /// If the current state of the sandbox has been captured in a snapshot,
     /// that snapshot is stored here.
-    snapshot: Option<Snapshot>,
+    snapshot: Option<Arc<Snapshot>>,
 }
 
 impl MultiUseSandbox {
@@ -163,7 +163,7 @@ impl MultiUseSandbox {
     /// # }
     /// ```
     #[instrument(err(Debug), skip_all, parent = Span::current())]
-    pub fn snapshot(&mut self) -> Result<Snapshot> {
+    pub fn snapshot(&mut self) -> Result<Arc<Snapshot>> {
         if self.poisoned {
             return Err(crate::HyperlightError::PoisonedSandbox);
         }
@@ -174,8 +174,7 @@ impl MultiUseSandbox {
         let mapped_regions_iter = self.vm.get_mapped_regions();
         let mapped_regions_vec: Vec<MemoryRegion> = mapped_regions_iter.cloned().collect();
         let memory_snapshot = self.mem_mgr.snapshot(self.id, mapped_regions_vec)?;
-        let inner = Arc::new(memory_snapshot);
-        let snapshot = Snapshot { inner };
+        let snapshot = Arc::new(memory_snapshot);
         self.snapshot = Some(snapshot.clone());
         Ok(snapshot)
     }
@@ -221,7 +220,7 @@ impl MultiUseSandbox {
     /// assert_eq!(value, 100);
     ///
     /// // Restore to previous state (same sandbox)
-    /// sandbox.restore(&snapshot)?;
+    /// sandbox.restore(snapshot)?;
     /// let restored_value: i32 = sandbox.call_guest_function_by_name("GetValue", ())?;
     /// assert_eq!(restored_value, 0); // Back to initial state
     /// # Ok(())
@@ -257,22 +256,22 @@ impl MultiUseSandbox {
     /// # }
     /// ```
     #[instrument(err(Debug), skip_all, parent = Span::current())]
-    pub fn restore(&mut self, snapshot: &Snapshot) -> Result<()> {
+    pub fn restore(&mut self, snapshot: Arc<Snapshot>) -> Result<()> {
         if let Some(snap) = &self.snapshot
-            && Arc::ptr_eq(&snap.inner, &snapshot.inner)
+            && Arc::ptr_eq(&snap, &snapshot)
         {
             // If the snapshot is already the current one, no need to restore
             return Ok(());
         }
 
-        if self.id != snapshot.inner.sandbox_id() {
+        if self.id != snapshot.sandbox_id() {
             return Err(SnapshotSandboxMismatch);
         }
 
-        self.mem_mgr.restore_snapshot(&snapshot.inner)?;
+        self.mem_mgr.restore_snapshot(&snapshot)?;
 
         let current_regions: HashSet<_> = self.vm.get_mapped_regions().cloned().collect();
-        let snapshot_regions: HashSet<_> = snapshot.inner.regions().iter().cloned().collect();
+        let snapshot_regions: HashSet<_> = snapshot.regions().iter().cloned().collect();
 
         let regions_to_unmap = current_regions.difference(&snapshot_regions);
         let regions_to_map = snapshot_regions.difference(&current_regions);
@@ -356,7 +355,7 @@ impl MultiUseSandbox {
         }
         let snapshot = self.snapshot()?;
         let res = self.call(func_name, args);
-        self.restore(&snapshot)?;
+        self.restore(snapshot)?;
         res
     }
 
@@ -963,7 +962,7 @@ mod tests {
         let res: i32 = sbox.call("GetStatic", ()).unwrap();
         assert_eq!(res, 5);
 
-        sbox.restore(&snapshot).unwrap();
+        sbox.restore(snapshot).unwrap();
         #[allow(deprecated)]
         let _ = sbox
             .call_guest_function_by_name::<i32>("AddToStatic", 5i32)
@@ -1027,7 +1026,7 @@ mod tests {
         let res: i32 = sbox.call("GetStatic", ()).unwrap();
         assert_eq!(res, 5);
 
-        sbox.restore(&snapshot).unwrap();
+        sbox.restore(snapshot).unwrap();
         let res: i32 = sbox.call("GetStatic", ()).unwrap();
         assert_eq!(res, 0);
     }
@@ -1244,12 +1243,12 @@ mod tests {
         assert_eq!(sbox.vm.get_mapped_regions().count(), 1);
 
         // 4. Restore to snapshot 1 (should unmap the region)
-        sbox.restore(&snapshot1).unwrap();
-        assert_eq!(sbox.vm.get_mapped_regions().count(), 0);
+        sbox.restore(snapshot1).unwrap();
+        assert_eq!(sbox.vm.get_mapped_regions().len(), 0);
 
         // 5. Restore forward to snapshot 2 (should remap the region)
-        sbox.restore(&snapshot2).unwrap();
-        assert_eq!(sbox.vm.get_mapped_regions().count(), 1);
+        sbox.restore(snapshot2).unwrap();
+        assert_eq!(sbox.vm.get_mapped_regions().len(), 1);
 
         // Verify the region is the same
         let mut restored_regions = sbox.vm.get_mapped_regions();
@@ -1282,7 +1281,7 @@ mod tests {
         assert_ne!(sandbox.id, sandbox2.id);
 
         let snapshot = sandbox.snapshot().unwrap();
-        let err = sandbox2.restore(&snapshot);
+        let err = sandbox2.restore(snapshot.clone());
         assert!(matches!(err, Err(HyperlightError::SnapshotSandboxMismatch)));
 
         let sandbox_id = sandbox.id;