Install DirtyPageTracker immediately after allocating shared memory. Uninstall DirtyPageTracker immediately before setting up VM and mapping memory into VM. Use dirty pages to create snapshots with new snapshotmanager.

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Move setting up page tables out of set_up_hypervisor_partition

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_host/src/hypervisor/mod.rs

@@ -565,11 +565,33 @@ pub(crate) mod tests {
             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 mut 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(&mut 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))
+        }?;
+
+        // we need to undo the mprotect(readonly) before mapping memory into vm.
+        // This is currently done by consuming the dirty page tracker by getting the dirty pages
         let _ = tracker.get_dirty_pages()?;
+
         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

@@ -24,6 +24,7 @@ use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue;
 use hyperlight_common::flatbuffer_wrappers::guest_error::GuestError;
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
 use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
+use hyperlight_common::mem::PAGES_IN_BLOCK;
 use tracing::{Span, instrument};
 
 use super::exe::ExeInfo;
@@ -33,8 +34,9 @@ use super::memory_region::{DEFAULT_GUEST_BLOB_MEM_FLAGS, MemoryRegion, MemoryReg
 use super::ptr::{GuestPtr, RawPtr};
 use super::ptr_offset::Offset;
 use super::shared_mem::{ExclusiveSharedMemory, GuestSharedMemory, HostSharedMemory, SharedMemory};
-use super::shared_mem_snapshot::SharedMemorySnapshot;
-use crate::HyperlightError::NoMemorySnapshot;
+use super::shared_memory_snapshot_manager::SharedMemorySnapshotManager;
+use crate::mem::bitmap::{bitmap_union, new_page_bitmap};
+use crate::mem::dirty_page_tracking::DirtyPageTracker;
 use crate::sandbox::SandboxConfiguration;
 use crate::sandbox::uninitialized::GuestBlob;
 use crate::{Result, log_then_return, new_error};
@@ -75,9 +77,8 @@ pub(crate) struct SandboxMemoryManager<S> {
     pub(crate) entrypoint_offset: Offset,
     /// How many memory regions were mapped after sandbox creation
     pub(crate) mapped_rgns: u64,
-    /// A vector of memory snapshots that can be used to save and  restore the state of the memory
-    /// This is used by the Rust Sandbox implementation (rather than the mem_snapshot field above which only exists to support current C API)
-    snapshots: Arc<Mutex<Vec<SharedMemorySnapshot>>>,
+    /// Shared memory snapshots that can be used to save and  restore the state of the memory
+    snapshot_manager: Arc<Mutex<Option<SharedMemorySnapshotManager>>>,
 }
 
 impl<S> SandboxMemoryManager<S>
@@ -98,7 +99,7 @@ where
             load_addr,
             entrypoint_offset,
             mapped_rgns: 0,
-            snapshots: Arc::new(Mutex::new(Vec::new())),
+            snapshot_manager: Arc::new(Mutex::new(None)),
         }
     }
 
@@ -107,17 +108,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: &mut [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 +122,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 +151,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
@@ -265,54 +260,100 @@ where
         }
     }
 
-    /// this function will create a memory snapshot and push it onto the stack of snapshots
-    /// It should be used when you want to save the state of the memory, for example, when evolving a sandbox to a new state
-    pub(crate) fn push_state(&mut self) -> Result<()> {
-        let snapshot = SharedMemorySnapshot::new(&mut self.shared_mem, self.mapped_rgns)?;
-        self.snapshots
+    /// this function will create an initial snapshot and then create the SnapshotManager
+    pub(crate) fn create_initial_snapshot(
+        &mut self,
+        vm_dirty_bitmap: &[u64],
+        host_dirty_page_idx: &[usize],
+        layout: &SandboxMemoryLayout,
+    ) -> Result<()> {
+        let mut existing_snapshot_manager = self
+            .snapshot_manager
             .try_lock()
-            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
-            .push(snapshot);
+            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
+
+        if existing_snapshot_manager.is_some() {
+            log_then_return!("Snapshot manager already initialized, not creating a new one");
+        }
+
+        // covert vec of page indices to bitmap
+        let mut res = new_page_bitmap(self.shared_mem.mem_size(), false)?;
+        for page_idx in host_dirty_page_idx {
+            let block_idx = page_idx / PAGES_IN_BLOCK;
+            let bit_idx = page_idx % PAGES_IN_BLOCK;
+            res[block_idx] |= 1 << bit_idx;
+        }
+
+        // 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,
+        )?;
+        existing_snapshot_manager.replace(snapshot_manager);
         Ok(())
     }
 
     /// this function restores a memory snapshot from the last snapshot in the list but does not pop the snapshot
     /// 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) -> Result<u64> {
-        let mut snapshots = self
-            .snapshots
+    pub(crate) fn restore_state_from_last_snapshot(&mut self, dirty_bitmap: &[u64]) -> Result<u64> {
+        let mut snapshot_manager = self
+            .snapshot_manager
             .try_lock()
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
-        let last = snapshots.last_mut();
-        if last.is_none() {
-            log_then_return!(NoMemorySnapshot);
+
+        match snapshot_manager.as_mut() {
+            None => {
+                log_then_return!("Snapshot manager not initialized");
+            }
+            Some(snapshot_manager) => {
+                snapshot_manager.restore_from_snapshot(&mut self.shared_mem, dirty_bitmap)
+            }
         }
-        #[allow(clippy::unwrap_used)] // We know that last is not None because we checked it above
-        let snapshot = last.unwrap();
-        let old_rgns = self.mapped_rgns;
-        self.mapped_rgns = snapshot.restore_from_snapshot(&mut self.shared_mem)?;
-        Ok(old_rgns - self.mapped_rgns)
     }
 
     /// this function pops the last snapshot off the stack and restores the memory to the previous state
     /// It should be used when you want to restore the state of the memory to a previous state and do not need to retain that state
     /// for example when devolving a sandbox to a previous state.
-    pub(crate) fn pop_and_restore_state_from_snapshot(&mut self) -> Result<u64> {
-        let last = self
-            .snapshots
+    pub(crate) fn pop_and_restore_state_from_snapshot(
+        &mut self,
+        dirty_bitmap: &[u64],
+    ) -> Result<u64> {
+        let mut snapshot_manager = self
+            .snapshot_manager
+            .try_lock()
+            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
+
+        match snapshot_manager.as_mut() {
+            None => {
+                log_then_return!("Snapshot manager not initialized");
+            }
+            Some(snapshot_manager) => snapshot_manager
+                .pop_and_restore_state_from_snapshot(&mut self.shared_mem, dirty_bitmap),
+        }
+    }
+
+    pub(crate) fn push_state(&mut self, dirty_bitmap: &[u64]) -> Result<()> {
+        let mut snapshot_manager = self
+            .snapshot_manager
             .try_lock()
-            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?
-            .pop();
-        if last.is_none() {
-            log_then_return!(NoMemorySnapshot);
+            .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
+
+        match snapshot_manager.as_mut() {
+            None => {
+                log_then_return!("Snapshot manager not initialized");
+            }
+            Some(snapshot_manager) => snapshot_manager.create_new_snapshot(
+                &mut self.shared_mem,
+                dirty_bitmap,
+                self.mapped_rgns,
+            ),
         }
-        self.restore_state_from_last_snapshot()
     }
 
     /// Sets `addr` to the correct offset in the memory referenced by
@@ -347,7 +388,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
         cfg: SandboxConfiguration,
         exe_info: &mut ExeInfo,
         guest_blob: Option<&GuestBlob>,
-    ) -> Result<Self> {
+    ) -> Result<(Self, DirtyPageTracker)> {
         let guest_blob_size = guest_blob.map(|b| b.data.len()).unwrap_or(0);
         let guest_blob_mem_flags = guest_blob.map(|b| b.permissions);
 
@@ -360,6 +401,7 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             guest_blob_mem_flags,
         )?;
         let mut shared_mem = ExclusiveSharedMemory::new(layout.get_memory_size()?)?;
+        let tracker = shared_mem.start_tracking_dirty_pages()?;
 
         let load_addr: RawPtr = RawPtr::try_from(layout.get_guest_code_address())?;
 
@@ -378,7 +420,10 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
             &mut shared_mem.as_mut_slice()[layout.get_guest_code_offset()..],
         )?;
 
-        Ok(Self::new(layout, shared_mem, load_addr, entrypoint_offset))
+        Ok((
+            Self::new(layout, shared_mem, load_addr, entrypoint_offset),
+            tracker,
+        ))
     }
 
     /// Writes host function details to memory
@@ -440,15 +485,15 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
                 mapped_rgns: 0,
-                snapshots: Arc::new(Mutex::new(Vec::new())),
+                snapshot_manager: Arc::new(Mutex::new(None)),
             },
             SandboxMemoryManager {
                 shared_mem: gshm,
                 layout: self.layout,
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
                 mapped_rgns: 0,
-                snapshots: Arc::new(Mutex::new(Vec::new())),
+                snapshot_manager: Arc::new(Mutex::new(None)),
             },
         )
     }

src/hyperlight_host/src/mem/mod.rs

@@ -43,9 +43,6 @@ pub mod ptr_offset;
 /// A wrapper around unsafe functionality to create and initialize
 /// a memory region for a guest running in a sandbox.
 pub mod shared_mem;
-/// A wrapper around a `SharedMemory` and a snapshot in time
-/// of the memory therein
-pub mod shared_mem_snapshot;
 /// Utilities for writing shared memory tests
 #[cfg(test)]
 pub(crate) mod shared_mem_tests;