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,18 +939,13 @@ 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 last_block_idx = sandbox_dirty_pages.len().saturating_sub(1);
+        let sandbox_pages = self.mem_regions_size / self.page_size;
         if let Some(last_block) = sandbox_dirty_pages.last_mut() {
-            let last_block_start_page = last_block_idx * 64;
-            let last_block_end_page = last_block_start_page + 64;
-
-            // If the last block extends beyond the sandbox, clear the invalid bits
-            if last_block_end_page > sandbox_pages {
-                let valid_bits_in_last_block = sandbox_pages - last_block_start_page;
-                let mask = (1u64 << valid_bits_in_last_block) - 1;
-                *last_block &= mask;
-            }
+            let mask = match sandbox_pages % 64 {
+                0 => u64::MAX,
+                tail_bits => (1u64 << tail_bits) - 1,
+            };
+            *last_block &= mask;
         }
         Ok(sandbox_dirty_pages)
     }

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 => {

src/hyperlight_host/src/mem/bitmap.rs

@@ -13,7 +13,6 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
-use std::cmp::Ordering;
 
 use hyperlight_common::mem::{PAGE_SIZE_USIZE, PAGES_IN_BLOCK};
 use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
@@ -47,27 +46,18 @@ pub fn new_page_bitmap(size_in_bytes: usize, init_dirty: bool) -> Result<Vec<u64
 
 /// Returns the union (bitwise OR) of two bitmaps. The resulting bitmap will have the same length
 /// as the longer of the two input bitmaps.
-pub(crate) fn bitmap_union(bitmap: &[u64], other_bitmap: &[u64]) -> Vec<u64> {
-    let min_len = bitmap.len().min(other_bitmap.len());
-    let max_len = bitmap.len().max(other_bitmap.len());
+pub(crate) fn bitmap_union(bitmap: &[u64], other_bitmap: &[u64]) -> Result<Vec<u64>> {
+    if bitmap.len() != other_bitmap.len() {
+        log_then_return!("Bitmaps must be of the same length to union them.");
+    }
 
-    let mut result = vec![0; max_len];
+    let mut result = vec![0; bitmap.len()];
 
-    for i in 0..min_len {
+    for i in 0..bitmap.len() {
         result[i] = bitmap[i] | other_bitmap[i];
     }
 
-    match bitmap.len().cmp(&other_bitmap.len()) {
-        Ordering::Greater => {
-            result[min_len..].copy_from_slice(&bitmap[min_len..]);
-        }
-        Ordering::Less => {
-            result[min_len..].copy_from_slice(&other_bitmap[min_len..]);
-        }
-        Ordering::Equal => {}
-    }
-
-    result
+    Ok(result)
 }
 
 // Used as a helper struct to implement an iterator on.
@@ -271,23 +261,14 @@ mod tests {
         let f = 0b100000000000000001010000000001010100000000000;
         let bitmap = vec![a, b, c];
         let other_bitmap = vec![d, e, f];
-        let union = bitmap_union(&bitmap, &other_bitmap);
+        let union = bitmap_union(&bitmap, &other_bitmap).unwrap();
         assert_eq!(union, vec![a | d, b | e, c | f]);
 
         // different length
-        let union = bitmap_union(&[a], &[d, e, f]);
-        assert_eq!(union, vec![a | d, e, f]);
-
-        let union = bitmap_union(&[a, b, c], &[d]);
-        assert_eq!(union, vec![a | d, b, c]);
-
-        let union = bitmap_union(&[], &[d, e]);
-        assert_eq!(union, vec![d, e]);
-
-        let union = bitmap_union(&[a, b, c], &[]);
-        assert_eq!(union, vec![a, b, c]);
+        bitmap_union(&[a], &[d, e, f]).unwrap_err();
 
-        let union = bitmap_union(&[], &[]);
+        // empty bitmaps
+        let union = bitmap_union(&[], &[]).unwrap();
         let empty: Vec<u64> = vec![];
         assert_eq!(union, empty);