Add signal-based dirty page tracker for tracking host-memory-writes into shared memory.

Add new SnapshotManager and Snapshot structs.

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

Cargo.lock

@@ -44,6 +44,7 @@ anyhow = "1.0"
 metrics = "0.24.2"
 serde_json = "1.0"
 elfcore = "2.0"
+lockfree ="0.5"
 
 [target.'cfg(windows)'.dependencies]
 windows = { version = "0.61", features = [

src/hyperlight_host/Cargo.toml

@@ -1183,7 +1183,7 @@ mod tests {
             return;
         }
         const MEM_SIZE: usize = 0x3000;
-        let gm = shared_mem_with_code(CODE.as_slice(), MEM_SIZE, 0).unwrap();
+        let mut gm = shared_mem_with_code(CODE.as_slice(), MEM_SIZE, 0).unwrap();
         let rsp_ptr = GuestPtr::try_from(0).unwrap();
         let pml4_ptr = GuestPtr::try_from(0).unwrap();
         let entrypoint_ptr = GuestPtr::try_from(0).unwrap();
@@ -1193,6 +1193,7 @@ mod tests {
             MemoryRegionFlags::READ | MemoryRegionFlags::WRITE | MemoryRegionFlags::EXECUTE,
             crate::mem::memory_region::MemoryRegionType::Code,
         );
+        gm.stop_tracking_dirty_pages().unwrap();
         let config: SandboxConfiguration = Default::default();
 
         super::HypervLinuxDriver::new(

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -513,6 +513,7 @@ pub(crate) mod tests {
     #[cfg(gdb)]
     use crate::hypervisor::DbgMemAccessHandlerCaller;
     use crate::mem::ptr::RawPtr;
+    use crate::mem::shared_mem::SharedMemory;
     use crate::sandbox::uninitialized::GuestBinary;
     #[cfg(any(crashdump, gdb))]
     use crate::sandbox::uninitialized::SandboxRuntimeConfig;
@@ -582,6 +583,9 @@ pub(crate) mod tests {
             GuestPtr::try_from(Offset::from(0))
         }?;
 
+        gshm.shared_mem
+            .with_exclusivity(|e| e.stop_tracking_dirty_pages())??;
+
         let mut vm = set_up_hypervisor_partition(
             &mut gshm,
             &config,

src/hyperlight_host/src/hypervisor/mod.rs

@@ -0,0 +1,59 @@
+/*
+Copyright 2025  The Hyperlight Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+use std::sync::Arc;
+
+use tracing::{Span, instrument};
+
+#[cfg(target_os = "linux")]
+pub use super::linux_dirty_page_tracker::LinuxDirtyPageTracker as PlatformDirtyPageTracker;
+use super::shared_mem::HostMapping;
+#[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 {
+    #[cfg(test)]
+    fn get_dirty_pages(&self) -> Result<Vec<usize>>;
+    fn uninstall(self) -> Result<Vec<usize>>;
+}
+
+/// Cross-platform dirty page tracker that delegates to platform-specific implementations
+#[derive(Debug)]
+pub struct DirtyPageTracker {
+    inner: PlatformDirtyPageTracker,
+}
+
+impl DirtyPageTracker {
+    /// Create a new dirty page tracker for the given shared memory
+    #[instrument(skip_all, parent = Span::current(), level = "Trace")]
+    pub fn new(mapping: Arc<HostMapping>) -> Result<Self> {
+        let inner = PlatformDirtyPageTracker::new(mapping)?;
+        Ok(Self { inner })
+    }
+}
+
+impl DirtyPageTracking for DirtyPageTracker {
+    fn uninstall(self) -> Result<Vec<usize>> {
+        self.inner.stop_tracking_and_get_dirty_pages()
+    }
+
+    #[cfg(test)]
+    fn get_dirty_pages(&self) -> Result<Vec<usize>> {
+        self.inner.get_dirty_pages()
+    }
+}

src/hyperlight_host/src/mem/dirty_page_tracking.rs

@@ -119,6 +119,8 @@ pub(crate) struct SandboxMemoryLayout {
     // other
     pub(crate) peb_address: usize,
     code_size: usize,
+    pub(super) input_size: usize,
+    pub(super) output_size: usize,
     // The total size of the page tables
     total_page_table_size: usize,
     // The offset in the sandbox memory where the code starts
@@ -246,7 +248,7 @@ impl SandboxMemoryLayout {
     pub(crate) const BASE_ADDRESS: usize = 0x0;
 
     // the offset into a sandbox's input/output buffer where the stack starts
-    const STACK_POINTER_SIZE_BYTES: u64 = 8;
+    pub(crate) const STACK_POINTER_SIZE_BYTES: u64 = 8;
 
     /// Create a new `SandboxMemoryLayout` with the given
     /// `SandboxConfiguration`, code size and stack/heap size.
@@ -295,11 +297,13 @@ impl SandboxMemoryLayout {
             input_data_buffer_offset + cfg.get_input_data_size(),
             PAGE_SIZE_USIZE,
         );
+        let input_size = output_data_buffer_offset - input_data_buffer_offset;
         // make sure heap buffer starts at 4K boundary
         let guest_heap_buffer_offset = round_up_to(
             output_data_buffer_offset + cfg.get_output_data_size(),
             PAGE_SIZE_USIZE,
         );
+        let output_size = guest_heap_buffer_offset - output_data_buffer_offset;
         // make sure guard page starts at 4K boundary
         let guard_page_offset = round_up_to(guest_heap_buffer_offset + heap_size, PAGE_SIZE_USIZE);
         let guest_user_stack_buffer_offset = guard_page_offset + PAGE_SIZE_USIZE;
@@ -324,7 +328,9 @@ impl SandboxMemoryLayout {
             code_size,
             host_function_definitions_buffer_offset,
             input_data_buffer_offset,
+            input_size,
             output_data_buffer_offset,
+            output_size,
             guest_heap_buffer_offset,
             guest_user_stack_buffer_offset,
             peb_address,