Add Sandbox::recreate to recover unrecoverable sandboxes

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 * `Sandbox::snapshot_with_options` takes a snapshot with `SnapshotOptions`. `SnapshotOptions::scratch_size` gives the sandbox, and later restores of the snapshot, a different scratch size. A guest can boot with the scratch its initialisation needs and run guest function calls from a snapshot with only the scratch they need, so each restore costs less. It can also give a guest more scratch than it booted with.
+* Add `Sandbox::recreate()`, which rebuilds a sandbox's VM and memory from a snapshot. It recovers unrecoverable sandboxes and keeps configuration, host functions, and interrupt handles by @ludfjig in https://github.com/hyperlight-dev/hyperlight/pull/1751
 
 ### Changed
 * Core dumps no longer record the guest binary path. The process name and command line are always `<unknown>`, as they already were for sandboxes created from snapshots by @ludfjig in https://github.com/hyperlight-dev/hyperlight/pull/1975

src/hyperlight_host/src/error.rs

@@ -193,8 +193,8 @@ pub enum HyperlightError {
     #[error("The sandbox was poisoned")]
     PoisonedSandbox,
 
-    /// The sandbox cannot safely perform further operations and must be discarded.
-    #[error("The sandbox is unrecoverable and must be discarded")]
+    /// The sandbox cannot safely perform further operations until it is recreated.
+    #[error("The sandbox is unrecoverable and must be recreated or discarded")]
     UnrecoverableSandbox,
 
     /// Raw pointer is less than base address

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

@@ -600,6 +600,21 @@ impl HyperlightVm {
         self.interrupt_handle.clone()
     }
 
+    /// Drops `old` and moves its interrupt handle to this VM, so handles
+    /// obtained from `old` control this VM.
+    pub(crate) fn adopt_interrupt_handle(&mut self, mut old: HyperlightVm) {
+        // Redirect kills before `old` tears down its partition.
+        #[cfg(target_os = "windows")]
+        old.interrupt_handle
+            .set_partition(self.vm.partition_handle());
+        // `old` marks this VM's unused handle dropped, so the adopted handle
+        // never reports dropped.
+        std::mem::swap(&mut self.interrupt_handle, &mut old.interrupt_handle);
+        #[cfg(hvf)]
+        self.vm.set_interrupt_handle(self.interrupt_handle.clone());
+        drop(old);
+    }
+
     pub(crate) fn clear_cancel(&self) {
         self.interrupt_handle.state().clear_cancel();
     }

src/hyperlight_host/src/hypervisor/hyperlight_vm/test_support.rs

@@ -277,6 +277,14 @@ impl VirtualMachine for FaultInjectingVirtualMachine {
     fn partition_handle(&self) -> windows::Win32::System::Hypervisor::WHV_PARTITION_HANDLE {
         self.inner().partition_handle()
     }
+
+    #[cfg(hvf)]
+    fn set_interrupt_handle(
+        &mut self,
+        handle: std::sync::Arc<dyn crate::hypervisor::InterruptHandleImpl>,
+    ) {
+        self.inner_mut().set_interrupt_handle(handle);
+    }
 }
 
 impl HyperlightVm {

src/hyperlight_host/src/hypervisor/mod.rs

@@ -129,6 +129,10 @@ pub(crate) trait InterruptHandleImpl: InterruptHandle {
 
     /// Mark the handle as dropped
     fn set_dropped(&self);
+
+    /// Points cancellation at a replacement partition.
+    #[cfg(target_os = "windows")]
+    fn set_partition(&self, partition: WHV_PARTITION_HANDLE);
 }
 
 pub(crate) trait InterruptHandleInternal {
@@ -334,6 +338,9 @@ trait SynchronousInterruptState: Debug + Send + Sync {
 
     #[cfg(hvf)]
     fn set_vcpu(&mut self, vcpu: Option<hv_vcpu_t>);
+
+    #[cfg(target_os = "windows")]
+    fn set_partition(&mut self, partition: WHV_PARTITION_HANDLE);
 }
 
 #[cfg(any(target_os = "windows", hvf))]
@@ -362,6 +369,17 @@ impl<T: SynchronousInterruptState> InterruptHandleImpl for SynchronousInterruptH
             }
         }
     }
+
+    #[cfg(target_os = "windows")]
+    fn set_partition(&self, partition: WHV_PARTITION_HANDLE) {
+        // The write lock waits for in-flight kills on the previous partition.
+        match self.dropped_state.write() {
+            Ok(mut guard) => guard.1.set_partition(partition),
+            Err(e) => {
+                tracing::error!("Failed to acquire partition_state write lock: {}", e);
+            }
+        }
+    }
 }
 
 #[cfg(any(target_os = "windows", hvf))]
@@ -425,6 +443,10 @@ impl SynchronousInterruptState for WHV_PARTITION_HANDLE {
         use windows::Win32::System::Hypervisor::WHvCancelRunVirtualProcessor;
         unsafe { WHvCancelRunVirtualProcessor(*self, 0, 0).is_ok() }
     }
+
+    fn set_partition(&mut self, partition: WHV_PARTITION_HANDLE) {
+        *self = partition;
+    }
 }
 
 #[cfg(hvf)]

src/hyperlight_host/src/hypervisor/virtual_machine/hvf/mod.rs

@@ -1233,4 +1233,8 @@ impl VirtualMachine for HvfVm {
         self.id.epoch += 1;
         Ok(())
     }
+
+    fn set_interrupt_handle(&mut self, handle: Arc<dyn InterruptHandleImpl>) {
+        self.interrupt_handle = handle;
+    }
 }

src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs

@@ -540,6 +540,13 @@ pub(crate) trait VirtualMachine: Debug + Send {
     /// Get partition handle
     #[cfg(target_os = "windows")]
     fn partition_handle(&self) -> windows::Win32::System::Hypervisor::WHV_PARTITION_HANDLE;
+
+    /// Replaces the interrupt handle that `run_vcpu` updates.
+    #[cfg(hvf)]
+    fn set_interrupt_handle(
+        &mut self,
+        handle: std::sync::Arc<dyn crate::hypervisor::InterruptHandleImpl>,
+    );
 }
 
 /// Why an MSR failed the read-then-write-back probe that restore relies on.

src/hyperlight_host/src/sandbox/host_funcs.rs

@@ -73,11 +73,6 @@ impl HostFunctions {
     pub(crate) fn inner_mut(&mut self) -> &mut FunctionRegistry {
         &mut self.0
     }
-
-    /// Borrow the inner registry immutably.
-    pub(crate) fn inner(&self) -> &FunctionRegistry {
-        &self.0
-    }
 }
 
 impl Default for HostFunctions {