From 85d03c4377df120ba87d67063ebe855650c27986 Mon Sep 17 00:00:00 2001 From: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:22:52 -0700 Subject: [PATCH] Retain page table root finder in memory snapshots Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com> --- CHANGELOG.md | 1 + src/hyperlight_host/src/mem/mgr.rs | 3 + .../src/sandbox/initialized_multi_use.rs | 103 +++++++++++++++++- .../src/sandbox/snapshot/file/mod.rs | 10 ++ .../src/sandbox/snapshot/file_tests.rs | 20 +++- .../src/sandbox/snapshot/mod.rs | 14 ++- 6 files changed, 142 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e82c6809d..80ce04474 100644 --- a/CHANGELOG.md +++ b/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` instead of `Into`. 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 diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index 8e4558770..0d9b13fab 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/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}; @@ -308,6 +309,7 @@ where #[cfg(target_arch = "x86_64")] msrs: Vec, next_action: NextAction, host_functions: HostFunctionDetails, + pt_root_finder: Option, ) -> Result { self.snapshot_count += 1; Snapshot::new( @@ -325,6 +327,7 @@ where self.original_entrypoint, self.snapshot_count, host_functions, + pt_root_finder, ) } } diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 2df1cafcc..b522e94fd 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -134,7 +134,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 Vec + Send>; +pub type PtRootFinder = Arc Vec + Send + Sync>; impl MultiUseSandbox { fn ensure_usable(&self) -> Result<()> { @@ -175,8 +175,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`], @@ -346,7 +350,8 @@ impl MultiUseSandbox { })?; } - let sbox = MultiUseSandbox::from_uninit(host_funcs, hshm, vm); + let mut sbox = MultiUseSandbox::from_uninit(host_funcs, hshm, vm); + sbox.pt_root_finder = snapshot.pt_root_finder().cloned(); Ok(sbox) } @@ -437,6 +442,7 @@ impl MultiUseSandbox { msrs, next_action, host_functions, + self.pt_root_finder.clone(), )?; let snapshot = Arc::new(memory_snapshot); self.snapshot = Some(snapshot.clone()); @@ -625,7 +631,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()); @@ -1189,6 +1195,7 @@ fn warn_on_layout_override( #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Barrier}; use std::thread; @@ -1202,6 +1209,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, @@ -1222,6 +1230,23 @@ mod tests { assert!(SandboxStatus::Unrecoverable.is_unrecoverable()); } + trait AmbiguousIfSync { + fn assert_not_sync() {} + } + + impl AmbiguousIfSync<()> for T {} + impl AmbiguousIfSync for T {} + + #[test] + fn snapshot_and_sandbox_thread_safety() { + fn assert_send() {} + fn assert_send_sync() {} + + assert_send::(); + let _ = >::assert_not_sync; + assert_send_sync::(); + } + #[test] fn poison() { let mut sbox: MultiUseSandbox = { @@ -2271,6 +2296,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() @@ -2279,8 +2306,7 @@ mod tests { assert_eq!(source.call::("StackAllocate", 256i32).unwrap(), 256); assert_eq!(target.call::("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(), @@ -2297,7 +2323,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::("StackAllocate", 512i32).unwrap(), 512); assert!(matches!( target.call::("GetStatic", ()), @@ -2308,6 +2337,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::("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 = diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs index 2a0f00d2f..773293d61 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs +++ b/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, }) } } diff --git a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs index 764e3ff4d..077fd475d 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs +++ b/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(); diff --git a/src/hyperlight_host/src/sandbox/snapshot/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/mod.rs index 8ab4dde31..22809a37c 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/mod.rs +++ b/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::(); @@ -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, } impl core::convert::AsRef 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, ) -> Result { let mut phys_seen = HashMap::::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();