Add test for creating sandbox from snapshot from sandbox

Previously, this code path was not tested, and would fail because
`UninitializedSandbox::evolve()` unconditionally tried to run the
guest initialisation function, which has already been run in this
case.  This simply ignores the initialisation stage if the snapshot
has no `preinitialise` entrypoint.

All that `UninitializedSandbox` is really good for at this point is
registering host functions, so it might make sense to further revisit
the API.

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

Author: syntactically

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -85,7 +85,7 @@ pub(crate) struct HyperlightVm {
     #[cfg(not(gdb))]
     vm: Box<dyn VirtualMachine>,
     page_size: usize,
-    entrypoint: u64,
+    entrypoint: Option<u64>, // only present if this vm has not yet been initialised
     orig_rsp: GuestPtr,
     interrupt_handle: Arc<dyn InterruptHandleImpl>,
 
@@ -350,7 +350,7 @@ impl HyperlightVm {
         snapshot_mem: GuestSharedMemory,
         scratch_mem: GuestSharedMemory,
         _pml4_addr: u64,
-        entrypoint: u64,
+        entrypoint: Option<u64>,
         rsp: u64,
         #[cfg_attr(target_os = "windows", allow(unused_variables))] config: &SandboxConfiguration,
         #[cfg(gdb)] gdb_conn: Option<DebugCommChannel<DebugResponse, DebugMsg>>,
@@ -450,11 +450,13 @@ impl HyperlightVm {
         #[cfg(gdb)]
         if ret.gdb_conn.is_some() {
             ret.send_dbg_msg(DebugResponse::InterruptHandle(ret.interrupt_handle.clone()))?;
-            // Add breakpoint to the entry point address
+            // Add breakpoint to the entry point address, if we are going to initialise
             ret.vm.set_debug(true).map_err(VmError::Debug)?;
-            ret.vm
-                .add_hw_breakpoint(entrypoint)
-                .map_err(CreateHyperlightVmError::AddHwBreakpoint)?;
+            if let Some(entrypoint) = entrypoint {
+                ret.vm
+                    .add_hw_breakpoint(entrypoint)
+                    .map_err(CreateHyperlightVmError::AddHwBreakpoint)?;
+            }
         }
 
         Ok(ret)
@@ -474,6 +476,10 @@ impl HyperlightVm {
         guest_max_log_level: Option<LevelFilter>,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<SandboxMemoryManager<HostSharedMemory>>>,
     ) -> std::result::Result<(), InitializeError> {
+        let Some(entrypoint) = self.entrypoint else {
+            return Ok(());
+        };
+
         self.page_size = page_size as usize;
 
         let guest_max_log_level: u64 = match guest_max_log_level {
@@ -482,7 +488,7 @@ impl HyperlightVm {
         };
 
         let regs = CommonRegisters {
-            rip: self.entrypoint,
+            rip: entrypoint,
             rsp: self
                 .orig_rsp
                 .absolute()
@@ -740,8 +746,12 @@ impl HyperlightVm {
                 #[cfg(gdb)]
                 Ok(VmExit::Debug { dr6, exception }) => {
                     // Handle debug event (breakpoints)
-                    let stop_reason =
-                        arch::vcpu_stop_reason(self.vm.as_mut(), dr6, self.entrypoint, exception)?;
+                    let stop_reason = arch::vcpu_stop_reason(
+                        self.vm.as_mut(),
+                        dr6,
+                        self.entrypoint.unwrap_or(0),
+                        exception,
+                    )?;
                     if let Err(e) = self.handle_debug(dbg_mem_access_fn.clone(), stop_reason) {
                         break Err(e.into());
                     }
@@ -1088,7 +1098,7 @@ impl HyperlightVm {
                 regions,
                 regs,
                 xsave.to_vec(),
-                self.entrypoint,
+                self.entrypoint.unwrap_or(0),
                 self.rt_cfg.binary_path.clone(),
                 filename,
             )))

src/hyperlight_host/src/sandbox/uninitialized.rs

@@ -1258,5 +1258,33 @@ mod tests {
 
             let _evolved: MultiUseSandbox = sandbox.evolve().expect("Failed to evolve sandbox");
         }
+
+        // Test 9: Create snapshot from existing sandbox
+        {
+            let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None);
+            let orig_snapshot = Arc::new(
+                Snapshot::from_env(env, Default::default())
+                    .expect("Failed to create snapshot with default config")
+            );
+            let orig_sandbox = UninitializedSandbox::from_snapshot(
+                orig_snapshot,
+                None,
+                #[cfg(crashdump)]
+                Some(binary_path.clone()),
+            )
+            .expect("Failed to create orig_sandbox");
+            let mut initialized_sandbox = orig_sandbox
+                .evolve()
+                .expect("Failed to evolve orig_sandbox");
+            let  new_snapshot = initialized_sandbox.snapshot().expect("Failed to create new_snapshot");
+            let new_sandbox = UninitializedSandbox::from_snapshot(
+                new_snapshot,
+                None,
+                #[cfg(crashdump)]
+                Some(binary_path.clone()),
+            )
+            .expect("Failed to create new_sandbox");
+            let _evolved = new_sandbox.evolve().expect("Failed to evolve new_sandbox");
+        }
     }
 }

src/hyperlight_host/src/sandbox/uninitialized_evolve.rs

@@ -119,11 +119,12 @@ pub(crate) fn set_up_hypervisor_partition(
     };
     let entrypoint_ptr = mgr
         .entrypoint_offset
-        .ok_or_else(|| new_error!("Entrypoint offset is None"))
-        .and_then(|x| {
+        .map(|x| {
             let entrypoint_total_offset = mgr.load_addr.clone() + x;
             GuestPtr::try_from(entrypoint_total_offset)
-        })?;
+                .and_then(|x| x.absolute())
+        })
+        .transpose()?;
 
     // Create gdb thread if gdb is enabled and the configuration is provided
     #[cfg(gdb)]
@@ -153,7 +154,7 @@ pub(crate) fn set_up_hypervisor_partition(
         mgr.shared_mem,
         mgr.scratch_mem,
         pml4_ptr.absolute()?,
-        entrypoint_ptr.absolute()?,
+        entrypoint_ptr,
         rsp_ptr.absolute()?,
         config,
         #[cfg(gdb)]