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

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_common/src/mem.rs

@@ -17,6 +17,8 @@ limitations under the License.
 pub const PAGE_SHIFT: u64 = 12;
 pub const PAGE_SIZE: u64 = 1 << 12;
 pub const PAGE_SIZE_USIZE: usize = 1 << 12;
+// The number of pages in 1 "block". A single u64 can be used as bitmap to keep track of all dirty pages in a block.
+pub const PAGES_IN_BLOCK: usize = 64;
 
 /// A memory region in the guest address space
 #[derive(Debug, Clone, Copy)]

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -29,6 +29,8 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 
 use log::{LevelFilter, error};
+#[cfg(mshv3)]
+use mshv_bindings::MSHV_GPAP_ACCESS_OP_CLEAR;
 #[cfg(mshv2)]
 use mshv_bindings::hv_message;
 use mshv_bindings::{
@@ -76,6 +78,9 @@ use crate::sandbox::SandboxConfiguration;
 use crate::sandbox::uninitialized::SandboxRuntimeConfig;
 use crate::{Result, log_then_return, new_error};
 
+#[cfg(mshv2)]
+const CLEAR_DIRTY_BIT_FLAG: u64 = 0b100;
+
 #[cfg(gdb)]
 mod debug {
     use std::sync::{Arc, Mutex};
@@ -351,6 +356,7 @@ impl HypervLinuxDriver {
             vm_fd.initialize()?;
             vm_fd
         };
+        vm_fd.enable_dirty_page_tracking()?;
 
         let mut vcpu_fd = vm_fd.create_vcpu(0)?;
 
@@ -391,13 +397,31 @@ impl HypervLinuxDriver {
             (None, None)
         };
 
+        let mut base_pfn = u64::MAX;
+        let mut total_size: usize = 0;
+
         mem_regions.iter().try_for_each(|region| {
-            let mshv_region = region.to_owned().into();
+            let mshv_region: mshv_user_mem_region = region.to_owned().into();
+            if base_pfn == u64::MAX {
+                base_pfn = mshv_region.guest_pfn;
+            }
+            total_size += mshv_region.size as usize;
             vm_fd.map_user_memory(mshv_region)
         })?;
 
         Self::setup_initial_sregs(&mut vcpu_fd, pml4_ptr.absolute()?)?;
 
+        // get/clear the dirty page bitmap, mshv sets all the bit dirty at initialization
+        // if we dont clear them then we end up taking a complete snapsot of memory page by page which gets
+        // progressively slower as the sandbox size increases
+        // 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)?;
+
         let interrupt_handle = Arc::new(LinuxInterruptHandle {
             running: AtomicU64::new(0),
             cancel_requested: AtomicBool::new(false),
@@ -863,6 +887,27 @@ impl Hypervisor for HypervLinuxDriver {
         self.interrupt_handle.clone()
     }
 
+    fn get_and_clear_dirty_pages(&mut self) -> Result<Vec<u64>> {
+        let first_mshv_region: mshv_user_mem_region = self
+            .mem_regions
+            .first()
+            .ok_or(new_error!(
+                "tried to get dirty page bitmap of 0-sized region"
+            ))?
+            .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(
+            first_mshv_region.guest_pfn,
+            total_size,
+            #[cfg(mshv2)]
+            CLEAR_DIRTY_BIT_FLAG,
+            #[cfg(mshv3)]
+            (MSHV_GPAP_ACCESS_OP_CLEAR as u8),
+        )?;
+        Ok(res)
+    }
+
     #[cfg(crashdump)]
     fn crashdump_context(&self) -> Result<Option<super::crashdump::CrashDumpContext>> {
         if self.rt_cfg.guest_core_dump {

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -55,6 +55,7 @@ use super::{
 use super::{HyperlightExit, Hypervisor, InterruptHandle, VirtualCPU};
 use crate::hypervisor::fpu::FP_CONTROL_WORD_DEFAULT;
 use crate::hypervisor::wrappers::WHvGeneralRegisters;
+use crate::mem::bitmap::new_page_bitmap;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::ptr::{GuestPtr, RawPtr};
 #[cfg(crashdump)]
@@ -606,6 +607,12 @@ impl Hypervisor for HypervWindowsDriver {
         Ok(())
     }
 
+    fn get_and_clear_dirty_pages(&mut self) -> Result<Vec<u64>> {
+        // For now we just mark all pages dirty which is the equivalent of taking a full snapshot
+        let total_size = self.mem_regions.iter().map(|r| r.guest_region.len()).sum();
+        new_page_bitmap(total_size, true)
+    }
+
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     unsafe fn map_region(&mut self, _rgn: &MemoryRegion) -> Result<()> {
         log_then_return!("Mapping host memory into the guest not yet supported on this platform");

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -21,7 +21,10 @@ use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 
-use kvm_bindings::{KVM_MEM_READONLY, kvm_fpu, kvm_regs, kvm_userspace_memory_region};
+use hyperlight_common::mem::{PAGE_SIZE_USIZE, PAGES_IN_BLOCK};
+use kvm_bindings::{
+    KVM_MEM_LOG_DIRTY_PAGES, KVM_MEM_READONLY, kvm_fpu, kvm_regs, kvm_userspace_memory_region,
+};
 use kvm_ioctls::Cap::UserMemory;
 use kvm_ioctls::{Kvm, VcpuExit, VcpuFd, VmFd};
 use log::LevelFilter;
@@ -43,7 +46,8 @@ use super::{
 use super::{HyperlightExit, Hypervisor, InterruptHandle, LinuxInterruptHandle, VirtualCPU};
 #[cfg(gdb)]
 use crate::HyperlightError;
-use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
+use crate::mem::bitmap::{bit_index_iterator, new_page_bitmap};
+use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
 use crate::mem::ptr::{GuestPtr, RawPtr};
 use crate::sandbox::SandboxConfiguration;
 #[cfg(crashdump)]
@@ -284,7 +288,7 @@ mod debug {
 /// A Hypervisor driver for KVM on Linux
 pub(crate) struct KVMDriver {
     _kvm: Kvm,
-    _vm_fd: VmFd,
+    vm_fd: VmFd,
     vcpu_fd: VcpuFd,
     entrypoint: u64,
     orig_rsp: GuestPtr,
@@ -329,7 +333,7 @@ impl KVMDriver {
                 userspace_addr: region.host_region.start as u64,
                 flags: match perm_flags {
                     MemoryRegionFlags::READ => KVM_MEM_READONLY,
-                    _ => 0, // normal, RWX
+                    _ => KVM_MEM_LOG_DIRTY_PAGES, // normal, RWX
                 },
             };
             unsafe { vm_fd.set_user_memory_region(kvm_region) }
@@ -378,7 +382,7 @@ impl KVMDriver {
         #[allow(unused_mut)]
         let mut hv = Self {
             _kvm: kvm,
-            _vm_fd: vm_fd,
+            vm_fd,
             vcpu_fd,
             entrypoint,
             orig_rsp: rsp_gp,
@@ -734,6 +738,45 @@ impl Hypervisor for KVMDriver {
         self.interrupt_handle.clone()
     }
 
+    fn get_and_clear_dirty_pages(&mut self) -> Result<Vec<u64>> {
+        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() {
+            let num_pages = mem_region.guest_region.len() / PAGE_SIZE_USIZE;
+            let bitmap = match mem_region.flags {
+                MemoryRegionFlags::READ => {
+                    // read-only page. It can never be dirty so return zero dirty pages.
+                    new_page_bitmap(mem_region.guest_region.len(), false)?
+                }
+                _ => {
+                    if mem_region.region_type == MemoryRegionType::GuardPage {
+                        //  Trying to get dirty pages for a guard page region results in a VMMSysError(2)
+                        new_page_bitmap(mem_region.guest_region.len(), false)?
+                    } else {
+                        // Get the dirty bitmap for the memory region
+                        self.vm_fd
+                            .get_dirty_log(i as u32, mem_region.guest_region.len())?
+                    }
+                }
+            };
+            for page_idx in bit_index_iterator(&bitmap) {
+                page_indices.push(current_page + page_idx);
+            }
+            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)?;
+        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;
+        }
+
+        Ok(res)
+    }
+
     #[cfg(crashdump)]
     fn crashdump_context(&self) -> Result<Option<crashdump::CrashDumpContext>> {
         if self.rt_cfg.guest_core_dump {

src/hyperlight_host/src/hypervisor/mod.rs

@@ -196,6 +196,11 @@ pub(crate) trait Hypervisor: Debug + Sync + Send {
         None
     }
 
+    /// 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>>;
+
     /// Get InterruptHandle to underlying VM
     fn interrupt_handle(&self) -> Arc<dyn InterruptHandle>;
 
@@ -506,6 +511,7 @@ pub(crate) mod tests {
     use super::handlers::{MemAccessHandler, OutBHandler};
     #[cfg(gdb)]
     use crate::hypervisor::DbgMemAccessHandlerCaller;
+    use crate::mem::dirty_page_tracking::DirtyPageTracking;
     use crate::mem::ptr::RawPtr;
     use crate::sandbox::uninitialized::GuestBinary;
     #[cfg(any(crashdump, gdb))]
@@ -557,7 +563,10 @@ pub(crate) mod tests {
         let rt_cfg: SandboxRuntimeConfig = Default::default();
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
+        let tracker = sandbox.tracker.unwrap();
         let (_hshm, mut gshm) = sandbox.mgr.build();
+        // we need to undo the mprotect(readonly) before mapping memory into vm, and that is done by getting dirty pages
+        let _ = tracker.get_dirty_pages()?;
         let mut vm = set_up_hypervisor_partition(
             &mut gshm,
             &config,