Avoid copying memory on snapshot restore, when possible

Whenever we are not running with certain features that still require
legacy writable access to the snapshot region (currently:
nanvix-unstable, gdb), map the snapshot region directly into the VM.
This lays the groundwork for sharing a single snapshot memory between
multiple VMs.

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

Author: syntactically

src/hyperlight_host/build.rs

@@ -105,6 +105,10 @@ fn main() -> Result<()> {
         crashdump: { all(feature = "crashdump", target_arch = "x86_64") },
         // print_debug feature is aliased with debug_assertions to make it only available in debug-builds.
         print_debug: { all(feature = "print_debug", debug_assertions) },
+        // the nanvix-unstable and gdb features both (only
+        // temporarily!) need to use writable/un-shared snapshot
+        // memories, and so can't share
+        unshared_snapshot_mem: { any(feature = "nanvix-unstable", feature = "gdb") },
     }
 
     #[cfg(feature = "build-metadata")]

src/hyperlight_host/src/error.rs

@@ -240,10 +240,6 @@ 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),
-
     /// Tried to restore snapshot to a sandbox that is not the same as the one the snapshot was taken from
     #[error("Snapshot was taken from a different sandbox")]
     SnapshotSandboxMismatch,
@@ -335,7 +331,6 @@ impl HyperlightError {
             | HyperlightError::PoisonedSandbox
             | HyperlightError::ExecutionAccessViolation(_)
             | HyperlightError::MemoryAccessViolation(_, _, _)
-            | HyperlightError::SnapshotSizeMismatch(_, _)
             | HyperlightError::MemoryRegionSizeMismatch(_, _, _)
             // HyperlightVmError::Restore is already handled manually in restore(), but we mark it
             // as poisoning here too for defense in depth.

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

@@ -21,7 +21,7 @@ mod x86_64_target;
 use std::io::{self, ErrorKind};
 use std::net::TcpListener;
 use std::sync::{Arc, Mutex};
-use std::{slice, thread};
+use std::thread;
 
 use crossbeam_channel::{Receiver, Sender, TryRecvError};
 use event_loop::event_loop_thread;
@@ -36,10 +36,10 @@ use super::regs::CommonRegisters;
 use crate::HyperlightError;
 use crate::hypervisor::regs::CommonFpu;
 use crate::hypervisor::virtual_machine::{HypervisorError, RegisterError, VirtualMachine};
-use crate::mem::layout::SandboxMemoryLayout;
+use crate::mem::layout::BaseGpaRegion;
 use crate::mem::memory_region::MemoryRegion;
 use crate::mem::mgr::SandboxMemoryManager;
-use crate::mem::shared_mem::{HostSharedMemory, SharedMemory};
+use crate::mem::shared_mem::HostSharedMemory;
 
 #[derive(Debug, Error)]
 pub enum GdbTargetError {
@@ -95,18 +95,11 @@ pub enum DebugMemoryAccessError {
     LockFailed(&'static str, u32, String),
     #[error("Failed to translate guest address {0:#x}")]
     TranslateGuestAddress(u64),
+    #[error("Failed to write to read-only region")]
+    WriteToReadOnly,
 }
 
 impl DebugMemoryAccess {
-    // TODO: There is a lot of common logic between both of these
-    // functions, as well as guest_page/access_gpa in snapshot.rs. It
-    // would be nice to factor that out at some point, but the
-    // snapshot versions deal with ExclusiveSharedMemory, since we
-    // never expect a guest to be running concurrent with a snapshot,
-    // and doesn't want to make unnecessary copies, since it runs over
-    // relatively large volumes of data, so it's not clear if it's
-    // terribly easy to combine them
-
     /// Reads memory from the guest's address space with a maximum length of a PAGE_SIZE
     ///
     /// # Arguments
@@ -120,74 +113,17 @@ impl DebugMemoryAccess {
         data: &mut [u8],
         gpa: u64,
     ) -> std::result::Result<(), DebugMemoryAccessError> {
-        let read_len = data.len();
-
-        let mem_offset = (gpa as usize)
-            .checked_sub(SandboxMemoryLayout::BASE_ADDRESS)
-            .ok_or_else(|| {
-                log::warn!(
-                    "gpa={:#X} causes subtract with underflow: \"gpa - BASE_ADDRESS={:#X}-{:#X}\"",
-                    gpa,
-                    gpa,
-                    SandboxMemoryLayout::BASE_ADDRESS
-                );
-                DebugMemoryAccessError::TranslateGuestAddress(gpa)
-            })?;
-
-        // First check the mapped memory regions to see if the address is within any of them
-        let mut region_found = false;
-        for reg in self.guest_mmap_regions.iter() {
-            if reg.guest_region.contains(&mem_offset) {
-                log::debug!("Found mapped region containing {:X}: {:#?}", gpa, reg);
-
-                // Region found - calculate the offset within the region
-                let region_offset = mem_offset.checked_sub(reg.guest_region.start).ok_or_else(|| {
-                    log::warn!(
-                        "Cannot calculate offset in memory region: mem_offset={:#X}, base={:#X}",
-                        mem_offset,
-                        reg.guest_region.start,
-                    );
-                    DebugMemoryAccessError::TranslateGuestAddress(mem_offset as u64)
-                })?;
-
-                let host_start_ptr = <_ as Into<usize>>::into(reg.host_region.start);
-                let bytes: &[u8] = unsafe {
-                    slice::from_raw_parts(host_start_ptr as *const u8, reg.guest_region.len())
-                };
-                data[..read_len].copy_from_slice(&bytes[region_offset..region_offset + read_len]);
-
-                region_found = true;
-                break;
-            }
-        }
-
-        if !region_found {
-            let mut mgr = self
-                .dbg_mem_access_fn
-                .try_lock()
-                .map_err(|e| DebugMemoryAccessError::LockFailed(file!(), line!(), e.to_string()))?;
-            let scratch_base =
-                hyperlight_common::layout::scratch_base_gpa(mgr.scratch_mem.mem_size());
-            let (mem, offset, name): (&mut HostSharedMemory, _, _) = if gpa >= scratch_base {
-                (
-                    &mut mgr.scratch_mem,
-                    (gpa - scratch_base) as usize,
-                    "scratch",
-                )
-            } else {
-                (&mut mgr.shared_mem, mem_offset, "snapshot")
-            };
-            log::debug!(
-                "No mapped region found containing {:X}. Trying {} memory at offset {:X} ...",
-                gpa,
-                name,
-                offset
-            );
-            mem.copy_to_slice(&mut data[..read_len], offset)
-                .map_err(|e| DebugMemoryAccessError::CopyFailed(Box::new(e)))?;
-        }
-
-        Ok(())
+        let mgr = self
+            .dbg_mem_access_fn
+            .try_lock()
+            .map_err(|e| DebugMemoryAccessError::LockFailed(file!(), line!(), e.to_string()))?;
+
+        mgr.layout
+            .resolve_gpa(gpa, &self.guest_mmap_regions)
+            .ok_or(DebugMemoryAccessError::TranslateGuestAddress(gpa))?
+            .with_memories(&mgr.shared_mem, &mgr.scratch_mem)
+            .copy_to_slice(data)
+            .map_err(|e| DebugMemoryAccessError::CopyFailed(Box::new(e)))
     }
 
     /// Writes memory from the guest's address space with a maximum length of a PAGE_SIZE
@@ -203,74 +139,30 @@ impl DebugMemoryAccess {
         data: &[u8],
         gpa: u64,
     ) -> std::result::Result<(), DebugMemoryAccessError> {
-        let write_len = data.len();
-
-        let mem_offset = (gpa as usize)
-            .checked_sub(SandboxMemoryLayout::BASE_ADDRESS)
-            .ok_or_else(|| {
-                log::warn!(
-                    "gpa={:#X} causes subtract with underflow: \"gpa - BASE_ADDRESS={:#X}-{:#X}\"",
-                    gpa,
-                    gpa,
-                    SandboxMemoryLayout::BASE_ADDRESS
-                );
-                DebugMemoryAccessError::TranslateGuestAddress(gpa)
-            })?;
-
-        // First check the mapped memory regions to see if the address is within any of them
-        let mut region_found = false;
-        for reg in self.guest_mmap_regions.iter() {
-            if reg.guest_region.contains(&mem_offset) {
-                log::debug!("Found mapped region containing {:X}: {:#?}", gpa, reg);
-
-                // Region found - calculate the offset within the region
-                let region_offset = mem_offset.checked_sub(reg.guest_region.start).ok_or_else(|| {
-                    log::warn!(
-                        "Cannot calculate offset in memory region: mem_offset={:#X}, base={:#X}",
-                        mem_offset,
-                        reg.guest_region.start,
-                    );
-                    DebugMemoryAccessError::TranslateGuestAddress(mem_offset as u64)
-                })?;
-
-                let host_start_ptr = <_ as Into<usize>>::into(reg.host_region.start);
-                let bytes: &mut [u8] = unsafe {
-                    slice::from_raw_parts_mut(host_start_ptr as *mut u8, reg.guest_region.len())
-                };
-                bytes[region_offset..region_offset + write_len].copy_from_slice(&data[..write_len]);
-
-                region_found = true;
-                break;
-            }
-        }
-
-        if !region_found {
-            let mut mgr = self
-                .dbg_mem_access_fn
-                .try_lock()
-                .map_err(|e| DebugMemoryAccessError::LockFailed(file!(), line!(), e.to_string()))?;
-            let scratch_base =
-                hyperlight_common::layout::scratch_base_gpa(mgr.scratch_mem.mem_size());
-            let (mem, offset, name): (&mut HostSharedMemory, _, _) = if gpa >= scratch_base {
-                (
-                    &mut mgr.scratch_mem,
-                    (gpa - scratch_base) as usize,
-                    "scratch",
-                )
-            } else {
-                (&mut mgr.shared_mem, mem_offset, "snapshot")
-            };
-            log::debug!(
-                "No mapped region found containing {:X}. Trying {} memory at offset {:X} ...",
-                gpa,
-                name,
-                offset
-            );
-            mem.copy_from_slice(&data[..write_len], offset)
-                .map_err(|e| DebugMemoryAccessError::CopyFailed(Box::new(e)))?;
+        let mgr = self
+            .dbg_mem_access_fn
+            .try_lock()
+            .map_err(|e| DebugMemoryAccessError::LockFailed(file!(), line!(), e.to_string()))?;
+
+        let resolved = mgr
+            .layout
+            .resolve_gpa(gpa, &self.guest_mmap_regions)
+            .ok_or(DebugMemoryAccessError::TranslateGuestAddress(gpa))?;
+
+        // We can only safely write (without causing UB in the host
+        // process) if the address is in the scratch region
+        match resolved.base {
+            #[cfg(unshared_snapshot_mem)]
+            BaseGpaRegion::Snapshot(()) => mgr
+                .shared_mem
+                .copy_from_slice(data, resolved.offset)
+                .map_err(|e| DebugMemoryAccessError::CopyFailed(Box::new(e))),
+            BaseGpaRegion::Scratch(()) => mgr
+                .scratch_mem
+                .copy_from_slice(data, resolved.offset)
+                .map_err(|e| DebugMemoryAccessError::CopyFailed(Box::new(e))),
+            _ => Err(DebugMemoryAccessError::WriteToReadOnly),
         }
-
-        Ok(())
     }
 }
 
@@ -490,6 +382,7 @@ mod tests {
         use hyperlight_testing::dummy_guest_as_string;
 
         use super::*;
+        use crate::mem::layout::SandboxMemoryLayout;
         use crate::mem::memory_region::{MemoryRegionFlags, MemoryRegionType};
         use crate::sandbox::UninitializedSandbox;
         use crate::sandbox::uninitialized::GuestBinary;

src/hyperlight_host/src/hypervisor/hyperlight_vm/mod.rs

@@ -47,7 +47,7 @@ use crate::hypervisor::virtual_machine::{
 };
 use crate::hypervisor::{InterruptHandle, InterruptHandleImpl};
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
-use crate::mem::mgr::SandboxMemoryManager;
+use crate::mem::mgr::{SandboxMemoryManager, SnapshotSharedMemory};
 use crate::mem::shared_mem::{GuestSharedMemory, HostSharedMemory, SharedMemory};
 use crate::metrics::{METRIC_ERRONEOUS_VCPU_KICKS, METRIC_GUEST_CANCELLATION};
 use crate::sandbox::host_funcs::FunctionRegistry;
@@ -375,7 +375,7 @@ pub(crate) struct HyperlightVm {
     pub(super) snapshot_slot: u32,
     // The current snapshot region, used to keep it alive as long as
     // it is used & when unmapping
-    pub(super) snapshot_memory: Option<GuestSharedMemory>,
+    pub(super) snapshot_memory: Option<SnapshotSharedMemory<GuestSharedMemory>>,
     pub(super) scratch_slot: u32, // The slot number used for the scratch region
     // The current scratch region, used to keep it alive as long as it
     // is used & when unmapping
@@ -460,7 +460,7 @@ impl HyperlightVm {
     /// Update the snapshot mapping to point to a new GuestSharedMemory
     pub(crate) fn update_snapshot_mapping(
         &mut self,
-        snapshot: GuestSharedMemory,
+        snapshot: SnapshotSharedMemory<GuestSharedMemory>,
     ) -> Result<(), UpdateRegionError> {
         let guest_base = crate::mem::layout::SandboxMemoryLayout::BASE_ADDRESS as u64;
         let rgn = snapshot.mapping_at(guest_base, MemoryRegionType::Snapshot);

src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs

@@ -74,7 +74,7 @@ impl HyperlightVm {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
-        snapshot_mem: GuestSharedMemory,
+        snapshot_mem: SnapshotSharedMemory<GuestSharedMemory>,
         scratch_mem: GuestSharedMemory,
         _pml4_addr: u64,
         entrypoint: NextAction,
@@ -885,7 +885,7 @@ mod tests {
     use crate::mem::memory_region::{GuestMemoryRegion, MemoryRegionFlags};
     use crate::mem::mgr::{GuestPageTableBuffer, SandboxMemoryManager};
     use crate::mem::ptr::RawPtr;
-    use crate::mem::shared_mem::ExclusiveSharedMemory;
+    use crate::mem::shared_mem::{ExclusiveSharedMemory, ReadonlySharedMemory};
     use crate::sandbox::SandboxConfiguration;
     use crate::sandbox::host_funcs::FunctionRegistry;
     #[cfg(any(crashdump, gdb))]
@@ -1479,20 +1479,22 @@ mod tests {
         layout.set_pt_size(pt_bytes.len()).unwrap();
 
         let mem_size = layout.get_memory_size().unwrap();
-        let mut eshm = ExclusiveSharedMemory::new(mem_size).unwrap();
+        let mut snapshot_contents = vec![0u8; mem_size];
         let snapshot_pt_start = mem_size - layout.get_pt_size();
-        eshm.copy_from_slice(&pt_bytes, snapshot_pt_start).unwrap();
-        eshm.copy_from_slice(code, layout.get_guest_code_offset())
-            .unwrap();
+        snapshot_contents[snapshot_pt_start..].copy_from_slice(&pt_bytes);
+        snapshot_contents
+            [layout.get_guest_code_offset()..layout.get_guest_code_offset() + code.len()]
+            .copy_from_slice(code);
+        layout.write_peb(&mut snapshot_contents).unwrap();
+        let ro_mem = ReadonlySharedMemory::from_bytes(&snapshot_contents).unwrap();
 
         let scratch_mem = ExclusiveSharedMemory::new(config.get_scratch_size()).unwrap();
-        let mut mem_mgr = SandboxMemoryManager::new(
+        let mem_mgr = SandboxMemoryManager::new(
             layout,
-            eshm,
+            ro_mem.to_mgr_snapshot_mem().unwrap(),
             scratch_mem,
             NextAction::Initialise(layout.get_guest_code_address() as u64),
         );
-        mem_mgr.write_memory_layout().unwrap();
 
         let (mut hshm, gshm) = mem_mgr.build().unwrap();