Fix some stuff

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

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -307,6 +307,7 @@ pub(crate) struct HypervLinuxDriver {
     vcpu_fd: VcpuFd,
     entrypoint: u64,
     mem_regions: Vec<MemoryRegion>,
+    n_initial_regions: usize,
     orig_rsp: GuestPtr,
     interrupt_handle: Arc<LinuxInterruptHandle>,
 
@@ -417,10 +418,20 @@ impl HypervLinuxDriver {
         // the downside of doing this here is that the call to get_dirty_log will takes longer as the number of pages increase
         // but for larger sandboxes its easily cheaper than copying all the pages
 
-        #[cfg(mshv2)]
-        vm_fd.get_dirty_log(base_pfn, total_size, CLEAR_DIRTY_BIT_FLAG)?;
-        #[cfg(mshv3)]
-        vm_fd.get_dirty_log(base_pfn, total_size, MSHV_GPAP_ACCESS_OP_CLEAR as u8)?;
+        // Clear dirty bits for each memory region separately since they may not be contiguous
+        for region in &mem_regions {
+            let mshv_region: mshv_user_mem_region = region.to_owned().into();
+            let region_size = region.guest_region.len();
+
+            #[cfg(mshv2)]
+            vm_fd.get_dirty_log(mshv_region.guest_pfn, region_size, CLEAR_DIRTY_BIT_FLAG)?;
+            #[cfg(mshv3)]
+            vm_fd.get_dirty_log(
+                mshv_region.guest_pfn,
+                region_size,
+                MSHV_GPAP_ACCESS_OP_CLEAR as u8,
+            )?;
+        }
 
         let interrupt_handle = Arc::new(LinuxInterruptHandle {
             running: AtomicU64::new(0),
@@ -452,6 +463,7 @@ impl HypervLinuxDriver {
             page_size: 0,
             vm_fd,
             vcpu_fd,
+            n_initial_regions: mem_regions.len(),
             mem_regions,
             entrypoint: entrypoint_ptr.absolute()?,
             orig_rsp: rsp_ptr,
@@ -887,7 +899,8 @@ impl Hypervisor for HypervLinuxDriver {
         self.interrupt_handle.clone()
     }
 
-    fn get_and_clear_dirty_pages(&mut self) -> Result<Vec<u64>> {
+    // TODO: Implement getting additional host-mapped dirty pages.
+    fn get_and_clear_dirty_pages(&mut self) -> Result<(Vec<u64>, Option<Vec<Vec<u64>>>)> {
         let first_mshv_region: mshv_user_mem_region = self
             .mem_regions
             .first()
@@ -896,16 +909,38 @@ impl Hypervisor for HypervLinuxDriver {
             ))?
             .to_owned()
             .into();
-        let total_size = self.mem_regions.iter().map(|r| r.guest_region.len()).sum();
-        let res = self.vm_fd.get_dirty_log(
+
+        let n_contiguous = self
+            .mem_regions
+            .windows(2)
+            .take_while(|window| window[0].guest_region.end == window[1].guest_region.start)
+            .count()
+            + 1; // +1 because windows(2) gives us n-1 pairs for n regions
+
+        if n_contiguous != self.n_initial_regions {
+            return Err(new_error!(
+                "get_and_clear_dirty_pages: not all regions are contiguous, expected {} but got {}",
+                self.n_initial_regions,
+                n_contiguous
+            ));
+        }
+
+        let sandbox_total_size = self
+            .mem_regions
+            .iter()
+            .take(n_contiguous)
+            .map(|r| r.guest_region.len())
+            .sum();
+
+        let sandbox_dirty_pages = self.vm_fd.get_dirty_log(
             first_mshv_region.guest_pfn,
-            total_size,
+            sandbox_total_size,
             #[cfg(mshv2)]
             CLEAR_DIRTY_BIT_FLAG,
             #[cfg(mshv3)]
             (MSHV_GPAP_ACCESS_OP_CLEAR as u8),
         )?;
-        Ok(res)
+        Ok((sandbox_dirty_pages, None))
     }
 
     #[cfg(crashdump)]

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -292,6 +292,7 @@ pub(crate) struct KVMDriver {
     entrypoint: u64,
     orig_rsp: GuestPtr,
     mem_regions: Vec<MemoryRegion>,
+    n_initial_regions: usize,
     interrupt_handle: Arc<LinuxInterruptHandle>,
 
     #[cfg(gdb)]
@@ -374,6 +375,7 @@ impl KVMDriver {
             vcpu_fd,
             entrypoint,
             orig_rsp: rsp_gp,
+            n_initial_regions: mem_regions.len(),
             mem_regions,
             interrupt_handle: interrupt_handle.clone(),
             #[cfg(gdb)]
@@ -752,11 +754,27 @@ impl Hypervisor for KVMDriver {
         self.interrupt_handle.clone()
     }
 
-    fn get_and_clear_dirty_pages(&mut self) -> Result<Vec<u64>> {
+    // TODO: Implement getting additional host-mapped dirty pages.
+    fn get_and_clear_dirty_pages(&mut self) -> Result<(Vec<u64>, Option<Vec<Vec<u64>>>)> {
+        let n_contiguous = self
+            .mem_regions
+            .windows(2)
+            .take_while(|window| window[0].guest_region.end == window[1].guest_region.start)
+            .count()
+            + 1; // +1 because windows(2) gives us n-1 pairs for n regions
+
+        if n_contiguous != self.n_initial_regions {
+            return Err(new_error!(
+                "get_and_clear_dirty_pages: not all regions are contiguous, expected {} but got {}",
+                self.n_initial_regions,
+                n_contiguous
+            ));
+        }
         let mut page_indices = vec![];
         let mut current_page = 0;
+
         // Iterate over all memory regions and get the dirty pages for each region ignoring guard pages which cannot be dirty
-        for (i, mem_region) in self.mem_regions.iter().enumerate() {
+        for (i, mem_region) in self.mem_regions.iter().take(n_contiguous).enumerate() {
             let num_pages = mem_region.guest_region.len() / PAGE_SIZE_USIZE;
             let bitmap = match mem_region.flags {
                 MemoryRegionFlags::READ => {
@@ -780,15 +798,15 @@ impl Hypervisor for KVMDriver {
             current_page += num_pages;
         }
 
-        // covert vec of page indices to vec of blocks
-        let mut res = new_page_bitmap(current_page * PAGE_SIZE_USIZE, false)?;
+        // convert vec of page indices to vec of blocks
+        let mut sandbox_dirty_pages = new_page_bitmap(current_page * PAGE_SIZE_USIZE, false)?;
         for page_idx in page_indices {
             let block_idx = page_idx / PAGES_IN_BLOCK;
             let bit_idx = page_idx % PAGES_IN_BLOCK;
-            res[block_idx] |= 1 << bit_idx;
+            sandbox_dirty_pages[block_idx] |= 1 << bit_idx;
         }
 
-        Ok(res)
+        Ok((sandbox_dirty_pages, None))
     }
 
     #[cfg(crashdump)]

src/hyperlight_host/src/hypervisor/mod.rs

@@ -199,7 +199,10 @@ pub(crate) trait Hypervisor: Debug + Sync + Send {
     /// Get dirty pages as a bitmap (Vec<u64>).
     /// Each bit in a u64 represents a page.
     /// This also clears the bitflags, marking the pages as non-dirty.
-    fn get_and_clear_dirty_pages(&mut self) -> Result<Vec<u64>>;
+    /// The Vec<u64> in the tuple is the bitmap of the first contiguous memory regions, which represents the sandbox itself.
+    /// The Vec<Vec<u64>> in the tuple are the host-mapped regions, which aren't necessarily contiguous, and not yet implemented
+    #[allow(clippy::type_complexity)]
+    fn get_and_clear_dirty_pages(&mut self) -> Result<(Vec<u64>, Option<Vec<Vec<u64>>>)>;
 
     /// Get InterruptHandle to underlying VM
     fn interrupt_handle(&self) -> Arc<dyn InterruptHandle>;

src/hyperlight_host/src/mem/mgr.rs

@@ -286,12 +286,8 @@ where
         // merge the host dirty page map into the dirty bitmap
         let merged = bitmap_union(&res, vm_dirty_bitmap);
 
-        let snapshot_manager = SharedMemorySnapshotManager::new(
-            &mut self.shared_mem,
-            &merged,
-            layout,
-            self.mapped_rgns,
-        )?;
+        let mut snapshot_manager = SharedMemorySnapshotManager::new(&mut self.shared_mem, layout)?;
+        snapshot_manager.create_new_snapshot(&mut self.shared_mem, &merged, self.mapped_rgns)?;
         existing_snapshot_manager.replace(snapshot_manager);
         Ok(())
     }
@@ -300,6 +296,10 @@ 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
+    ///
+    /// 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_state_from_last_snapshot(&mut self, dirty_bitmap: &[u64]) -> Result<u64> {
         let mut snapshot_manager = self
             .snapshot_manager
@@ -311,7 +311,10 @@ where
                 log_then_return!("Snapshot manager not initialized");
             }
             Some(snapshot_manager) => {
-                snapshot_manager.restore_from_snapshot(&mut self.shared_mem, dirty_bitmap)
+                let old_rgns = self.mapped_rgns;
+                self.mapped_rgns =
+                    snapshot_manager.restore_from_snapshot(&mut self.shared_mem, dirty_bitmap)?;
+                Ok(old_rgns - self.mapped_rgns)
             }
         }
     }

src/hyperlight_host/src/mem/shared_memory_snapshot_manager.rs

@@ -45,15 +45,8 @@ impl SharedMemorySnapshotManager {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     pub(super) fn new<S: SharedMemory>(
         shared_mem: &mut S,
-        dirty_page_bitmap: &[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_bitmap, 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();
         let output_data_size_offset = layout.get_output_data_size_offset();
@@ -76,20 +69,20 @@ impl SharedMemorySnapshotManager {
         })??;
 
         Ok(Self {
-            snapshots: vec![diff],
+            snapshots: vec![],
             input_data_size,
             output_data_size,
             output_data_buffer_offset,
             input_data_buffer_offset,
         })
     }
 
-    fn build_snapshot_from_dirty_page_map<S: SharedMemory>(
+    pub(super) fn create_new_snapshot<S: SharedMemory>(
+        &mut self,
         shared_mem: &mut S,
         dirty_page_bitmap: &[u64],
         mapped_rgns: u64,
-    ) -> Result<PageSnapshot> {
-        // If there is no dirty page map, return an empty snapshot
+    ) -> Result<()> {
         if dirty_page_bitmap.is_empty() {
             return Err(new_error!(
                 "Tried to build snapshot from empty dirty page bitmap"
@@ -143,18 +136,6 @@ impl SharedMemorySnapshotManager {
 
         // Create the snapshot with the pre-allocated buffer
         let snapshot = PageSnapshot::with_pages_and_buffer(dirty_pages, buffer, mapped_rgns);
-
-        Ok(snapshot)
-    }
-
-    pub(super) fn create_new_snapshot<S: SharedMemory>(
-        &mut self,
-        shared_mem: &mut S,
-        dirty_page_map: &[u64],
-        mapped_rgns: u64,
-    ) -> Result<()> {
-        let snapshot =
-            Self::build_snapshot_from_dirty_page_map(shared_mem, dirty_page_map, mapped_rgns)?;
         self.snapshots.push(snapshot);
         Ok(())
     }
@@ -383,8 +364,11 @@ mod tests {
 
         // Create snapshot
         let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, &dirty_pages, &layout, 0)
-                .unwrap();
+            super::SharedMemorySnapshotManager::new(&mut shared_mem, &layout).unwrap();
+
+        snapshot_manager
+            .create_new_snapshot(&mut shared_mem, &dirty_pages, 0)
+            .unwrap();
 
         // Modify memory
         let modified_data = vec![0xBB; PAGE_SIZE_USIZE];
@@ -454,8 +438,11 @@ mod tests {
 
         // Create initial snapshot (State 1)
         let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, &dirty_pages, &layout, 0)
-                .unwrap();
+            super::SharedMemorySnapshotManager::new(&mut shared_mem, &layout).unwrap();
+
+        snapshot_manager
+            .create_new_snapshot(&mut shared_mem, &dirty_pages, 0)
+            .unwrap();
 
         // State 2: Modify and create second snapshot
         let state2_data = vec![0x22; PAGE_SIZE_USIZE];
@@ -570,8 +557,10 @@ mod tests {
 
         // Create snapshot
         let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, &dirty_pages, &layout, 0)
-                .unwrap();
+            super::SharedMemorySnapshotManager::new(&mut shared_mem, &layout).unwrap();
+        snapshot_manager
+            .create_new_snapshot(&mut shared_mem, &dirty_pages, 0)
+            .unwrap();
 
         // Modify first and third pages
         let modified_data = [vec![0x11; PAGE_SIZE_USIZE], vec![0x22; PAGE_SIZE_USIZE]];
@@ -651,8 +640,10 @@ mod tests {
         }
 
         let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, &dirty_pages, &layout, 0)
-                .unwrap();
+            super::SharedMemorySnapshotManager::new(&mut shared_mem, &layout).unwrap();
+        snapshot_manager
+            .create_new_snapshot(&mut shared_mem, &dirty_pages, 0)
+            .unwrap();
 
         // Cycle 2: Modify and snapshot
         let cycle2_page0 = vec![0x02; PAGE_SIZE_USIZE];
@@ -768,13 +759,11 @@ mod tests {
             }
         }
 
-        let mut snapshot_manager = super::SharedMemorySnapshotManager::new(
-            &mut shared_mem,
-            &dirty_pages_snapshot,
-            &layout,
-            0,
-        )
-        .unwrap();
+        let mut snapshot_manager =
+            super::SharedMemorySnapshotManager::new(&mut shared_mem, &layout).unwrap();
+        snapshot_manager
+            .create_new_snapshot(&mut shared_mem, &dirty_pages_snapshot, 0)
+            .unwrap();
 
         // Modify pages in init_data area
         let page0_offset = init_data_offset;
@@ -921,8 +910,10 @@ mod tests {
 
         // Create initial checkpoint
         let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, &dirty_pages, &layout, 0)
-                .unwrap();
+            super::SharedMemorySnapshotManager::new(&mut shared_mem, &layout).unwrap();
+        snapshot_manager
+            .create_new_snapshot(&mut shared_mem, &dirty_pages, 0)
+            .unwrap();
 
         // Simulate function call 1: modify pages 0 and 2
         let func1_page0 = vec![0x10; PAGE_SIZE_USIZE];
@@ -1168,8 +1159,10 @@ mod tests {
 
         // Create snapshot
         let mut snapshot_manager =
-            super::SharedMemorySnapshotManager::new(&mut shared_mem, &dirty_pages, &layout, 0)
-                .unwrap();
+            super::SharedMemorySnapshotManager::new(&mut shared_mem, &layout).unwrap();
+        snapshot_manager
+            .create_new_snapshot(&mut shared_mem, &dirty_pages, 0)
+            .unwrap();
 
         // Modify only the dirty pages
         let modified_patterns = [