fixup! Avoid copying memory on snapshot restore, when possible

Address review comments

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/hyperlight_host/src/hypervisor/gdb/mod.rs

@@ -96,7 +96,7 @@ pub enum DebugMemoryAccessError {
     #[error("Failed to translate guest address {0:#x}")]
     TranslateGuestAddress(u64),
     #[error("Failed to write to read-only region")]
-    WriteToReadOnly(),
+    WriteToReadOnly,
 }
 
 impl DebugMemoryAccess {
@@ -161,7 +161,7 @@ impl DebugMemoryAccess {
                 .scratch_mem
                 .copy_from_slice(data, resolved.offset)
                 .map_err(|e| DebugMemoryAccessError::CopyFailed(Box::new(e))),
-            _ => Err(DebugMemoryAccessError::WriteToReadOnly()),
+            _ => Err(DebugMemoryAccessError::WriteToReadOnly),
         }
     }
 }

src/hyperlight_host/src/mem/layout.rs

@@ -132,6 +132,10 @@ impl<'a> BaseGpaRegion<&'a [u8], &'a [u8]> {
 }
 impl<'a> ResolvedGpa<&'a [u8], &'a [u8]> {
     pub(crate) fn as_ref<'b>(&'b self) -> &'a [u8] {
+        let base = self.base.as_ref();
+        if self.offset > base.len() {
+            return &[];
+        }
         &self.base.as_ref()[self.offset..]
     }
 }
@@ -158,7 +162,7 @@ mod coherence_hack {
 #[cfg(any(gdb, feature = "mem_profile"))]
 impl<T: coherence_hack::SharedMemoryAsRefMarker> ReadableSharedMemory for T {
     fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> {
-        let ss: &[u8] = &self.as_ref()[offset..];
+        let ss: &[u8] = self.as_ref();
         let end = offset + slice.len();
         if end > ss.len() {
             return Err(new_error!(
@@ -184,6 +188,14 @@ impl<Sn: ReadableSharedMemory, Sc: ReadableSharedMemory> ResolvedGpa<Sn, Sc> {
                 #[allow(clippy::useless_conversion)]
                 let host_region_end: usize = r.host_region.end.into();
                 let len = host_region_end - host_region_base;
+                // Safety: it's a documented invariant of MemoryRegion
+                // that the memory must remain alive as long as the
+                // sandbox is alive, and the way this code is used,
+                // the lifetimes of the snapshot and scratch memories
+                // ensure that the sandbox is still alive. This could
+                // perhaps be cleaned up/improved/made harder to
+                // misuse significantly, but it would require a much
+                // larger rework.
                 let ss = std::slice::from_raw_parts(host_region_base as *const u8, len);
                 let end = self.offset + slice.len();
                 if end > ss.len() {
@@ -230,7 +242,7 @@ pub(crate) struct SandboxMemoryLayout {
     // The size of the scratch region in physical memory; note that
     // this will appear under the top of physical memory.
     scratch_size: usize,
-    // The size of the scratch region in physical memory; note that
+    // The size of the snapshot region in physical memory; note that
     // this will appear somewhere near the base of physical memory.
     snapshot_size: usize,
 }
@@ -653,10 +665,10 @@ impl SandboxMemoryLayout {
         Ok(())
     }
 
-    /// Write the finished memory layout to `shared_mem` and return
-    /// `Ok` if successful.
+    /// Write the finished memory layout to `mem` and return `Ok` if
+    /// successful.
     ///
-    /// Note: `shared_mem` may have been modified, even if `Err` was returned
+    /// Note: `mem` may have been modified, even if `Err` was returned
     /// from this function.
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(crate) fn write_peb(&self, mem: &mut [u8]) -> Result<()> {

src/hyperlight_host/src/mem/mgr.rs

@@ -428,6 +428,13 @@ impl SandboxMemoryManager<HostSharedMemory> {
         Option<GuestSharedMemory>,
     )> {
         let gsnapshot = if *snapshot.memory() == self.shared_mem {
+            // If the snapshot memory is already the correct memory,
+            // which is readonly, don't bother with restoring it,
+            // since its contents must be the same.  Note that in the
+            // #[cfg(unshared_snapshot_mem)] case, this condition will
+            // never be true, since even immediately after a restore,
+            // self.shared_mem is a (writable) copy, not the original
+            // shared_mem.
             None
         } else {
             let new_snapshot_mem = snapshot.memory().to_mgr_snapshot_mem()?;

src/hyperlight_host/src/mem/shared_mem.rs

@@ -1998,6 +1998,16 @@ mod tests {
 pub struct ReadonlySharedMemory {
     region: Arc<HostMapping>,
 }
+// Safety: HostMapping is only non-Send/Sync (causing
+// ReadonlySharedMemory to not be automatically Send/Sync) because raw
+// pointers are not ("as a lint", as the Rust docs say). We don't want
+// to mark HostMapping Send/Sync immediately, because that could
+// socially imply that it's "safe" to use unsafe accesses from
+// multiple threads at once in more cases, including ones that don't
+// actually ensure immutability/synchronisation. Since
+// ReadonlySharedMemory can only be accessed by reading, and reading
+// concurrently from multiple threads is not racy,
+// ReadonlySharedMemory can be Send and Sync.
 unsafe impl Send for ReadonlySharedMemory {}
 unsafe impl Sync for ReadonlySharedMemory {}
 
@@ -2032,11 +2042,11 @@ impl ReadonlySharedMemory {
         guest_base: u64,
         region_type: MemoryRegionType,
     ) -> MemoryRegion {
+        #[allow(clippy::panic)]
+        // This will not ever actually panic: the only place this is
+        // called is HyperlightVm::update_snapshot_mapping, which
+        // always calls it with the Snapshot region type.
         if region_type != MemoryRegionType::Snapshot {
-            #[allow(clippy::panic)]
-            // This will not ever actually panic: the only place this
-            // is called is HyperlightVm::update_snapshot_mapping,
-            // which always calls it with the Snapshot region type.
             panic!("ReadonlySharedMemory::mapping_at should only be used for Snapshot regions");
         }
         mapping_at(

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -287,7 +287,10 @@ unsafe fn guest_page<'a>(
     let resolved = layout
         .resolve_gpa(gpa, regions)?
         .with_memories(snap, scratch);
-    Some(&resolved.as_ref()[0..PAGE_SIZE])
+    if resolved.as_ref().len() < PAGE_SIZE {
+        return None;
+    }
+    Some(&resolved.as_ref()[..PAGE_SIZE])
 }
 
 fn map_specials(pt_buf: &GuestPageTableBuffer, scratch_size: usize) {
@@ -666,7 +669,7 @@ mod tests {
         // Restore snapshot B
         mgr.restore_snapshot(&snapshot_b).unwrap();
         mgr.shared_mem
-            .with_contents(|contents| assert_eq!(&contents[0..pattern_a.len()], &pattern_b[..]))
+            .with_contents(|contents| assert_eq!(&contents[0..pattern_b.len()], &pattern_b[..]))
             .unwrap();
     }
 }