Ensure sandboxes can only restore to snapshots taken from the same sandbox (#746)

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

Author: ludfjig

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),

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<MemoryRegion>,
     ) -> Result<SharedMemorySnapshot> {
-        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.

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<u8>,
     /// The memory regions that were mapped when this snapshot was taken (excluding initial sandbox regions)
     regions: Vec<MemoryRegion>,
@@ -35,11 +38,16 @@ impl SharedMemorySnapshot {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(super) fn new<S: SharedMemory>(
         shared_mem: &mut S,
+        sandbox_id: u64,
         regions: Vec<MemoryRegion>,
     ) -> Result<Self> {
         // 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::<Vec<u8>>();
         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

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<Mutex<FunctionRegistry>>,
     pub(crate) mem_mgr: MemMgrWrapper<HostSharedMemory>,
@@ -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<Snapshot> {
         let mapped_regions_iter = self.vm.get_mapped_regions();
         let mapped_regions_vec: Vec<MemoryRegion> = 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);
+    }
 }