fixup! Add scratch region

Pass the scratch region to HyperlightVm::new, allowing it to set up
the scratch region when the sandbox is created for the first time,
before a snapshot restore method is called. This turns out to be
slightly more convenient.

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

Author: syntactically

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -275,6 +275,8 @@ pub enum CreateHyperlightVmError {
     SendDbgMsg(#[from] SendDbgMsgError),
     #[error("VM operation error: {0}")]
     Vm(#[from] VmError),
+    #[error("Set scratch error: {0}")]
+    UpdateScratch(#[from] UpdateScratchError),
 }
 
 /// Errors that can occur during debug exit handling
@@ -334,6 +336,7 @@ impl HyperlightVm {
     #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
         mem_regions: Vec<MemoryRegion>,
+        scratch_mem: GuestSharedMemory,
         _pml4_addr: u64,
         entrypoint: u64,
         rsp: u64,
@@ -436,6 +439,8 @@ impl HyperlightVm {
             rt_cfg,
         };
 
+        ret.update_scratch_mapping(scratch_mem)?;
+
         // Send the interrupt handle to the GDB thread if debugging is enabled
         // This is used to allow the GDB thread to stop the vCPU
         #[cfg(gdb)]

src/hyperlight_host/src/hypervisor/mod.rs

@@ -525,9 +525,9 @@ pub(crate) mod tests {
         let rt_cfg: SandboxRuntimeConfig = Default::default();
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
-        let (mut mem_mgr, mut gshm) = sandbox.mgr.build();
+        let (mut mem_mgr, gshm) = sandbox.mgr.build();
         let mut vm = set_up_hypervisor_partition(
-            &mut gshm,
+            gshm,
             &config,
             #[cfg(any(crashdump, gdb))]
             &rt_cfg,

src/hyperlight_host/src/sandbox/uninitialized_evolve.rs

@@ -39,9 +39,9 @@ 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, mut gshm) = u_sbox.mgr.build();
+    let (mut hshm, gshm) = u_sbox.mgr.build();
     let mut vm = set_up_hypervisor_partition(
-        &mut gshm,
+        gshm,
         &u_sbox.config,
         #[cfg(any(crashdump, gdb))]
         &u_sbox.rt_cfg,
@@ -53,7 +53,7 @@ pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result<Mult
         rng.random::<u64>()
     };
     let peb_addr = {
-        let peb_u64 = u64::try_from(gshm.layout.peb_address)?;
+        let peb_u64 = u64::try_from(hshm.layout.peb_address)?;
         RawPtr::from(peb_u64)
     };
 
@@ -98,7 +98,7 @@ pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result<Mult
 }
 
 pub(crate) fn set_up_hypervisor_partition(
-    mgr: &mut SandboxMemoryManager<GuestSharedMemory>,
+    mgr: SandboxMemoryManager<GuestSharedMemory>,
     #[cfg_attr(target_os = "windows", allow(unused_variables))] config: &SandboxConfiguration,
     #[cfg(any(crashdump, gdb))] rt_cfg: &SandboxRuntimeConfig,
     _load_info: LoadInfo,
@@ -153,6 +153,7 @@ pub(crate) fn set_up_hypervisor_partition(
 
     Ok(HyperlightVm::new(
         regions,
+        mgr.scratch_mem,
         pml4_ptr.absolute()?,
         entrypoint_ptr.absolute()?,
         rsp_ptr.absolute()?,