expose snapshots to the public API

Signed-off-by: Jorge Prendes <[email protected]>

Author: jprendes

src/hyperlight_host/src/mem/mgr.rs

@@ -265,10 +265,33 @@ where
         }
     }
 
+    pub(crate) fn snapshot(&mut self) -> Result<SharedMemorySnapshot> {
+        SharedMemorySnapshot::new(&mut self.shared_mem, self.mapped_rgns)
+    }
+
+    
+    /// This function restores a memory snapshot the given snapshot.
+    ///
+    /// Returns the number of memory regions mapped into the sandbox
+    /// that need to be unmapped in order for the restore to be
+    /// completed.
+    pub(crate) fn restore_snapshot(&mut self, snapshot: &SharedMemorySnapshot) -> Result<u64> {
+        if self.shared_mem.mem_size() != snapshot.mem_size() {
+            return Err(new_error!(
+                "Snapshot size does not match current memory size: {} != {}",
+                self.shared_mem.raw_mem_size(),
+                snapshot.mem_size()
+            ));
+        }
+        let old_rgns = self.mapped_rgns;
+        self.mapped_rgns = snapshot.restore_from_snapshot(&mut self.shared_mem)?;
+        Ok(old_rgns - self.mapped_rgns)
+    }
+
     /// this function will create a memory snapshot and push it onto the stack of snapshots
     /// It should be used when you want to save the state of the memory, for example, when evolving a sandbox to a new state
     pub(crate) fn push_state(&mut self) -> Result<()> {
-        let snapshot = SharedMemorySnapshot::new(&mut self.shared_mem, self.mapped_rgns)?;
+        let snapshot = self.snapshot()?;
         self.snapshots
             .try_lock()
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
@@ -285,16 +308,14 @@ where
     /// that need to be unmapped in order for the restore to be
     /// completed.
     pub(crate) fn restore_state_from_last_snapshot(&mut self) -> Result<u64> {
-        let mut snapshots = self
+        let snapshots = self
             .snapshots
             .try_lock()
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
-        let last = snapshots.last_mut();
-        if last.is_none() {
+        let last = snapshots.last();
+        let Some(snapshot) = last else {
             log_then_return!(NoMemorySnapshot);
-        }
-        #[allow(clippy::unwrap_used)] // We know that last is not None because we checked it above
-        let snapshot = last.unwrap();
+        };
         let old_rgns = self.mapped_rgns;
         self.mapped_rgns = snapshot.restore_from_snapshot(&mut self.shared_mem)?;
         Ok(old_rgns - self.mapped_rgns)

src/hyperlight_host/src/mem/shared_mem_snapshot.rs

@@ -22,7 +22,7 @@ use crate::Result;
 /// A wrapper around a `SharedMemory` reference and a snapshot
 /// of the memory therein
 #[derive(Clone)]
-pub(super) struct SharedMemorySnapshot {
+pub(crate) struct SharedMemorySnapshot {
     snapshot: Vec<u8>,
     /// How many non-main-RAM regions were mapped when this snapshot was taken?
     mapped_rgns: u64,
@@ -54,12 +54,18 @@ impl SharedMemorySnapshot {
     /// into the internally-stored `SharedMemory`
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(super) fn restore_from_snapshot<S: SharedMemory>(
-        &mut self,
+        &self,
         shared_mem: &mut S,
     ) -> Result<u64> {
         shared_mem.with_exclusivity(|e| e.copy_from_slice(self.snapshot.as_slice(), 0))??;
         Ok(self.mapped_rgns)
     }
+
+    /// Return the size of the snapshot in bytes.
+    #[instrument(skip_all, parent = Span::current(), level= "Trace")]
+    pub(super) fn mem_size(&self) -> usize {
+        self.snapshot.len()
+    }
 }
 
 #[cfg(test)]

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -28,6 +28,7 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{
 use tracing::{Span, instrument};
 
 use super::host_funcs::FunctionRegistry;
+use super::snapshot::Snapshot;
 use super::{MemMgrWrapper, WrapperGetter};
 use crate::func::call_ctx::MultiUseGuestCallContext;
 use crate::func::guest_err::check_for_guest_error;
@@ -163,6 +164,23 @@ impl MultiUseSandbox {
         MultiUseGuestCallContext::start(self)
     }
 
+    /// Create a snapshot of the current state of the sandbox's memory.
+    #[instrument(err(Debug), skip_all, parent = Span::current())]
+    pub fn snapshot(&mut self) -> Result<Snapshot> {
+        let snapshot = self.mem_mgr.unwrap_mgr_mut().snapshot()?;
+        Ok(Snapshot { inner: snapshot })
+    }
+
+    /// 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<()> {
+        let rgns_to_unmap = self.mem_mgr
+            .unwrap_mgr_mut()
+            .restore_snapshot(&snapshot.inner)?;
+        unsafe { self.vm.unmap_regions(rgns_to_unmap)? };
+        Ok(())
+    }
+
     /// Call a guest function by name, with the given return type and arguments.
     #[instrument(err(Debug), skip(self, args), parent = Span::current())]
     pub fn call_guest_function_by_name<Output: SupportedReturnType>(

src/hyperlight_host/src/sandbox/mod.rs

@@ -37,6 +37,9 @@ pub mod uninitialized;
 /// initialized `Sandbox`es.
 pub(crate) mod uninitialized_evolve;
 
+/// Representation of a snapshot of a `Sandbox`.
+pub mod snapshot;
+
 /// Trait used by the macros to paper over the differences between hyperlight and hyperlight-wasm
 mod callable;
 

src/hyperlight_host/src/sandbox/snapshot.rs

@@ -0,0 +1,8 @@
+use crate::mem::shared_mem_snapshot::SharedMemorySnapshot;
+
+/// A snapshot capturing the state of the memory in a `MultiUseSandbox`.
+#[derive(Clone)]
+pub struct Snapshot {
+    /// TODO: Use Arc<SharedMemorySnapshot>
+    pub(crate) inner: SharedMemorySnapshot,
+}