fixup! Snapshot only contains pages dirties since last snapshot was taken

Author: ludfjig

src/hyperlight_host/src/mem/shared_mem.rs

@@ -680,7 +680,7 @@ impl ExclusiveSharedMemory {
     /// 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<()> {
+    pub unsafe fn copy_from_slice_no_dirty(&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);
@@ -692,7 +692,7 @@ impl ExclusiveSharedMemory {
     /// 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<()> {
+    pub(crate) unsafe fn zero_fill_no_dirty(&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);

src/hyperlight_host/src/mem/shared_mem_snapshot.rs

@@ -85,66 +85,30 @@ impl SharedMemorySnapshot {
         }
 
         let mut pages_to_restore = HashSet::new();
+
+        // Unconditionally restore all current dirty pages
         for page_num in bit_index_iterator(current_dirty_pages) {
             pages_to_restore.insert(page_num);
         }
 
         if let Some(most_recent) = most_recent_snapshot {
-            // Check if we're restoring to the same snapshot state the sandbox was most recently in
             if Arc::ptr_eq(self, most_recent) {
-                // No need to collect additional pages - just use current dirty pages
-            }
-            // Check if we're "rolling back" to a previous state, i.e. `self` is older than `most_recent` (`self` is ancestor of `most_recent`)
-            // We need to restore all pages that exist in `most_recent` and its ancestors up to `self` (exclusive)
-            else if self.is_ancestor_of(most_recent) {
-                let mut current_snapshot = Some(most_recent);
-
-                while let Some(snapshot) = current_snapshot {
-                    if Arc::ptr_eq(snapshot, self) {
-                        break;
-                    }
-
-                    for page_num in snapshot.data.keys() {
-                        pages_to_restore.insert(*page_num);
-                    }
-
-                    current_snapshot = snapshot.parent.as_ref();
-                }
-            // Check if we're "fast forwarding" to a "newer" state, i.e. `self` is newer than `most_recent` (`most_recent` is ancestor of `self`)
-            // We need to restore all pages that exist in `self` and its ancestors up to `most_recent` (exclusive)
+                // Restoring to same snapshot sandbox was most recently in: only restore current dirty pages
+            } else if self.is_ancestor_of(most_recent) {
+                // Rolling back: collect pages from most_recent back to self (exclusive)
+                Self::collect_pages_from_chain(most_recent, Some(self), &mut pages_to_restore);
             } else if most_recent.is_ancestor_of(self) {
-                let mut current_snapshot = Some(self);
-
-                while let Some(snapshot) = current_snapshot {
-                    if Arc::ptr_eq(snapshot, most_recent) {
-                        break;
-                    }
-
-                    for page_num in snapshot.data.keys() {
-                        pages_to_restore.insert(*page_num);
-                    }
-
-                    current_snapshot = snapshot.parent.as_ref();
-                }
+                // Fast-forwarding: collect pages from self back to most_recent (exclusive)
+                Self::collect_pages_from_chain(self, Some(most_recent), &mut pages_to_restore);
             } else {
-                // Neither is ancestor of the other - they're on different branches
-                // This is not supported for now
+                // Different branches not supported for now
                 return Err(crate::new_error!(
                     "Cannot restore between snapshots on different branches"
                 ));
             }
         } else {
-            // No previous snapshots exist, meaning we're restoring a fresh sandbox to this snapshot.
-            // We need to restore pages that exist in this snapshot and all its ancestors.
-            let mut current_snapshot = Some(self);
-
-            while let Some(snapshot) = current_snapshot {
-                for page_num in snapshot.data.keys() {
-                    pages_to_restore.insert(*page_num);
-                }
-
-                current_snapshot = snapshot.parent.as_ref();
-            }
+            // Sandbox has no previous snapshots: restore all pages from this snapshot and its ancestors
+            Self::collect_pages_from_chain(self, None, &mut pages_to_restore);
         }
 
         // Restore all collected pages
@@ -156,11 +120,11 @@ impl SharedMemorySnapshot {
                 if let Some(page_data) = self.find_page_in_snapshots(page_num) {
                     // Restore from snapshot
                     // # Safety: We don't want to dirty the pages we restore
-                    unsafe { e.restore_copy_from_slice(page_data, offset)? };
+                    unsafe { e.copy_from_slice_no_dirty(page_data, offset)? };
                 } else {
                     // Zero the page (return to initial state)
                     // # Safety: We don't want to dirty the pages we restore
-                    unsafe { e.restore_zero_fill(offset, PAGE_SIZE_USIZE)? };
+                    unsafe { e.zero_fill_no_dirty(offset, PAGE_SIZE_USIZE)? };
                 }
             }
             Ok(())
@@ -170,7 +134,7 @@ impl SharedMemorySnapshot {
     }
 
     /// Check if this snapshot is an ancestor of the other snapshot
-    /// (i.e., the other snapshot was taken after this one in the same chain)
+    /// (i.e. the other snapshot was taken after this one in the same chain)
     fn is_ancestor_of(self: &Arc<Self>, other: &Arc<SharedMemorySnapshot>) -> bool {
         let mut current = other.parent.as_ref();
 
@@ -204,6 +168,30 @@ impl SharedMemorySnapshot {
     pub(super) fn sandbox_size(&self) -> usize {
         self.sandbox_size
     }
+
+    /// Collect all page numbers from snapshots in a chain, starting from `start_snapshot`
+    /// and stopping when `stop_snapshot` is reached (exclusive) or when the chain ends.
+    fn collect_pages_from_chain(
+        start_snapshot: &Arc<SharedMemorySnapshot>,
+        stop_snapshot: Option<&Arc<SharedMemorySnapshot>>,
+        pages_to_restore: &mut HashSet<usize>,
+    ) {
+        let mut current_snapshot = Some(start_snapshot);
+
+        while let Some(snapshot) = current_snapshot {
+            if let Some(stop) = stop_snapshot {
+                if Arc::ptr_eq(snapshot, stop) {
+                    break;
+                }
+            }
+
+            for page_num in snapshot.data.keys() {
+                pages_to_restore.insert(*page_num);
+            }
+
+            current_snapshot = snapshot.parent.as_ref();
+        }
+    }
 }
 
 #[cfg(test)]

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -103,7 +103,7 @@ impl MultiUseSandbox {
             .with_exclusivity(|e| e.get_and_clear_dirty_pages())??;
         let vm_dirty_pages = self.vm.get_and_clear_dirty_pages()?;
 
-        let dirty_pages_bitmap = bitmap_union(&vm_dirty_pages, &host_dirty_pages);
+        let dirty_pages_bitmap = bitmap_union(&vm_dirty_pages, &host_dirty_pages)?;
 
         let snapshot = self
             .mem_mgr
@@ -123,7 +123,7 @@ impl MultiUseSandbox {
             .with_exclusivity(|e| e.get_and_clear_dirty_pages())??;
         let vm_dirty_pages = self.vm.get_and_clear_dirty_pages()?;
 
-        let dirty_pages_bitmap = bitmap_union(&vm_dirty_pages, &host_dirty_pages);
+        let dirty_pages_bitmap = bitmap_union(&vm_dirty_pages, &host_dirty_pages)?;
 
         let rgns_to_unmap = self
             .mem_mgr