Snapshot only contains pages dirties since last snapshot was taken

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/src/mem/mgr.rs

@@ -15,6 +15,7 @@ limitations under the License.
 */
 
 use std::cmp::Ordering;
+use std::sync::Arc;
 
 use hyperlight_common::flatbuffer_wrappers::function_call::{
     FunctionCall, validate_guest_function_call_buffer,
@@ -73,6 +74,8 @@ pub(crate) struct SandboxMemoryManager<S> {
     pub(crate) entrypoint_offset: Offset,
     /// How many memory regions were mapped after sandbox creation
     pub(crate) mapped_rgns: u64,
+    /// Most recent snapshot taken, in other words, the most recent state that `self` has been in (disregarding currently dirty pages)
+    pub(crate) most_recent_snapshot: Option<Arc<SharedMemorySnapshot>>,
 }
 
 impl<S> SandboxMemoryManager<S>
@@ -93,6 +96,7 @@ where
             load_addr,
             entrypoint_offset,
             mapped_rgns: 0,
+            most_recent_snapshot: None,
         }
     }
 
@@ -259,25 +263,40 @@ where
         }
     }
 
-    pub(crate) fn snapshot(&mut self) -> Result<SharedMemorySnapshot> {
-        SharedMemorySnapshot::new(&mut self.shared_mem, self.mapped_rgns)
+    pub(crate) fn snapshot(
+        &mut self,
+        dirty_pages_bitmap: &[u64],
+    ) -> Result<Arc<SharedMemorySnapshot>> {
+        let snapshot = Arc::new(SharedMemorySnapshot::new(
+            &mut self.shared_mem,
+            dirty_pages_bitmap,
+            self.mapped_rgns,
+            self.most_recent_snapshot.clone(),
+        )?);
+        self.most_recent_snapshot = Some(snapshot.clone());
+        Ok(snapshot)
     }
 
     /// This function restores a memory snapshot from a 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()
-            ));
-        }
+    pub(crate) fn restore_snapshot(
+        &mut self,
+        snapshot: &Arc<SharedMemorySnapshot>,
+        dirty_pages_bitmap: &[u64],
+    ) -> Result<u64> {
         let old_rgns = self.mapped_rgns;
-        self.mapped_rgns = snapshot.restore_from_snapshot(&mut self.shared_mem)?;
+        self.mapped_rgns = snapshot.restore_from_snapshot(
+            &mut self.shared_mem,
+            dirty_pages_bitmap,
+            &self.most_recent_snapshot,
+        )?;
+
+        // Update the most recent snapshot to the one we just restored to
+        self.most_recent_snapshot = Some(snapshot.clone());
+
         Ok(old_rgns - self.mapped_rgns)
     }
 
@@ -407,13 +426,15 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
                 mapped_rgns: 0,
+                most_recent_snapshot: self.most_recent_snapshot.clone(),
             },
             SandboxMemoryManager {
                 shared_mem: gshm,
                 layout: self.layout,
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
                 mapped_rgns: 0,
+                most_recent_snapshot: self.most_recent_snapshot.clone(),
             },
         )
     }

src/hyperlight_host/src/mem/shared_mem.rs

@@ -677,6 +677,30 @@ impl ExclusiveSharedMemory {
         Ok(())
     }
 
+    /// Same as `copy_from_slice` but doesn't dirty the pages.
+    /// # Safety
+    /// This function is unsafe because it does not mark the pages as dirty.
+    /// Only use this if you are certain that the pages do not need to be marked as dirty.
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub unsafe fn restore_copy_from_slice(&mut self, src: &[u8], offset: usize) -> Result<()> {
+        let data = self.as_mut_slice();
+        bounds_check!(offset, src.len(), data.len());
+        data[offset..offset + src.len()].copy_from_slice(src);
+        Ok(())
+    }
+
+    /// Same as `zero_fill` but doesn't dirty the pages.
+    /// # Safety
+    /// This function is unsafe because it does not mark the pages as dirty.
+    /// Only use this if you are certain that the pages do not need to be marked as dirty.
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub(crate) unsafe fn restore_zero_fill(&mut self, offset: usize, len: usize) -> Result<()> {
+        bounds_check!(offset, len, self.mem_size());
+        let data = self.as_mut_slice();
+        data[offset..offset + len].fill(0);
+        Ok(())
+    }
+
     generate_reader!(read_u8, u8);
     generate_reader!(read_i8, i8);
     generate_reader!(read_u16, u16);