fixup! Modify guest physical page allocator to allocate from the scratch region

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

Author: syntactically

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

@@ -497,7 +497,7 @@ mod tests {
                 .inspect_err(|_| unsafe {
                     libc::munmap(mapped_mem, size);
                 })?;
-            let (mem_mgr, _) = sandbox.mgr.build();
+            let (mem_mgr, _) = sandbox.mgr.build()?;
 
             // Create the memory access struct
             let mem_access = DebugMemoryAccess {

src/hyperlight_host/src/hypervisor/mod.rs

@@ -525,7 +525,7 @@ pub(crate) mod tests {
         let rt_cfg: SandboxRuntimeConfig = Default::default();
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
-        let (mut mem_mgr, gshm) = sandbox.mgr.build();
+        let (mut mem_mgr, gshm) = sandbox.mgr.build().unwrap();
         let mut vm = set_up_hypervisor_partition(
             gshm,
             &config,

src/hyperlight_host/src/mem/mgr.rs

@@ -235,12 +235,21 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
     }
 
     /// Wraps ExclusiveSharedMemory::build
+    // Morally, this should not have to be a Result: this operation is
+    // infallible. The source of the Result is
+    // update_scratch_bookkeeping(), which calls functions that can
+    // fail due to bounds checks (which are statically known to be ok
+    // in this situation) or due to failing to take the scratch shared
+    // memory lock, but the scratch shared memory is built in this
+    // function, its lock does not escape before the end of the
+    // function, and the lock is taken by no other code path, so we
+    // know it is not contended.
     pub fn build(
         self,
-    ) -> (
+    ) -> Result<(
         SandboxMemoryManager<HostSharedMemory>,
         SandboxMemoryManager<GuestSharedMemory>,
-    ) {
+    )> {
         let (hshm, gshm) = self.shared_mem.build();
         let (hscratch, gscratch) = self.scratch_mem.build();
         let mut host_mgr = SandboxMemoryManager {
@@ -265,8 +274,8 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
         };
         host_mgr.update_scratch_bookkeeping(
             (SandboxMemoryLayout::BASE_ADDRESS + self.layout.get_pt_offset()) as u64,
-        );
-        (host_mgr, guest_mgr)
+        )?;
+        Ok((host_mgr, guest_mgr))
     }
 }
 
@@ -424,39 +433,30 @@ impl SandboxMemoryManager<HostSharedMemory> {
 
             Some(gscratch)
         };
-        self.update_scratch_bookkeeping(snapshot.root_pt_gpa());
+        self.update_scratch_bookkeeping(snapshot.root_pt_gpa())?;
         Ok((gsnapshot, gscratch))
     }
 
-    fn update_scratch_bookkeeping(&mut self, snapshot_pt_base_gpa: u64) {
+    fn update_scratch_bookkeeping(&mut self, snapshot_pt_base_gpa: u64) -> Result<()> {
         let scratch_size = self.scratch_mem.mem_size();
+
         let size_offset =
             scratch_size - hyperlight_common::layout::SCRATCH_TOP_SIZE_OFFSET as usize;
-        // The only way that write can fail is if the offset is
-        // outside of the memory, which would be sufficiently much of
-        // an invariant violation that panicking is probably
-        // sensible...
-        #[allow(clippy::unwrap_used)]
         self.scratch_mem
-            .write::<u64>(size_offset, scratch_size as u64)
-            .unwrap();
+            .write::<u64>(size_offset, scratch_size as u64)?;
+
         let alloc_offset =
             scratch_size - hyperlight_common::layout::SCRATCH_TOP_ALLOCATOR_OFFSET as usize;
-        // See above comment about unwrap() on write
-        #[allow(clippy::unwrap_used)]
-        self.scratch_mem
-            .write::<u64>(
-                alloc_offset,
-                hyperlight_common::layout::scratch_base_gpa(scratch_size),
-            )
-            .unwrap();
+        self.scratch_mem.write::<u64>(
+            alloc_offset,
+            hyperlight_common::layout::scratch_base_gpa(scratch_size),
+        )?;
+
         let snapshot_pt_base_gpa_offset = scratch_size
             - hyperlight_common::layout::SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET as usize;
-        // See above comment about unwrap() on write
-        #[allow(clippy::unwrap_used)]
         self.scratch_mem
-            .write::<u64>(snapshot_pt_base_gpa_offset, snapshot_pt_base_gpa)
-            .unwrap();
+            .write::<u64>(snapshot_pt_base_gpa_offset, snapshot_pt_base_gpa)?;
+        Ok(())
     }
 }
 

src/hyperlight_host/src/sandbox/outb.rs

@@ -287,7 +287,7 @@ mod tests {
             layout
                 .write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
                 .unwrap();
-            let (hmgr, _) = mgr.build();
+            let (hmgr, _) = mgr.build().unwrap();
             hmgr
         };
         {
@@ -399,7 +399,7 @@ mod tests {
                 layout
                     .write(shared_mem, SandboxMemoryLayout::BASE_ADDRESS, mem_size)
                     .unwrap();
-                let (hmgr, _) = mgr.build();
+                let (hmgr, _) = mgr.build().unwrap();
                 hmgr
             };
 

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -613,7 +613,7 @@ mod tests {
             None,
             [0u8; 16],
         );
-        let (mgr, _) = mgr.build();
+        let (mgr, _) = mgr.build().unwrap();
         (mgr, pt_base as u64)
     }
 

src/hyperlight_host/src/sandbox/uninitialized_evolve.rs

@@ -39,7 +39,7 @@ use crate::{MultiUseSandbox, Result, UninitializedSandbox, new_error};
 
 #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
 pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result<MultiUseSandbox> {
-    let (mut hshm, gshm) = u_sbox.mgr.build();
+    let (mut hshm, gshm) = u_sbox.mgr.build()?;
     let mut vm = set_up_hypervisor_partition(
         gshm,
         &u_sbox.config,