diff --git a/src/hyperlight_host/src/error.rs b/src/hyperlight_host/src/error.rs index df8ee8a37..951066c53 100644 --- a/src/hyperlight_host/src/error.rs +++ b/src/hyperlight_host/src/error.rs @@ -221,6 +221,10 @@ pub enum HyperlightError { #[cfg(all(feature = "seccomp", target_os = "linux"))] SeccompFilterError(#[from] seccompiler::Error), + /// Tried to restore snapshot to a sandbox that is not the same as the one the snapshot was taken from + #[error("Snapshot was taken from a different sandbox")] + SnapshotSandboxMismatch, + /// SystemTimeError #[error("SystemTimeError {0:?}")] SystemTimeError(#[from] SystemTimeError), diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index 526190d4c..e33c4b08d 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -263,9 +263,10 @@ where /// Create a snapshot with the given mapped regions pub(crate) fn snapshot( &mut self, + sandbox_id: u64, mapped_regions: Vec, ) -> Result { - SharedMemorySnapshot::new(&mut self.shared_mem, mapped_regions) + SharedMemorySnapshot::new(&mut self.shared_mem, sandbox_id, mapped_regions) } /// This function restores a memory snapshot from a given snapshot. diff --git a/src/hyperlight_host/src/mem/shared_mem_snapshot.rs b/src/hyperlight_host/src/mem/shared_mem_snapshot.rs index 8f55a8b25..49af2d99d 100644 --- a/src/hyperlight_host/src/mem/shared_mem_snapshot.rs +++ b/src/hyperlight_host/src/mem/shared_mem_snapshot.rs @@ -24,6 +24,9 @@ use crate::Result; /// of the memory therein #[derive(Clone)] pub(crate) struct SharedMemorySnapshot { + // Unique ID of the sandbox this snapshot was taken from + sandbox_id: u64, + // Memory of the sandbox at the time this snapshot was taken snapshot: Vec, /// The memory regions that were mapped when this snapshot was taken (excluding initial sandbox regions) regions: Vec, @@ -35,11 +38,16 @@ impl SharedMemorySnapshot { #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] pub(super) fn new( shared_mem: &mut S, + sandbox_id: u64, regions: Vec, ) -> Result { // TODO: Track dirty pages instead of copying entire memory let snapshot = shared_mem.with_exclusivity(|e| e.copy_all_to_vec())??; - Ok(Self { snapshot, regions }) + Ok(Self { + sandbox_id, + snapshot, + regions, + }) } /// Take another snapshot of the internally-stored `SharedMemory`, @@ -59,6 +67,11 @@ impl SharedMemorySnapshot { Ok(()) } + /// The id of the sandbox this snapshot was taken from. + pub(crate) fn sandbox_id(&self) -> u64 { + self.sandbox_id + } + /// Get the mapped regions from this snapshot pub(crate) fn regions(&self) -> &[MemoryRegion] { &self.regions @@ -84,7 +97,7 @@ mod tests { let data2 = data1.iter().map(|b| b + 1).collect::>(); let mut gm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap(); gm.copy_from_slice(data1.as_slice(), 0).unwrap(); - let mut snap = super::SharedMemorySnapshot::new(&mut gm, Vec::new()).unwrap(); + let mut snap = super::SharedMemorySnapshot::new(&mut gm, 0, Vec::new()).unwrap(); { // after the first snapshot is taken, make sure gm has the equivalent // of data1 diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index b21ab6361..6b925f264 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -20,6 +20,7 @@ use std::os::fd::AsRawFd; #[cfg(unix)] use std::os::linux::fs::MetadataExt; use std::path::Path; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; @@ -31,6 +32,7 @@ use tracing::{Span, instrument}; use super::host_funcs::FunctionRegistry; use super::snapshot::Snapshot; use super::{Callable, MemMgrWrapper, WrapperGetter}; +use crate::HyperlightError::SnapshotSandboxMismatch; use crate::func::guest_err::check_for_guest_error; use crate::func::{ParameterTuple, SupportedReturnType}; #[cfg(gdb)] @@ -44,6 +46,9 @@ use crate::mem::shared_mem::HostSharedMemory; use crate::metrics::maybe_time_and_emit_guest_call; use crate::{HyperlightError, Result, log_then_return}; +/// Global counter for assigning unique IDs to sandboxes +static SANDBOX_ID_COUNTER: AtomicU64 = AtomicU64::new(0); + /// A sandbox that supports being used Multiple times. /// The implication of being used multiple times is two-fold: /// @@ -53,6 +58,8 @@ use crate::{HyperlightError, Result, log_then_return}; /// 2. A MultiUseGuestCallContext can be created from the sandbox and used to make multiple guest function calls to the Sandbox. /// in this case the state of the sandbox is not reset until the context is finished and the `MultiUseSandbox` is returned. pub struct MultiUseSandbox { + /// Unique identifier for this sandbox instance + id: u64, // We need to keep a reference to the host functions, even if the compiler marks it as unused. The compiler cannot detect our dynamic usages of the host function in `HyperlightFunction::call`. pub(super) _host_funcs: Arc>, pub(crate) mem_mgr: MemMgrWrapper, @@ -77,6 +84,7 @@ impl MultiUseSandbox { #[cfg(gdb)] dbg_mem_access_fn: DbgMemAccessHandlerWrapper, ) -> MultiUseSandbox { Self { + id: SANDBOX_ID_COUNTER.fetch_add(1, Ordering::Relaxed), _host_funcs: host_funcs, mem_mgr: mgr, vm, @@ -91,7 +99,10 @@ impl MultiUseSandbox { pub fn snapshot(&mut self) -> Result { let mapped_regions_iter = self.vm.get_mapped_regions(); let mapped_regions_vec: Vec = mapped_regions_iter.cloned().collect(); - let memory_snapshot = self.mem_mgr.unwrap_mgr_mut().snapshot(mapped_regions_vec)?; + let memory_snapshot = self + .mem_mgr + .unwrap_mgr_mut() + .snapshot(self.id, mapped_regions_vec)?; Ok(Snapshot { inner: memory_snapshot, }) @@ -100,6 +111,10 @@ impl MultiUseSandbox { /// Restore the sandbox's memory to the state captured in the given snapshot. #[instrument(err(Debug), skip_all, parent = Span::current())] pub fn restore(&mut self, snapshot: &Snapshot) -> Result<()> { + if self.id != snapshot.inner.sandbox_id() { + return Err(SnapshotSandboxMismatch); + } + self.mem_mgr .unwrap_mgr_mut() .restore_snapshot(&snapshot.inner)?; @@ -765,4 +780,36 @@ mod tests { err ); } + + #[test] + fn snapshot_different_sandbox() { + let mut sandbox = { + let path = simple_guest_as_string().unwrap(); + let u_sbox = UninitializedSandbox::new(GuestBinary::FilePath(path), None).unwrap(); + u_sbox.evolve().unwrap() + }; + + let mut sandbox2 = { + let path = simple_guest_as_string().unwrap(); + let u_sbox = UninitializedSandbox::new(GuestBinary::FilePath(path), None).unwrap(); + u_sbox.evolve().unwrap() + }; + assert_ne!(sandbox.id, sandbox2.id); + + let snapshot = sandbox.snapshot().unwrap(); + let err = sandbox2.restore(&snapshot); + assert!(matches!(err, Err(HyperlightError::SnapshotSandboxMismatch))); + + let sandbox_id = sandbox.id; + drop(sandbox); + drop(sandbox2); + drop(snapshot); + + let sandbox3 = { + let path = simple_guest_as_string().unwrap(); + let u_sbox = UninitializedSandbox::new(GuestBinary::FilePath(path), None).unwrap(); + u_sbox.evolve().unwrap() + }; + assert_ne!(sandbox3.id, sandbox_id); + } }