Retain page table root finder in memory snapshots

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

Author: ludfjig

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 * **Breaking:** Filesystem paths are now represented using `PathBuf`. `GuestBinary::FilePath` now stores a `PathBuf` instead of a `String`, and `MultiUseSandbox::generate_crashdump_to_dir` accepts `Into<PathBuf>` instead of `Into<String>`. Callers passing a `String` to `GuestBinary::FilePath` must convert it using `.into()`.
 * Deprecate `MultiUseSandbox::poisoned` in favor of `MultiUseSandbox::status().is_poisoned()`.
 * `MultiUseSandbox::restore` has been made more flexible and now accepts snapshots from any guest binary or memory layout when host functions are compatible.
+* **Breaking:** `PtRootFinder` now uses `Arc` and requires `Sync`.
 
 ### Removed
 

src/hyperlight_host/src/mem/mgr.rs

@@ -34,6 +34,7 @@ use crate::hypervisor::regs::CommonSpecialRegisters;
 use crate::mem::memory_region::MemoryRegion;
 #[cfg(crashdump)]
 use crate::mem::memory_region::{CrashDumpRegion, MemoryRegionFlags, MemoryRegionType};
+use crate::sandbox::PtRootFinder;
 use crate::sandbox::snapshot::{NextAction, Snapshot};
 use crate::{Result, new_error};
 
@@ -359,6 +360,7 @@ where
         #[cfg(target_arch = "x86_64")] msrs: Vec<crate::hypervisor::regs::MsrEntry>,
         next_action: NextAction,
         host_functions: HostFunctionDetails,
+        pt_root_finder: Option<PtRootFinder>,
     ) -> Result<Snapshot> {
         self.snapshot_count += 1;
         Snapshot::new(
@@ -376,6 +378,7 @@ where
             self.original_entrypoint,
             self.snapshot_count,
             host_functions,
+            pt_root_finder,
         )
     }
 }

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -136,7 +136,7 @@ pub struct MultiUseSandbox {
 ///
 /// Returns a list of root page table GPAs to walk. If the list is
 /// empty, only `root_pt_gpa` is used.
-pub type PtRootFinder = Box<dyn Fn(&[u8], &[u8], u64) -> Vec<u64> + Send>;
+pub type PtRootFinder = Arc<dyn Fn(&[u8], &[u8], u64) -> Vec<u64> + Send + Sync>;
 
 impl MultiUseSandbox {
     fn ensure_usable(&self) -> Result<()> {
@@ -180,8 +180,12 @@ impl MultiUseSandbox {
     /// Set a callback that discovers page table roots from guest memory.
     /// The callback receives (snapshot_mem, scratch_mem, cr3) and returns
     /// the list of root GPAs to walk during snapshot creation.
+    ///
+    /// In-memory snapshots retain the finder across restore. The finder is not
+    /// serialized.
     pub fn set_pt_root_finder(&mut self, finder: PtRootFinder) {
         self.pt_root_finder = Some(finder);
+        self.snapshot = None;
     }
 
     /// Create a `MultiUseSandbox` directly from a [`Snapshot`],
@@ -365,13 +369,14 @@ impl MultiUseSandbox {
         #[cfg(gdb)]
         let dbg_mem_wrapper = Arc::new(Mutex::new(hshm.clone()));
 
-        let sbox = MultiUseSandbox::from_uninit(
+        let mut sbox = MultiUseSandbox::from_uninit(
             host_funcs,
             hshm,
             vm,
             #[cfg(gdb)]
             dbg_mem_wrapper,
         );
+        sbox.pt_root_finder = snapshot.pt_root_finder().cloned();
         Ok(sbox)
     }
 
@@ -462,6 +467,7 @@ impl MultiUseSandbox {
             msrs,
             next_action,
             host_functions,
+            self.pt_root_finder.clone(),
         )?;
         let snapshot = Arc::new(memory_snapshot);
         self.snapshot = Some(snapshot.clone());
@@ -682,7 +688,7 @@ impl MultiUseSandbox {
             self.vm.clear_crashdump_binary_path();
         }
 
-        self.pt_root_finder = None;
+        self.pt_root_finder = snapshot.pt_root_finder().cloned();
 
         // The restored snapshot is now our most current snapshot
         self.snapshot = Some(snapshot.clone());
@@ -1249,6 +1255,7 @@ fn warn_on_layout_override(
 
 #[cfg(test)]
 mod tests {
+    use std::sync::atomic::{AtomicUsize, Ordering};
     use std::sync::{Arc, Barrier};
     use std::thread;
 
@@ -1262,6 +1269,7 @@ mod tests {
     use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
     use crate::mem::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, SharedMemory as _};
     use crate::sandbox::SandboxConfiguration;
+    use crate::sandbox::snapshot::Snapshot;
     use crate::sandbox::uninitialized::{GuestBlob, GuestEnvironment};
     use crate::{
         GuestBinary, HyperlightError, MultiUseSandbox, Result, SandboxStatus, UninitializedSandbox,
@@ -1282,6 +1290,23 @@ mod tests {
         assert!(SandboxStatus::Unrecoverable.is_unrecoverable());
     }
 
+    trait AmbiguousIfSync<Marker> {
+        fn assert_not_sync() {}
+    }
+
+    impl<T: ?Sized> AmbiguousIfSync<()> for T {}
+    impl<T: ?Sized + Sync> AmbiguousIfSync<u8> for T {}
+
+    #[test]
+    fn snapshot_and_sandbox_thread_safety() {
+        fn assert_send<T: Send>() {}
+        fn assert_send_sync<T: Send + Sync>() {}
+
+        assert_send::<MultiUseSandbox>();
+        let _ = <MultiUseSandbox as AmbiguousIfSync<_>>::assert_not_sync;
+        assert_send_sync::<Snapshot>();
+    }
+
     #[test]
     fn poison() {
         let mut sbox: MultiUseSandbox = {
@@ -2590,6 +2615,8 @@ mod tests {
             .unwrap()
             .evolve()
             .unwrap();
+        let source_finder: crate::sandbox::PtRootFinder = Arc::new(|_, _, root| vec![root]);
+        source.set_pt_root_finder(source_finder.clone());
         let mut target =
             UninitializedSandbox::new(GuestBinary::FilePath(simple_guest_as_pathbuf()), None)
                 .unwrap()
@@ -2598,8 +2625,7 @@ mod tests {
 
         assert_eq!(source.call::<i32>("StackAllocate", 256i32).unwrap(), 256);
         assert_eq!(target.call::<i32>("AddToStatic", 17i32).unwrap(), 17);
-        target.set_pt_root_finder(Box::new(|_, _, root| vec![root]));
-        assert!(target.pt_root_finder.is_some());
+        target.set_pt_root_finder(Arc::new(|_, _, _| Vec::new()));
 
         assert_ne!(
             source.mem_mgr.layout.code_size(),
@@ -2616,7 +2642,10 @@ mod tests {
 
         let snapshot = source.snapshot().unwrap();
         target.restore(snapshot).unwrap();
-        assert!(target.pt_root_finder.is_none());
+        assert!(Arc::ptr_eq(
+            target.pt_root_finder.as_ref().unwrap(),
+            &source_finder
+        ));
         assert_eq!(target.call::<i32>("StackAllocate", 512i32).unwrap(), 512);
         assert!(matches!(
             target.call::<i32>("GetStatic", ()),
@@ -2627,6 +2656,68 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn snapshot_restore_clears_absent_pt_root_finder() {
+        let path = simple_guest_as_pathbuf();
+        let mut source = UninitializedSandbox::new(GuestBinary::FilePath(path), None)
+            .unwrap()
+            .evolve()
+            .unwrap();
+        let snapshot = source.snapshot().unwrap();
+        assert!(snapshot.pt_root_finder().is_none());
+
+        let path = simple_guest_as_pathbuf();
+        let mut target = UninitializedSandbox::new(GuestBinary::FilePath(path), None)
+            .unwrap()
+            .evolve()
+            .unwrap();
+        target.set_pt_root_finder(Arc::new(|_, _, root| vec![root]));
+
+        target.restore(snapshot).unwrap();
+        assert!(target.pt_root_finder.is_none());
+    }
+
+    #[test]
+    fn snapshot_restore_uses_retained_pt_root_finder() {
+        let source_calls = Arc::new(AtomicUsize::new(0));
+        let source_calls_in_finder = source_calls.clone();
+        let source_finder: crate::sandbox::PtRootFinder = Arc::new(move |_, _, _| {
+            source_calls_in_finder.fetch_add(1, Ordering::Relaxed);
+            Vec::new()
+        });
+        let path = simple_guest_as_pathbuf();
+        let mut source = UninitializedSandbox::new(GuestBinary::FilePath(path), None)
+            .unwrap()
+            .evolve()
+            .unwrap();
+        source.set_pt_root_finder(source_finder);
+        let snapshot = source.snapshot().unwrap();
+
+        let target_calls = Arc::new(AtomicUsize::new(0));
+        let target_calls_in_finder = target_calls.clone();
+        let target_finder: crate::sandbox::PtRootFinder = Arc::new(move |_, _, root| {
+            target_calls_in_finder.fetch_add(1, Ordering::Relaxed);
+            vec![root]
+        });
+        let path = simple_guest_as_pathbuf();
+        let mut target = UninitializedSandbox::new(GuestBinary::FilePath(path), None)
+            .unwrap()
+            .evolve()
+            .unwrap();
+        target.set_pt_root_finder(target_finder);
+        target.restore(snapshot).unwrap();
+
+        let source_calls_before = source_calls.load(Ordering::Relaxed);
+        target.call::<i32>("GetStatic", ()).unwrap();
+        target.snapshot().unwrap();
+
+        assert_eq!(
+            source_calls.load(Ordering::Relaxed),
+            source_calls_before + 1
+        );
+        assert_eq!(target_calls.load(Ordering::Relaxed), 0);
+    }
+
     #[test]
     fn snapshot_restore_replaces_c_guest_with_rust_guest() {
         let mut source =

src/hyperlight_host/src/sandbox/snapshot/file/mod.rs

@@ -334,6 +334,10 @@ impl Snapshot {
     /// guest is running. Any release that breaks the format is called
     /// out in the Hyperlight changelog.
     ///
+    /// A [`PtRootFinder`](crate::sandbox::PtRootFinder) configured with
+    /// [`set_pt_root_finder`](crate::MultiUseSandbox::set_pt_root_finder) is not
+    /// serialized. Set it again on any sandbox created from the loaded snapshot.
+    ///
     /// # Examples
     ///
     /// ```no_run
@@ -668,6 +672,11 @@ impl Snapshot {
     /// guest is running. Any release that breaks the format is called
     /// out in the Hyperlight changelog.
     ///
+    /// If the source sandbox used
+    /// [`MultiUseSandbox::set_pt_root_finder`](crate::MultiUseSandbox::set_pt_root_finder),
+    /// set the finder again on the sandbox created from this snapshot. The finder
+    /// is not serialized.
+    ///
     /// # Verification
     ///
     /// This method does not check the manifest, config, or snapshot
@@ -909,6 +918,7 @@ impl Snapshot {
             original_entrypoint: cfg.original_entrypoint_addr,
             snapshot_generation,
             host_functions,
+            pt_root_finder: None,
         })
     }
 }

src/hyperlight_host/src/sandbox/snapshot/file_tests.rs

@@ -26,6 +26,7 @@ use sha2::{Digest as _, Sha256};
 
 use crate::func::Registerable;
 use crate::mem::layout::SandboxMemoryLayout;
+use crate::sandbox::PtRootFinder;
 use crate::sandbox::snapshot::{OciDigest, OciReference, OciTag, Snapshot};
 use crate::{GuestBinary, HostFunctions, MultiUseSandbox, UninitializedSandbox};
 
@@ -95,9 +96,21 @@ fn find_snapshot_blob(oci_dir: &std::path::Path) -> std::path::PathBuf {
 
 #[test]
 fn from_snapshot_already_initialized_in_memory() {
-    let snapshot = create_snapshot();
+    let mut source = create_test_sandbox();
+    let initial_snapshot = source.snapshot().unwrap();
+    let finder: PtRootFinder = Arc::new(|_, _, root| vec![root]);
+    source.set_pt_root_finder(finder.clone());
+    let snapshot = source.snapshot().unwrap();
+    assert!(!Arc::ptr_eq(&initial_snapshot, &snapshot));
+    assert!(Arc::ptr_eq(snapshot.pt_root_finder().unwrap(), &finder));
+
     let mut sbox2 =
         MultiUseSandbox::from_snapshot(snapshot, HostFunctions::default(), None).unwrap();
+    let restored_snapshot = sbox2.snapshot().unwrap();
+    assert!(Arc::ptr_eq(
+        restored_snapshot.pt_root_finder().unwrap(),
+        &finder
+    ));
     let result: i32 = sbox2.call("GetStatic", ()).unwrap();
     assert_eq!(result, 0);
 }
@@ -119,7 +132,9 @@ fn from_snapshot_in_memory_pre_init() {
 
 #[test]
 fn round_trip_save_load_call() {
-    let snapshot = create_snapshot();
+    let mut source = create_test_sandbox();
+    source.set_pt_root_finder(Arc::new(|_, _, root| vec![root]));
+    let snapshot = source.snapshot().unwrap();
 
     let dir = tempfile::tempdir().unwrap();
     let oci = dir.path().join("snap");
@@ -128,6 +143,7 @@ fn round_trip_save_load_call() {
         .unwrap();
 
     let loaded = Snapshot::checked_load(&oci, OciTag::new("latest").unwrap()).unwrap();
+    assert!(loaded.pt_root_finder().is_none());
     let mut sbox2 =
         MultiUseSandbox::from_snapshot(Arc::new(loaded), HostFunctions::default(), None).unwrap();
 

src/hyperlight_host/src/sandbox/snapshot/mod.rs

@@ -39,8 +39,8 @@ use crate::mem::layout::SandboxMemoryLayout;
 use crate::mem::memory_region::{GuestMemoryRegion, MemoryRegion, MemoryRegionFlags};
 use crate::mem::mgr::{GuestPageTableBuffer, SnapshotSharedMemory};
 use crate::mem::shared_mem::{ReadonlySharedMemory, SharedMemory};
-use crate::sandbox::SandboxConfiguration;
 use crate::sandbox::uninitialized::{GuestBinary, GuestEnvironment};
+use crate::sandbox::{PtRootFinder, SandboxConfiguration};
 
 const PTE_SIZE: usize = size_of::<vmem::PageTableEntry>();
 
@@ -123,6 +123,9 @@ pub struct Snapshot {
     /// `HostFunctions` set that is missing required functions or
     /// has mismatched signatures.
     host_functions: HostFunctionDetails,
+
+    /// Runtime-only page-table root finder retained by in-memory snapshots.
+    pt_root_finder: Option<PtRootFinder>,
 }
 impl core::convert::AsRef<Snapshot> for Snapshot {
     fn as_ref(&self) -> &Self {
@@ -406,6 +409,7 @@ impl Snapshot {
             host_functions: HostFunctionDetails {
                 host_functions: None,
             },
+            pt_root_finder: None,
         })
     }
 
@@ -432,6 +436,7 @@ impl Snapshot {
         original_entrypoint: u64,
         snapshot_generation: u64,
         host_functions: HostFunctionDetails,
+        pt_root_finder: Option<PtRootFinder>,
     ) -> Result<Self> {
         let mut phys_seen = HashMap::<u64, usize>::new();
         let scratch_gva = scratch_base_gva(layout.get_scratch_size());
@@ -588,6 +593,7 @@ impl Snapshot {
             original_entrypoint,
             snapshot_generation,
             host_functions,
+            pt_root_finder,
         })
     }
 
@@ -596,6 +602,10 @@ impl Snapshot {
         self.snapshot_generation
     }
 
+    pub(crate) fn pt_root_finder(&self) -> Option<&PtRootFinder> {
+        self.pt_root_finder.as_ref()
+    }
+
     /// Return the main memory contents of the snapshot
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub(crate) fn memory(&self) -> &ReadonlySharedMemory {
@@ -785,6 +795,7 @@ mod tests {
             0,
             1,
             HostFunctionDetails::default(),
+            None,
         )
         .unwrap();
 
@@ -805,6 +816,7 @@ mod tests {
             0,
             2,
             HostFunctionDetails::default(),
+            None,
         )
         .unwrap();