feat: restore snapshots with less scratch

CHANGELOG.md

@@ -5,8 +5,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 ## [Prerelease] - Unreleased
 
 ### Added
+* `Sandbox::snapshot_with_scratch_size` takes a snapshot that gives a restored sandbox a different scratch size. A guest can boot with the scratch its initialisation needs and run guest function calls from a snapshot with only the scratch they need, so each restore costs less. It cannot grow scratch.
 
 ### Changed
+* **Breaking:** The transport arena sits one free page below the reserved scratch-top pages, at an address independent of the scratch size, and the page tables at the scratch base. The guest's page allocator stops at the arena GPA the host publishes in the scratch-top metadata, so guests must be rebuilt. `hyperlight_common::layout::scratch_allocator_limit_gpa` is renamed `scratch_reserved_top_gpa`, the start of the reserved pages. An allocator that stops there runs into the arena. The minimum scratch size grows by that free page. Snapshot ABI 7: existing snapshots must be regenerated.
 
 ### Removed
 

docs/paging-development-notes.md

@@ -128,9 +128,13 @@ scratch region that will be used when this snapshot is next restored
 into a sandbox. This will then be baked into the guest page tables
 created in the snapshot.
 
-TODO: add support, if found to be useful operationally, for either
-dynamically growing the scratch region, or changing its size between
-taking a snapshot and restoring it.
+`Sandbox::snapshot_with_scratch_size` takes a snapshot whose scratch is
+smaller than the sandbox's. Its page tables map the smaller window, and
+restoring it replaces the scratch region. A guest can boot with the
+scratch its initialisation needs and run guest function calls from a
+snapshot with only the scratch they need. A snapshot cannot have more
+scratch than the sandbox it is taken from. Restoring a larger snapshot
+gives a sandbox its scratch back.
 
 ### Call descriptors
 
@@ -139,14 +143,18 @@ calls, i.e. there may be no calls in flight at the time of
 snapshotting. This is not enforced, but odd things may happen if it is
 violated.
 
-Host and guest calls use two virtqueues in a fixed transport arena at
-the bottom of scratch. The arena contains both rings and their
-fixed-slot buffer pools. Copied page tables follow the arena. Dynamic
-scratch allocations begin after the copied page tables.
-
-The minimum scratch size is calculated by `min_scratch_size()` in the
-architecture-specific layout modules under `hyperlight_common`; see
-that function for the detailed breakdown of required overhead.
+Host function calls and guest function calls use two virtqueues in a
+fixed transport arena below the reserved scratch-top pages, with one
+free page between them for an exception stack overrun. The arena's
+address does not depend on the scratch size. The arena contains both
+rings and their fixed-slot buffer pools. Copied page tables start at the
+scratch base. Dynamic scratch allocations begin after the copied page
+tables and stop at the arena.
+
+The minimum scratch size is the architecture's fixed overhead
+(`min_scratch_size()` in the layout modules under `hyperlight_common`,
+which breaks it down) plus the transport arena and the free page above
+it (`SandboxMemoryLayout::min_scratch`), plus the copied page tables.
 
 ## Creating a fresh guest
 

docs/snapshot-versioning.md

@@ -41,6 +41,10 @@ A snapshot carries four independently evolvable version markers:
 ABI 6 requires checkpoint readiness after guest initialization and permits
 H2G prefill to omit retained slots. Earlier snapshots must be regenerated.
 
+ABI 7 places the transport arena below the reserved scratch-top pages, at an
+address independent of the scratch size, and the page tables at the scratch
+base. Earlier snapshots must be regenerated.
+
 The `OCI_LAYOUT_VERSION` constant is pinned by the OCI image-layout
 spec at `1.0.0`.
 

docs/virtio-host-guest-communication.md

@@ -42,16 +42,17 @@ optional writable response capacity.
 
 ## Transport arena
 
-Both rings, the transport mailbox, and both pools occupy one fixed prefix of
-guest scratch memory.
+Both rings, the transport mailbox, and both pools occupy one fixed range of
+guest scratch memory, below the reserved scratch-top pages with one free page
+between them. Its address does not depend on the scratch size. Guest page
+allocations stop at its base.
 
 ```text
- scratch base
-     |
-     v
- +----------+-----+----------+-----+-----+-----+----------+----------+
- | G2H ring | pad | H2G ring | pad | mbx | pad | G2H pool | H2G pool |
- +----------+-----+----------+-----+-----+-----+----------+----------+
+                                                           reserved top pages
+                                                                            |
+ +----------+-----+----------+-----+-----+-----+----------+----------+------+
+ | G2H ring | pad | H2G ring | pad | mbx | pad | G2H pool | H2G pool | free |
+ +----------+-----+----------+-----+-----+-----+----------+----------+------+
 ```
 
 The host derives this layout from `SandboxConfiguration`. Ring starts follow
@@ -525,8 +526,8 @@ both consumers.
 ## Placement and relocation limitations
 
 The fixed-pool runtime places both rings, the mailbox, and both pools in one
-host-owned arena at the scratch base. The guest reconstructs that layout from
-host metadata.
+host-owned arena below the reserved scratch-top pages. The guest reconstructs
+that layout from host metadata.
 
 Descriptors, pool owners, and producer state contain absolute GVAs. Restore
 adopts the snapshot's scratch size, queue geometry, and transport addresses,

src/hyperlight_common/src/layout.rs

@@ -113,11 +113,33 @@ const _: () = {
     assert!((SCRATCH_TOP_GPA as u64) < VIRTQ_BUFFER_GVA_END - VIRTQ_BUFFER_GVA_START);
 };
 
-/// Exclusive upper GPA boundary for dynamic scratch allocations.
-pub const fn scratch_allocator_limit_gpa() -> u64 {
+/// Free pages between the transport arena and the reserved top pages. A
+/// small overrun of the exception stack lands in them.
+pub const SCRATCH_TRANSPORT_GAP_PAGES: usize = 1;
+
+/// Start of the reserved pages at the top of scratch. The transport arena
+/// sits [`SCRATCH_TRANSPORT_GAP_PAGES`] below it. A guest's page allocator
+/// stops at the arena GPA the host publishes.
+pub const fn scratch_reserved_top_gpa() -> u64 {
     (SCRATCH_TOP_GPA + 1 - SCRATCH_TOP_RESERVED_PAGES * crate::vmem::PAGE_SIZE) as u64
 }
 
+/// [`scratch_reserved_top_gpa`] in the scratch mapping.
+pub const fn scratch_reserved_top_gva() -> u64 {
+    SCRATCH_TOP_GVA as u64 + 1 - (SCRATCH_TOP_RESERVED_PAGES * crate::vmem::PAGE_SIZE) as u64
+}
+
+/// Exclusive end of the transport arena: the reserved top pages, less the
+/// free pages between them.
+pub const fn scratch_transport_end_gpa() -> u64 {
+    scratch_reserved_top_gpa() - (SCRATCH_TRANSPORT_GAP_PAGES * crate::vmem::PAGE_SIZE) as u64
+}
+
+/// [`scratch_transport_end_gpa`] in the scratch mapping.
+pub const fn scratch_transport_end_gva() -> u64 {
+    scratch_reserved_top_gva() - (SCRATCH_TRANSPORT_GAP_PAGES * crate::vmem::PAGE_SIZE) as u64
+}
+
 pub fn scratch_base_gpa(size: usize) -> u64 {
     (SCRATCH_TOP_GPA - size + 1) as u64
 }
@@ -127,8 +149,9 @@ pub const fn scratch_base_gva(size: usize) -> u64 {
 
 /// Compute fixed scratch overhead for a sandbox, saturating at [`usize::MAX`].
 ///
-/// `transport_len` covers rings and buffer pools. Snapshot page tables and
-/// runtime allocations need additional scratch.
+/// `transport_len` covers the transport arena (rings, mailbox and buffer
+/// pools) and the [`SCRATCH_TRANSPORT_GAP_PAGES`] above it. Snapshot page
+/// tables and runtime allocations need additional scratch.
 pub fn min_scratch_size(transport_len: usize) -> usize {
     arch::min_scratch_size()
         .and_then(|fixed| fixed.checked_add(transport_len))

src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs

@@ -2,7 +2,7 @@
 // Copyright 2025 The Hyperlight Authors.
 
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
-use hyperlight_common::{layout, vmem};
+use hyperlight_common::vmem;
 
 // There are no notable architecture-specific safety considerations
 // here, and the general conditions are documented in the
@@ -22,7 +22,7 @@ pub unsafe fn alloc_phys_pages(n: u64) -> u64 {
             prev_base = out(reg) prev_base,
         );
     }
-    let limit = layout::scratch_allocator_limit_gpa();
+    let limit = crate::layout::allocator_limit_gpa();
     if prev_base.checked_add(nbytes).is_none_or(|end| end > limit) {
         unsafe {
             crate::exit::abort_with_code_and_message(

src/hyperlight_guest/src/arch/amd64/prim_alloc.rs

@@ -2,7 +2,7 @@
 // Copyright 2025 The Hyperlight Authors.
 
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
-use hyperlight_common::{layout, vmem};
+use hyperlight_common::vmem;
 
 // There are no notable architecture-specific safety considerations
 // here, and the general conditions are documented in the
@@ -19,7 +19,7 @@ pub unsafe fn alloc_phys_pages(n: u64) -> u64 {
             x = inout(reg) x
         );
     }
-    let limit = layout::scratch_allocator_limit_gpa();
+    let limit = crate::layout::allocator_limit_gpa();
     if x.checked_add(nbytes).is_none_or(|end| end > limit) {
         unsafe {
             crate::exit::abort_with_code_and_message(

src/hyperlight_guest/src/layout.rs

@@ -52,3 +52,30 @@ pub fn guest_log_level_gva() -> *mut u64 {
     (SCRATCH_TOP_GVA as u64 - SCRATCH_TOP_GUEST_LOG_LEVEL_OFFSET + 1) as *mut u64
 }
 pub use arch::{scratch_base_gpa, scratch_base_gva};
+
+/// Exclusive upper GPA for the scratch page allocator: the transport arena,
+/// which sits below the reserved top pages.
+///
+/// Aborts if the host has not published an arena between the scratch base
+/// and the end of the transport range, so allocations never reach the
+/// rings.
+pub fn allocator_limit_gpa() -> u64 {
+    // SAFETY: the host writes the arena GPA into scratch metadata before entry.
+    let arena = unsafe { transport_arena_gpa_gva().read_volatile() };
+    if arena <= scratch_base_gpa()
+        || arena >= hyperlight_common::layout::scratch_transport_end_gpa()
+    {
+        // SAFETY: aborting is always sound.
+        unsafe {
+            crate::exit::abort_with_code_and_message(
+                &[
+                    hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode::MallocFailed
+                        as u8,
+                ],
+                c"Transport arena not published between the scratch base and the transport end"
+                    .as_ptr(),
+            )
+        }
+    }
+    arena
+}

src/hyperlight_guest/src/transport/context.rs

@@ -12,7 +12,7 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{
     FunctionCallResult, ParameterValue, ReturnType, ReturnValue,
 };
 use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity;
-use hyperlight_common::layout::SCRATCH_TOP_GVA;
+use hyperlight_common::layout::scratch_transport_end_gva;
 use hyperlight_common::outb::OutBAction;
 use hyperlight_common::transport::{
     EncodedMessage, ExternalValues, MailboxValue, MsgHeader, MsgKind,
@@ -119,23 +119,27 @@ impl GuestContext {
     /// The caller must run in an initialized single-vCPU Hyperlight guest and
     /// exclusively own both queues, their pools, and the writable `u64` mailbox
     /// at `mbx_gva`. The mailbox must lie in scratch, disjoint from both queues
-    /// and pools. Scratch must remain mapped while the context or any returned
-    /// buffer views exist.
+    /// and pools. Both rings and pools must lie in the transport arena, which
+    /// starts at the G2H ring and ends below the reserved scratch-top pages,
+    /// with the pools after the rings. Scratch must remain mapped while the
+    /// context or any returned buffer views exist.
     pub unsafe fn new(g2h: QueueConfig, h2g: QueueConfig, mbx_gva: u64) -> Result<Self> {
         // SAFETY: Generation has been initialized by host.
         let generation = unsafe { layout::snapshot_generation_gva().read_volatile() };
 
         let g2h_pool = g2h_pool(g2h.pool_gva, g2h.pool_pages, g2h.buffer_size)
             .with_context(|| "failed to create G2H pool")?;
+        // The G2H ring starts the transport arena.
+        let arena_gva = g2h.layout.desc_table_addr();
         // SAFETY: The caller owns the pool's scratch pages.
-        let mem = unsafe { pool_mem(&g2h_pool) }
+        let mem = unsafe { pool_mem(&g2h_pool, arena_gva) }
             .with_context(|| "failed to initialize G2H pool access")?;
         let g2h_producer = VirtqProducer::new(g2h.layout, mem, OutbNotifier, g2h_pool);
 
         let h2g_pool = h2g_pool(h2g.pool_gva, h2g.pool_pages, h2g.buffer_size)
             .with_context(|| "failed to create H2G slot pool")?;
         // SAFETY: The caller owns this disjoint pool.
-        let mem = unsafe { pool_mem(&h2g_pool) }
+        let mem = unsafe { pool_mem(&h2g_pool, arena_gva) }
             .with_context(|| "failed to initialize H2G pool access")?;
         // H2G prefill supplies writable buffers for host-initiated messages.
         let h2g_producer = VirtqProducer::new(h2g.layout, mem, NoopNotifier, h2g_pool);
@@ -549,24 +553,28 @@ fn pool_len(pages: usize) -> result::Result<usize, AllocError> {
         .ok_or(AllocError::Overflow)
 }
 
-/// Validate pool placement and bind scratch access.
+/// Validate pool placement and bind access to the arena up to the pool's end.
+///
+/// The bounds start at the arena, whose address does not depend on the
+/// scratch size. The pools follow the rings, so a producer's rings and pool
+/// lie in [arena, end of its pool), which must end below the free page
+/// under the reserved scratch-top pages.
 ///
 /// # Safety
 ///
 /// Run on the initialized guest vCPU and own the pool's scratch.
-unsafe fn pool_mem(pool: &SlotPool) -> result::Result<GuestMemOps, AllocError> {
-    let scratch_gva = layout::scratch_base_gva();
-    let scratch_end = SCRATCH_TOP_GVA as u64 + 1;
+unsafe fn pool_mem(pool: &SlotPool, arena_gva: u64) -> result::Result<GuestMemOps, AllocError> {
+    let limit = scratch_transport_end_gva();
 
     let base = pool.base_addr();
     let end = base + pool.byte_len() as u64;
 
-    if !base.is_multiple_of(PAGE_SIZE as u64) || base < scratch_gva || end > scratch_end {
+    if !base.is_multiple_of(PAGE_SIZE as u64) || base < arena_gva || end > limit {
         return Err(AllocError::InvalidArg);
     }
 
-    // SAFETY: The validated pool lies in live scratch on the serialized guest vCPU.
-    unsafe { Ok(GuestMemOps::new(scratch_gva, scratch_end)) }
+    // SAFETY: The validated pool lies in the arena on the serialized guest vCPU.
+    unsafe { Ok(GuestMemOps::new(arena_gva, end)) }
 }
 
 /// Build the uniform H2G pool.
@@ -612,6 +620,22 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn pool_mem_binds_the_arena_below_the_reserved_top() {
+        let limit = scratch_transport_end_gva();
+        let page = PAGE_SIZE as u64;
+        let arena = limit - 8 * page;
+        let pool = |base, pages| h2g_pool(base, pages, 256).unwrap();
+        // SAFETY: only bounds are computed. No memory is touched.
+        unsafe {
+            assert!(pool_mem(&pool(arena + 4 * page, 4), arena).is_ok());
+            // Into the free page or the reserved pages.
+            assert!(pool_mem(&pool(arena + 4 * page, 5), arena).is_err());
+            // Below the arena.
+            assert!(pool_mem(&pool(arena - page, 1), arena).is_err());
+        }
+    }
+
     fn encoded_message(external: &[u8]) -> EncodedMessage<'_> {
         let mut values = ExternalValues::new();
         values.push_bytes(external).unwrap();

src/hyperlight_guest/src/transport/mem.rs

@@ -17,13 +17,14 @@ use super::backing;
 // Mark types not [`Send`]
 type NotThreadSafe = PhantomData<*mut ()>;
 
-/// Bounded scratch memory access.
+/// Memory access bounded to a range of scratch: the transport arena, for
+/// the producers.
 ///
-/// Copies reject buffers that overlap the accessed scratch range.
+/// Copies reject buffers that overlap the accessed range.
 #[derive(Clone, Copy, Debug)]
 pub(crate) struct GuestMemOps {
-    scratch_gva: u64,
-    scratch_end: u64,
+    start: u64,
+    end: u64,
     _nts: NotThreadSafe,
 }
 
@@ -32,24 +33,25 @@ pub(crate) struct GuestMemOps {
 pub struct GuestMemError;
 
 impl GuestMemOps {
-    /// Bind memory access to live scratch bounds.
+    /// Bind memory access to `[start, end)`, a live part of
+    /// scratch (the transport arena, for the producers).
     ///
     /// # Safety
     ///
-    /// Scratch must be live, initialized, and writable.
+    /// The range must be live, initialized, and writable scratch.
     /// All access and destruction must stay on the serialized guest vCPU.
     /// Serialize paging and share alias leaf entries across accessing roots.
-    pub(super) unsafe fn new(scratch_gva: u64, scratch_end: u64) -> Self {
+    pub(super) unsafe fn new(start: u64, end: u64) -> Self {
         Self {
-            scratch_gva,
-            scratch_end,
+            start,
+            end,
             _nts: PhantomData,
         }
     }
 
     fn ptr(&self, addr: u64, len: usize) -> Result<*mut u8, GuestMemError> {
         let end = addr.checked_add(len as u64).ok_or(GuestMemError)?;
-        if addr < self.scratch_gva || end > self.scratch_end {
+        if addr < self.start || end > self.end {
             return Err(GuestMemError);
         }
         Ok(addr as *mut u8)
@@ -119,7 +121,7 @@ impl AsRef<[u8]> for GuestMapping {
     }
 }
 
-// SAFETY: Every address is restricted to the scratch mapping. Payload
+// SAFETY: Every address is restricted to the bound range of scratch. Payload
 // references rely on descriptor ownership, and ring flags use aligned atomics.
 unsafe impl MemOps for GuestMemOps {
     type Error = GuestMemError;