Move PEB setup to snapshot creation

For historical reasons, PEB setup still happened when an uninitialised
sandbox was created, instead of happening at snapshot time, which is
far more sensible. This removes the need for the snapshot memory to be
writable during initialisation.

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

Author: syntactically

src/hyperlight_host/src/error.rs

@@ -104,10 +104,6 @@ pub enum HyperlightError {
     #[error("Unsupported type: {0}")]
     GuestInterfaceUnsupportedType(String),
 
-    /// The guest offset is invalid.
-    #[error("The guest offset {0} is invalid.")]
-    GuestOffsetIsInvalid(usize),
-
     /// The guest binary was built with a different hyperlight-guest-bin version than the host expects.
     /// Hyperlight currently provides no backwards compatibility guarantees for guest binaries,
     /// so the guest and host versions must match exactly. This might change in the future.
@@ -369,7 +365,6 @@ impl HyperlightError {
             | HyperlightError::GuestExecutionHungOnHostFunctionCall()
             | HyperlightError::GuestFunctionCallAlreadyInProgress()
             | HyperlightError::GuestInterfaceUnsupportedType(_)
-            | HyperlightError::GuestOffsetIsInvalid(_)
             | HyperlightError::HostFunctionNotFound(_)
             | HyperlightError::HyperlightVmError(HyperlightVmError::Create(_))
             | HyperlightError::HyperlightVmError(HyperlightVmError::Initialize(_))

src/hyperlight_host/src/mem/layout.rs

@@ -72,9 +72,7 @@ use super::memory_region::{
     MemoryRegionVecBuilder,
 };
 use super::shared_mem::{ExclusiveSharedMemory, SharedMemory};
-use crate::error::HyperlightError::{
-    GuestOffsetIsInvalid, MemoryRequestTooBig, MemoryRequestTooSmall,
-};
+use crate::error::HyperlightError::{MemoryRequestTooBig, MemoryRequestTooSmall};
 use crate::sandbox::SandboxConfiguration;
 use crate::{Result, new_error};
 
@@ -584,68 +582,70 @@ impl SandboxMemoryLayout {
     /// Note: `shared_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(
-        &self,
-        shared_mem: &mut ExclusiveSharedMemory,
-        guest_offset: usize,
-        //TODO: Unused remove
-        _size: usize,
-    ) -> Result<()> {
+    pub(crate) fn write_peb(&self, mem: &mut [u8]) -> Result<()> {
+        let guest_offset = SandboxMemoryLayout::BASE_ADDRESS;
+
+        fn write_u64(mem: &mut [u8], offset: usize, value: u64) -> Result<()> {
+            if offset + 8 > mem.len() {
+                return Err(new_error!(
+                    "Cannot write to offset {} in slice of len {}",
+                    offset,
+                    mem.len()
+                ));
+            }
+            mem[offset..offset + 8].copy_from_slice(&u64::to_ne_bytes(value));
+            Ok(())
+        }
+
         macro_rules! get_address {
             ($something:ident) => {
                 u64::try_from(guest_offset + self.$something)?
             };
         }
 
-        if guest_offset != SandboxMemoryLayout::BASE_ADDRESS
-            && guest_offset != shared_mem.base_addr()
-        {
-            return Err(GuestOffsetIsInvalid(guest_offset));
-        }
-
         // Start of setting up the PEB. The following are in the order of the PEB fields
 
-        // Skip guest_dispatch_function_ptr_offset because it is set by the guest
-
-        // Skip code, is set when loading binary
-        // skip outb and outb context, is set when running in_proc
-
         // Set up input buffer pointer
-        shared_mem.write_u64(
+        write_u64(
+            mem,
             self.get_input_data_size_offset(),
             self.sandbox_memory_config
                 .get_input_data_size()
                 .try_into()?,
         )?;
-        shared_mem.write_u64(
+        write_u64(
+            mem,
             self.get_input_data_pointer_offset(),
             self.get_input_data_buffer_gva(),
         )?;
 
         // Set up output buffer pointer
-        shared_mem.write_u64(
+        write_u64(
+            mem,
             self.get_output_data_size_offset(),
             self.sandbox_memory_config
                 .get_output_data_size()
                 .try_into()?,
         )?;
-        shared_mem.write_u64(
+        write_u64(
+            mem,
             self.get_output_data_pointer_offset(),
             self.get_output_data_buffer_gva(),
         )?;
 
         // Set up init data pointer
-        shared_mem.write_u64(
+        write_u64(
+            mem,
             self.get_init_data_size_offset(),
             (self.get_unaligned_memory_size() - self.init_data_offset).try_into()?,
         )?;
         let addr = get_address!(init_data_offset);
-        shared_mem.write_u64(self.get_init_data_pointer_offset(), addr)?;
+        write_u64(mem, self.get_init_data_pointer_offset(), addr)?;
 
         // Set up heap buffer pointer
         let addr = get_address!(guest_heap_buffer_offset);
-        shared_mem.write_u64(self.get_heap_size_offset(), self.heap_size.try_into()?)?;
-        shared_mem.write_u64(self.get_heap_pointer_offset(), addr)?;
+        write_u64(mem, self.get_heap_size_offset(), self.heap_size.try_into()?)?;
+        write_u64(mem, self.get_heap_pointer_offset(), addr)?;
 
         // Set up the file_mappings descriptor in the PEB.
         // - The `size` field holds the number of valid FileMappingInfo

src/hyperlight_host/src/mem/mgr.rs

@@ -303,17 +303,6 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
         Ok(Self::new(layout, shared_mem, scratch_mem, entrypoint))
     }
 
-    /// Write memory layout
-    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
-    pub(crate) fn write_memory_layout(&mut self) -> Result<()> {
-        let mem_size = self.shared_mem.mem_size();
-        self.layout.write(
-            &mut self.shared_mem,
-            SandboxMemoryLayout::BASE_ADDRESS,
-            mem_size,
-        )
-    }
-
     /// Wraps ExclusiveSharedMemory::build
     // Morally, this should not have to be a Result: this operation is
     // infallible. The source of the Result is

src/hyperlight_host/src/sandbox/outb.rs

@@ -274,13 +274,7 @@ mod tests {
         let new_mgr = || {
             let bin = GuestBinary::FilePath(simple_guest_as_string().unwrap());
             let snapshot = crate::sandbox::snapshot::Snapshot::from_env(bin, sandbox_cfg).unwrap();
-            let mut mgr = SandboxMemoryManager::from_snapshot(&snapshot).unwrap();
-            let mem_size = mgr.get_shared_mem_mut().mem_size();
-            let layout = mgr.layout;
-            let shared_mem = mgr.get_shared_mem_mut();
-            layout
-                .write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
-                .unwrap();
+            let mgr = SandboxMemoryManager::from_snapshot(&snapshot).unwrap();
             let (hmgr, _) = mgr.build().unwrap();
             hmgr
         };
@@ -386,13 +380,7 @@ mod tests {
                 let bin = GuestBinary::FilePath(simple_guest_as_string().unwrap());
                 let snapshot =
                     crate::sandbox::snapshot::Snapshot::from_env(bin, sandbox_cfg).unwrap();
-                let mut mgr = SandboxMemoryManager::from_snapshot(&snapshot).unwrap();
-                let mem_size = mgr.get_shared_mem_mut().mem_size();
-                let layout = mgr.layout;
-                let shared_mem = mgr.get_shared_mem_mut();
-                layout
-                    .write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
-                    .unwrap();
+                let mgr = SandboxMemoryManager::from_snapshot(&snapshot).unwrap();
                 let (hmgr, _) = mgr.build().unwrap();
                 hmgr
             };

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -382,6 +382,8 @@ impl Snapshot {
             &mut memory[layout.get_guest_code_offset()..],
         )?;
 
+        layout.write_peb(&mut memory)?;
+
         blob.map(|x| layout.write_init_data(&mut memory, x.data))
             .transpose()?;
 

src/hyperlight_host/src/sandbox/uninitialized.rs

@@ -362,11 +362,9 @@ impl UninitializedSandbox {
             }
         };
 
-        let mut mem_mgr_wrapper =
+        let mem_mgr_wrapper =
             SandboxMemoryManager::<ExclusiveSharedMemory>::from_snapshot(snapshot.as_ref())?;
 
-        mem_mgr_wrapper.write_memory_layout()?;
-
         let host_funcs = Arc::new(Mutex::new(FunctionRegistry::default()));
 
         let mut sandbox = Self {