more updates

Signed-off-by: Simon Davies <[email protected]>

Author: simongdavies

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::{
@@ -79,9 +81,6 @@ use crate::{Result, log_then_return, new_error};
 #[cfg(mshv2)]
 const CLEAR_DIRTY_BIT_FLAG: u64 = 0b100;
 
-#[cfg(mshv3)]
-const CLEAR_DIRTY_BIT_FLAG: u8 = 0b100;
-
 #[cfg(gdb)]
 mod debug {
     use std::sync::{Arc, Mutex};
@@ -397,12 +396,30 @@ 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),
@@ -427,7 +444,7 @@ impl HypervLinuxDriver {
             sig_rt_min_offset: config.get_interrupt_vcpu_sigrtmin_offset(),
             dropped: AtomicBool::new(false),
         });
-
+        
         #[allow(unused_mut)]
         let mut hv = Self {
             _mshv: mshv,
@@ -846,7 +863,10 @@ impl Hypervisor for HypervLinuxDriver {
         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)
     }

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")]
     fn dispatch_call_from_host(
         &mut self,

src/hyperlight_host/src/mem/bitmap.rs

@@ -20,6 +20,7 @@ pub fn new_page_bitmap(size_in_bytes: usize, init_dirty: bool) -> Result<Vec<u64
             let mut bitmap = vec![!0u64; num_blocks]; // all pages are dirty
             let num_unused_bits = num_blocks * PAGES_IN_BLOCK - num_pages;
             // set the unused bits to 0, could cause problems otherwise
+            #[allow(clippy::unwrap_used)]
             let last_block = bitmap.last_mut().unwrap(); // unwrap is safe since size_in_bytes>0
             *last_block >>= num_unused_bits;
             Ok(bitmap)

src/hyperlight_host/src/mem/dirty_page_tracking.rs

@@ -14,47 +14,16 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-#![allow(dead_code, unused_variables, unused_imports)]
-use std::sync::Arc;
-
 use tracing::{Span, instrument};
 
-#[cfg(target_os = "windows")]
-pub use super::WindowsDirtyPageTracker as PlatformDirtyPageTracker;
 #[cfg(target_os = "linux")]
 pub use super::linux_dirty_page_tracker::LinuxDirtyPageTracker as PlatformDirtyPageTracker;
+#[cfg(target_os = "windows")]
+pub use super::windows_dirty_page_tracker::WindowsDirtyPageTracker as PlatformDirtyPageTracker;
 use crate::Result;
 
 /// Trait defining the interface for dirty page tracking implementations
 pub trait DirtyPageTracking {
-    // /// Start tracking dirty pages
-    // fn start_tracking(&self) -> Result<()>;
-
-    // /// Stop tracking dirty pages
-    // fn stop_tracking(&self) -> Result<()>;
-
-    // /// Take a snapshot of only the dirty pages
-    // fn take_new_snapshot(&self) -> Result<()>;
-
-    // /// Restore memory from the snapshot of dirty pages
-    // fn restore_from_snapshot(&self) -> Result<()>;
-
-    // /// Check if tracking is currently active
-    // fn is_tracking_active(&self) -> bool;
-
-    // /// Get the base address of the tracked region
-    // fn base_addr(&self) -> usize;
-
-    // /// Get the size of the tracked region
-    // fn size(&self) -> usize;
-
-    // /// Get the number of pages in the tracked region
-    // fn num_pages(&self) -> usize;
-
-    // /// Add pages to the last snapshot
-    // /// This is used to update the snapshot with new dirty pages from hypervisor dirty page tracking
-    // fn add_pages_to_last_snapshot(&self, pages: &[usize]) -> Result<()>;
-
     fn get_dirty_pages(self) -> Vec<usize>;
 }
 
@@ -66,61 +35,13 @@ pub struct DirtyPageTracker {
 impl DirtyPageTracker {
     /// Create a new dirty page tracker for the given memory region
     #[instrument(skip_all, parent = Span::current(), level = "Trace")]
-    #[cfg(target_os = "linux")]
     pub fn new(base_addr: usize, size: usize) -> Result<Self> {
         let inner = PlatformDirtyPageTracker::new(base_addr, size)?;
         Ok(Self { inner })
     }
-
-    /// Create a new tracker with a Windows file mapping handle (Windows only)
-    #[cfg(target_os = "windows")]
-    #[instrument(skip_all, parent = Span::current(), level = "Trace")]
-    pub fn new_with_handle(
-        base_addr: usize,
-        size: usize,
-        handle: windows::Win32::Foundation::HANDLE,
-    ) -> Result<Arc<Self>> {
-        let inner = PlatformDirtyPageTracker::new_with_handle(base_addr, size, handle)?;
-        Ok(Arc::new(Self { inner }))
-    }
 }
 
 impl DirtyPageTracking for DirtyPageTracker {
-    // fn start_tracking(&self) -> Result<()> {
-    //     self.inner.start_tracking()
-    // }
-
-    // fn stop_tracking(&self) -> Result<()> {
-    //     self.inner.stop_tracking()
-    // }
-
-    // fn take_new_snapshot(&self) -> Result<()> {
-    //     self.inner.take_new_snapshot()
-    // }
-
-    // fn restore_from_snapshot(&self) -> Result<()> {
-    //     self.inner.restore_from_snapshot()
-    // }
-
-    // fn is_tracking_active(&self) -> bool {
-    //     self.inner.is_tracking_active()
-    // }
-
-    // fn base_addr(&self) -> usize {
-    //     self.inner.base_addr()
-    // }
-
-    // fn size(&self) -> usize {
-    //     self.inner.size()
-    // }
-
-    // fn num_pages(&self) -> usize {
-    //     self.inner.num_pages()
-    // }
-
-    // fn add_pages_to_last_snapshot(&self, pages: &[usize]) -> Result<()> {
-    //     self.inner.add_pages_to_last_snapshot(pages)
-    // }
     fn get_dirty_pages(self) -> Vec<usize> {
         self.inner.get_dirty_pages()
     }