diff --git a/src/hyperlight_host/src/sandbox/initialized.rs b/src/hyperlight_host/src/sandbox/initialized.rs index a74d9631b..f07f71bf8 100644 --- a/src/hyperlight_host/src/sandbox/initialized.rs +++ b/src/hyperlight_host/src/sandbox/initialized.rs @@ -1437,6 +1437,102 @@ mod tests { } } + /// A held byte chunk keeps a slot in the transport arena live, mapped + /// through an alias page. A resize moves the scratch base but not the + /// arena. Held bytes, their release, and slot reuse must work after + /// snapshots taken at other scratch sizes. + #[test] + fn held_buffers_survive_resized_snapshots() { + let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf()) + .scratch_size(8 << 20) + // Pin the default pool sizes so 32 holds always outnumber the slots. + .h2g_pool_pages(8) + .g2h_pool_pages(12) + .host_function("HostEchoByteChunks", |v: Vec| v) + .host_function("HostNoOp", || {}) + .build() + .unwrap(); + // Bytes differ between pages, chunks and holds, so a read from a + // shifted or stale page fails. + let chunks = |seed: u8| { + let mut first: Bytes = (0..11 * 1024usize) + .map(|i| (i ^ (i >> 8)) as u8 ^ seed) + .collect(); + let second = first.split_off(6 * 1024); + vec![first, second] + }; + // Chunk boundaries follow transport buffers, so compare bytes. + let read_back = + |sbox: &mut Sandbox, read: &str| sbox.call::>(read, ()).unwrap().concat(); + let base = sbox.snapshot().unwrap(); + + for (seed, hold, read, release) in [ + ( + 0x00, + "RetainGuestByteChunks", + "ReadRetainedGuestByteChunks", + "ReleaseGuestByteChunks", + ), + ( + 0x40, + "RetainHostByteChunks", + "ReadRetainedHostByteChunks", + "ReleaseHostByteChunks", + ), + ] { + sbox.restore(base.clone()).unwrap(); + let held = chunks(seed); + let len: i32 = sbox.call(hold, held.clone()).unwrap(); + assert_eq!(len, 11 * 1024); + assert_eq!(read_back(&mut sbox, read), held.concat()); + + // Shrink from the 8 MiB boot size, then grow from 2 MiB, with + // the chunks held in each snapshot. + for size in [2 << 20, 16 << 20] { + let snapshot = sbox.snapshot_with_options(&scratch(size)).unwrap(); + // Other chunks are held at the restore, so only the snapshot + // can bring the first ones back, and the restore must drop + // the live ones. + assert_eq!(sbox.call::(release, ()).unwrap(), len); + assert_eq!(sbox.call::(hold, chunks(!seed)).unwrap(), len); + sbox.restore(snapshot).unwrap(); + assert_eq!(sbox.mem_mgr.scratch_mem.mem_size(), size); + + // Bytes through both pools, which recycle the held slots after + // a restore. The held chunks read from their captured copy, so + // this traffic must not change them. + let other = chunks(seed ^ 0x80); + let echoed: Vec = + sbox.call("RoundTripHostByteChunks", other.clone()).unwrap(); + assert_eq!(echoed.concat(), other.concat()); + assert_eq!(read_back(&mut sbox, read), held.concat()); + } + assert_eq!(sbox.call::(release, ()).unwrap(), len); + // The restores dropped the swapped-in chunks, so nothing is left. + assert!(read_back(&mut sbox, read).is_empty()); + + // Hold and release more times than either pool has slots, so a + // release that leaks its slot runs the pool dry. + for round in 1..=32 { + let next = chunks(seed + round); + assert_eq!(sbox.call::(hold, next.clone()).unwrap(), len); + assert_eq!(read_back(&mut sbox, read), next.concat()); + assert_eq!(sbox.call::(release, ()).unwrap(), len); + } + + // After the churn the sandbox can still resize, back to 8 MiB, + // and hold chunks there. + let back = sbox.snapshot_with_options(&scratch(8 << 20)).unwrap(); + assert!(!sbox.transport_dirty); + sbox.restore(back).unwrap(); + assert_eq!(sbox.mem_mgr.scratch_mem.mem_size(), 8 << 20); + let last = chunks(seed + 0x21); + assert_eq!(sbox.call::(hold, last.clone()).unwrap(), len); + assert_eq!(read_back(&mut sbox, read), last.concat()); + assert_eq!(sbox.call::(release, ()).unwrap(), len); + } + } + /// The guest's page allocator stops at the transport arena. A resized /// sandbox that runs out of scratch aborts once an allocation reaches the /// arena, with nothing handed out inside it, and a restore recovers it.