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 <[email protected]>

Author: ludfjig

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)),
         }
     }
 
@@ -265,54 +266,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.raw_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))?
-            .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
+                .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))?;
+
+        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 +394,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 +407,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 +426,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 +491,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;

src/hyperlight_host/src/mem/shared_mem_snapshot.rs

@@ -270,8 +270,9 @@ impl MultiUseSandbox {
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     pub(crate) fn restore_state(&mut self) -> Result<()> {
         let mem_mgr = self.mem_mgr.unwrap_mgr_mut();
-        let rgns_to_unmap = mem_mgr.restore_state_from_last_snapshot()?;
-        unsafe { self.vm.unmap_regions(rgns_to_unmap)? };
+        let dirty_pages = self.vm.get_and_clear_dirty_pages()?;
+        let rgns_to_umap = mem_mgr.restore_state_from_last_snapshot(&dirty_pages)?;
+        unsafe { self.vm.unmap_regions(rgns_to_umap)? };
         Ok(())
     }
 
@@ -354,10 +355,11 @@ impl DevolvableSandbox<MultiUseSandbox, MultiUseSandbox, Noop<MultiUseSandbox, M
     /// The devolve can be used to return the MultiUseSandbox to the state before the code was loaded. Thus avoiding initialisation overhead
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
     fn devolve(mut self, _tsn: Noop<MultiUseSandbox, MultiUseSandbox>) -> Result<MultiUseSandbox> {
+        let dirty_pages = self.vm.get_and_clear_dirty_pages()?;
         let rgns_to_unmap = self
             .mem_mgr
             .unwrap_mgr_mut()
-            .pop_and_restore_state_from_snapshot()?;
+            .pop_and_restore_state_from_snapshot(&dirty_pages)?;
         unsafe { self.vm.unmap_regions(rgns_to_unmap)? };
         Ok(self)
     }
@@ -389,7 +391,8 @@ where
         let mut ctx = self.new_call_context();
         transition_func.call(&mut ctx)?;
         let mut sbox = ctx.finish_no_reset();
-        sbox.mem_mgr.unwrap_mgr_mut().push_state()?;
+        let vm_dirty_pages = sbox.vm.get_and_clear_dirty_pages()?;
+        sbox.mem_mgr.unwrap_mgr_mut().push_state(&vm_dirty_pages)?;
         Ok(sbox)
     }
 }

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -241,7 +241,7 @@ mod tests {
 
         let new_mgr = || {
             let mut exe_info = simple_guest_exe_info().unwrap();
-            let mut mgr = SandboxMemoryManager::load_guest_binary_into_memory(
+            let (mut mgr, _) = SandboxMemoryManager::load_guest_binary_into_memory(
                 sandbox_cfg,
                 &mut exe_info,
                 None,
@@ -356,7 +356,7 @@ mod tests {
         tracing::subscriber::with_default(subscriber.clone(), || {
             let new_mgr = || {
                 let mut exe_info = simple_guest_exe_info().unwrap();
-                let mut mgr = SandboxMemoryManager::load_guest_binary_into_memory(
+                let (mut mgr, _) = SandboxMemoryManager::load_guest_binary_into_memory(
                     sandbox_cfg,
                     &mut exe_info,
                     None,