diff --git a/CHANGELOG.md b/CHANGELOG.md index 21ecfcb42..b88d07b5b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -71,6 +71,21 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). will be extended in the future. * MSHV: vCPU creation reads the partition's XSAVE format once. This removes one hypercall from every restore. +* On KVM, restore zeroes the scratch pages that hold data and keeps them + mapped, so the next run takes no faults on them. Pages holding only zeros + are kept as well, and every 16th restore gives them back (64 runs of them + at most), so what stays resident follows what recent runs wrote. On Linux + 6.7 or newer, a restore costs time in how much of scratch is in use rather + than in its size. +* On WHP, and on MSHV on x86_64, restore zeroes only the scratch pages written + since the last restore, from the hypervisor's dirty-page log and a log of host + writes. WHP no longer replaces the scratch mapping on each restore, so the + pages a guest writes stay resident between restores. MSHV tracks scratch of + 2 MiB or more. After a run writes more than a tenth of it (a quarter from 32 + MiB), where zeroing all of it costs less, it stops tracking for 16 restores. +* When Windows cannot tell which scratch pages were written, restore zeroes + scratch of up to 16 MiB in place, which is faster than replacing it below + that size, and replaces larger scratch. ### Removed * `RunPool` and the run-specific `AllocError::InvalidAlign` variant. diff --git a/bench_report.toml b/bench_report.toml index dbdfa6001..6670652ac 100644 --- a/bench_report.toml +++ b/bench_report.toml @@ -4,6 +4,21 @@ # nothing fails the report. allowlist = [ "^function_call_codec/decode_vec_bytes_copy$", + "^guest_calls/call_with_restore/default$", + "^guest_calls/call_with_restore/scratch_64mib$", + "^guest_calls/call_with_restore/scratch_64mib_write_1mib$", + "^guest_calls/call_with_restore/scratch_256mib$", + "^guest_calls/call_with_restore/scratch_256mib_write_1mib$", + "^guest_calls/call_with_restore/scratch_256mib_write_16mib$", + "^guest_calls/call_with_restore/32_after_32mib_scratch_64mib$", + "^guest_calls/call_after_restore/default$", + "^guest_calls/call_after_restore/scratch_256mib$", + "^guest_calls/call_after_restore/scratch_256mib_write_1mib$", + "^guest_calls/call_after_restore/scratch_256mib_write_16mib$", + "^snapshots/restore/default$", + "^snapshots/restore/scratch_256mib$", + "^snapshots/restore/scratch_256mib_write_1mib$", + "^snapshots/restore/scratch_256mib_write_16mib$", "^function_call_codec/encode_control/byte_chunks$", "^function_call_codec/encode_control/vec_bytes$", "^payload_allocation/slot_pool_segmented/262144$", diff --git a/src/hyperlight_host/benches/benchmarks.rs b/src/hyperlight_host/benches/benchmarks.rs index 791018ae6..40e24231a 100644 --- a/src/hyperlight_host/benches/benchmarks.rs +++ b/src/hyperlight_host/benches/benchmarks.rs @@ -131,6 +131,103 @@ fn bench_guest_call_with_restore(b: &mut criterion::Bencher, size: SandboxSize) }); } +/// Sandboxes with large scratch, where resetting scratch dominates a +/// restore: the name, MiB of scratch, and bytes of heap the guest writes a +/// call (0 for only a few pages, an Echo). +const LARGE_SCRATCH: [(&str, usize, u64); 5] = [ + ("scratch_64mib", 64, 0), + ("scratch_64mib_write_1mib", 64, 1 << 20), + ("scratch_256mib", 256, 0), + ("scratch_256mib_write_1mib", 256, 1 << 20), + ("scratch_256mib_write_16mib", 256, 16 << 20), +]; + +/// A sandbox with `mib` MiB of scratch whose heap can allocate `write` +/// bytes. +fn large_scratch_sandbox(mib: usize, write: u64) -> Sandbox { + SandboxBuilder::from_file(simple_guest_as_pathbuf()) + .heap_size((4 << 20).max(4 * write)) + .scratch_size(mib << 20) + .build() + .unwrap() +} + +/// A guest call that writes `write` bytes of heap, or an Echo at 0. +fn call_writing(sbox: &mut Sandbox, write: u64) { + if write == 0 { + sbox.call::("Echo", "hello\n".to_string()).unwrap(); + } else { + sbox.call::<()>("AllocAndWritePattern", write).unwrap(); + } +} + +fn bench_large_scratch_call_with_restore(b: &mut criterion::Bencher, mib: usize, write: u64) { + let mut sbox = large_scratch_sandbox(mib, write); + let snapshot = sbox.snapshot().unwrap(); + + b.iter(|| { + call_writing(&mut sbox, write); + sbox.restore(snapshot.clone()).unwrap(); + }); +} + +/// After one call that wrote `heavy` bytes, the next 32 calls that write +/// little and their restores, as a pool of sandboxes sees after one +/// large request. An iteration times those 32. +fn bench_call_with_restore_after_heavy(b: &mut criterion::Bencher, mib: usize, heavy: u64) { + let mut sbox = large_scratch_sandbox(mib, heavy); + let snapshot = sbox.snapshot().unwrap(); + + b.iter_custom(|iters| { + let mut total = Duration::ZERO; + for _ in 0..iters { + call_writing(&mut sbox, heavy); + sbox.restore(snapshot.clone()).unwrap(); + let start = Instant::now(); + for _ in 0..32 { + call_writing(&mut sbox, 0); + sbox.restore(snapshot.clone()).unwrap(); + } + total += start.elapsed(); + } + total + }); +} + +/// Only the call after a restore, which pays for any page the restore +/// left the guest to fault in again. +fn bench_call_after_restore(b: &mut criterion::Bencher, mut sbox: Sandbox, write: u64) { + let snapshot = sbox.snapshot().unwrap(); + + b.iter_custom(|iters| { + let mut total = Duration::ZERO; + for _ in 0..iters { + sbox.restore(snapshot.clone()).unwrap(); + let start = Instant::now(); + call_writing(&mut sbox, write); + total += start.elapsed(); + } + total + }); +} + +/// Only the restore after a call. +fn bench_restore_after_call(b: &mut criterion::Bencher, mib: usize, write: u64) { + let mut sbox = large_scratch_sandbox(mib, write); + let snapshot = sbox.snapshot().unwrap(); + + b.iter_custom(|iters| { + let mut total = Duration::ZERO; + for _ in 0..iters { + call_writing(&mut sbox, write); + let start = Instant::now(); + sbox.restore(snapshot.clone()).unwrap(); + total += start.elapsed(); + } + total + }); +} + fn bench_guest_call_with_host_function(b: &mut criterion::Bencher, size: SandboxSize) { let mut multiuse_sandbox = size .builder() @@ -238,6 +335,29 @@ fn guest_calls_benchmark(c: &mut Criterion) { }); } + for (name, mib, write) in LARGE_SCRATCH { + group.bench_function(format!("call_with_restore/{name}"), |b| { + bench_large_scratch_call_with_restore(b, mib, write) + }); + } + + group.bench_function("call_with_restore/32_after_32mib_scratch_64mib", |b| { + bench_call_with_restore_after_heavy(b, 64, 32 << 20) + }); + + group.bench_function("call_after_restore/default", |b| { + bench_call_after_restore( + b, + create_multiuse_sandbox_with_size(SandboxSize::Default), + 0, + ) + }); + for (name, mib, write) in LARGE_SCRATCH.into_iter().skip(2) { + group.bench_function(format!("call_after_restore/{name}"), |b| { + bench_call_after_restore(b, large_scratch_sandbox(mib, write), write) + }); + } + for size in SandboxSize::all() { group.bench_function(format!("call_with_host_function/{}", size.name()), |b| { bench_guest_call_with_host_function(b, size) @@ -338,6 +458,12 @@ fn snapshots_benchmark(c: &mut Criterion) { }); } + for (name, mib, write) in LARGE_SCRATCH.into_iter().skip(2) { + group.bench_function(format!("restore/{name}"), |b| { + bench_restore_after_call(b, mib, write) + }); + } + group.finish(); } diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm/aarch64.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm/aarch64.rs index 0c6870f23..29ec4219f 100644 --- a/src/hyperlight_host/src/hypervisor/hyperlight_vm/aarch64.rs +++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm/aarch64.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use super::{ AccessPageTableError, CreateHyperlightVmError, DispatchGuestCallError, HyperlightVm, - InitializeError, + InitializeError, ScratchDirtyLog, }; #[cfg(hvf)] use crate::hypervisor::HvfInterruptHandle; @@ -100,6 +100,7 @@ impl HyperlightVm { snapshot_memory: None, scratch_slot, scratch_memory: None, + scratch_dirty: ScratchDirtyLog::default(), mmap_regions: Vec::new(), diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm/dirty_log.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm/dirty_log.rs new file mode 100644 index 000000000..6e26c103c --- /dev/null +++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm/dirty_log.rs @@ -0,0 +1,539 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Which scratch pages the guest wrote since the last restore. +//! +//! A restore zeroes scratch. Zeroing all of it costs time in its size, +//! however little the guest wrote, and mapping fresh memory instead +//! costs a fault on every page the guest then touches. Where the +//! hypervisor logs the guest's writes, a restore zeroes only the pages +//! in the log (see [`HostSharedMemory::zero_written`]). +//! +//! Tracking starts when scratch is mapped, before the guest runs, so +//! every restore, the first included, has the log, unless MSHV tracking +//! is paused (see below). On MSHV the log is not read until the first +//! restore: until then every page reads as written, so the guest sets up +//! without tracking faults and the first restore zeroes all of scratch. +//! +//! On WHP, scratch is mapped with tracking, which costs the guest +//! nothing measurable on bare metal, so the log is always used. Scratch +//! is no longer replaced on each restore, so the pages the guest writes +//! stay resident between restores: releasing them would cost a fault on +//! each the next run touches. The section is charged to commit in full +//! when it is created, as before. +//! +//! On MSHV, tracking is switched on for the whole VM. While it is on, +//! the guest's first write to each page after a read faults to the +//! hypervisor, and reading the log costs time in the size of scratch. +//! Measured on MSHV, that costs more than zeroing all of scratch when +//! scratch is under [`MIN_TRACKED_SCRATCH`], or when a run writes more +//! of it than [`writes_much`] allows. Small scratch is not tracked, and +//! tracking stops after such a run for [`RESUME_AFTER`] restores, which +//! zero all of scratch. +//! +//! [`HostSharedMemory::zero_written`]: crate::mem::shared_mem::HostSharedMemory::zero_written + +use tracing::{debug, warn}; + +use crate::hypervisor::virtual_machine::{DirtyLog, DirtyTracking, HypervisorError}; +use crate::mem::shared_mem::{DIRTY_PAGE_SIZE, DirtyRuns}; + +/// The smallest scratch tracked on MSHV. At 1 MiB, tracking a guest +/// that writes nothing costs as much as zeroing all of scratch; from 2 +/// MiB on it costs less. +const MIN_TRACKED_SCRATCH: usize = 2 << 20; + +/// Scratch from which zeroing all of it costs about three times more per +/// MiB on MSHV, so tracking pays off for runs that write more of it. +const LARGE_SCRATCH: usize = 32 << 20; + +/// On MSHV, whether a run that wrote `written` of `pages` costs more +/// tracked than zeroing all of scratch would. Measured, tracking costs +/// more past about 5-10% of scratch written below [`LARGE_SCRATCH`] and +/// past about 25% from it on. +fn writes_much(written: usize, pages: usize) -> bool { + let share = if pages * DIRTY_PAGE_SIZE < LARGE_SCRATCH { + 10 + } else { + 4 + }; + written * share > pages +} + +/// On MSHV, the restores after a run that wrote much that zero all of +/// scratch before tracking resumes. A guest that keeps writing much pays +/// for tracking one run in this many; one that wrote much once is tracked +/// again soon after. +const RESUME_AFTER: usize = 16; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum State { + /// Not tracking: scratch is reset in full. + Off, + /// Tracking since the last restore. + On, + /// The hypervisor failed; scratch is reset in full. + Failed, +} + +/// The guest's writes to scratch since the last restore, where the +/// hypervisor logs them. See the module docs. +#[derive(Debug)] +pub(crate) struct ScratchDirtyLog { + state: State, + /// The last log read, reused across restores where the hypervisor + /// API fills a caller's buffer. + bitmap: Vec, + /// Tracking is switched on in the hypervisor. + enabled: bool, + /// No restore has read the log since scratch was mapped. + since_mapped: bool, + /// Restores left before tracking resumes after a run that wrote much. + paused: usize, + /// Restores that had the log. + #[cfg(test)] + pub(crate) logs: usize, +} + +impl Default for ScratchDirtyLog { + fn default() -> Self { + Self { + state: State::Off, + bitmap: Vec::new(), + enabled: false, + since_mapped: false, + paused: 0, + #[cfg(test)] + logs: 0, + } + } +} + +impl ScratchDirtyLog { + /// Scratch at `[gpa, gpa + size)` was just mapped, before the guest + /// runs on it. Start tracking it where that pays off. + pub(crate) fn mapped(&mut self, vm: &mut (impl DirtyLog + ?Sized), gpa: u64, size: usize) { + if self.state == State::Failed { + return; + } + self.since_mapped = true; + self.paused = 0; + let result = match vm.dirty_tracking() { + DirtyTracking::Switched if size >= MIN_TRACKED_SCRATCH => { + vm.enable_dirty_tracking().map(|()| self.enabled = true) + } + // A range mapped with tracking starts clean: the first restore + // zeroes only what was written since, leaving the rest of + // scratch untouched (not resident). + DirtyTracking::Mapped => Ok(()), + _ => { + self.state = State::Off; + return; + } + }; + self.state = match result { + Ok(()) => State::On, + Err(e) => self.fail(vm, gpa, size, e), + }; + } + + /// Scratch at `[gpa, gpa + size)` is about to be unmapped. Tracking + /// is switched off while the range is still mapped, since MSHV stops + /// only once the bits of every page it read are set again. + pub(crate) fn unmapping(&mut self, vm: &mut (impl DirtyLog + ?Sized), gpa: u64, size: usize) { + if self.enabled { + match vm.disable_dirty_tracking(gpa, size) { + Ok(()) => self.enabled = false, + Err(e) => self.state = failed(e), + } + } + if self.state != State::Failed { + self.state = State::Off; + } + self.paused = 0; + } + + /// Called once per restore, before scratch at `[gpa, gpa + size)` + /// is reset. Returns the pages the guest wrote since the last + /// restore, one bit per page of [`DIRTY_PAGE_SIZE`], or `None` when + /// they are not known and all of scratch must be reset. The log read + /// is cleared: a restore that then fails leaves the sandbox + /// unrecoverable, so nothing restores it on a log missing those pages. + pub(crate) fn take( + &mut self, + vm: &mut (impl DirtyLog + ?Sized), + gpa: u64, + size: usize, + ) -> Option<&mut Vec> { + if self.state == State::Off && self.paused > 0 { + self.paused -= 1; + if self.paused == 0 { + self.resume(vm, gpa, size); + } + return None; + } + if self.state != State::On { + return None; + } + if let Err(e) = vm.read_dirty_log(gpa, size, &mut self.bitmap) { + self.state = self.fail(vm, gpa, size, e); + return None; + } + // The first log after mapping holds what the guest wrote setting + // up, before its first snapshot, which is not a run. + let first = std::mem::take(&mut self.since_mapped); + if vm.dirty_tracking() == DirtyTracking::Switched && !first { + let pages = size / DIRTY_PAGE_SIZE; + let written = written(&self.bitmap, pages); + if writes_much(written, pages) { + debug!( + "Scratch dirty tracking paused for {RESUME_AFTER} restores: a run wrote {written} of {pages} pages" + ); + // The log just read is right either way. + self.state = match vm.disable_dirty_tracking(gpa, size) { + Ok(()) => { + self.enabled = false; + self.paused = RESUME_AFTER; + State::Off + } + Err(e) => failed(e), + }; + } + } + #[cfg(test)] + { + self.logs += 1; + } + Some(&mut self.bitmap) + } + + /// Track again from the next run, from this restore on, which zeroes + /// all of scratch. The log is read now so the next restore reads that + /// run's writes alone. + fn resume(&mut self, vm: &mut (impl DirtyLog + ?Sized), gpa: u64, size: usize) { + let result = vm.enable_dirty_tracking().and_then(|()| { + self.enabled = true; + vm.read_dirty_log(gpa, size, &mut self.bitmap) + }); + self.state = match result { + Ok(()) => State::On, + Err(e) => self.fail(vm, gpa, size, e), + }; + } + + fn fail( + &mut self, + vm: &mut (impl DirtyLog + ?Sized), + gpa: u64, + size: usize, + e: HypervisorError, + ) -> State { + // Best effort. Left on, tracking costs each page one fault at + // most: pages fault only after a read clears them. + if self.enabled && vm.disable_dirty_tracking(gpa, size).is_ok() { + self.enabled = false; + } + failed(e) + } +} + +/// The log failed after `e`: scratch is reset in full from now on. +fn failed(e: HypervisorError) -> State { + warn!("Scratch dirty log failed, restores zero all of scratch from now on: {e}"); + State::Failed +} + +/// The set bits below `pages` in `bitmap`. +fn written(bitmap: &[u64], pages: usize) -> usize { + DirtyRuns::new(bitmap, pages).map(|run| run.len()).sum() +} + +#[cfg(test)] +mod tests { + use super::*; + + const GPA: u64 = 0x1_0000_0000; + const SIZE: usize = 16 << 20; + const PAGES: usize = SIZE / DIRTY_PAGE_SIZE; + + /// A VM whose guest writes the first `written` pages between reads. + /// As on MSHV, every page reads as written until the first read after + /// switching tracking on. + #[derive(Debug, Default, PartialEq, Eq)] + struct FakeVm { + tracking: Option, + written: usize, + all_written: bool, + fail: bool, + fail_enable: bool, + fail_disable: bool, + enables: u32, + disables: u32, + reads: u32, + } + + impl DirtyLog for FakeVm { + fn dirty_tracking(&self) -> DirtyTracking { + self.tracking.unwrap_or(DirtyTracking::None) + } + fn enable_dirty_tracking(&mut self) -> Result<(), HypervisorError> { + if self.fail_enable { + return Err(HypervisorError::Injected); + } + self.enables += 1; + self.all_written = true; + Ok(()) + } + fn disable_dirty_tracking(&mut self, _: u64, _: usize) -> Result<(), HypervisorError> { + if self.fail_disable { + return Err(HypervisorError::Injected); + } + self.disables += 1; + Ok(()) + } + fn read_dirty_log( + &mut self, + gpa: u64, + size: usize, + bitmap: &mut Vec, + ) -> Result<(), HypervisorError> { + assert_eq!((gpa, size % DIRTY_PAGE_SIZE), (GPA, 0)); + if self.fail { + return Err(HypervisorError::Injected); + } + self.reads += 1; + *bitmap = vec![0; (size / DIRTY_PAGE_SIZE).div_ceil(64)]; + let written = if std::mem::take(&mut self.all_written) { + size / DIRTY_PAGE_SIZE + } else { + self.written + }; + for page in 0..written { + bitmap[page / 64] |= 1 << (page % 64); + } + Ok(()) + } + } + + fn vm(tracking: DirtyTracking, written: usize) -> FakeVm { + FakeVm { + tracking: Some(tracking), + written, + ..FakeVm::default() + } + } + + fn restore(log: &mut ScratchDirtyLog, vm: &mut FakeVm) -> Option { + log.take(vm, GPA, SIZE).map(|bitmap| written(bitmap, PAGES)) + } + + #[test] + fn written_counts_only_pages_in_range() { + assert_eq!(written(&[u64::MAX, u64::MAX], 70), 70); + assert_eq!(written(&[0b1011], 3), 2); + assert_eq!(written(&[], 10), 0); + } + + #[test] + fn untracked_vms_have_no_log() { + let mut vm = vm(DirtyTracking::None, 1); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + assert_eq!(restore(&mut log, &mut vm), None); + assert_eq!((vm.enables, vm.disables, vm.reads), (0, 0, 0)); + } + + /// Every restore, the first included, zeroes what the log says, + /// whatever the guest writes. + #[test] + fn mapped_tracking_is_always_read() { + let mut vm = vm(DirtyTracking::Mapped, PAGES / 2); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + for _ in 0..10 { + assert_eq!(restore(&mut log, &mut vm), Some(PAGES / 2)); + } + assert_eq!((vm.enables, vm.disables, vm.reads), (0, 0, 10)); + } + + /// Enabled when mapped and not read: the first log reports every page, + /// later ones what was written. + #[test] + fn switched_tracking_starts_when_mapped() { + let mut vm = vm(DirtyTracking::Switched, 1); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + assert_eq!((vm.enables, vm.reads), (1, 0)); + assert_eq!(restore(&mut log, &mut vm), Some(PAGES)); + assert_eq!(vm.disables, 0); + vm.written = 1; + for _ in 0..1000 { + assert_eq!(restore(&mut log, &mut vm), Some(1)); + } + assert_eq!(vm.disables, 0); + } + + #[test] + fn small_scratch_is_not_tracked() { + let size = MIN_TRACKED_SCRATCH - DIRTY_PAGE_SIZE; + let mut vm = vm(DirtyTracking::Switched, 1); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, size); + assert!(log.take(&mut vm, GPA, size).is_none()); + assert_eq!((vm.enables, vm.reads), (0, 0)); + } + + /// A run that writes much of scratch stops tracking for + /// [`RESUME_AFTER`] restores, and its own log is still used. Setting + /// up does not count. Tracking then resumes, from the next run on, and + /// stops again if that run writes much. + #[test] + fn a_run_writing_much_pauses_tracking() { + let mut vm = vm(DirtyTracking::Switched, PAGES); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + assert_eq!(restore(&mut log, &mut vm), Some(PAGES)); + assert_eq!(vm.disables, 0); + for round in 1..=2 { + vm.written = PAGES / 10 + 1; + assert_eq!(restore(&mut log, &mut vm), Some(PAGES / 10 + 1)); + assert_eq!((vm.enables, vm.disables), (round, round)); + vm.written = 1; + for _ in 0..RESUME_AFTER { + assert_eq!(restore(&mut log, &mut vm), None); + } + assert_eq!(vm.enables, round + 1); + assert_eq!(restore(&mut log, &mut vm), Some(1)); + } + } + + /// Tracking that fails to resume resets in full from then on, and is + /// switched off if it was switched on. + #[test] + fn a_failed_resume_resets_in_full() { + for fail_read in [false, true] { + let mut vm = vm(DirtyTracking::Switched, PAGES); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + restore(&mut log, &mut vm); + restore(&mut log, &mut vm); + assert_eq!(vm.disables, 1); + vm.written = 1; + vm.fail_enable = !fail_read; + vm.fail = fail_read; + for _ in 0..RESUME_AFTER * 2 { + assert_eq!(restore(&mut log, &mut vm), None); + } + assert_eq!(log.state, State::Failed); + assert!(!log.enabled); + assert_eq!(vm.disables, 1 + u32::from(fail_read)); + } + } + + /// Tracking that fails to stop resets in full from then on. + #[test] + fn a_failed_stop_resets_in_full() { + let mut vm = vm(DirtyTracking::Switched, PAGES); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + restore(&mut log, &mut vm); + vm.fail_disable = true; + assert_eq!(restore(&mut log, &mut vm), Some(PAGES)); + vm.written = 1; + for _ in 0..RESUME_AFTER * 2 { + assert_eq!(restore(&mut log, &mut vm), None); + } + assert_eq!((log.state, vm.enables), (State::Failed, 1)); + } + + #[test] + fn larger_scratch_tolerates_more_writes() { + let large = LARGE_SCRATCH / DIRTY_PAGE_SIZE; + assert!(!writes_much(large / 4, large)); + assert!(writes_much(large / 4 + 1, large)); + assert!(!writes_much(PAGES / 10, PAGES)); + assert!(writes_much(PAGES / 10 + 1, PAGES)); + } + + /// Unmapping switches tracking off over the range still mapped, and + /// the next mapping starts over, tracking again if it is large enough. + #[test] + fn remapping_stops_and_starts_over() { + let mut vm = vm(DirtyTracking::Switched, 1); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + restore(&mut log, &mut vm); + log.unmapping(&mut vm, GPA, SIZE); + assert_eq!(vm.disables, 1); + log.mapped(&mut vm, GPA, MIN_TRACKED_SCRATCH - DIRTY_PAGE_SIZE); + assert!( + log.take(&mut vm, GPA, MIN_TRACKED_SCRATCH - DIRTY_PAGE_SIZE) + .is_none() + ); + log.unmapping(&mut vm, GPA, MIN_TRACKED_SCRATCH - DIRTY_PAGE_SIZE); + assert_eq!(vm.disables, 1); + log.mapped(&mut vm, GPA, SIZE); + assert_eq!(restore(&mut log, &mut vm), Some(PAGES)); + assert_eq!(restore(&mut log, &mut vm), Some(1)); + assert_eq!(vm.enables, 2); + } + + /// A pause ends with the mapping: scratch mapped again while tracking + /// is paused is tracked, or not, as any new mapping. + #[test] + fn remapping_ends_a_pause() { + let small = MIN_TRACKED_SCRATCH - DIRTY_PAGE_SIZE; + for (unmap, size) in [(true, SIZE), (true, small), (false, SIZE), (false, small)] { + let mut vm = vm(DirtyTracking::Switched, PAGES); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + restore(&mut log, &mut vm); + restore(&mut log, &mut vm); + assert_eq!(log.paused, RESUME_AFTER); + if unmap { + log.unmapping(&mut vm, GPA, SIZE); + for _ in 0..RESUME_AFTER * 2 { + assert!(log.take(&mut vm, GPA, SIZE).is_none()); + } + assert_eq!(vm.enables, 1); + } + let enables = vm.enables; + log.mapped(&mut vm, GPA, size); + vm.written = 1; + for _ in 0..RESUME_AFTER * 2 { + log.take(&mut vm, GPA, size); + } + let tracked = size == SIZE; + assert_eq!(vm.enables - enables, u32::from(tracked)); + assert_eq!(log.take(&mut vm, GPA, size).is_some(), tracked); + } + } + + /// A mapping that loses tracking has no log. + #[test] + fn a_mapping_without_tracking_has_no_log() { + let mut vm = vm(DirtyTracking::Mapped, 1); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + log.unmapping(&mut vm, GPA, SIZE); + vm.tracking = Some(DirtyTracking::None); + log.mapped(&mut vm, GPA, SIZE); + assert!(log.take(&mut vm, GPA, SIZE).is_none()); + } + + #[test] + fn a_failed_read_resets_in_full_from_then_on() { + let mut vm = vm(DirtyTracking::Switched, 1); + let mut log = ScratchDirtyLog::default(); + log.mapped(&mut vm, GPA, SIZE); + vm.fail = true; + assert_eq!(restore(&mut log, &mut vm), None); + assert_eq!(vm.disables, 1); + vm.fail = false; + log.mapped(&mut vm, GPA, SIZE); + for _ in 0..10 { + assert_eq!(restore(&mut log, &mut vm), None); + } + assert_eq!((vm.enables, vm.reads), (1, 0)); + } +} diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm/mod.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm/mod.rs index 9d6157364..b1ca9cd45 100644 --- a/src/hyperlight_host/src/hypervisor/hyperlight_vm/mod.rs +++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm/mod.rs @@ -7,6 +7,8 @@ mod x86_64; #[cfg(target_arch = "aarch64")] mod aarch64; +mod dirty_log; + #[cfg(all(test, not(gdb)))] pub(crate) mod test_support; @@ -18,6 +20,7 @@ use std::sync::{Arc, Mutex}; use hyperlight_common::log_level::GuestLogFilter; use tracing_core::LevelFilter; +use self::dirty_log::ScratchDirtyLog; use crate::HyperlightError; #[cfg(gdb)] use crate::hypervisor::gdb::DebuggableVm; @@ -378,6 +381,8 @@ pub(crate) struct HyperlightVm { // The current scratch region, used to keep it alive as long as it // is used & when unmapping pub(super) scratch_memory: Option, + /// What the guest wrote to scratch since the last restore. + pub(super) scratch_dirty: ScratchDirtyLog, pub(super) mmap_regions: Vec<(u32, MemoryRegion)>, // Later mapped regions (slot number, region) @@ -535,6 +540,23 @@ impl HyperlightVm { Ok(()) } + /// The scratch pages the guest wrote since the last restore, where + /// the hypervisor logs them. Called once per restore, before + /// scratch is reset. See [`ScratchDirtyLog`]. + pub(crate) fn scratch_dirty_pages(&mut self) -> Option<&mut Vec> { + let size = self.scratch_memory.as_ref()?.mem_size(); + let gpa = hyperlight_common::layout::scratch_base_gpa(size); + self.scratch_dirty.take(&mut *self.vm, gpa, size) + } + + /// The restores that had the guest's writes to scratch from the + /// hypervisor, or `None` where it does not log them. + #[cfg(test)] + pub(crate) fn scratch_dirty_logs(&self) -> Option { + use crate::hypervisor::virtual_machine::DirtyTracking; + (self.vm.dirty_tracking() != DirtyTracking::None).then_some(self.scratch_dirty.logs) + } + /// Update the scratch mapping to point to a new GuestSharedMemory pub(crate) fn update_scratch_mapping( &mut self, @@ -547,11 +569,16 @@ impl HyperlightVm { if let Some(old_scratch) = self.scratch_memory.as_ref() { let old_base = hyperlight_common::layout::scratch_base_gpa(old_scratch.mem_size()); let old_rgn = old_scratch.mapping_at(old_base, MemoryRegionType::Scratch); + self.scratch_dirty + .unmapping(&mut *self.vm, old_base, old_rgn.guest_region.len()); self.vm.unmap_memory((self.scratch_slot, &old_rgn))?; } self.scratch_memory = None; unsafe { self.vm.map_memory((self.scratch_slot, &rgn))? }; self.scratch_memory = Some(scratch); + // Track the guest's writes to it from before the guest runs. + self.scratch_dirty + .mapped(&mut *self.vm, guest_base, rgn.guest_region.len()); Ok(()) } diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm/test_support.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm/test_support.rs index e410d6266..d8ff2ff6e 100644 --- a/src/hyperlight_host/src/hypervisor/hyperlight_vm/test_support.rs +++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm/test_support.rs @@ -11,7 +11,9 @@ use crate::hypervisor::regs::{ }; #[cfg(target_arch = "x86_64")] use crate::hypervisor::virtual_machine::CreateVmError; -use crate::hypervisor::virtual_machine::{HypervisorError, VirtualMachine}; +use crate::hypervisor::virtual_machine::{ + DirtyLog, DirtyTracking, HypervisorError, VirtualMachine, +}; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(crate) enum VmOperation { @@ -103,6 +105,33 @@ impl FaultInjectingVirtualMachine { } } +impl DirtyLog for FaultInjectingVirtualMachine { + fn dirty_tracking(&self) -> DirtyTracking { + self.inner().dirty_tracking() + } + + fn enable_dirty_tracking(&mut self) -> std::result::Result<(), HypervisorError> { + self.inner_mut().enable_dirty_tracking() + } + + fn disable_dirty_tracking( + &mut self, + gpa: u64, + size: usize, + ) -> std::result::Result<(), HypervisorError> { + self.inner_mut().disable_dirty_tracking(gpa, size) + } + + fn read_dirty_log( + &mut self, + gpa: u64, + size: usize, + bitmap: &mut Vec, + ) -> std::result::Result<(), HypervisorError> { + self.inner_mut().read_dirty_log(gpa, size, bitmap) + } +} + impl VirtualMachine for FaultInjectingVirtualMachine { unsafe fn map_memory( &mut self, diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs index 8ca3024fc..24e07bf7e 100644 --- a/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs +++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs @@ -122,6 +122,7 @@ impl HyperlightVm { snapshot_memory: None, scratch_slot, scratch_memory: None, + scratch_dirty: ScratchDirtyLog::default(), mmap_regions: Vec::new(), diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/hvf/mod.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/hvf/mod.rs index 9756c3953..bf38e63d4 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/hvf/mod.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/hvf/mod.rs @@ -1095,6 +1095,8 @@ impl MemorySpace { } } +impl crate::hypervisor::virtual_machine::DirtyLog for HvfVm {} + impl VirtualMachine for HvfVm { unsafe fn map_memory( &mut self, diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/kvm/aarch64.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/kvm/aarch64.rs index fdd8d4963..341561bcc 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/kvm/aarch64.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/kvm/aarch64.rs @@ -105,6 +105,8 @@ impl KvmVm { } } +impl crate::hypervisor::virtual_machine::DirtyLog for KvmVm {} + impl VirtualMachine for KvmVm { unsafe fn map_memory( &mut self, diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/kvm/x86_64.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/kvm/x86_64.rs index e352ab20c..f27d57256 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/kvm/x86_64.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/kvm/x86_64.rs @@ -393,6 +393,8 @@ impl KvmVm { } } +impl crate::hypervisor::virtual_machine::DirtyLog for KvmVm {} + impl VirtualMachine for KvmVm { unsafe fn map_memory( &mut self, diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs index e1922a1ec..8f3cb91d8 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs @@ -18,6 +18,7 @@ use crate::hypervisor::regs::{ #[cfg(target_arch = "x86_64")] use crate::hypervisor::regs::{MsrEntry, is_resettable_msr}; use crate::mem::memory_region::MemoryRegion; +use crate::mem::shared_mem::DIRTY_PAGE_SIZE; #[cfg(feature = "trace_guest")] use crate::sandbox::trace::TraceContext as SandboxTraceContext; @@ -429,9 +430,88 @@ pub enum HvfSyncError { SyncInvariant(String), } +/// How a VM reports the pages the guest writes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum DirtyTracking { + /// It does not. + None, + /// For the scratch region, from when it is mapped (WHP). + #[cfg_attr(not(target_os = "windows"), allow(dead_code))] + Mapped, + /// For the whole VM, switched on and off (MSHV). While on, the + /// guest's first write to a page after each read faults to the + /// hypervisor. + #[cfg_attr(not(all(mshv3, target_arch = "x86_64")), allow(dead_code))] + Switched, +} + +/// The pages the guest writes, where a VM logs them. Restores use it to +/// zero only those (see `hyperlight_vm::dirty_log`). +pub(crate) trait DirtyLog { + /// How this VM reports the pages the guest writes. The dirty-log + /// methods below are only called when this is not + /// [`DirtyTracking::None`]. + fn dirty_tracking(&self) -> DirtyTracking { + DirtyTracking::None + } + + /// Start tracking the pages the guest writes, for + /// [`DirtyTracking::Switched`]. Until the first + /// [`read_dirty_log`](Self::read_dirty_log) after this, every page + /// may read as written. + fn enable_dirty_tracking(&mut self) -> std::result::Result<(), HypervisorError> { + Ok(()) + } + + /// Stop tracking, for [`DirtyTracking::Switched`]. `[gpa, gpa + + /// size)` is the range read since tracking was enabled. + fn disable_dirty_tracking( + &mut self, + _gpa: u64, + _size: usize, + ) -> std::result::Result<(), HypervisorError> { + Ok(()) + } + + /// Set `bitmap` to the pages, of [`DIRTY_PAGE_SIZE`], in `[gpa, gpa + + /// size)` the guest wrote since the last read, one bit per page from + /// bit 0 of word 0, and clear them. Bits past the range may be set. + fn read_dirty_log( + &mut self, + _gpa: u64, + size: usize, + bitmap: &mut Vec, + ) -> std::result::Result<(), HypervisorError> { + // A backend that tracks but does not read reports every page, so + // a restore zeroes all of scratch rather than too little. + bitmap.clear(); + bitmap.resize(size.div_ceil(DIRTY_PAGE_SIZE * 64), u64::MAX); + Ok(()) + } +} + +#[cfg(test)] +mod dirty_log_tests { + use super::*; + + /// A backend that tracks but does not read its log reports every + /// page written, so a restore zeroes all of scratch. + #[test] + fn the_default_log_reports_every_page() { + struct Untracked; + impl DirtyLog for Untracked {} + let mut bitmap = vec![0; 1]; + Untracked + .read_dirty_log(0, 130 * DIRTY_PAGE_SIZE, &mut bitmap) + .unwrap(); + assert_eq!(bitmap, vec![u64::MAX; 3]); + assert_eq!(Untracked.dirty_tracking(), DirtyTracking::None); + } +} + /// Trait for single-vCPU VMs. Provides a common interface for basic VM operations. /// Abstracts over differences between KVM, MSHV and WHP implementations. -pub(crate) trait VirtualMachine: Debug + Send { +pub(crate) trait VirtualMachine: Debug + Send + DirtyLog { /// Map memory region into this VM /// /// # Safety diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/mshv/x86_64.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/mshv/x86_64.rs index 1407fceaf..2a7818a98 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/mshv/x86_64.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/mshv/x86_64.rs @@ -13,8 +13,9 @@ use mshv_bindings::LapicState; #[cfg(gdb)] use mshv_bindings::{DebugRegisters, hv_message_type_HVMSG_X64_EXCEPTION_INTERCEPT}; use mshv_bindings::{ - FloatingPointUnit, HV_X64_REGISTER_CLASS_IP, SpecialRegisters, StandardRegisters, XSave, Xcrs, - hv_message_type, hv_message_type_HVMSG_GPA_INTERCEPT, hv_message_type_HVMSG_UNMAPPED_GPA, + FloatingPointUnit, HV_X64_REGISTER_CLASS_IP, MSHV_GPAP_ACCESS_OP_CLEAR, + MSHV_GPAP_ACCESS_OP_SET, SpecialRegisters, StandardRegisters, XSave, Xcrs, hv_message_type, + hv_message_type_HVMSG_GPA_INTERCEPT, hv_message_type_HVMSG_UNMAPPED_GPA, hv_message_type_HVMSG_X64_HALT, hv_message_type_HVMSG_X64_IO_PORT_INTERCEPT, hv_partition_property_code_HV_PARTITION_PROPERTY_SYNTHETIC_PROC_FEATURES, hv_partition_synthetic_processor_features, hv_register_assoc, @@ -47,10 +48,11 @@ use crate::hypervisor::virtual_machine::XSAVE_BUFFER_SIZE; #[cfg(feature = "hw-interrupts")] use crate::hypervisor::virtual_machine::x86_64::hw_interrupts::TimerThread; use crate::hypervisor::virtual_machine::{ - CreateVmError, MapMemoryError, RegisterError, RunVcpuError, UnmapMemoryError, VirtualMachine, - VmExit, + CreateVmError, DirtyLog, DirtyTracking, HypervisorError, MapMemoryError, RegisterError, + RunVcpuError, UnmapMemoryError, VirtualMachine, VmExit, }; use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags}; +use crate::mem::shared_mem::DIRTY_PAGE_SIZE; #[cfg(feature = "trace_guest")] use crate::sandbox::trace::TraceContext as SandboxTraceContext; @@ -230,6 +232,44 @@ impl MshvVm { } } +impl DirtyLog for MshvVm { + fn dirty_tracking(&self) -> DirtyTracking { + DirtyTracking::Switched + } + + fn enable_dirty_tracking(&mut self) -> std::result::Result<(), HypervisorError> { + Ok(self.vm_fd.enable_dirty_page_tracking()?) + } + + fn disable_dirty_tracking( + &mut self, + gpa: u64, + size: usize, + ) -> std::result::Result<(), HypervisorError> { + // MSHV refuses to stop tracking while any page's bit is clear. + self.vm_fd.get_dirty_log( + gpa / DIRTY_PAGE_SIZE as u64, + size, + MSHV_GPAP_ACCESS_OP_SET as u8, + )?; + Ok(self.vm_fd.disable_dirty_page_tracking()?) + } + + fn read_dirty_log( + &mut self, + gpa: u64, + size: usize, + bitmap: &mut Vec, + ) -> std::result::Result<(), HypervisorError> { + *bitmap = self.vm_fd.get_dirty_log( + gpa / DIRTY_PAGE_SIZE as u64, + size, + MSHV_GPAP_ACCESS_OP_CLEAR as u8, + )?; + Ok(()) + } +} + impl VirtualMachine for MshvVm { unsafe fn map_memory( &mut self, diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/whp/aarch64.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/whp/aarch64.rs index bfff3ecd7..e279723ac 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/whp/aarch64.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/whp/aarch64.rs @@ -330,6 +330,10 @@ pub(crate) struct WhpVm { /// Tracks host-side file mappings for cleanup. file_mappings: Vec<(HandleWrapper, *mut c_void)>, _no_surrogate_guard: Option, + /// The scratch region is mapped with dirty-page tracking. + scratch_dirty_tracked: bool, + /// Mapping with dirty-page tracking failed, so it is not tried again. + dirty_tracking_failed: bool, } // Safety: same reasoning as x86_64 WhpVm — raw pointers are kernel resource handles, @@ -406,6 +410,8 @@ impl WhpVm { surrogate_process: None, file_mappings: Vec::new(), _no_surrogate_guard: no_surrogate_guard, + scratch_dirty_tracked: false, + dirty_tracking_failed: false, }; if !no_surrogate { @@ -457,6 +463,26 @@ impl WhpVm { } } +impl crate::hypervisor::virtual_machine::DirtyLog for WhpVm { + fn dirty_tracking(&self) -> crate::hypervisor::virtual_machine::DirtyTracking { + use crate::hypervisor::virtual_machine::DirtyTracking; + if self.scratch_dirty_tracked { + DirtyTracking::Mapped + } else { + DirtyTracking::None + } + } + + fn read_dirty_log( + &mut self, + gpa: u64, + size: usize, + bitmap: &mut Vec, + ) -> Result<(), super::super::HypervisorError> { + super::read_dirty_bitmap(self.partition, gpa, size, bitmap) + } +} + impl VirtualMachine for WhpVm { unsafe fn map_memory( &mut self, @@ -502,22 +528,45 @@ impl VirtualMachine for WhpVm { } }; - let result = unsafe { - (self.map_gpa_range2)( - self.partition, - process_handle, - host_addr, - region.guest_region.start as u64, - region.guest_region.len() as u64, - flags, - ) + let map = |flags| { + // SAFETY: the caller of `map_memory` keeps the region valid + // while it is mapped. + let result = unsafe { + (self.map_gpa_range2)( + self.partition, + process_handle, + host_addr, + region.guest_region.start as u64, + region.guest_region.len() as u64, + flags, + ) + }; + if result.is_err() { + return Err(MapMemoryError::Hypervisor( + super::super::HypervisorError::WindowsError( + windows_result::Error::from_hresult(result), + ), + )); + } + Ok(()) }; - if result.is_err() { - return Err(MapMemoryError::Hypervisor( - super::super::HypervisorError::WindowsError(windows_result::Error::from_hresult( - result, - )), - )); + if region.region_type == MemoryRegionType::Scratch { + self.scratch_dirty_tracked = false; + // Tracking lets a restore zero only the pages the guest wrote. + // Without it, scratch is mapped as before. + if self.dirty_tracking_failed { + map(flags)?; + } else if let Err(e) = map(flags | WHvMapGpaRangeFlagTrackDirtyPages) { + // Expected where Windows does not track dirty pages, so not + // a warning, unlike a failed read of the log. + tracing::debug!("Mapping scratch with dirty tracking failed: {e}"); + self.dirty_tracking_failed = true; + map(flags)?; + } else { + self.scratch_dirty_tracked = true; + } + } else { + map(flags)?; } if region.region_type == MemoryRegionType::MappedFile { diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/whp/mod.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/whp/mod.rs index 6b730526e..fa4b60ce4 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/whp/mod.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/whp/mod.rs @@ -18,6 +18,35 @@ mod aarch64; #[cfg(target_arch = "aarch64")] pub(crate) use aarch64::*; +/// Set `bitmap` to the pages of [`DIRTY_PAGE_SIZE`] in `[gpa, gpa + size)` +/// the guest wrote since the last read, and clear them. +/// +/// [`DIRTY_PAGE_SIZE`]: crate::mem::shared_mem::DIRTY_PAGE_SIZE +fn read_dirty_bitmap( + partition: windows::Win32::System::Hypervisor::WHV_PARTITION_HANDLE, + gpa: u64, + size: usize, + bitmap: &mut Vec, +) -> Result<(), super::HypervisorError> { + use crate::mem::shared_mem::DIRTY_PAGE_SIZE; + // Every word is written: resize only. + bitmap.resize((size / DIRTY_PAGE_SIZE).div_ceil(64), 0); + let len = u32::try_from(bitmap.len() * size_of::()).map_err(|_| { + windows_result::Error::from_hresult(windows::Win32::Foundation::E_INVALIDARG) + })?; + // SAFETY: `bitmap` holds `len` bytes, which the call fills. + unsafe { + windows::Win32::System::Hypervisor::WHvQueryGpaRangeDirtyBitmap( + partition, + gpa, + size as u64, + Some(bitmap.as_mut_ptr()), + len, + )? + }; + Ok(()) +} + fn release_file_mapping(view_base: *mut c_void, mapping_handle: HandleWrapper) { unsafe { if let Err(error) = UnmapViewOfFile(MEMORY_MAPPED_VIEW_ADDRESS { Value: view_base }) { diff --git a/src/hyperlight_host/src/hypervisor/virtual_machine/whp/x86_64.rs b/src/hyperlight_host/src/hypervisor/virtual_machine/whp/x86_64.rs index 57342a532..f28af33ea 100644 --- a/src/hyperlight_host/src/hypervisor/virtual_machine/whp/x86_64.rs +++ b/src/hyperlight_host/src/hypervisor/virtual_machine/whp/x86_64.rs @@ -35,8 +35,8 @@ use crate::hypervisor::surrogate_process_manager::{ #[cfg(feature = "hw-interrupts")] use crate::hypervisor::virtual_machine::x86_64::hw_interrupts::{self, TimerThread}; use crate::hypervisor::virtual_machine::{ - CreateVmError, HypervisorError, MapMemoryError, RegisterError, RunVcpuError, UnmapMemoryError, - VirtualMachine, VmExit, XSAVE_MIN_SIZE, + CreateVmError, DirtyLog, DirtyTracking, HypervisorError, MapMemoryError, RegisterError, + RunVcpuError, UnmapMemoryError, VirtualMachine, VmExit, XSAVE_MIN_SIZE, }; use crate::hypervisor::wrappers::HandleWrapper; use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType}; @@ -249,6 +249,10 @@ pub(crate) struct WhpVm { /// Handle to the background timer (if started). #[cfg(feature = "hw-interrupts")] timer: Option, + /// The scratch region is mapped with dirty-page tracking. + scratch_dirty_tracked: bool, + /// Mapping with dirty-page tracking failed, so it is not tried again. + dirty_tracking_failed: bool, } // Safety: `WhpVm` is !Send because it holds `Option` which @@ -317,6 +321,8 @@ impl WhpVm { _no_surrogate_guard: no_surrogate_guard, #[cfg(feature = "hw-interrupts")] timer: None, + scratch_dirty_tracked: false, + dirty_tracking_failed: false, }) } @@ -359,58 +365,21 @@ impl WhpVm { } } -impl VirtualMachine for WhpVm { - unsafe fn map_memory( - &mut self, - (_slot, region): (u32, &MemoryRegion), +impl WhpVm { + /// Map `region` into the partition, from `surrogate_addr` in the + /// surrogate process when there is one. + /// + /// # Safety + /// As for [`VirtualMachine::map_memory`]: the memory `region` describes, + /// and its view at `surrogate_addr`, must stay valid while it is mapped. + unsafe fn map_gpa_range( + &self, + region: &MemoryRegion, + surrogate_addr: Option<*mut c_void>, + flags: WHV_MAP_GPA_RANGE_FLAGS, ) -> Result<(), MapMemoryError> { - let flags = region - .flags - .iter() - .map(|flag| match flag { - MemoryRegionFlags::NONE => Ok(WHvMapGpaRangeFlagNone), - MemoryRegionFlags::READ => Ok(WHvMapGpaRangeFlagRead), - MemoryRegionFlags::WRITE => Ok(WHvMapGpaRangeFlagWrite), - MemoryRegionFlags::EXECUTE => Ok(WHvMapGpaRangeFlagExecute), - _ => Err(MapMemoryError::InvalidFlags(format!( - "Invalid memory region flag: {:?}", - flag - ))), - }) - .collect::, MapMemoryError>>()? - .iter() - .fold(WHvMapGpaRangeFlagNone, |acc, flag| acc | *flag); - - match &mut self.surrogate_process { - None => { - let host_addr = (region.host_region.start.handle_base - + region.host_region.start.offset) - as *const c_void; - let res = unsafe { - WHvMapGpaRange( - self.partition, - host_addr, - region.guest_region.start as u64, - region.guest_region.len() as u64, - flags, - ) - }; - if let Err(e) = res { - return Err(MapMemoryError::Hypervisor(HypervisorError::WindowsError(e))); - } - } - Some(surrogate) => { - // Calculate the surrogate process address for this region - let surrogate_base = surrogate - .map( - region.host_region.start.from_handle, - region.host_region.start.handle_base, - region.host_region.start.handle_size, - ®ion.region_type.surrogate_mapping(), - ) - .map_err(|e| MapMemoryError::SurrogateProcess(e.to_string()))?; - let surrogate_addr = surrogate_base.wrapping_add(region.host_region.start.offset); - + match (&self.surrogate_process, surrogate_addr) { + (Some(surrogate), Some(surrogate_addr)) => { // This function dynamically loads the WHvMapGpaRange2 function from the winhvplatform.dll // WHvMapGpaRange2 only available on Windows 11 or Windows Server 2022 and later // we do things this way to allow a user trying to load hyperlight on an older version of windows to @@ -418,6 +387,8 @@ impl VirtualMachine for WhpVm { // with an error about a missing entrypoint // This function should always succeed since before we get here we have already checked that the hypervisor is present and // that we are on a supported version of windows. + // SAFETY: loads a function from winhvplatform.dll, whose + // signature `try_load_whv_map_gpa_range2` declares. let whvmapgparange2_func = unsafe { match try_load_whv_map_gpa_range2() { Ok(func) => func, @@ -430,6 +401,8 @@ impl VirtualMachine for WhpVm { } }; + // SAFETY: the caller keeps the region valid while it is + // mapped, and `surrogate_addr` is its view in the surrogate. let res = unsafe { whvmapgparange2_func( self.partition, @@ -446,6 +419,116 @@ impl VirtualMachine for WhpVm { ))); } } + (None, None) => { + let host_addr = (region.host_region.start.handle_base + + region.host_region.start.offset) + as *const c_void; + // SAFETY: the caller keeps the region at `host_addr` valid + // while it is mapped. + let res = unsafe { + WHvMapGpaRange( + self.partition, + host_addr, + region.guest_region.start as u64, + region.guest_region.len() as u64, + flags, + ) + }; + if let Err(e) = res { + return Err(MapMemoryError::Hypervisor(HypervisorError::WindowsError(e))); + } + } + _ => { + return Err(MapMemoryError::SurrogateProcess( + "surrogate address given without a surrogate, or not given with one".into(), + )); + } + } + Ok(()) + } +} + +impl DirtyLog for WhpVm { + fn dirty_tracking(&self) -> DirtyTracking { + if self.scratch_dirty_tracked { + DirtyTracking::Mapped + } else { + DirtyTracking::None + } + } + + fn read_dirty_log( + &mut self, + gpa: u64, + size: usize, + bitmap: &mut Vec, + ) -> Result<(), HypervisorError> { + super::read_dirty_bitmap(self.partition, gpa, size, bitmap) + } +} + +impl VirtualMachine for WhpVm { + unsafe fn map_memory( + &mut self, + (_slot, region): (u32, &MemoryRegion), + ) -> Result<(), MapMemoryError> { + let flags = region + .flags + .iter() + .map(|flag| match flag { + MemoryRegionFlags::NONE => Ok(WHvMapGpaRangeFlagNone), + MemoryRegionFlags::READ => Ok(WHvMapGpaRangeFlagRead), + MemoryRegionFlags::WRITE => Ok(WHvMapGpaRangeFlagWrite), + MemoryRegionFlags::EXECUTE => Ok(WHvMapGpaRangeFlagExecute), + _ => Err(MapMemoryError::InvalidFlags(format!( + "Invalid memory region flag: {:?}", + flag + ))), + }) + .collect::, MapMemoryError>>()? + .iter() + .fold(WHvMapGpaRangeFlagNone, |acc, flag| acc | *flag); + + // The surrogate's view of the region, when there is a surrogate. + let surrogate_addr = match &mut self.surrogate_process { + None => None, + Some(surrogate) => { + // Calculate the surrogate process address for this region + let surrogate_base = surrogate + .map( + region.host_region.start.from_handle, + region.host_region.start.handle_base, + region.host_region.start.handle_size, + ®ion.region_type.surrogate_mapping(), + ) + .map_err(|e| MapMemoryError::SurrogateProcess(e.to_string()))?; + Some(surrogate_base.wrapping_add(region.host_region.start.offset)) + } + }; + + if region.region_type == MemoryRegionType::Scratch { + self.scratch_dirty_tracked = false; + // Tracking lets a restore zero only the pages the guest wrote. + // Without it, scratch is mapped as before. + let tracked = flags | WHvMapGpaRangeFlagTrackDirtyPages; + // SAFETY: the caller of `map_memory` keeps the region valid + // while it is mapped. + unsafe { + if self.dirty_tracking_failed { + self.map_gpa_range(region, surrogate_addr, flags)?; + } else if let Err(e) = self.map_gpa_range(region, surrogate_addr, tracked) { + // Expected where Windows does not track dirty pages, so + // not a warning, unlike a failed read of the log. + tracing::debug!("Mapping scratch with dirty tracking failed: {e}"); + self.dirty_tracking_failed = true; + self.map_gpa_range(region, surrogate_addr, flags)?; + } else { + self.scratch_dirty_tracked = true; + } + } + } else { + // SAFETY: as above. + unsafe { self.map_gpa_range(region, surrogate_addr, flags)? }; } // Track host-side file mappings for cleanup on unmap or drop. diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index 970525211..f749df8ca 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -15,6 +15,7 @@ use hyperlight_common::vmem::{BasicMapping, MappingKind}; use tracing::{Span, instrument}; use super::layout::SandboxMemoryLayout; +use super::scratch_reset::ScratchReset; use super::shared_mem::{ ExclusiveSharedMemory, GuestSharedMemory, HostSharedMemory, ReadonlySharedMemory, SharedMemory, }; @@ -174,6 +175,8 @@ pub(crate) struct SandboxMemoryManager { pub(crate) h2g_consumer: Option, /// Correlation ID sequence survives consumer replacement and manager cloning. next_guest_cid: u32, + /// How scratch resets keep or drop backed pages. + pub(crate) scratch_reset: ScratchReset, } impl Clone for SandboxMemoryManager { @@ -190,6 +193,8 @@ impl Clone for SandboxMemoryManager { g2h_consumer: None, h2g_consumer: None, next_guest_cid: self.next_guest_cid, + // Counts its own resets. + scratch_reset: ScratchReset::default(), } } } @@ -331,6 +336,7 @@ where g2h_consumer: None, h2g_consumer: None, next_guest_cid: 1, + scratch_reset: ScratchReset::default(), } } @@ -344,7 +350,7 @@ impl SandboxMemoryManager { pub(crate) fn from_snapshot(s: &Snapshot) -> Result { let layout = *s.layout(); let shared_mem = s.memory().to_mgr_snapshot_mem()?; - let scratch_mem = ExclusiveSharedMemory::new(s.layout().get_scratch_size())?; + let scratch_mem = ExclusiveSharedMemory::new_scratch(s.layout().get_scratch_size())?; let next_action = s.next_action(); let mut mgr = Self::new(layout, shared_mem, scratch_mem, next_action); mgr.original_entrypoint = s.original_entrypoint(); @@ -386,6 +392,7 @@ impl SandboxMemoryManager { g2h_consumer: None, h2g_consumer: None, next_guest_cid: self.next_guest_cid, + scratch_reset: self.scratch_reset, }; let guest_mgr = SandboxMemoryManager { shared_mem: gshm, @@ -399,6 +406,7 @@ impl SandboxMemoryManager { g2h_consumer: None, h2g_consumer: None, next_guest_cid: self.next_guest_cid, + scratch_reset: ScratchReset::default(), }; host_mgr.update_scratch_bookkeeping()?; @@ -689,9 +697,14 @@ impl SandboxMemoryManager { } /// Restore base memory after the caller checks snapshot compatibility. + /// + /// `guest_written` is the scratch pages the guest wrote since the + /// last restore, one bit per 4 KiB page, when the hypervisor logs + /// them. Only those, and what the host wrote, are zeroed then. pub(crate) fn restore_snapshot( &mut self, snapshot: &Snapshot, + guest_written: Option<&mut Vec>, ) -> Result<( Option>, Option, @@ -728,11 +741,17 @@ impl SandboxMemoryManager { }; let new_scratch_size = snapshot.layout().get_scratch_size(); let gscratch = if new_scratch_size == self.scratch_mem.mem_size() { - // zero_or_replace picks the fastest zeroing strategy for - // the current platform (see SharedMemory::zero_or_replace). - self.scratch_mem.zero_or_replace()? + match guest_written { + Some(written) => { + self.scratch_mem.zero_written(written)?; + None + } + // zero_or_replace picks the fastest zeroing strategy for + // the current platform (see SharedMemory::zero_or_replace). + None => self.scratch_mem.zero_or_replace(&mut self.scratch_reset)?, + } } else { - let new_scratch_mem = ExclusiveSharedMemory::new(new_scratch_size)?; + let new_scratch_mem = ExclusiveSharedMemory::new_scratch(new_scratch_size)?; let (hscratch, gscratch) = new_scratch_mem.build(); // Even though this destroys the reference to the host // side of the old scratch mapping, the VM should still @@ -793,6 +812,13 @@ impl SandboxMemoryManager { ) } + /// Where [`update_scratch_bookkeeping`](Self::update_scratch_bookkeeping) + /// copies the snapshot's page tables into scratch. + pub(crate) fn scratch_pt_range(&self) -> std::ops::Range { + let start = self.layout.get_pt_base_scratch_offset(); + start..start + self.layout.get_pt_size().next_multiple_of(page_size::get()) + } + fn update_scratch_bookkeeping(&mut self) -> Result<()> { use hyperlight_common::layout::*; let scratch_size = self.scratch_mem.mem_size(); @@ -856,22 +882,22 @@ impl SandboxMemoryManager { // overlapping with `map_file_cow` regions installed // immediately after the snapshot in the guest PA space. let snapshot_pt_end = self.shared_mem.mem_size(); + #[cfg(unshared_snapshot_mem)] let snapshot_pt_size = self.layout.get_pt_size(); - let snapshot_pt_start = - snapshot_pt_end - snapshot_pt_size.next_multiple_of(page_size::get()); - self.scratch_mem.with_exclusivity(|scratch| { - #[cfg(not(unshared_snapshot_mem))] - let bytes = &self.shared_mem.as_slice()[snapshot_pt_start..snapshot_pt_end]; - #[cfg(unshared_snapshot_mem)] - let bytes = { - let mut bytes = vec![0u8; snapshot_pt_size]; - self.shared_mem - .copy_to_slice(&mut bytes, snapshot_pt_start)?; - bytes - }; - #[allow(clippy::needless_borrow)] - scratch.copy_from_slice(&bytes, self.layout.get_pt_base_scratch_offset()) - })??; + let snapshot_pt_start = snapshot_pt_end - self.scratch_pt_range().len(); + // A logged write, so a reset that zeroes only written pages sees it. + #[cfg(not(unshared_snapshot_mem))] + let bytes = &self.shared_mem.as_slice()[snapshot_pt_start..snapshot_pt_end]; + #[cfg(unshared_snapshot_mem)] + let bytes = { + let mut bytes = vec![0u8; snapshot_pt_size]; + self.shared_mem + .copy_to_slice(&mut bytes, snapshot_pt_start)?; + bytes + }; + #[allow(clippy::needless_borrow)] + self.scratch_mem + .copy_from_slice(&bytes, self.layout.get_pt_base_scratch_offset())?; Ok(()) } diff --git a/src/hyperlight_host/src/mem/mod.rs b/src/hyperlight_host/src/mem/mod.rs index aa26437f0..cdb9f79ef 100644 --- a/src/hyperlight_host/src/mem/mod.rs +++ b/src/hyperlight_host/src/mem/mod.rs @@ -19,9 +19,19 @@ pub mod ptr; pub(super) mod ptr_addr_space; /// Structures to represent an offset into a memory space pub mod ptr_offset; +#[cfg(all(kvm, not(miri)))] +pub(crate) mod scratch_reset; /// A wrapper around unsafe functionality to create and initialize /// a memory region for a guest running in a sandbox. pub mod shared_mem; +/// Without KVM, scratch is reset by zeroing what was written, filling or +/// replacing it (see `HostSharedMemory::zero_written` and +/// `zero_or_replace`): nothing to keep between resets. +#[cfg(not(all(kvm, not(miri))))] +pub(crate) mod scratch_reset { + #[derive(Debug, Default)] + pub(crate) struct ScratchReset {} +} /// Utilities for writing shared memory tests #[cfg(all(test, not(miri)))] // uses proptest which isn't miri-compatible pub(crate) mod shared_mem_tests; diff --git a/src/hyperlight_host/src/mem/scratch_reset.rs b/src/hyperlight_host/src/mem/scratch_reset.rs new file mode 100644 index 000000000..09f1c6082 --- /dev/null +++ b/src/hyperlight_host/src/mem/scratch_reset.rs @@ -0,0 +1,1052 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Resetting a sandbox's scratch region to zeros on restore, on KVM. +//! +//! Dropping a page (`MADV_DONTNEED`) makes the guest's next touch of it +//! fault twice, in the host for a zeroed page and in KVM to map it, which +//! costs a few microseconds a page on a nested hypervisor. Zeroing in +//! place the pages a run wrote keeps them mapped, which costs a fraction +//! of that, so every reset zeroes those and keeps them. +//! +//! `/proc/self/pagemap` reports each page's page-table entry, and every +//! reset reads it for the whole region. A page with memory of its own +//! that holds data is zeroed and kept. One that holds only zeros is kept +//! as it is: it may be one the last run wrote zeros to, which a drop would +//! make the next run fault on again. Every [`GIVE_BACK_EVERY`] resets, +//! pages holding only zeros are dropped instead, up to [`MAX_DROPS`] runs +//! of them, so memory a run once used is given back without a guest that +//! writes zeros faulting on every run. A page in swap, or mapped to +//! memory it shares that holds data (a page shared with a forked child), +//! is dropped. One that shares only zeros (the zero page a read maps) +//! holds no data, and is left like an unwritten page of its own. A page +//! with an empty entry reads as zero and is left alone. The kernel reports +//! any page that holds data as present or swapped, so this leaves no data +//! behind. Pagemap is opened once per process (and once per sandbox in a +//! forked child), and used only after it reports a page the process +//! wrote. +//! +//! From Linux 6.7, the `PAGEMAP_SCAN` ioctl reports pagemap as ranges of +//! pages alike, skipping what was never touched, so a reset costs time +//! in what scratch holds rather than in its size. It does not tell a page +//! shared with a forked child from one of our own: one holding data is +//! zeroed then, which copies it first, as safe as dropping it. Where the +//! ioctl fails (an older kernel, a filter), entries are read one a page. + +use std::fs::File; +use std::ops::Range; +use std::os::unix::fs::FileExt; +use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; + +use tracing::{debug, warn}; + +use super::shared_mem::{ExclusiveSharedMemory, SharedMemory}; + +/// The most separate drops of each kind a reset makes. A guest that +/// scatters shared pages so a reset would drop them one by one has the +/// whole region dropped instead. Runs of pages holding only zeros past +/// it stay, for a later reset that gives them back. +const MAX_DROPS: usize = 64; + +/// How often a reset gives back the pages holding only zeros. A page +/// dropped costs a fault when the guest next touches it, which costs +/// about as much as checking it for data on 16 resets. +const GIVE_BACK_EVERY: usize = 16; + +/// Pagemap entries read at once. +const CHUNK_PAGES: usize = 512; + +/// Ranges one `PAGEMAP_SCAN` call reports. +const SCAN_RANGES: usize = 256; + +/// `PAGEMAP_SCAN` failed in this process, so resets read entries one a +/// page. Whether it works is the kernel's, not a sandbox's. +static NO_RANGES: AtomicBool = AtomicBool::new(false); + +const PM_PRESENT: u64 = 1 << 63; +const PM_SWAPPED: u64 = 1 << 62; +const PM_EXCLUSIVE: u64 = 1 << 56; + +/// The process's pagemap, with the process that opened it, set once and +/// never freed. A forked child inherits its parent's, which shows the +/// parent's pages. Published with an atomic rather than a lock, which a +/// fork could leave held. +static PAGEMAP: AtomicPtr<(u32, Pagemap)> = AtomicPtr::new(std::ptr::null_mut()); + +/// The process's pagemap, once a reset opened it. +fn process_pagemap() -> Option<&'static (u32, Pagemap)> { + // SAFETY: set only from a leaked `Box`, never freed or changed after. + unsafe { PAGEMAP.load(Ordering::Acquire).as_ref() } +} + +/// That pagemap is unusable was logged. +static WARNED: AtomicBool = AtomicBool::new(false); + +/// How a scratch region is reset in place. See the module docs. +#[derive(Debug, Default)] +pub(crate) struct ScratchReset { + /// Resets since pages holding only zeros were last given back. + resets: usize, + /// The pagemap of a forked child, with the child. + forked: Option<(u32, Pagemap)>, + /// The process for which pagemap could not be opened, so it is not + /// tried again. + unusable: Option, + /// Read entries one a page, to test that. + #[cfg(test)] + no_ranges: bool, + /// The last scan read ranges. + #[cfg(test)] + read_ranges: bool, + /// Pagemap fails, to test that. + #[cfg(test)] + fail_pagemap: bool, +} + +/// This process's `/proc/self/pagemap`. +#[derive(Debug)] +struct Pagemap { + file: File, +} + +impl Pagemap { + /// Open pagemap, and check that it reports a page this process wrote + /// as present and its own. Pagemap that cannot see that (an + /// emulated `/proc`, say) is not used. A fork in between shares the + /// page again, so the check is tried a few times. + fn open() -> std::io::Result { + let pagemap = Self { + file: File::open("/proc/self/pagemap")?, + }; + let mut probe = Box::new(0u64); + for _ in 0..3 { + // SAFETY: a valid, aligned `u64` of our own. Volatile, so + // the write that backs the page happens. + unsafe { std::ptr::write_volatile(&mut *probe, 1) }; + let mut entry = [0u64]; + pagemap.read(&*probe as *const u64 as usize, &mut entry)?; + if kept(entry[0]) { + return Ok(pagemap); + } + } + Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "pagemap does not report this process's pages", + )) + } + + /// Read the entries of the pages from `addr` on into `entries`, at + /// most [`CHUNK_PAGES`]. + fn read(&self, addr: usize, entries: &mut [u64]) -> std::io::Result<()> { + // SAFETY: the bytes of `entries`, which any bit pattern fills. + let bytes = unsafe { + std::slice::from_raw_parts_mut(entries.as_mut_ptr().cast::(), entries.len() * 8) + }; + let offset = (addr / page_size::get()) as u64 * 8; + self.file.read_exact_at(bytes, offset) + } +} + +/// [`Pagemap::open`], saying once per process when it fails. +fn open_warned() -> std::io::Result { + Pagemap::open().inspect_err(|e| { + if !WARNED.swap(true, Ordering::Relaxed) { + warn!("scratch is dropped on restore while pagemap is unusable: {e}"); + } + }) +} + +/// Backed by memory of its own: zero it and keep it. +fn kept(entry: u64) -> bool { + entry & (PM_PRESENT | PM_EXCLUSIVE) == PM_PRESENT | PM_EXCLUSIVE +} + +/// Backed: present, or in swap. +#[cfg(test)] +fn held(entry: u64) -> bool { + entry & (PM_PRESENT | PM_SWAPPED) != 0 +} + +/// The region a reset works on, borrowed exclusively. +struct Region<'a> { + mem: &'a mut ExclusiveSharedMemory, + pages: usize, +} + +impl Region<'_> { + fn addr(&self, page: usize) -> usize { + self.mem.base_ptr() as usize + page * page_size::get() + } + + /// Drop `pages`. They read as zero and refill on demand. + fn drop_pages(&mut self, pages: Range) -> std::io::Result<()> { + assert!(pages.start <= pages.end && pages.end <= self.pages); + if pages.is_empty() { + return Ok(()); + } + // SAFETY: the pages lie in the region, a private anonymous + // mapping borrowed exclusively. + let ret = unsafe { + libc::madvise( + self.addr(pages.start) as *mut libc::c_void, + pages.len() * page_size::get(), + libc::MADV_DONTNEED, + ) + }; + if ret == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } + } + + /// Holds only zeros. Stops at the first word that is not, checking + /// the last first, where data often sits (a stack grows down from the + /// end of its page). + fn is_zero(&mut self, page: usize) -> bool { + let size = page_size::get(); + let bytes = &self.mem.as_mut_slice()[page * size..(page + 1) * size]; + // SAFETY: any bytes are valid `u64`s; the aligned middle is + // checked as words, the rest as bytes. + let (head, words, tail) = unsafe { bytes.align_to::() }; + words.last().is_none_or(|&w| w == 0) + && head.iter().all(|&b| b == 0) + && tail.iter().all(|&b| b == 0) + && words.iter().all(|&w| w == 0) + } + + fn zero_pages(&mut self, pages: Range) { + let page = page_size::get(); + self.mem.as_mut_slice()[pages.start * page..pages.end * page].fill(0); + } +} + +impl ScratchReset { + /// Reset `mem` to zeros, keeping mapped the pages that are its own, + /// as the module docs say. On an error the region may be partly + /// reset, and the caller must reset it some other way. + pub(crate) fn reset(&mut self, mem: &mut ExclusiveSharedMemory) -> std::io::Result<()> { + let pages = mem.mem_size() / page_size::get(); + let mut region = Region { mem, pages }; + self.resets += 1; + let give_back = self.resets == GIVE_BACK_EVERY; + if give_back { + self.resets = 0; + } + if self.scan(&mut region, give_back)? > MAX_DROPS { + region.drop_pages(0..pages)?; + } + Ok(()) + } + + /// This process's pagemap. Where it cannot be opened, this sandbox + /// does not try again. + fn pagemap(&mut self) -> std::io::Result<&Pagemap> { + let pid = std::process::id(); + if self.unusable == Some(pid) { + return Err(std::io::ErrorKind::Unsupported.into()); + } + if let Err(e) = self.open_pagemap(pid) { + self.unusable = Some(pid); + return Err(e); + } + match process_pagemap() { + Some((opener, pagemap)) if *opener == pid => Ok(pagemap), + _ => self + .forked + .as_ref() + .map(|(_, pagemap)| pagemap) + .ok_or_else(|| std::io::ErrorKind::Unsupported.into()), + } + } + + /// Open the pagemap of process `pid` if it is not open: the process's + /// own, or in a forked child, the child's. + fn open_pagemap(&mut self, pid: u32) -> std::io::Result<()> { + #[cfg(test)] + if self.fail_pagemap { + return Err(std::io::ErrorKind::Unsupported.into()); + } + if process_pagemap().is_none() { + let opened = Box::into_raw(Box::new((pid, open_warned()?))); + // Another thread may set it first; either is this process's. + if PAGEMAP + .compare_exchange( + std::ptr::null_mut(), + opened, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_err() + { + // SAFETY: from `Box::into_raw` above, and not published. + drop(unsafe { Box::from_raw(opened) }); + } + } + let own = matches!(process_pagemap(), Some((opener, _)) if *opener == pid); + if !own && !matches!(&self.forked, Some((child, _)) if *child == pid) { + self.forked = Some((pid, open_warned()?)); + } + Ok(()) + } + + /// Zero the pages with memory of their own that hold data, and drop + /// the rest that hold data, one drop per run between kept pages (see + /// [`Run::end`]), and with `give_back` those holding only zeros too. + /// Returns the drops of runs that had to go, and stops past + /// [`MAX_DROPS`] of them for the caller to drop it all. + fn scan(&mut self, region: &mut Region<'_>, give_back: bool) -> std::io::Result { + #[cfg(test)] + let ranges = !self.no_ranges; + #[cfg(not(test))] + let ranges = true; + if ranges + && !NO_RANGES.load(Ordering::Relaxed) + && let Some(drops) = self.scan_ranges(region, give_back)? + { + #[cfg(test)] + { + self.read_ranges = true; + } + return Ok(drops); + } + self.scan_entries(region, give_back) + } + + /// [`scan`](Self::scan) from pagemap entries, one a page. + fn scan_entries(&mut self, region: &mut Region<'_>, give_back: bool) -> std::io::Result { + let mut walk = Walk::new(give_back); + let mut entries = [0u64; CHUNK_PAGES]; + let pagemap = self.pagemap()?; + let mut chunk = 0; + while chunk < region.pages { + let n = CHUNK_PAGES.min(region.pages - chunk); + pagemap.read(region.addr(chunk), &mut entries[..n])?; + for (i, &entry) in entries[..n].iter().enumerate() { + let page = chunk + i; + let own = kept(entry); + let present = entry & PM_PRESENT != 0; + let zero = present && region.is_zero(page); + // A shared page holding only zeros (the zero page a read + // maps) holds no data; one holding data, or a page in + // swap, must go. + let shared = !own && (entry & PM_SWAPPED != 0 || present && !zero); + if walk.page(region, page, own && !zero, shared, present && zero)? { + return Ok(walk.drops.shared); + } + } + chunk += n; + } + walk.finish(region) + } + + /// [`scan`](Self::scan) from `PAGEMAP_SCAN`'s ranges of present or + /// swapped pages; what lies between them is empty. `None` when the + /// first call fails, before anything is reset, for the caller to read + /// entries instead. A later call failing is an error, as the region is + /// partly reset. Either way, later scans read entries. + fn scan_ranges( + &mut self, + region: &mut Region<'_>, + give_back: bool, + ) -> std::io::Result> { + let size = page_size::get(); + let base = region.addr(0) as u64; + let pagemap = self.pagemap()?; + let mut ranges = [PageRegion::default(); SCAN_RANGES]; + let mut arg = PmScanArg::new(base, base + (region.pages * size) as u64); + let mut walk = Walk::new(give_back); + let mut next = 0; + loop { + let n = match pagemap_scan(&pagemap.file, &mut arg, &mut ranges) { + Ok(n) => n, + Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue, + Err(e) => { + // Before Linux 6.7 there is no such ioctl; anything else + // is unexpected. + if e.raw_os_error() == Some(libc::ENOTTY) { + debug!("no PAGEMAP_SCAN, reading pagemap one entry a page"); + } else { + warn!("PAGEMAP_SCAN failed, reading pagemap one entry a page: {e}"); + } + NO_RANGES.store(true, Ordering::Relaxed); + // Past the first call, part of the region is reset: the + // caller drops it all. + return if arg.start == base { Ok(None) } else { Err(e) }; + } + }; + for range in ranges.iter().take(n) { + let first = (range.start - base) as usize / size; + let end = (range.end - base) as usize / size; + if first > next && walk.page(region, next, false, false, false)? { + return Ok(Some(walk.drops.shared)); + } + let swapped = range.categories & PAGE_IS_SWAPPED != 0; + let zero_page = range.categories & PAGE_IS_PFNZERO != 0; + for page in first..end { + let zero = !swapped && (zero_page || region.is_zero(page)); + if walk.page(region, page, !swapped && !zero, swapped, zero)? { + return Ok(Some(walk.drops.shared)); + } + } + next = end; + } + if arg.walk_end >= arg.end { + break; + } + arg.start = arg.walk_end; + } + if next < region.pages && walk.page(region, next, false, false, false)? { + return Ok(Some(walk.drops.shared)); + } + walk.finish(region).map(Some) + } +} + +/// Run `PAGEMAP_SCAN` on `pagemap` into `ranges`. Returns the ranges it +/// filled. +fn pagemap_scan( + pagemap: &File, + arg: &mut PmScanArg, + ranges: &mut [PageRegion], +) -> std::io::Result { + use std::os::fd::AsRawFd; + arg.vec = ranges.as_mut_ptr() as u64; + arg.vec_len = ranges.len() as u64; + // SAFETY: `arg` points at `ranges`, borrowed for the call, which the + // kernel fills with at most `vec_len` entries. The kernel writes no + // other memory of ours but `arg`. + let n = unsafe { libc::ioctl(pagemap.as_raw_fd(), PAGEMAP_SCAN, &mut *arg) }; + usize::try_from(n).map_err(|_| std::io::Error::last_os_error()) +} + +/// `PAGEMAP_SCAN`: `_IOWR('f', 16, struct pm_scan_arg)`. The request is +/// a `c_ulong` on glibc and a `c_int` on musl. +const PAGEMAP_SCAN: libc::Ioctl = 0xC060_6610u32 as libc::Ioctl; +const PAGE_IS_PRESENT: u64 = 1 << 3; +const PAGE_IS_SWAPPED: u64 = 1 << 4; +const PAGE_IS_PFNZERO: u64 = 1 << 5; + +/// `struct pm_scan_arg`. +#[repr(C)] +#[derive(Default)] +struct PmScanArg { + size: u64, + flags: u64, + start: u64, + end: u64, + walk_end: u64, + vec: u64, + vec_len: u64, + max_pages: u64, + category_inverted: u64, + category_mask: u64, + category_anyof_mask: u64, + return_mask: u64, +} + +impl PmScanArg { + /// Report the present or swapped pages in `[start, end)`, flagging + /// those mapped to the zero page. [`pagemap_scan`] says where. + fn new(start: u64, end: u64) -> Self { + Self { + size: size_of::() as u64, + start, + end, + category_anyof_mask: PAGE_IS_PRESENT | PAGE_IS_SWAPPED, + return_mask: PAGE_IS_PRESENT | PAGE_IS_SWAPPED | PAGE_IS_PFNZERO, + ..Self::default() + } + } +} + +/// `struct page_region`. +#[repr(C)] +#[derive(Clone, Copy, Default)] +struct PageRegion { + start: u64, + end: u64, + categories: u64, +} + +/// The runs a scan makes from pages fed in order. +struct Walk { + run: Run, + drops: Drops, +} + +impl Walk { + /// With `give_back`, runs holding only zeros are dropped. + fn new(give_back: bool) -> Self { + Self { + run: Run { + start: 0, + keep: false, + shared: false, + zero: false, + }, + drops: Drops { + give_back, + ..Drops::default() + }, + } + } + + /// Page `page` is kept, or not and holds shared or swapped data, or + /// zeros. Every page up to the next one fed is as this one. True once + /// shared drops pass [`MAX_DROPS`]. + fn page( + &mut self, + region: &mut Region<'_>, + page: usize, + keep: bool, + shared: bool, + zero: bool, + ) -> std::io::Result { + if keep != self.run.keep { + if self.run.end(region, page, &mut self.drops)? { + return Ok(true); + } + self.run = Run { + start: page, + keep, + shared: false, + zero: false, + }; + } + self.run.shared |= shared; + self.run.zero |= zero; + Ok(false) + } + + /// End the last run at the end of the region. + fn finish(mut self, region: &mut Region<'_>) -> std::io::Result { + self.run.end(region, region.pages, &mut self.drops)?; + Ok(self.drops.shared) + } +} + +/// Drops made by a reset. +#[derive(Default)] +struct Drops { + /// Of runs holding shared or swapped pages, which must go. + shared: usize, + /// Of runs holding only zeros, dropped to give their memory back. + zero: usize, + /// Runs holding only zeros are dropped. + give_back: bool, +} + +/// A run of pages kept, or not, between pages that are the other. +struct Run { + start: usize, + keep: bool, + shared: bool, + zero: bool, +} + +impl Run { + /// Zero the run if kept. Otherwise drop it, empty entries and all, + /// so the guest cannot make a reset drop page by page: always when it + /// holds shared or swapped pages, and when it holds pages holding only + /// zeros on a reset that gives them back, up to [`MAX_DROPS`] of those; + /// past that they stay, zero. + /// True once shared drops pass [`MAX_DROPS`]. + fn end(&self, region: &mut Region<'_>, end: usize, drops: &mut Drops) -> std::io::Result { + let pages = self.start..end; + if self.keep { + region.zero_pages(pages); + } else if self.shared { + region.drop_pages(pages)?; + drops.shared += 1; + } else if self.zero && drops.give_back && drops.zero < MAX_DROPS { + region.drop_pages(pages)?; + drops.zero += 1; + } + Ok(drops.shared > MAX_DROPS) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn page() -> usize { + page_size::get() + } + + /// Backed 4 KiB at a time, as KVM scratch is, whatever the host's + /// hypervisor. + fn region(pages: usize) -> ExclusiveSharedMemory { + let mem = ExclusiveSharedMemory::new(pages * page()).unwrap(); + // SAFETY: the test's own mapping. + unsafe { + libc::madvise( + mem.base_ptr() as *mut libc::c_void, + mem.mem_size(), + libc::MADV_NOHUGEPAGE, + ) + }; + mem + } + + fn write(mem: &mut ExclusiveSharedMemory, page_index: usize, byte: u8) { + mem.as_mut_slice()[page_index * page() + 7] = byte; + } + + fn entry(mem: &ExclusiveSharedMemory, page_index: usize) -> u64 { + let pagemap = Pagemap { + file: File::open("/proc/self/pagemap").unwrap(), + }; + let mut e = [0u64]; + pagemap + .read(mem.base_ptr() as usize + page_index * page(), &mut e) + .unwrap(); + e[0] + } + + /// Reading maps the zero page on dropped pages, which a reset that + /// gives back drops again. + fn all_zero(mem: &mut ExclusiveSharedMemory) -> bool { + mem.as_mut_slice().iter().all(|b| *b == 0) + } + + fn reset(state: &mut ScratchReset, mem: &mut ExclusiveSharedMemory) { + state.reset(mem).unwrap(); + } + + /// A reset that gives back the pages holding only zeros. + fn give_back(state: &mut ScratchReset, mem: &mut ExclusiveSharedMemory) { + state.resets = GIVE_BACK_EVERY - 1; + reset(state, mem); + assert_eq!(state.resets, 0); + } + + /// Every reset leaves the region zero, whatever the run wrote: the + /// usual pages, others far away, and nothing at all. + #[test] + fn every_reset_leaves_the_region_zero() { + for mut state in states() { + let mut mem = region(2048); + for round in 0..100usize { + match round % 7 { + 3 => {} + 5 => { + write(&mut mem, 1500 + round % 300, 0xa5); + write(&mut mem, 40 + round % 8, 0x5a); + } + _ => { + for p in 32..96 { + write(&mut mem, p, round as u8 | 1); + } + } + } + reset(&mut state, &mut mem); + assert!(all_zero(&mut mem), "round {round}"); + } + } + } + + /// A reset reading `PAGEMAP_SCAN` ranges, and one reading entries a + /// page. + fn states() -> [ScratchReset; 2] { + [ + ScratchReset::default(), + ScratchReset { + no_ranges: true, + ..ScratchReset::default() + }, + ] + } + + /// Both ways of reading pagemap leave every page alike. + #[test] + fn ranges_and_entries_agree() { + let pages = 4096; + let run = |state: &mut ScratchReset| { + let mut mem = region(pages); + for p in 0..1000 { + write(&mut mem, p, 1); + } + reset(state, &mut mem); + for p in (0..100).step_by(3) { + write(&mut mem, p, 1); + } + for p in 2000..2100 { + assert_eq!(mem.as_mut_slice()[p * page()], 0); + } + write(&mut mem, pages - 1, 1); + reset(state, &mut mem); + assert!(all_zero(&mut mem)); + (0..pages) + .map(|p| { + let e = entry(&mem, p); + (kept(e), held(e)) + }) + .collect::>() + }; + let [mut ranges, mut entries] = states(); + assert_eq!(run(&mut ranges), run(&mut entries)); + } + + /// From Linux 6.7, resets read ranges, not entries a page: the ioctl + /// is only missing (ENOTTY) before that, and any other failure is a + /// wrong call. + #[test] + fn ranges_are_read_where_supported() { + let mut mem = region(64); + write(&mut mem, 3, 1); + let file = File::open("/proc/self/pagemap").unwrap(); + let mut ranges = [PageRegion::default(); 4]; + let base = mem.base_ptr() as u64; + let mut arg = PmScanArg::new(base, base + 64 * page() as u64); + match pagemap_scan(&file, &mut arg, &mut ranges) { + Err(e) if e.raw_os_error() == Some(libc::ENOTTY) => { + eprintln!("no PAGEMAP_SCAN: ranges_are_read_where_supported skipped"); + return; + } + result => assert_eq!(result.unwrap(), 1), + } + assert_eq!(ranges[0].start, base + 3 * page() as u64); + assert_eq!(ranges[0].categories & PAGE_IS_PRESENT, PAGE_IS_PRESENT); + let [mut ranges, mut entries] = states(); + reset(&mut ranges, &mut mem); + assert!(ranges.read_ranges && !NO_RANGES.load(Ordering::Relaxed)); + reset(&mut entries, &mut mem); + assert!(!entries.read_ranges); + } + + /// A scan goes on where a call that filled its ranges stopped. + #[test] + fn a_scan_continues_past_full_ranges() { + let pages = SCAN_RANGES * 8; + let run = |state: &mut ScratchReset| { + let mut mem = region(pages); + // Every other page, so each is a range of its own. + for p in (0..pages).step_by(2) { + write(&mut mem, p, 1); + } + reset(state, &mut mem); + assert!(all_zero(&mut mem)); + (0..pages).map(|p| kept(entry(&mem, p))).collect::>() + }; + let [mut ranges, mut entries] = states(); + let kept = run(&mut ranges); + assert!(kept.iter().step_by(2).all(|&k| k)); + assert_eq!(kept, run(&mut entries)); + } + + #[test] + fn pagemap_bits() { + assert!(!kept(0) && !held(0)); + assert!(kept(PM_PRESENT | PM_EXCLUSIVE)); + assert!(!kept(PM_PRESENT) && held(PM_PRESENT)); + assert!(!kept(PM_SWAPPED) && held(PM_SWAPPED)); + } + + /// The pages a run wrote stay backed from the first reset on, however + /// many there are. + #[test] + fn written_pages_stay_backed() { + for mut state in states() { + let mut mem = region(8192); + for p in 0..8192 { + write(&mut mem, p, 1); + } + reset(&mut state, &mut mem); + assert!((0..8192).all(|p| kept(entry(&mem, p)))); + assert!(all_zero(&mut mem)); + } + } + + /// Pages holding only zeros stay backed until a reset gives them + /// back: a guest that writes zeros to a page every run would + /// otherwise fault on it every run. What stays resident then follows + /// the last run, not the most any run wrote. + #[test] + fn zero_pages_stay_until_given_back() { + for mut state in states() { + let mut mem = region(2048); + for p in 0..512 { + write(&mut mem, p, 1); + } + reset(&mut state, &mut mem); + assert!(kept(entry(&mem, 300))); + write(&mut mem, 0, 1); + write(&mut mem, 301, 0); + for _ in 2..GIVE_BACK_EVERY { + reset(&mut state, &mut mem); + assert!((0..512).all(|p| kept(entry(&mem, p)))); + } + assert_eq!(state.resets, GIVE_BACK_EVERY - 1); + write(&mut mem, 0, 1); + reset(&mut state, &mut mem); + assert!(kept(entry(&mem, 0))); + assert!((1..512).all(|p| !held(entry(&mem, p)))); + assert!(all_zero(&mut mem)); + } + } + + /// Sparse writes across an earlier, larger footprint leave many runs + /// of zero pages. Those are dropped up to the limit and otherwise + /// left, but never make the reset drop everything. + #[test] + fn sparse_writes_keep_their_pages() { + for mut state in states() { + let mut mem = region(4096); + for p in 0..4096 { + write(&mut mem, p, 1); + } + reset(&mut state, &mut mem); + for p in (0..4096).step_by(16) { + write(&mut mem, p, 1); + } + give_back(&mut state, &mut mem); + assert!((0..4096).step_by(16).all(|p| kept(entry(&mem, p)))); + assert!(all_zero(&mut mem)); + } + } + + /// The zero page a read maps holds no data: it is left like an + /// unwritten page, and given back like one, never as a run that must + /// go. + #[test] + fn zero_page_reads_are_left_as_unwritten() { + for mut state in states() { + let mut mem = region(512); + write(&mut mem, 10, 1); + write(&mut mem, 200, 1); + reset(&mut state, &mut mem); + for p in (11..200).step_by(2) { + assert_eq!(mem.as_mut_slice()[p * page()], 0); + } + for give_back in [false, true] { + write(&mut mem, 10, 1); + write(&mut mem, 200, 1); + assert_eq!( + state + .scan( + &mut Region { + pages: 512, + mem: &mut mem + }, + give_back + ) + .unwrap(), + 0 + ); + assert_eq!(held(entry(&mem, 11)), !give_back); + } + assert!(kept(entry(&mem, 10)) && kept(entry(&mem, 200))); + assert!(all_zero(&mut mem)); + } + } + + /// A run not kept is one drop however many pagemap chunks it spans: + /// more separate drops than [`MAX_DROPS`] would leave some of it. + #[test] + fn a_run_across_chunks_is_one_drop() { + for mut state in states() { + let chunks = MAX_DROPS + 4; + let pages = CHUNK_PAGES * chunks; + let mut mem = region(pages); + for c in 0..chunks { + write(&mut mem, c * CHUNK_PAGES + 7, 1); + } + write(&mut mem, 0, 1); + write(&mut mem, pages - 1, 1); + reset(&mut state, &mut mem); + // Only the ends written again: what lies between is left zero. + write(&mut mem, 0, 1); + write(&mut mem, pages - 1, 1); + give_back(&mut state, &mut mem); + assert!((0..chunks).all(|c| !held(entry(&mem, c * CHUNK_PAGES + 7)))); + assert!(kept(entry(&mem, 0)) && kept(entry(&mem, pages - 1))); + } + } + + /// Runs that must go always drop, and past [`MAX_DROPS`] of them the + /// reset drops everything. Runs holding only zeros drop only on a + /// reset that gives them back, up to their own budget, then stay. + #[test] + fn drop_limits() { + let mut mem = region(16); + let mut region = Region { + pages: 16, + mem: &mut mem, + }; + let run = |shared, zero| Run { + start: 0, + keep: false, + shared, + zero, + }; + let mut drops = Drops::default(); + for _ in 0..MAX_DROPS { + assert!(!run(true, false).end(&mut region, 1, &mut drops).unwrap()); + } + assert!(run(true, false).end(&mut region, 1, &mut drops).unwrap()); + for give_back in [false, true] { + let mut drops = Drops { + give_back, + ..Drops::default() + }; + for _ in 0..MAX_DROPS + 10 { + assert!(!run(false, true).end(&mut region, 1, &mut drops).unwrap()); + } + let zero = if give_back { MAX_DROPS } else { 0 }; + assert_eq!((drops.shared, drops.zero), (0, zero)); + } + } + + /// A page in swap holds the guest's data. The reset drops it. + #[test] + fn a_swapped_page_is_dropped() { + for mut state in states() { + let mut mem = region(512); + write(&mut mem, 10, 0x77); + // SAFETY: the test's own mapping. + unsafe { + libc::madvise( + mem.base_ptr().add(10 * page()) as *mut libc::c_void, + page(), + libc::MADV_PAGEOUT, + ) + }; + if entry(&mem, 10) & PM_SWAPPED == 0 { + eprintln!("no swap: a_swapped_page_is_dropped skipped"); + return; + } + reset(&mut state, &mut mem); + assert!(all_zero(&mut mem)); + } + } + + /// A forked child resets its own pages, not the ones its parent's + /// pagemap shows. After the fork the parent writes pages 10 and 20 + /// again, so its pagemap shows them as its own and page 15, which the + /// child wrote, as empty: a reset reading it would leave page 15. + /// Page 40 holds data shared with the child at the fork, which the + /// child's reset must not leave; the child keeps page 39 so page 40 + /// is not in a run dropped anyway. + /// Ignored: it forks, so `forked_child_shim` runs it alone in a + /// process of its own, away from the pages and locks of tests running + /// in parallel. + #[test] + #[ignore] + fn a_forked_child_resets_its_own_pages() { + for state in states() { + forked_child_resets(state); + } + } + + fn forked_child_resets(mut state: ScratchReset) { + let mut mem = region(512); + write(&mut mem, 10, 1); + write(&mut mem, 20, 1); + reset(&mut state, &mut mem); + assert_eq!( + process_pagemap().map(|(pid, _)| *pid), + Some(std::process::id()) + ); + assert!(state.forked.is_none()); + write(&mut mem, 40, 0x33); + let mut fds = [0; 2]; + // SAFETY: two valid fds for the pipe. + assert_eq!(unsafe { libc::pipe(fds.as_mut_ptr()) }, 0); + // SAFETY: the process runs this test alone (`forked_child_shim`). + // The child waits for the parent, writes its copy of `mem`, resets + // it (opening pagemap, which allocates), and exits without + // unwinding. + let pid = unsafe { libc::fork() }; + assert!(pid >= 0); + if pid == 0 { + let mut byte = 0u8; + // SAFETY: reads one byte into `byte` from the pipe. + unsafe { libc::read(fds[0], (&mut byte as *mut u8).cast(), 1) }; + let ok = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + write(&mut mem, 15, 0x5a); + write(&mut mem, 39, 0x5a); + state.reset(&mut mem).is_ok() + && all_zero(&mut mem) + && state.forked.as_ref().map(|(pid, _)| *pid) == Some(std::process::id()) + })); + // SAFETY: ends the child. + unsafe { libc::_exit(if matches!(ok, Ok(true)) { 0 } else { 1 }) }; + } + write(&mut mem, 10, 2); + write(&mut mem, 20, 2); + assert!(kept(entry(&mem, 10)) && !held(entry(&mem, 15))); + // SAFETY: writes one byte to the pipe, then waits for the child. + let mut status = 0; + unsafe { + libc::write(fds[1], [1u8].as_ptr().cast(), 1); + libc::waitpid(pid, &mut status, 0); + } + assert!(libc::WIFEXITED(status) && libc::WEXITSTATUS(status) == 0); + } + + #[test] + fn forked_child_shim() { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--ignored", + "--exact", + "--test-threads=1", + "mem::scratch_reset::tests::a_forked_child_resets_its_own_pages", + ]) + .stdin(std::process::Stdio::null()) + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed"), + "{stdout}{}", + String::from_utf8_lossy(&output.stderr) + ); + } + + /// Without pagemap, a reset fails, and the caller drops the region + /// instead. Pagemap is not tried again for the sandbox. + #[test] + fn without_pagemap_a_reset_fails() { + let mut mem = region(512); + let mut state = ScratchReset { + fail_pagemap: true, + ..ScratchReset::default() + }; + write(&mut mem, 10, 1); + assert!(state.reset(&mut mem).is_err()); + assert_eq!(state.unusable, Some(std::process::id())); + state.fail_pagemap = false; + assert!(state.reset(&mut mem).is_err()); + } + + /// Pages that must go (here, in swap) scattered so that a reset would + /// drop them one by one are dropped all at once. + #[test] + fn scattered_swapped_pages_are_dropped_wholesale() { + for mut state in states() { + let mut mem = region(2048); + for p in (0..400).step_by(2) { + write(&mut mem, p, 1); + write(&mut mem, p + 1, 2); + // SAFETY: the test's own mapping. + unsafe { + libc::madvise( + mem.base_ptr().add((p + 1) * page()) as *mut libc::c_void, + page(), + libc::MADV_PAGEOUT, + ) + }; + } + if entry(&mem, 1) & PM_SWAPPED == 0 { + eprintln!("no swap: scattered_swapped_pages_are_dropped_wholesale skipped"); + return; + } + reset(&mut state, &mut mem); + assert!(!kept(entry(&mem, 0))); + assert!(all_zero(&mut mem)); + } + } +} diff --git a/src/hyperlight_host/src/mem/shared_mem.rs b/src/hyperlight_host/src/mem/shared_mem.rs index b9e803ca8..c895642ae 100644 --- a/src/hyperlight_host/src/mem/shared_mem.rs +++ b/src/hyperlight_host/src/mem/shared_mem.rs @@ -4,9 +4,10 @@ use std::ffi::c_void; use std::io::Error; use std::mem::{align_of, size_of}; +use std::ops::Range; #[cfg(unix)] use std::ptr::null_mut; -use std::sync::atomic::Ordering; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, RwLock}; use bytemuck::Pod; @@ -30,6 +31,7 @@ use windows::core::PCSTR; use super::memory_region::{ HostGuestMemoryRegion, MemoryRegion, MemoryRegionFlags, MemoryRegionKind, MemoryRegionType, }; +use super::scratch_reset::ScratchReset; use crate::log_then_return; type Result = core::result::Result; @@ -938,6 +940,7 @@ impl ExclusiveSharedMemory { let hshm = HostSharedMemory { region: self.region.clone(), lock: lock.clone(), + host_writes: Arc::new(HostWrites::new(self.mem_size())), }; ( hshm, @@ -955,6 +958,33 @@ impl ExclusiveSharedMemory { } } +impl ExclusiveSharedMemory { + /// New memory for a sandbox's scratch region. On KVM it is backed 4 + /// KiB at a time: KVM maps scratch 4 KiB at a time, since the guard + /// page before it puts its host address 4 KiB off its guest address + /// within 2 MiB, so a huge page would only make a touch, and a reset + /// that zeroes the touched pages (see `ScratchReset`), zero 2 MiB. + pub(crate) fn new_scratch(size: usize) -> Result { + let mem = Self::new(size)?; + #[cfg(all(kvm, not(miri)))] + if matches!( + crate::hypervisor::virtual_machine::get_available_hypervisor(), + Some(crate::hypervisor::virtual_machine::HypervisorType::Kvm) + ) { + // SAFETY: the region is a private anonymous mapping of our + // own. Failure leaves the default page size, so it is ignored. + unsafe { + libc::madvise( + mem.base_ptr() as *mut libc::c_void, + mem.mem_size(), + libc::MADV_NOHUGEPAGE, + ) + }; + } + Ok(mem) + } +} + impl SharedMemory for ExclusiveSharedMemory { fn region(&self) -> &HostMapping { &self.region @@ -1202,6 +1232,8 @@ impl SharedMemory for GuestSharedMemory { pub struct HostSharedMemory { region: Arc, lock: Arc>, + /// The pages written through any clone of this. + host_writes: Arc, } unsafe impl Send for HostSharedMemory {} @@ -1275,6 +1307,7 @@ impl HostSharedMemory { // AtomicAccess is sealed to integer atomics, whose bit patterns are valid. let atomic = unsafe { &*ptr.cast::() }; atomic.store(value, ordering); + self.host_writes.note(offset, size_of::()); Ok(()) } @@ -1370,6 +1403,7 @@ impl HostSharedMemory { i += 1; } + self.host_writes.note(offset, len); drop(guard); Ok(()) } @@ -1415,68 +1449,393 @@ impl HostSharedMemory { i += 1; } + self.host_writes.note(offset, len); drop(guard); Ok(()) } } impl HostSharedMemory { - /// Reset this memory region to all-zeros, choosing the fastest - /// strategy for the current platform and hypervisor configuration. + /// Reset this memory region to all-zeros given the pages written + /// since the last reset: `written` (the guest's writes, one bit per + /// 4 KiB page, bits past the end ignored) and the pages host writes + /// through this memory touched. Pages no one wrote are left alone, + /// so the cost follows the pages written, not the size. /// - /// On Linux/KVM (without mshv3), uses `MADV_DONTNEED` for lazy - /// zeroing. On Linux/mshv3, falls through to `fill(0)`. + /// `written` comes from the hypervisor's dirty log, which sees only + /// the guest's writes. Every host write is noted, page by page or, + /// with exclusive access, all of the memory at once. + pub(crate) fn zero_written(&mut self, written: &mut Vec) -> Result<()> { + let pages = self.mem_size() / DIRTY_PAGE_SIZE; + written.resize(written.len().max(pages.div_ceil(64)), 0); + let host_writes = self.host_writes.clone(); + self.with_exclusivity_unnoted(|e| { + // Under the lock, so no host write is noted after this. + host_writes.take_into(written); + let mem = e.as_mut_slice(); + for run in DirtyRuns::new(written, pages) { + mem[run.start * DIRTY_PAGE_SIZE..run.end * DIRTY_PAGE_SIZE].fill(0); + } + // All of memory, so only in tests: on WHP it also makes pages + // never touched resident. + #[cfg(test)] + assert!( + mem.chunks(DIRTY_PAGE_SIZE) + .all(|page| page == &ZERO_PAGE[..page.len()]), + "a scratch page written since the last reset was not logged" + ); + }) + } + + /// Reset this memory region to all-zeros, choosing the fastest + /// strategy for the current platform and hypervisor, without + /// knowing which pages were written (see + /// [`zero_written`](Self::zero_written) for when they are known). /// - /// On Windows, zeroing via `fill(0)` is prohibitively expensive - /// for large regions (e.g. 448 MiB scratch). Instead, the - /// mapping is replaced with a fresh demand-zero allocation. - /// Returns `Some(GuestSharedMemory)` when the mapping was - /// replaced (the caller must update the VM mapping), or `None` - /// when zeroed in place. + /// On KVM, `state` resets the region in place. It zeroes its own pages + /// that hold data, keeping them mapped, drops pages holding data that + /// are not its own, and gives back now and then those holding only + /// zeros (see [`ScratchReset`]). If that + /// fails, the whole region is dropped with `MADV_DONTNEED` for lazy + /// zeroing. On MSHV, whose + /// mappings must stay in sync with userspace, it is zeroed with + /// `fill(0)`. /// - // TODO: Find the break-even point between zero-in-place and - // replace for each hypervisor and use a size-based heuristic - // instead of a compile-time platform check. - pub(crate) fn zero_or_replace(&mut self) -> Result> { + /// On Windows, scratch up to [`FILL_IN_PLACE_MAX`] is zeroed with + /// `fill(0)`. Larger scratch is replaced with a fresh demand-zero + /// allocation, since `fill(0)` takes long and makes all of it + /// resident (e.g. 448 MiB scratch). Returns `Some(GuestSharedMemory)` + /// when the mapping was replaced (the caller must update the VM + /// mapping), or `None` when zeroed in place. + pub(crate) fn zero_or_replace( + &mut self, + #[cfg_attr(not(all(kvm, not(miri))), allow(unused_variables))] state: &mut ScratchReset, + ) -> Result> { + // Zeroing in place makes the whole region resident; above the cap, + // Windows maps fresh memory instead (#1765). #[cfg(target_os = "windows")] - { - let new_mem = ExclusiveSharedMemory::new(self.mem_size())?; + if self.mem_size() > FILL_IN_PLACE_MAX { + let new_mem = ExclusiveSharedMemory::new_scratch(self.mem_size())?; let (hscratch, gscratch) = new_mem.build(); *self = hscratch; - Ok(Some(gscratch)) - } - #[cfg(not(target_os = "windows"))] - { - self.with_exclusivity(|e| { - #[allow(unused_mut)] - let mut do_copy = true; - // TODO: Find a similar lazy zeroing approach that works on MSHV. - // (See Note [Keeping mappings in sync between userspace and the guest]) - #[cfg(all(feature = "kvm", not(any(feature = "mshv3"))))] - unsafe { - let ret = libc::madvise( - e.region.ptr() as *mut libc::c_void, - e.region.size(), - libc::MADV_DONTNEED, - ); - if ret == 0 { - do_copy = false; - } + return Ok(Some(gscratch)); + } + let host_writes = self.host_writes.clone(); + self.with_exclusivity_unnoted(|e| { + // Nothing written before this reset is left to zero. Under the + // lock, so no host write is noted before it ends. + host_writes.clear(); + // TODO: Find a similar lazy zeroing approach that works on MSHV. + // (See Note [Keeping mappings in sync between userspace and the guest]) + #[cfg(all(kvm, not(miri)))] + if matches!( + crate::hypervisor::virtual_machine::get_available_hypervisor(), + Some(crate::hypervisor::virtual_machine::HypervisorType::Kvm) + ) { + if state.reset(e).is_ok() { + return; } - if do_copy { - e.as_mut_slice().fill(0); + // SAFETY: the region is a private anonymous mapping, + // held exclusively. + let ret = unsafe { + libc::madvise( + e.base_ptr() as *mut libc::c_void, + e.mem_size(), + libc::MADV_DONTNEED, + ) + }; + if ret == 0 { + return; } - })?; - Ok(None) + } + e.as_mut_slice().fill(0); + })?; + Ok(None) + } +} + +#[cfg(test)] +impl HostSharedMemory { + /// The pages host writes touched since the last reset, one bit per + /// page, taken. + pub(crate) fn take_host_writes(&self) -> Vec { + self.host_writes.take() + } +} + +/// On Windows without a dirty log, the largest scratch zeroed in place. +/// Zeroing in place makes all of scratch resident, where a fresh mapping +/// holds only what the next run touches (#1765), so larger scratch is +/// replaced. Below it, zeroing in place is about 2x faster than a fresh +/// mapping, and over 5x when the guest writes a megabyte or more, since +/// a fresh mapping faults on every page touched. +#[cfg(target_os = "windows")] +const FILL_IN_PLACE_MAX: usize = 16 << 20; + +/// The page size of dirty-page bitmaps: hypervisor dirty logs and +/// [`HostWrites`]. +pub(crate) const DIRTY_PAGE_SIZE: usize = 4096; + +#[cfg(test)] +static ZERO_PAGE: [u8; DIRTY_PAGE_SIZE] = [0; DIRTY_PAGE_SIZE]; + +/// The pages host writes through a [`HostSharedMemory`] touched since +/// the last reset, one bit per page. A hypervisor's dirty log sees only +/// the guest's writes, and the host writes wherever the guest points +/// it (virtqueue buffers), so a reset that zeroes only written pages +/// needs these too. +#[derive(Debug)] +struct HostWrites { + size: usize, + /// Allocated on the first write, since most memory other than + /// scratch is never written by the host. + words: std::sync::OnceLock>, +} + +impl HostWrites { + fn new(size: usize) -> Self { + Self { + size, + words: std::sync::OnceLock::new(), + } + } + + fn words(&self) -> &[AtomicU64] { + self.words.get().map_or(&[], |words| words) + } + + /// Note a write of `len` bytes at `offset`, already bounds checked. + /// Called with the lock held for reading, so a reset, which holds + /// it for writing, sees it. + fn note(&self, offset: usize, len: usize) { + if len == 0 { + return; + } + let words = self.words.get_or_init(|| { + let words = self.size.div_ceil(DIRTY_PAGE_SIZE * 64); + (0..words).map(|_| AtomicU64::new(0)).collect() + }); + for page in offset / DIRTY_PAGE_SIZE..=(offset + len - 1) / DIRTY_PAGE_SIZE { + let bit = 1 << (page % 64); + let word = &words[page / 64]; + // Most writes land on pages already noted; skip the RMW. + if word.load(Ordering::Relaxed) & bit == 0 { + word.fetch_or(bit, Ordering::Relaxed); + } + } + } + + /// Move the noted pages into `bitmap`, at least as long as these. + fn take_into(&self, bitmap: &mut [u64]) { + for (word, out) in self.words().iter().zip(bitmap) { + *out |= word.swap(0, Ordering::Relaxed); + } + } + + /// Take the noted pages, to test what was noted. + #[cfg(test)] + fn take(&self) -> Vec { + let mut bitmap = vec![0; self.words().len()]; + self.take_into(&mut bitmap); + bitmap + } + + fn clear(&self) { + for word in self.words() { + word.store(0, Ordering::Relaxed); } } } -impl SharedMemory for HostSharedMemory { - fn region(&self) -> &HostMapping { - &self.region +/// The runs of set bits below `pages` in a page bitmap. +pub(crate) struct DirtyRuns<'a> { + bitmap: &'a [u64], + pages: usize, + next: usize, +} + +impl<'a> DirtyRuns<'a> { + pub(crate) fn new(bitmap: &'a [u64], pages: usize) -> Self { + Self { + bitmap, + pages, + next: 0, + } } - fn with_exclusivity T>( + + /// The first page at or after `from` whose bit is `set`, or `pages`. + fn seek(&self, from: usize, set: bool) -> usize { + let mut page = from; + while page < self.pages { + let word = self.bitmap.get(page / 64).copied().unwrap_or(0); + let word = (if set { word } else { !word }) >> (page % 64); + if word != 0 { + return (page + word.trailing_zeros() as usize).min(self.pages); + } + page = (page / 64 + 1) * 64; + } + self.pages + } +} + +impl Iterator for DirtyRuns<'_> { + type Item = Range; + + fn next(&mut self) -> Option> { + let start = self.seek(self.next, true); + if start >= self.pages { + return None; + } + let end = self.seek(start, false); + self.next = end; + Some(start..end) + } +} + +#[cfg(test)] +mod zero_written_tests { + use super::*; + + const PAGE: usize = DIRTY_PAGE_SIZE; + + fn bits(pages: &[usize]) -> Vec { + let mut bitmap = vec![0; 4]; + for &page in pages { + bitmap[page / 64] |= 1 << (page % 64); + } + bitmap + } + + fn runs(bitmap: &[u64], pages: usize) -> Vec> { + DirtyRuns::new(bitmap, pages).collect() + } + + #[test] + fn runs_span_words_and_stop_at_the_end() { + assert!(runs(&bits(&[]), 200).is_empty()); + assert_eq!(runs(&bits(&[0, 1, 2, 5]), 200), vec![0..3, 5..6]); + assert_eq!( + runs(&bits(&[62, 63, 64, 65, 130]), 200), + vec![62..66, 130..131] + ); + assert_eq!(runs(&[u64::MAX; 4], 200), vec![0..200]); + // Bits past the end, and a bitmap shorter than the region. + assert_eq!(runs(&bits(&[10, 150]), 100), vec![10..11]); + assert_eq!(runs(&[u64::MAX], 100), vec![0..64]); + } + + #[test] + fn host_writes_note_every_page_touched() { + let log = HostWrites::new(200 * PAGE); + log.note(PAGE - 1, 2); + log.note(70 * PAGE, 1); + log.note(5 * PAGE, 0); + let mut bitmap = vec![0; 4]; + log.take_into(&mut bitmap); + assert_eq!(runs(&bitmap, 200), vec![0..2, 70..71]); + // Taken, so cleared. + let mut bitmap = vec![0; 4]; + log.take_into(&mut bitmap); + assert!(runs(&bitmap, 200).is_empty()); + } + + fn scratch(pages: usize) -> HostSharedMemory { + ExclusiveSharedMemory::new(pages * PAGE).unwrap().build().0 + } + + /// Write `byte` at `page` unnoted, as the guest's writes are. + fn guest_write(mem: &mut HostSharedMemory, page: usize, byte: u8) { + mem.with_exclusivity_unnoted(|e| e.as_mut_slice()[page * PAGE + 7] = byte) + .unwrap(); + } + + fn nonzero_pages(mem: &mut HostSharedMemory) -> Vec { + mem.with_contents(|contents| { + contents + .chunks(PAGE) + .enumerate() + .filter(|(_, page)| page.iter().any(|&b| b != 0)) + .map(|(i, _)| i) + .collect() + }) + .unwrap() + } + + #[test] + fn zeroes_the_guest_and_host_pages() { + let mut mem = scratch(100); + guest_write(&mut mem, 3, 1); + guest_write(&mut mem, 64, 2); + mem.write::(40 * PAGE, u64::MAX).unwrap(); + let mut written = bits(&[3, 64]); + mem.zero_written(&mut written).unwrap(); + assert!(nonzero_pages(&mut mem).is_empty()); + assert_eq!(runs(&written, 100), vec![3..4, 40..41, 64..65]); + } + + /// Host writes with exclusive access are not logged page by page, so + /// they count as writing every page; reads count as none. + #[test] + fn exclusive_access_writes_every_page() { + let mut mem = scratch(100); + mem.with_contents(|_| ()).unwrap(); + let mut written = bits(&[]); + mem.host_writes.take_into(&mut written); + assert!(runs(&written, 100).is_empty()); + mem.with_exclusivity(|e| e.as_mut_slice()[90 * PAGE] = 3) + .unwrap(); + mem.zero_written(&mut written).unwrap(); + assert!(nonzero_pages(&mut mem).is_empty()); + assert_eq!(runs(&written, 100), vec![0..100]); + } + + /// A page no one noted is left alone, which tests catch. + #[test] + fn leaves_pages_no_one_wrote() { + let mut mem = scratch(100); + guest_write(&mut mem, 9, 1); + guest_write(&mut mem, 10, 1); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + mem.zero_written(&mut bits(&[10])) + })); + assert!(result.is_err()); + } + + /// Without a dirty log, Windows zeroes small scratch in place and + /// replaces large scratch. + #[test] + #[cfg(target_os = "windows")] + fn windows_replaces_only_large_scratch() { + use super::FILL_IN_PLACE_MAX; + let mut state = crate::mem::scratch_reset::ScratchReset::default(); + let mut small = scratch(FILL_IN_PLACE_MAX / PAGE); + guest_write(&mut small, 3, 1); + small.write::(5 * PAGE, 1).unwrap(); + assert!(small.zero_or_replace(&mut state).unwrap().is_none()); + assert!(nonzero_pages(&mut small).is_empty()); + let mut bitmap = vec![0; FILL_IN_PLACE_MAX / PAGE / 64]; + small.host_writes.take_into(&mut bitmap); + assert!(bitmap.iter().all(|&w| w == 0)); + let mut large = scratch(FILL_IN_PLACE_MAX / PAGE + 1); + assert!(large.zero_or_replace(&mut state).unwrap().is_some()); + } + + #[test] + fn host_writes_are_cleared_by_a_full_reset() { + let mut mem = scratch(16); + mem.write::(3 * PAGE, 1).unwrap(); + let mut state = crate::mem::scratch_reset::ScratchReset::default(); + assert!(mem.zero_or_replace(&mut state).unwrap().is_none()); + let mut bitmap = vec![0; 1]; + mem.host_writes.take_into(&mut bitmap); + assert_eq!(bitmap, vec![0]); + } +} + +impl HostSharedMemory { + /// [`SharedMemory::with_exclusivity`] without noting the writes `f` + /// makes, for reads and resets. + fn with_exclusivity_unnoted T>( &mut self, f: F, ) -> Result { @@ -1491,6 +1850,31 @@ impl SharedMemory for HostSharedMemory { } } +impl SharedMemory for HostSharedMemory { + fn region(&self) -> &HostMapping { + &self.region + } + /// Notes all of the memory as written: what `f` writes is not logged + /// page by page, so a reset that zeroes only written pages zeroes + /// all of it. The logged writes, such as + /// [`copy_from_slice`](HostSharedMemory::copy_from_slice), avoid that. + fn with_exclusivity T>( + &mut self, + f: F, + ) -> Result { + let size = self.mem_size(); + let host_writes = self.host_writes.clone(); + self.with_exclusivity_unnoted(|e| { + // First, so a panic in `f` leaves its writes noted. + host_writes.note(0, size); + f(e) + }) + } + fn with_contents T>(&mut self, f: F) -> Result { + self.with_exclusivity_unnoted(|e| f(e.as_slice())) + } +} + /// A ReadonlySharedMemory is a different kind of shared memory, /// separate from the exclusive/host/guest lifecycle, used to /// represent read-only mappings of snapshot pages into the guest diff --git a/src/hyperlight_host/src/sandbox/initialized.rs b/src/hyperlight_host/src/sandbox/initialized.rs index 37ab42084..997855785 100644 --- a/src/hyperlight_host/src/sandbox/initialized.rs +++ b/src/hyperlight_host/src/sandbox/initialized.rs @@ -487,7 +487,8 @@ impl Sandbox { } fn restore_memory_and_mappings(&mut self, snapshot: &Snapshot) -> Result<()> { - let (snapshot_mem, scratch_mem) = self.mem_mgr.restore_snapshot(snapshot)?; + let guest_written = self.vm.scratch_dirty_pages(); + let (snapshot_mem, scratch_mem) = self.mem_mgr.restore_snapshot(snapshot, guest_written)?; if let Some(snapshot_mem) = snapshot_mem { self.vm .update_snapshot_mapping(snapshot_mem) @@ -1499,6 +1500,94 @@ mod tests { } } + /// A restore copies the page tables into scratch with a logged write, + /// page by page, so the next restore zeroes them, and only them, before + /// the snapshot's own go in. + #[test] + fn restore_logs_the_page_table_copy() { + let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf()) + .scratch_size(4 << 20) + .build() + .unwrap(); + let snapshot = sbox.snapshot().unwrap(); + sbox.restore(snapshot).unwrap(); + let written = sbox.mem_mgr.scratch_mem.take_host_writes(); + let range = sbox.mem_mgr.scratch_pt_range(); + assert!(!range.is_empty()); + let noted = |page: usize| written[page / 64] & 1 << (page % 64) != 0; + assert!((range.start / 4096..range.end / 4096).all(noted)); + // Free scratch, past the page tables and below the bookkeeping. + let free = (4 << 20) / 4096 / 2; + assert!(free > range.end / 4096 && !noted(free)); + } + + /// Restores that zero only the scratch pages written since the last + /// one leave nothing behind: from the first restore on, and after the + /// guest writes much of scratch (which pauses MSHV tracking). + #[test] + fn restore_zeroes_every_scratch_page_written() { + // After a restore, scratch holds only what the host wrote: the + // rings at the bottom, the page tables after the transport pools, + // and the bookkeeping page at the top. The pools held the last + // run's messages. + fn assert_reset(sbox: &mut Sandbox, i: u64) { + const ZERO: [u8; 4096] = [0; 4096]; + let arena = sbox.mem_mgr.layout.get_transport_arena(); + let pools = arena.ring_span_len()..arena.size(); + let free = sbox.mem_mgr.scratch_pt_range().end + ..sbox.mem_mgr.scratch_mem.mem_size() - ZERO.len(); + for range in [pools, free] { + let left = sbox + .mem_mgr + .scratch_mem + .with_contents(|scratch| { + scratch[range.clone()] + .chunks(ZERO.len()) + .position(|page| page != &ZERO[..page.len()]) + }) + .unwrap(); + assert_eq!(left, None, "restore {i} left scratch in {range:#x?}"); + } + } + + const SCRATCH: usize = 4 << 20; + let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf()) + .heap_size(8 << 20) + .scratch_size(SCRATCH) + .build() + .unwrap(); + let snapshot = sbox.snapshot().unwrap(); + for i in 0..300u64 { + // A few pages, with a stretch writing a quarter of scratch. + let len = if (100..110).contains(&i) { + SCRATCH as u64 / 4 + } else { + 4096 * (i % 5) + }; + sbox.call::<()>("AllocAndWritePattern", len).unwrap(); + sbox.restore(snapshot.clone()).unwrap(); + assert_reset(&mut sbox, i); + let pattern: Vec = sbox.call("ReadPattern", ()).unwrap(); + assert!(pattern.is_empty(), "restore {i} kept the pattern"); + sbox.restore(snapshot.clone()).unwrap(); + assert_reset(&mut sbox, i); + } + // Where the hypervisor logs the guest's writes, most restores + // used the log: tracking pauses only after the stretch. + #[cfg(all(mshv3, target_arch = "x86_64"))] + let mshv = matches!( + crate::hypervisor::virtual_machine::get_available_hypervisor(), + Some(crate::hypervisor::virtual_machine::HypervisorType::Mshv) + ); + #[cfg(not(all(mshv3, target_arch = "x86_64")))] + let mshv = false; + // WHP may lack tracking, which the VM then reports. + let logs = sbox.vm.scratch_dirty_logs(); + if mshv || logs.is_some() { + assert!(logs.is_some_and(|logs| logs > 500), "{logs:?}"); + } + } + /// Tests that evolving from Sandbox to Sandbox creates a new state /// and restoring a snapshot from before evolving restores the previous state #[test] @@ -2091,8 +2180,9 @@ mod tests { let new_mappings = sandbox.vm.base_mapping_state(); // Snapshot mapping must be identical (no remap). assert_eq!(new_mappings.0, mappings.0); - // On Windows, scratch is freshly allocated each restore so the - // base address may change, but the size must stay the same. + // Where scratch is replaced on restore (Windows without dirty + // tracking) the base address may change, but the size must stay + // the same. assert_eq!(new_mappings.1.map(|m| m.1), mappings.1.map(|m| m.1)); assert!(!fault_plan.is_consumed()); assert_eq!(sandbox.call::("GetStatic", ()).unwrap(), 0); diff --git a/src/hyperlight_host/src/sandbox/snapshot/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/mod.rs index 197edadf0..0707ff0a8 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/mod.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/mod.rs @@ -945,13 +945,13 @@ mod tests { .unwrap(); // Restore snapshot A - mgr.restore_snapshot(&snapshot_a).unwrap(); + mgr.restore_snapshot(&snapshot_a, None).unwrap(); mgr.shared_mem .with_contents(|contents| assert_eq!(&contents[0..pattern_a.len()], &pattern_a[..])) .unwrap(); // Restore snapshot B - mgr.restore_snapshot(&snapshot_b).unwrap(); + mgr.restore_snapshot(&snapshot_b, None).unwrap(); mgr.shared_mem .with_contents(|contents| assert_eq!(&contents[0..pattern_b.len()], &pattern_b[..])) .unwrap();