Fix merge errors from rebasing

Signed-off-by: Simon Davies <simongdavies@users.noreply.github.com>

Author: simongdavies

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -245,7 +245,7 @@ impl MultiUseSandbox {
     /// if result.is_err() {
     ///     if sandbox.poisoned() {
     ///         // Restore from snapshot to clear poison
-    ///         sandbox.restore(&snapshot)?;
+    ///         sandbox.restore(snapshot.clone())?;
     ///         assert!(!sandbox.poisoned());
     ///         
     ///         // Sandbox is now usable again
@@ -429,7 +429,7 @@ impl MultiUseSandbox {
     ///     
     ///     if sandbox.poisoned() {
     ///         eprintln!("Sandbox was poisoned, restoring from snapshot");
-    ///         sandbox.restore(&snapshot)?;
+    ///         sandbox.restore(snapshot.clone())?;
     ///     }
     /// }
     /// # Ok(())
@@ -856,7 +856,7 @@ mod tests {
         assert!(matches!(res, HyperlightError::PoisonedSandbox));
 
         // restore to non-poisoned snapshot should work and clear poison
-        sbox.restore(&snapshot).unwrap();
+        sbox.restore(snapshot.clone()).unwrap();
         assert!(!sbox.poisoned());
 
         // guest calls should work again after restore
@@ -874,7 +874,7 @@ mod tests {
         assert!(sbox.poisoned());
 
         // restore to non-poisoned snapshot should work again
-        sbox.restore(&snapshot).unwrap();
+        sbox.restore(snapshot.clone()).unwrap();
         assert!(!sbox.poisoned());
 
         // guest calls should work again
@@ -1243,12 +1243,12 @@ mod tests {
         assert_eq!(sbox.vm.get_mapped_regions().count(), 1);
 
         // 4. Restore to snapshot 1 (should unmap the region)
-        sbox.restore(snapshot1).unwrap();
-        assert_eq!(sbox.vm.get_mapped_regions().len(), 0);
+        sbox.restore(snapshot1.clone()).unwrap();
+        assert_eq!(sbox.vm.get_mapped_regions().count(), 0);
 
         // 5. Restore forward to snapshot 2 (should remap the region)
-        sbox.restore(snapshot2).unwrap();
-        assert_eq!(sbox.vm.get_mapped_regions().len(), 1);
+        sbox.restore(snapshot2.clone()).unwrap();
+        assert_eq!(sbox.vm.get_mapped_regions().count(), 1);
 
         // Verify the region is the same
         let mut restored_regions = sbox.vm.get_mapped_regions();

src/hyperlight_host/tests/integration_test.rs

@@ -68,7 +68,7 @@ fn interrupt_host_call() {
     assert!(sandbox.poisoned());
 
     // Restore from snapshot to clear poison
-    sandbox.restore(&snapshot).unwrap();
+    sandbox.restore(snapshot.clone()).unwrap();
     assert!(!sandbox.poisoned());
 
     thread.join().unwrap();
@@ -98,7 +98,7 @@ fn interrupt_in_progress_guest_call() {
     assert!(sbox1.poisoned());
 
     // Restore from snapshot to clear poison
-    sbox1.restore(&snapshot).unwrap();
+    sbox1.restore(snapshot.clone()).unwrap();
     assert!(!sbox1.poisoned());
 
     barrier.wait();
@@ -191,7 +191,7 @@ fn interrupt_same_thread() {
             _ => panic!("Unexpected return"),
         };
         if sbox2.poisoned() {
-            sbox2.restore(&snapshot2).unwrap();
+            sbox2.restore(snapshot2.clone()).unwrap();
         }
         sbox3
             .call::<String>("Echo", "hello".to_string())
@@ -238,7 +238,7 @@ fn interrupt_same_thread_no_barrier() {
             _ => panic!("Unexpected return"),
         };
         if sbox2.poisoned() {
-            sbox2.restore(&snapshot2).unwrap();
+            sbox2.restore(snapshot2.clone()).unwrap();
         }
         sbox3
             .call::<String>("Echo", "hello".to_string())
@@ -267,7 +267,7 @@ fn interrupt_moved_sandbox() {
         let res = sbox1.call::<i32>("Spin", ()).unwrap_err();
         assert!(matches!(res, HyperlightError::ExecutionCanceledByHost()));
         assert!(sbox1.poisoned());
-        sbox1.restore(&snapshot1).unwrap();
+        sbox1.restore(snapshot1.clone()).unwrap();
         assert!(!sbox1.poisoned());
     });
 
@@ -327,7 +327,7 @@ fn interrupt_custom_signal_no_and_retry_delay() {
         assert!(sbox1.poisoned());
         // immediately reenter another guest function call after having being cancelled,
         // so that the vcpu is running again before the interruptor-thread has a chance to see that the vcpu is not running
-        sbox1.restore(&snapshot1).unwrap();
+        sbox1.restore(snapshot1.clone()).unwrap();
         assert!(!sbox1.poisoned());
     }
     thread.join().expect("Thread should finish");
@@ -906,7 +906,7 @@ fn interrupt_random_kill_stress_test() {
     // Wrapper to hold a sandbox and its snapshot together
     struct SandboxWithSnapshot {
         sandbox: MultiUseSandbox,
-        snapshot: Snapshot,
+        snapshot: Arc<Snapshot>,
     }
 
     use std::collections::VecDeque;
@@ -1128,7 +1128,10 @@ fn interrupt_random_kill_stress_test() {
                         assert!(sandbox_wrapper.sandbox.poisoned());
 
                         // Try to restore the snapshot
-                        if let Err(e) = sandbox_wrapper.sandbox.restore(&sandbox_wrapper.snapshot) {
+                        if let Err(e) = sandbox_wrapper
+                            .sandbox
+                            .restore(sandbox_wrapper.snapshot.clone())
+                        {
                             error!(
                                 "CRITICAL: Thread {} iteration {}: Failed to restore snapshot: {:?}",
                                 thread_id, iteration, e
@@ -1442,7 +1445,7 @@ fn interrupt_infinite_loop_stress_test() {
                 }
 
                 // Restore the sandbox for the next iteration
-                sandbox.restore(&snapshot).unwrap();
+                sandbox.restore(snapshot.clone()).unwrap();
             }
         }));
     }

src/hyperlight_host/tests/sandbox_host_tests.rs

@@ -365,7 +365,7 @@ fn host_function_error() -> Result<()> {
             );
             // C guest panics in rust guest lib when host function returns error, which will poison the sandbox
             if init_sandbox.poisoned() {
-                init_sandbox.restore(&snapshot)?;
+                init_sandbox.restore(snapshot.clone())?;
             }
         }
     }