Remove Asserts on snapshot/memory/region size checks

Make asserts errors and poision sandbox
Fix hasher inputs

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

Author: simongdavies

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

@@ -291,13 +291,6 @@ where
 
     /// This function restores a memory snapshot from a given snapshot.
     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()
-            ));
-        }
         self.shared_mem.restore_from_snapshot(snapshot)?;
         Ok(())
     }

src/hyperlight_host/src/mem/shared_mem.rs

@@ -37,6 +37,7 @@ 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;
@@ -679,7 +680,9 @@ pub trait SharedMemory {
 
     /// 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());
+        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))?
     }
 }

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -16,6 +16,7 @@ limitations under the License.
 
 use tracing::{Span, instrument};
 
+use crate::HyperlightError::MemoryRegionSizeMismatch;
 use crate::Result;
 use crate::mem::memory_region::MemoryRegion;
 use crate::mem::shared_mem::SharedMemory;
@@ -39,19 +40,21 @@ pub struct Snapshot {
     hash: [u8; 32],
 }
 
-fn hash(memory: &[u8], regions: &[MemoryRegion]) -> [u8; 32] {
+fn hash(memory: &[u8], regions: &[MemoryRegion]) -> Result<[u8; 32]> {
     let mut hasher = blake3::Hasher::new();
     hasher.update(memory);
     for rgn in regions {
         hasher.update(&usize::to_le_bytes(rgn.guest_region.start));
         let guest_len = rgn.guest_region.end - rgn.guest_region.start;
-        hasher.update(&usize::to_le_bytes(rgn.guest_region.start));
+        hasher.update(&usize::to_le_bytes(rgn.host_region.start));
         let host_len = rgn.host_region.end - rgn.host_region.start;
-        assert!(guest_len == host_len);
+        if guest_len != host_len {
+            return Err(MemoryRegionSizeMismatch(host_len, guest_len, rgn.clone()));
+        }
         hasher.update(&usize::to_le_bytes(guest_len));
         hasher.update(&u32::to_le_bytes(rgn.flags.bits()));
     }
-    hasher.finalize().into()
+    Ok(hasher.finalize().into())
 }
 
 impl Snapshot {
@@ -65,7 +68,7 @@ impl Snapshot {
     ) -> Result<Self> {
         // TODO: Track dirty pages instead of copying entire memory
         let memory = shared_mem.with_exclusivity(|e| e.copy_all_to_vec())??;
-        let hash = hash(&memory, &regions);
+        let hash = hash(&memory, &regions)?;
         Ok(Self {
             sandbox_id,
             memory,