Integrate DirtyPageTracker into UninitializedSandbox starting host page tracking after memory is allocated and stopping it once the uninitialized sandbox is evolved

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

Author: simongdavies

src/hyperlight_host/src/sandbox/uninitialized.rs

@@ -29,6 +29,7 @@ use crate::func::host_functions::{HostFunction, register_host_function};
 use crate::func::{ParameterTuple, SupportedReturnType};
 #[cfg(feature = "build-metadata")]
 use crate::log_build_details;
+use crate::mem::dirty_page_tracking::DirtyPageTracker;
 use crate::mem::exe::ExeInfo;
 use crate::mem::memory_region::{DEFAULT_GUEST_BLOB_MEM_FLAGS, MemoryRegionFlags};
 use crate::mem::mgr::{STACK_COOKIE_LEN, SandboxMemoryManager};
@@ -80,6 +81,7 @@ pub struct UninitializedSandbox {
     pub(crate) config: SandboxConfiguration,
     #[cfg(any(crashdump, gdb))]
     pub(crate) rt_cfg: SandboxRuntimeConfig,
+    pub(crate) tracker: Option<DirtyPageTracker>,
 }
 
 impl crate::sandbox_state::sandbox::UninitializedSandbox for UninitializedSandbox {
@@ -250,17 +252,15 @@ impl UninitializedSandbox {
             }
         };
 
-        let mut mem_mgr_wrapper = {
-            let mut mgr = UninitializedSandbox::load_guest_binary(
-                sandbox_cfg,
-                &guest_binary,
-                guest_blob.as_ref(),
-            )?;
+        let (mut mgr, tracker) = UninitializedSandbox::load_guest_binary(
+            sandbox_cfg,
+            &guest_binary,
+            guest_blob.as_ref(),
+        )?;
 
-            let stack_guard = Self::create_stack_guard();
-            mgr.set_stack_guard(&stack_guard)?;
-            MemMgrWrapper::new(mgr, stack_guard)
-        };
+        let stack_guard = Self::create_stack_guard();
+        mgr.set_stack_guard(&stack_guard)?;
+        let mut mem_mgr_wrapper = MemMgrWrapper::new(mgr, stack_guard);
 
         mem_mgr_wrapper.write_memory_layout()?;
 
@@ -278,6 +278,7 @@ impl UninitializedSandbox {
             config: sandbox_cfg,
             #[cfg(any(crashdump, gdb))]
             rt_cfg,
+            tracker: Some(tracker),
         };
 
         // If we were passed a writer for host print register it otherwise use the default.
@@ -308,7 +309,10 @@ impl UninitializedSandbox {
         cfg: SandboxConfiguration,
         guest_binary: &GuestBinary,
         guest_blob: Option<&GuestBlob>,
-    ) -> Result<SandboxMemoryManager<ExclusiveSharedMemory>> {
+    ) -> Result<(
+        SandboxMemoryManager<ExclusiveSharedMemory>,
+        DirtyPageTracker,
+    )> {
         let mut exe_info = match guest_binary {
             GuestBinary::FilePath(bin_path_str) => ExeInfo::from_file(bin_path_str)?,
             GuestBinary::Buffer(buffer) => ExeInfo::from_buf(buffer)?,
@@ -396,6 +400,7 @@ impl UninitializedSandbox {
         Ok(())
     }
 }
+
 // Check to see if the current version of Windows is supported
 // Hyperlight is only supported on Windows 11 and Windows Server 2022 and later
 #[cfg(target_os = "windows")]

src/hyperlight_host/src/sandbox/uninitialized_evolve.rs

@@ -28,6 +28,7 @@ use super::uninitialized::SandboxRuntimeConfig;
 use crate::HyperlightError::NoHypervisorFound;
 use crate::hypervisor::Hypervisor;
 use crate::hypervisor::handlers::{MemAccessHandlerCaller, OutBHandlerCaller};
+use crate::mem::dirty_page_tracking::DirtyPageTracker;
 use crate::mem::layout::SandboxMemoryLayout;
 use crate::mem::mgr::SandboxMemoryManager;
 use crate::mem::ptr::{GuestPtr, RawPtr};
@@ -73,11 +74,13 @@ where
     ) -> Result<ResSandbox>,
 {
     let (hshm, mut gshm) = u_sbox.mgr.build();
+
     let mut vm = set_up_hypervisor_partition(
         &mut gshm,
         &u_sbox.config,
         #[cfg(any(crashdump, gdb))]
         &u_sbox.rt_cfg,
+        u_sbox.tracker,
     )?;
     let outb_hdl = outb_handler_wrapper(hshm.clone(), u_sbox.host_funcs.clone());
 
@@ -129,9 +132,12 @@ where
 pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result<MultiUseSandbox> {
     evolve_impl(
         u_sbox,
-        |hf, mut hshm, vm, out_hdl, mem_hdl, dispatch_ptr| {
+        |hf, mut hshm, mut vm, out_hdl, mem_hdl, dispatch_ptr| {
             {
-                hshm.as_mut().push_state()?;
+                let vm_dirty_pages = vm.get_and_clear_dirty_pages()?;
+                let layout = hshm.unwrap_mgr().layout;
+                hshm.as_mut()
+                    .create_initial_snapshot(Some(&vm_dirty_pages), &layout)?;
             }
             Ok(MultiUseSandbox::from_uninit(
                 hf,
@@ -151,6 +157,7 @@ pub(crate) fn set_up_hypervisor_partition(
     mgr: &mut SandboxMemoryManager<GuestSharedMemory>,
     #[cfg_attr(target_os = "windows", allow(unused_variables))] config: &SandboxConfiguration,
     #[cfg(any(crashdump, gdb))] rt_cfg: &SandboxRuntimeConfig,
+    tracker: Option<DirtyPageTracker>,
 ) -> Result<Box<dyn Hypervisor>> {
     #[cfg(feature = "init-paging")]
     let rsp_ptr = {
@@ -163,6 +170,10 @@ pub(crate) fn set_up_hypervisor_partition(
     #[cfg(not(feature = "init-paging"))]
     let rsp_ptr = GuestPtr::try_from(Offset::from(0))?;
     let regions = mgr.layout.get_memory_regions(&mgr.shared_mem)?;
+
+    let tracker = tracker.ok_or_else(|| new_error!("No Dirty page tracker found"))?;
+    mgr.shared_mem
+        .with_exclusivity(|e| e.stop_tracking_dirty_pages(tracker))??;
     let base_ptr = GuestPtr::try_from(Offset::from(0))?;
     let pml4_ptr = {
         let pml4_offset_u64 = u64::try_from(SandboxMemoryLayout::PML4_OFFSET)?;