Update snapshots

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/error.rs

@@ -32,7 +32,7 @@ use thiserror::Error;
 
 #[cfg(target_os = "windows")]
 use crate::hypervisor::wrappers::HandleWrapper;
-use crate::mem::memory_region::MemoryRegionFlags;
+use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::ptr::RawPtr;
 
 /// The error type for Hyperlight operations
@@ -148,6 +148,10 @@ pub enum HyperlightError {
     #[error("Memory Protection Failed with OS Error {0:?}.")]
     MemoryProtectionFailed(Option<i32>),
 
+    /// Memory region size mismatch
+    #[error("Memory region size mismatch: host size {0:?}, guest size {1:?} region {2:?}")]
+    MemoryRegionSizeMismatch(usize, usize, MemoryRegion),
+
     /// The memory request exceeds the maximum size allowed
     #[error("Memory requested {0} exceeds maximum size allowed {1}")]
     MemoryRequestTooBig(usize, usize),
@@ -222,6 +226,10 @@ pub enum HyperlightError {
     #[error("Failed To Convert Return Value {0:?} to {1:?}")]
     ReturnValueConversionFailure(ReturnValue, &'static str),
 
+    /// Attempted to process a snapshot but the snapshot size does not match the current memory size
+    #[error("Snapshot Size Mismatch: Memory Size {0:?} Snapshot Size {1:?}")]
+    SnapshotSizeMismatch(usize, usize),
+
     /// Stack overflow detected in guest
     #[error("Stack overflow detected")]
     StackOverflow(),
@@ -322,7 +330,9 @@ impl HyperlightError {
             | HyperlightError::PoisonedSandbox
             | HyperlightError::ExecutionAccessViolation(_)
             | HyperlightError::StackOverflow()
-            | HyperlightError::MemoryAccessViolation(_, _, _) => true,
+            | HyperlightError::MemoryAccessViolation(_, _, _)
+            | HyperlightError::SnapshotSizeMismatch(_, _)
+            | HyperlightError::MemoryRegionSizeMismatch(_, _, _) => true,
 
             // All other errors do not poison the sandbox.
             HyperlightError::AnyhowError(_)

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,13 @@ 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<()> {
-        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)?;
+    pub(crate) fn restore_snapshot(&mut self, snapshot: &Snapshot) -> Result<()> {
+        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

@@ -37,8 +37,10 @@ use windows::core::PCSTR;
 
 #[cfg(target_os = "windows")]
 use crate::HyperlightError::MemoryAllocationFailed;
+use crate::HyperlightError::SnapshotSizeMismatch;
 #[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 +677,14 @@ 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<()> {
+        if snapshot.memory().len() != self.mem_size() {
+            return Err(SnapshotSizeMismatch(self.mem_size(), snapshot.mem_size()));
+        }
+        self.with_exclusivity(|e| e.copy_from_slice(snapshot.memory(), 0))?
+    }
 }
 
 impl SharedMemory for ExclusiveSharedMemory {