Move setting up page-tables out of set_up_hypervisor_partition

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

Author: ludfjig

src/hyperlight_host/src/hypervisor/mod.rs

@@ -564,9 +564,29 @@ pub(crate) mod tests {
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
         let (_hshm, mut gshm) = sandbox.mgr.build();
+
+        let regions = gshm.layout.get_memory_regions(&gshm.shared_mem)?;
+
+        // Set up shared memory to calculate rsp_ptr
+        #[cfg(feature = "init-paging")]
+        let rsp_ptr = {
+            use crate::mem::ptr::GuestPtr;
+            let rsp_u64 = gshm.set_up_page_tables(&regions)?;
+            let rsp_raw = RawPtr::from(rsp_u64);
+            GuestPtr::try_from(rsp_raw)
+        }?;
+        #[cfg(not(feature = "init-paging"))]
+        let rsp_ptr = {
+            use crate::mem::ptr::GuestPtr;
+            use crate::mem::ptr_offset::Offset;
+            GuestPtr::try_from(Offset::from(0))
+        }?;
+
         let mut vm = set_up_hypervisor_partition(
             &mut gshm,
             &config,
+            rsp_ptr,
+            regions,
             #[cfg(any(crashdump, gdb))]
             &rt_cfg,
         )?;

src/hyperlight_host/src/mem/mgr.rs

@@ -107,17 +107,12 @@ where
         &mut self.shared_mem
     }
 
-    /// Set up the hypervisor partition in the given `SharedMemory` parameter
-    /// `shared_mem`, with the given memory size `mem_size`
+    /// Set up the page tables in the shared memory
     // TODO: This should perhaps happen earlier and use an
     // ExclusiveSharedMemory from the beginning.
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
     #[cfg(feature = "init-paging")]
-    pub(crate) fn set_up_shared_memory(
-        &mut self,
-        mem_size: u64,
-        regions: &mut [MemoryRegion],
-    ) -> Result<u64> {
+    pub(crate) fn set_up_page_tables(&mut self, regions: &[MemoryRegion]) -> Result<u64> {
         let rsp: u64 = self.layout.get_top_of_user_stack_offset() as u64
             + SandboxMemoryLayout::BASE_ADDRESS as u64
             + self.layout.stack_size as u64
@@ -126,6 +121,7 @@ where
             // test from `sandbox_host_tests` fails. We should investigate this further.
             // See issue #498 for more details.
             - 0x28;
+        let mem_size = self.shared_mem.mem_size();
 
         self.shared_mem.with_exclusivity(|shared_mem| {
             // Create PDL4 table with only 1 PML4E
@@ -154,8 +150,6 @@ where
             // We can use the memory size to calculate the number of PTs we need
             // We round up mem_size/2MB
 
-            let mem_size = usize::try_from(mem_size)?;
-
             let num_pages: usize = mem_size.div_ceil(AMOUNT_OF_MEMORY_PER_PT);
 
             // Create num_pages PT with 512 PTEs

src/hyperlight_host/src/sandbox/uninitialized_evolve.rs

@@ -73,12 +73,7 @@ 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,
-    )?;
+
     let outb_hdl = outb_handler_wrapper(hshm.clone(), u_sbox.host_funcs.clone());
 
     let seed = {
@@ -99,6 +94,32 @@ where
     #[cfg(target_os = "linux")]
     setup_signal_handlers(&u_sbox.config)?;
 
+    let regions = gshm.layout.get_memory_regions(&gshm.shared_mem)?;
+
+    // Set up shared memory before stopping dirty page tracking to ensure page table setup is tracked
+    #[cfg(feature = "init-paging")]
+    let rsp_ptr = {
+        let rsp_u64 = gshm.set_up_page_tables(&regions)?;
+        let rsp_raw = RawPtr::from(rsp_u64);
+        GuestPtr::try_from(rsp_raw)
+    }?;
+    #[cfg(not(feature = "init-paging"))]
+    let rsp_ptr = GuestPtr::try_from(Offset::from(0))?;
+
+    // before entering VM (and before mapping memory into VM), stop tracking dirty pages from the host side
+    let dirty_host_pages_idx = gshm
+        .get_shared_mem_mut()
+        .with_exclusivity(|e| e.stop_tracking_dirty_pages())??;
+
+    let mut vm = set_up_hypervisor_partition(
+        &mut gshm,
+        &u_sbox.config,
+        rsp_ptr,
+        regions,
+        #[cfg(any(crashdump, gdb))]
+        &u_sbox.rt_cfg,
+    )?;
+
     vm.initialise(
         peb_addr,
         seed,
@@ -150,19 +171,10 @@ pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result<Mult
 pub(crate) fn set_up_hypervisor_partition(
     mgr: &mut SandboxMemoryManager<GuestSharedMemory>,
     #[cfg_attr(target_os = "windows", allow(unused_variables))] config: &SandboxConfiguration,
+    rsp_ptr: GuestPtr,
+    regions: Vec<MemoryRegion>,
     #[cfg(any(crashdump, gdb))] rt_cfg: &SandboxRuntimeConfig,
 ) -> Result<Box<dyn Hypervisor>> {
-    #[cfg(feature = "init-paging")]
-    let rsp_ptr = {
-        let mut regions = mgr.layout.get_memory_regions(&mgr.shared_mem)?;
-        let mem_size = u64::try_from(mgr.shared_mem.mem_size())?;
-        let rsp_u64 = mgr.set_up_shared_memory(mem_size, &mut regions)?;
-        let rsp_raw = RawPtr::from(rsp_u64);
-        GuestPtr::try_from(rsp_raw)
-    }?;
-    #[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 base_ptr = GuestPtr::try_from(Offset::from(0))?;
     let pml4_ptr = {
         let pml4_offset_u64 = u64::try_from(SandboxMemoryLayout::PML4_OFFSET)?;