Update MultiUseSandbox to use dirty page tracking in restore and devolve methods

Signed-off-by: Simon Davies <[email protected]>

Author: simongdavies

src/hyperlight_host/src/mem/mgr.rs

@@ -77,9 +77,6 @@ pub(crate) struct SandboxMemoryManager<S> {
     pub(crate) entrypoint_offset: Offset,
     /// How many memory regions were mapped after sandbox creation
     pub(crate) mapped_rgns: u64,
-    /// A vector of memory snapshots that can be used to save and  restore the state of the memory
-    /// This is used by the Rust Sandbox implementation (rather than the mem_snapshot field above which only exists to support current C API)
-    snapshots: Arc<Mutex<Vec<SharedMemorySnapshot>>>,
     /// Shared memory snapshots that can be used to save and  restore the state of the memory
     snapshot_manager: Arc<Mutex<Option<SharedMemorySnapshotManager>>>,
 }
@@ -102,7 +99,6 @@ where
             load_addr,
             entrypoint_offset,
             mapped_rgns: 0,
-            snapshots: Arc::new(Mutex::new(Vec::new())),
             snapshot_manager: Arc::new(Mutex::new(None)),
         }
     }
@@ -311,8 +307,12 @@ where
             dirty_bitmap
         };
 
-        let snapshot_manager =
-            SharedMemorySnapshotManager::new(&mut self.shared_mem, dirty_page_map, layout)?;
+        let snapshot_manager = SharedMemorySnapshotManager::new(
+            &mut self.shared_mem,
+            dirty_page_map,
+            layout,
+            self.mapped_rgns,
+        )?;
         existing_snapshot_manager.replace(snapshot_manager);
         Ok(())
     }
@@ -321,7 +321,7 @@ where
     /// off the stack
     /// It should be used when you want to restore the state of the memory to a previous state but still want to
     /// retain that state, for example after calling a function in the guest
-    pub(crate) fn restore_state_from_last_snapshot(&mut self, dirty_bitmap: &[u64]) -> Result<()> {
+    pub(crate) fn restore_state_from_last_snapshot(&mut self, dirty_bitmap: &[u64]) -> Result<u64> {
         let mut snapshot_manager = self
             .snapshot_manager
             .try_lock()
@@ -343,7 +343,7 @@ where
     pub(crate) fn pop_and_restore_state_from_snapshot(
         &mut self,
         dirty_bitmap: &[u64],
-    ) -> Result<()> {
+    ) -> Result<u64> {
         let mut snapshot_manager = self
             .snapshot_manager
             .try_lock()
@@ -368,9 +368,11 @@ where
             None => {
                 log_then_return!("Snapshot manager not initialized");
             }
-            Some(snapshot_manager) => {
-                snapshot_manager.create_new_snapshot(&mut self.shared_mem, dirty_bitmap)
-            }
+            Some(snapshot_manager) => snapshot_manager.create_new_snapshot(
+                &mut self.shared_mem,
+                dirty_bitmap,
+                self.mapped_rgns,
+            ),
         }
     }
 
@@ -504,7 +506,6 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
                 mapped_rgns: 0,
-                snapshots: Arc::new(Mutex::new(Vec::new())),
                 snapshot_manager: Arc::new(Mutex::new(None)),
             },
             SandboxMemoryManager {
@@ -513,7 +514,6 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
                 mapped_rgns: 0,
-                snapshots: Arc::new(Mutex::new(Vec::new())),
                 snapshot_manager: Arc::new(Mutex::new(None)),
             },
         )

src/hyperlight_host/src/mem/page_snapshot.rs

@@ -36,6 +36,8 @@ pub(super) struct PageSnapshot {
     page_index: HashMap<usize, usize>, // page_number -> buffer_offset_in_pages
     /// Contiguous buffer containing all the page data
     buffer: Vec<u8>,
+    /// How many non-main-RAM regions were mapped when this snapshot was taken?
+    mapped_rgns: u64,
 }
 
 impl PageSnapshot {
@@ -44,11 +46,16 @@ impl PageSnapshot {
         Self {
             page_index: HashMap::new(),
             buffer: Vec::new(),
+            mapped_rgns: 0,
         }
     }
 
     /// Create a snapshot from a list of page numbers with pre-allocated buffer
-    pub(super) fn with_pages_and_buffer(page_numbers: Vec<usize>, buffer: Vec<u8>) -> Self {
+    pub(super) fn with_pages_and_buffer(
+        page_numbers: Vec<usize>,
+        buffer: Vec<u8>,
+        mapped_rgns: u64,
+    ) -> Self {
         let page_count = page_numbers.len();
         let mut page_index = HashMap::with_capacity(page_count);
 
@@ -57,7 +64,11 @@ impl PageSnapshot {
             page_index.insert(page_num, buffer_offset);
         }
 
-        Self { page_index, buffer }
+        Self {
+            page_index,
+            buffer,
+            mapped_rgns,
+        }
     }
 
     /// Get page data by page number, returns None if page is not in snapshot
@@ -78,4 +89,9 @@ impl PageSnapshot {
     pub(super) fn max_page(&self) -> Option<usize> {
         self.page_index.keys().max().copied()
     }
+
+    /// Get the number of mapped regions when this snapshot was taken
+    pub(super) fn mapped_rgns(&self) -> u64 {
+        self.mapped_rgns
+    }
 }

src/hyperlight_host/src/mem/shared_memory_snapshot_manager.rs

@@ -48,10 +48,12 @@ impl SharedMemorySnapshotManager {
         shared_mem: &mut S,
         dirty_page_map: Option<&Vec<u64>>,
         layout: &SandboxMemoryLayout,
+        mapped_rgns: u64,
     ) -> Result<Self> {
         // Build a snapshot of memory from the dirty_page_map
 
-        let diff = Self::build_snapshot_from_dirty_page_map(shared_mem, dirty_page_map)?;
+        let diff =
+            Self::build_snapshot_from_dirty_page_map(shared_mem, dirty_page_map, mapped_rgns)?;
 
         // Get the input output buffer details from the layout so that they can be reset to their initial state
         let input_data_size_offset = layout.get_input_data_size_offset();
@@ -86,6 +88,7 @@ impl SharedMemorySnapshotManager {
     fn build_snapshot_from_dirty_page_map<S: SharedMemory>(
         shared_mem: &mut S,
         dirty_page_map: Option<&Vec<u64>>,
+        mapped_rgns: u64,
     ) -> Result<PageSnapshot> {
         // If there is no dirty page map, return an empty snapshot
         if dirty_page_map.is_none() {
@@ -146,7 +149,7 @@ impl SharedMemorySnapshotManager {
         }
 
         // Create the snapshot with the pre-allocated buffer
-        let snapshot = PageSnapshot::with_pages_and_buffer(dirty_pages, buffer);
+        let snapshot = PageSnapshot::with_pages_and_buffer(dirty_pages, buffer, mapped_rgns);
 
         Ok(snapshot)
     }
@@ -155,8 +158,10 @@ impl SharedMemorySnapshotManager {
         &mut self,
         shared_mem: &mut S,
         dirty_page_map: Option<&Vec<u64>>,
+        mapped_rgns: u64,
     ) -> Result<()> {
-        let snapshot = Self::build_snapshot_from_dirty_page_map(shared_mem, dirty_page_map)?;
+        let snapshot =
+            Self::build_snapshot_from_dirty_page_map(shared_mem, dirty_page_map, mapped_rgns)?;
         self.snapshots.push(snapshot);
         Ok(())
     }
@@ -168,7 +173,7 @@ impl SharedMemorySnapshotManager {
         &mut self,
         shared_mem: &mut S,
         dirty_bitmap: &[u64],
-    ) -> Result<()> {
+    ) -> Result<u64> {
         // check the each index in the dirty bitmap and restore only the corresponding pages from the snapshots vector
         // starting at the last snapshot look for the page in each snapshot if it exists and restore it
         // if it does not exist set the page to zero
@@ -232,14 +237,16 @@ impl SharedMemorySnapshotManager {
                 self.output_data_buffer_offset,
                 SandboxMemoryLayout::STACK_POINTER_SIZE_BYTES,
             )
-        })?
+        })??;
+
+        Ok(self.snapshots.last().unwrap().mapped_rgns())
     }
 
     pub(super) fn pop_and_restore_state_from_snapshot<S: SharedMemory>(
         &mut self,
         shared_mem: &mut S,
         dirty_bitmap: &[u64],
-    ) -> Result<()> {
+    ) -> Result<u64> {
         // Check that there is a snapshot to restore from
         if self.snapshots.is_empty() {
             return Err(crate::HyperlightError::NoMemorySnapshot);
@@ -256,9 +263,7 @@ impl SharedMemorySnapshotManager {
         }
 
         // restore the state from the last snapshot
-        self.restore_from_snapshot(shared_mem, &merged_bitmap)?;
-
-        Ok(())
+        self.restore_from_snapshot(shared_mem, &merged_bitmap)
     }
 
     fn get_bitmap_from_snapshot(&self, snapshot_index: usize) -> Vec<u64> {
@@ -390,9 +395,13 @@ mod tests {
         }
 
         // Create snapshot
-        let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, Some(&dirty_pages), &layout)
-                .unwrap();
+        let mut snapshot_manager = super::SharedMemorySnapshotManager::new(
+            &mut shared_mem,
+            Some(&dirty_pages),
+            &layout,
+            0,
+        )
+        .unwrap();
 
         // Modify memory
         let modified_data = vec![0xBB; PAGE_SIZE_USIZE];
@@ -468,9 +477,13 @@ mod tests {
         }
 
         // Create initial snapshot (State 1)
-        let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, Some(&dirty_pages), &layout)
-                .unwrap();
+        let mut snapshot_manager = super::SharedMemorySnapshotManager::new(
+            &mut shared_mem,
+            Some(&dirty_pages),
+            &layout,
+            0,
+        )
+        .unwrap();
 
         // State 2: Modify and create second snapshot
         let tracker2 = shared_mem.start_tracking_dirty_pages().unwrap();
@@ -494,7 +507,7 @@ mod tests {
         }
 
         snapshot_manager
-            .create_new_snapshot(&mut shared_mem, Some(&dirty_pages2))
+            .create_new_snapshot(&mut shared_mem, Some(&dirty_pages2), 0)
             .unwrap();
 
         // State 3: Modify again
@@ -594,9 +607,13 @@ mod tests {
         }
 
         // Create snapshot
-        let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, Some(&dirty_pages), &layout)
-                .unwrap();
+        let mut snapshot_manager = super::SharedMemorySnapshotManager::new(
+            &mut shared_mem,
+            Some(&dirty_pages),
+            &layout,
+            0,
+        )
+        .unwrap();
 
         // Modify first and third pages
         let modified_data = [vec![0x11; PAGE_SIZE_USIZE], vec![0x22; PAGE_SIZE_USIZE]];
@@ -680,9 +697,13 @@ mod tests {
             }
         }
 
-        let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, Some(&dirty_pages), &layout)
-                .unwrap();
+        let mut snapshot_manager = super::SharedMemorySnapshotManager::new(
+            &mut shared_mem,
+            Some(&dirty_pages),
+            &layout,
+            0,
+        )
+        .unwrap();
 
         // Cycle 2: Modify and snapshot
         let tracker2 = shared_mem.start_tracking_dirty_pages().unwrap();
@@ -713,7 +734,7 @@ mod tests {
         }
 
         snapshot_manager
-            .create_new_snapshot(&mut shared_mem, Some(&dirty_pages2))
+            .create_new_snapshot(&mut shared_mem, Some(&dirty_pages2), 0)
             .unwrap();
 
         // Cycle 3: Modify again
@@ -814,6 +835,7 @@ mod tests {
             &mut shared_mem,
             Some(&dirty_pages_snapshot),
             &layout,
+            0,
         )
         .unwrap();
 
@@ -966,9 +988,13 @@ mod tests {
         }
 
         // Create initial checkpoint
-        let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, Some(&dirty_pages), &layout)
-                .unwrap();
+        let mut snapshot_manager = super::SharedMemorySnapshotManager::new(
+            &mut shared_mem,
+            Some(&dirty_pages),
+            &layout,
+            0,
+        )
+        .unwrap();
 
         // Simulate function call 1: modify pages 0 and 2
         let tracker1 = shared_mem.start_tracking_dirty_pages().unwrap();
@@ -998,7 +1024,7 @@ mod tests {
 
         // Checkpoint after function 1
         snapshot_manager
-            .create_new_snapshot(&mut shared_mem, Some(&dirty_pages1))
+            .create_new_snapshot(&mut shared_mem, Some(&dirty_pages1), 0)
             .unwrap();
 
         // Simulate function call 2: modify pages 1 and 3
@@ -1029,7 +1055,7 @@ mod tests {
 
         // Checkpoint after function 2
         snapshot_manager
-            .create_new_snapshot(&mut shared_mem, Some(&dirty_pages2))
+            .create_new_snapshot(&mut shared_mem, Some(&dirty_pages2), 0)
             .unwrap();
 
         // Simulate function call 3: modify all pages
@@ -1228,9 +1254,13 @@ mod tests {
         }
 
         // Create snapshot
-        let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, Some(&dirty_pages), &layout)
-                .unwrap();
+        let mut snapshot_manager = super::SharedMemorySnapshotManager::new(
+            &mut shared_mem,
+            Some(&dirty_pages),
+            &layout,
+            0,
+        )
+        .unwrap();
 
         // Modify only the dirty pages
         let modified_patterns = [

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -270,8 +270,9 @@ impl MultiUseSandbox {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(crate) fn restore_state(&mut self) -> Result<()> {
         let mem_mgr = self.mem_mgr.unwrap_mgr_mut();
-        let rgns_to_unmap = mem_mgr.restore_state_from_last_snapshot()?;
-        unsafe { self.vm.unmap_regions(rgns_to_unmap)? };
+        let dirty_pages = self.vm.get_and_clear_dirty_pages()?;
+        let rgns_to_umap = mem_mgr.restore_state_from_last_snapshot(&dirty_pages)?;
+        unsafe { self.vm.unmap_regions(rgns_to_umap)? };
         Ok(())
     }
 
@@ -354,10 +355,11 @@ impl DevolvableSandbox<MultiUseSandbox, MultiUseSandbox, Noop<MultiUseSandbox, M
     /// The devolve can be used to return the MultiUseSandbox to the state before the code was loaded. Thus avoiding initialisation overhead
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     fn devolve(mut self, _tsn: Noop<MultiUseSandbox, MultiUseSandbox>) -> Result<MultiUseSandbox> {
+        let dirty_pages = self.vm.get_and_clear_dirty_pages()?;
         let rgns_to_unmap = self
             .mem_mgr
             .unwrap_mgr_mut()
-            .pop_and_restore_state_from_snapshot()?;
+            .pop_and_restore_state_from_snapshot(&dirty_pages)?;
         unsafe { self.vm.unmap_regions(rgns_to_unmap)? };
         Ok(self)
     }
@@ -389,7 +391,10 @@ where
         let mut ctx = self.new_call_context();
         transition_func.call(&mut ctx)?;
         let mut sbox = ctx.finish_no_reset();
-        sbox.mem_mgr.unwrap_mgr_mut().push_state()?;
+        let vm_dirty_pages = sbox.vm.get_and_clear_dirty_pages()?;
+        sbox.mem_mgr
+            .unwrap_mgr_mut()
+            .push_state(Some(&vm_dirty_pages))?;
         Ok(sbox)
     }
 }