Revert "Be more careful about poisoning"

This reverts commit cd50b1b.

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -287,64 +287,35 @@ impl MultiUseSandbox {
             return Err(SnapshotSandboxMismatch);
         }
 
-        // A prior failed restore may have swapped mem_mgr.shared_mem
-        // without updating the hypervisor's GPA mapping. Track this
-        // so we can force a remap below if needed.
-        let was_poisoned = self.poisoned;
-
-        // Assume poisoned: any failure below leaves the sandbox
-        // inconsistent. Only cleared at the end on success.
-        self.poisoned = true;
-
-        // Failure here may partially reset vCPU state. The sandbox is
-        // poisoned, and the next restore() retry will call
-        // reset_vm_state() again which is idempotent.
-        self.vm
-            .reset_vm_state()
-            .map_err(HyperlightVmError::Restore)?;
+        // Reset VM/vCPU state before memory restore.
+        self.vm.reset_vm_state().map_err(|e| {
+            self.poisoned = true;
+            HyperlightVmError::Restore(e)
+        })?;
 
-        // Failure here may leave mem_mgr with partially swapped
-        // shared_mem/scratch_mem. The sandbox is poisoned, and if
-        // shared_mem was swapped without updating the hypervisor, the
-        // was_poisoned path below handles it on the retry.
         let (gsnapshot, gscratch) = self.mem_mgr.restore_snapshot(&snapshot)?;
         if let Some(gsnapshot) = gsnapshot {
-            // Failure here leaves the hypervisor pointing at the old
-            // (possibly freed) snapshot. The sandbox is poisoned, and
-            // the next restore() will remap via the was_poisoned path.
-            self.vm
-                .update_snapshot_mapping(gsnapshot)
-                .map_err(|e| HyperlightError::HyperlightVmError(e.into()))?;
-        } else if was_poisoned {
-            // restore_snapshot() skipped the swap because shared_mem
-            // already matched, but the hypervisor may still point at
-            // the old backing from before the failed restore. Force
-            // a remap. Failure leaves us still poisoned for the same
-            // reason as above.
-            let snap_mem = snapshot.memory().to_mgr_snapshot_mem()?;
-            let (_, gsnapshot) = snap_mem.build();
             self.vm
                 .update_snapshot_mapping(gsnapshot)
                 .map_err(|e| HyperlightError::HyperlightVmError(e.into()))?;
         }
         if let Some(gscratch) = gscratch {
-            // Failure here leaves the hypervisor pointing at old
-            // scratch. The sandbox is poisoned; scratch size is fixed
-            // per sandbox so retry will zero the same allocation and
-            // remap here again.
             self.vm
                 .update_scratch_mapping(gscratch)
                 .map_err(|e| HyperlightError::HyperlightVmError(e.into()))?;
         }
 
-        // Failure here means sregs (CR3, etc.) are not set. The
-        // sandbox is poisoned; retry will call restore_sregs() again.
         let sregs = snapshot.sregs().ok_or_else(|| {
             HyperlightError::Error("snapshot from running sandbox should have sregs".to_string())
         })?;
+        // TODO (ludfjig): Go through the rest of possible errors in this `MultiUseSandbox::restore` function
+        // and determine if they should also poison the sandbox.
         self.vm
             .restore_sregs(snapshot.root_pt_gpa(), sregs)
-            .map_err(HyperlightVmError::Restore)?;
+            .map_err(|e| {
+                self.poisoned = true;
+                HyperlightVmError::Restore(e)
+            })?;
 
         self.vm.set_stack_top(snapshot.stack_top_gva());
         self.vm.set_entrypoint(snapshot.entrypoint());
@@ -356,19 +327,12 @@ impl MultiUseSandbox {
         let regions_to_map = snapshot_regions.difference(&current_regions);
 
         for region in regions_to_unmap {
-            // Failure here leaves extra regions mapped. The sandbox is
-            // poisoned; retry will attempt the unmap again since the
-            // region will still be in the current set.
             self.vm
                 .unmap_region(region)
                 .map_err(HyperlightVmError::UnmapRegion)?;
         }
 
         for region in regions_to_map {
-            // Failure here leaves regions missing. The sandbox is
-            // poisoned; retry will attempt the map again since the
-            // region will still be in the snapshot set.
-            //
             // Safety: The region has been mapped before, and at that point the caller promised that the memory region is valid
             // in their call to `MultiUseSandbox::map_region`
             unsafe { self.vm.map_region(region) }.map_err(HyperlightVmError::MapRegion)?;
@@ -377,8 +341,9 @@ impl MultiUseSandbox {
         // The restored snapshot is now our most current snapshot
         self.snapshot = Some(snapshot.clone());
 
-        // Clear poison state: restore completed successfully.
+        // Clear poison state when successfully restoring from snapshot.
         //
+        // # Safety:
         // This is safe because:
         // 1. Snapshots can only be taken from non-poisoned sandboxes (verified at snapshot creation)
         // 2. Restoration completely replaces all memory state, eliminating: