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)]
@@ -1158,7 +1193,8 @@ mod tests {
             return;
         }
         const MEM_SIZE: usize = 0x3000;
-        let gm = shared_mem_with_code(CODE.as_slice(), MEM_SIZE, 0).unwrap();
+        let mut gm = shared_mem_with_code(CODE.as_slice(), MEM_SIZE, 0).unwrap();
+        gm.stop_tracking_dirty_pages().unwrap();
         let rsp_ptr = GuestPtr::try_from(0).unwrap();
         let pml4_ptr = GuestPtr::try_from(0).unwrap();
         let entrypoint_ptr = GuestPtr::try_from(0).unwrap();

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -607,19 +607,21 @@ impl Hypervisor for HypervWindowsDriver {
         Ok(())
     }
 
-    fn get_and_clear_dirty_pages(&mut self) -> Result<Vec<u64>> {
+    fn get_and_clear_dirty_pages(&mut self) -> Result<(Vec<u64>, Option<Vec<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)
+        Ok((new_page_bitmap(total_size, true)?, None))
     }
 
     #[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");
+        // TODO: when adding support, also update `get_and_clear_dirty_pages`, see kvm/mshv for details
+        log_then_return!("Mapping host memory into the guest not yet supported on this platform.");
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     unsafe fn unmap_regions(&mut self, n: u64) -> Result<()> {
+        // TODO: when adding support, also update `get_and_clear_dirty_pages`, see kvm/mshv for details
         if n > 0 {
             log_then_return!(
                 "Mapping host memory into the guest not yet supported on this platform"

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/memory_region.rs

@@ -30,9 +30,8 @@ use bitflags::bitflags;
 #[cfg(mshv)]
 use hyperlight_common::mem::PAGE_SHIFT;
 use hyperlight_common::mem::PAGE_SIZE_USIZE;
-use kvm_bindings::KVM_MEM_LOG_DIRTY_PAGES;
 #[cfg(kvm)]
-use kvm_bindings::{KVM_MEM_READONLY, kvm_userspace_memory_region};
+use kvm_bindings::{KVM_MEM_LOG_DIRTY_PAGES, KVM_MEM_READONLY, kvm_userspace_memory_region};
 #[cfg(mshv2)]
 use mshv_bindings::{
     HV_MAP_GPA_EXECUTABLE, HV_MAP_GPA_PERMISSIONS_NONE, HV_MAP_GPA_READABLE, HV_MAP_GPA_WRITABLE,

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_mem.rs

@@ -524,21 +524,28 @@ impl ExclusiveSharedMemory {
             log_then_return!(WindowsAPIError(e.clone()));
         }
 
+        // HostMapping is only non-Send/Sync because raw pointers
+        // are not ("as a lint", as the Rust docs say). We don't
+        // want to mark HostMapping Send/Sync immediately, because
+        // that could socially imply that it's "safe" to use
+        // unsafe accesses from multiple threads at once. Instead, we
+        // directly impl Send and Sync on this type. Since this
+        // type does have Send and Sync manually impl'd, the Arc
+        // is not pointless as the lint suggests.
+        #[allow(clippy::arc_with_non_send_sync)]
+        let host_mapping = Arc::new(HostMapping {
+            ptr: addr.Value as *mut u8,
+            size: total_size,
+            handle,
+        });
+
+        let dirty_page_tracker = Arc::new(Mutex::new(Some(DirtyPageTracker::new(Arc::clone(
+            &host_mapping,
+        ))?)));
+
         Ok(Self {
-            // HostMapping is only non-Send/Sync because raw pointers
-            // are not ("as a lint", as the Rust docs say). We don't
-            // want to mark HostMapping Send/Sync immediately, because
-            // that could socially imply that it's "safe" to use
-            // unsafe accesses from multiple threads at once. Instead, we
-            // directly impl Send and Sync on this type. Since this
-            // type does have Send and Sync manually impl'd, the Arc
-            // is not pointless as the lint suggests.
-            #[allow(clippy::arc_with_non_send_sync)]
-            region: Arc::new(HostMapping {
-                ptr: addr.Value as *mut u8,
-                size: total_size,
-                handle,
-            }),
+            region: host_mapping,
+            signal_dirty_bitmap_tracker: dirty_page_tracker,
         })
     }