fixup! Add API for getting VM's dirty pages, and add some bitmap utility functions.

Author: ludfjig

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -306,8 +306,12 @@ pub(crate) struct HypervLinuxDriver {
     vm_fd: VmFd,
     vcpu_fd: VcpuFd,
     entrypoint: u64,
+    // Regions part of the original sandbox
     mem_regions: Vec<MemoryRegion>,
-    n_initial_regions: usize,
+    // Size of the (contigous) sandbox mem_regions
+    mem_regions_size: usize,
+    // Regions that are mapped after sandbox creation
+    mmap_regions: Vec<MemoryRegion>,
     orig_rsp: GuestPtr,
     interrupt_handle: Arc<LinuxInterruptHandle>,
 
@@ -453,8 +457,9 @@ impl HypervLinuxDriver {
             page_size: 0,
             vm_fd,
             vcpu_fd,
-            n_initial_regions: mem_regions.len(),
             mem_regions,
+            mem_regions_size: total_size,
+            mmap_regions: Vec::new(),
             entrypoint: entrypoint_ptr.absolute()?,
             orig_rsp: rsp_ptr,
             interrupt_handle: interrupt_handle.clone(),
@@ -623,16 +628,14 @@ impl Hypervisor for HypervLinuxDriver {
         }
         let mshv_region: mshv_user_mem_region = rgn.to_owned().into();
         self.vm_fd.map_user_memory(mshv_region)?;
-        self.mem_regions.push(rgn.to_owned());
+        self.mmap_regions.push(rgn.to_owned());
         Ok(())
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     unsafe fn unmap_regions(&mut self, n: u64) -> Result<()> {
-        for rgn in self
-            .mem_regions
-            .split_off(self.mem_regions.len() - n as usize)
-        {
+        let n_keep = self.mmap_regions.len() - n as usize;
+        for rgn in self.mmap_regions.split_off(n_keep) {
             let mshv_region: mshv_user_mem_region = rgn.to_owned().into();
             self.vm_fd.unmap_user_memory(mshv_region)?;
         }
@@ -922,31 +925,9 @@ impl Hypervisor for HypervLinuxDriver {
             .to_owned()
             .into();
 
-        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 mut sandbox_dirty_pages = self.vm_fd.get_dirty_log(
             first_mshv_region.guest_pfn,
-            sandbox_total_size,
+            self.mem_regions_size,
             #[cfg(mshv2)]
             CLEAR_DIRTY_BIT_FLAG,
             #[cfg(mshv3)]
@@ -958,7 +939,7 @@ impl Hypervisor for HypervLinuxDriver {
         // TODO: remove this once bug in mshv is fixed. The bug makes it possible
         // for non-mapped memory to incorrectly be marked dirty. To fix this, we just zero out
         // any bits that are not within the sandbox size.
-        let sandbox_pages = sandbox_total_size / self.page_size;
+        let sandbox_pages = self.mem_regions_size / self.page_size;
         let last_block_idx = sandbox_dirty_pages.len().saturating_sub(1);
         if let Some(last_block) = sandbox_dirty_pages.last_mut() {
             let last_block_start_page = last_block_idx * 64;

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -291,8 +291,10 @@ pub(crate) struct KVMDriver {
     vcpu_fd: VcpuFd,
     entrypoint: u64,
     orig_rsp: GuestPtr,
+    // Regions part of the original sandbox
     mem_regions: Vec<MemoryRegion>,
-    n_initial_regions: usize,
+    // Regions that are mapped after sandbox creation
+    mmap_regions: Vec<MemoryRegion>,
     interrupt_handle: Arc<LinuxInterruptHandle>,
 
     #[cfg(gdb)]
@@ -375,8 +377,8 @@ impl KVMDriver {
             vcpu_fd,
             entrypoint,
             orig_rsp: rsp_gp,
-            n_initial_regions: mem_regions.len(),
             mem_regions,
+            mmap_regions: Vec::new(),
             interrupt_handle: interrupt_handle.clone(),
             #[cfg(gdb)]
             debug,
@@ -507,18 +509,19 @@ impl Hypervisor for KVMDriver {
         }
 
         let mut kvm_region: kvm_userspace_memory_region = region.clone().into();
-        kvm_region.slot = self.mem_regions.len() as u32;
+        kvm_region.slot = (self.mem_regions.len() + self.mmap_regions.len()) as u32;
         unsafe { self.vm_fd.set_user_memory_region(kvm_region) }?;
-        self.mem_regions.push(region.to_owned());
+        self.mmap_regions.push(region.to_owned());
         Ok(())
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     unsafe fn unmap_regions(&mut self, n: u64) -> Result<()> {
-        let n_keep = self.mem_regions.len() - n as usize;
-        for (k, region) in self.mem_regions.split_off(n_keep).iter().enumerate() {
+        let n_keep = self.mmap_regions.len() - n as usize;
+        let n_sandbox_regions = self.mem_regions.len();
+        for (k, region) in self.mmap_regions.split_off(n_keep).iter().enumerate() {
             let mut kvm_region: kvm_userspace_memory_region = region.clone().into();
-            kvm_region.slot = (n_keep + k) as u32;
+            kvm_region.slot = (n_sandbox_regions + n_keep + k) as u32;
             // Setting memory_size to 0 unmaps the slot's region
             // From https://docs.kernel.org/virt/kvm/api.html
             // > Deleting a slot is done by passing zero for memory_size.
@@ -756,25 +759,11 @@ impl Hypervisor for KVMDriver {
 
     // TODO: Implement getting additional host-mapped dirty pages.
     fn get_and_clear_dirty_pages(&mut self) -> Result<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().take(n_contiguous).enumerate() {
+        for (i, mem_region) in self.mem_regions.iter().enumerate() {
             let num_pages = mem_region.guest_region.len() / PAGE_SIZE_USIZE;
             let bitmap = match mem_region.flags {
                 MemoryRegionFlags::READ => {