From 5c04ce379e78c7b825fc2c15977371339371f7eb Mon Sep 17 00:00:00 2001
From: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>
Date: Fri, 9 Oct 2026 16:54:28 -0700
Subject: [PATCH] Write initial snapshot memory in place

Snapshot::from_env writes the guest code, PEB and init data straight
into shared memory and freezes it. The host populates the remaining
pages on first guest touch, so sandbox creation cost stays independent
of heap size.

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>
---
 CHANGELOG.md                                  |  3 ++
 src/hyperlight_host/src/mem/shared_mem.rs     |  1 +
 .../src/sandbox/snapshot/mod.rs               | 38 +++++++++----------
 3 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1284bf79f..3eca76ff9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -27,6 +27,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   tables become unreachable at the snapshot instead of at the next restore.
 * `vmem::walk_va_spaces` merges adjacent mappings of the same kind into one
   mapping.
+* A new sandbox writes only the guest code, PEB and init data into its memory.
+  The host populates other pages on first guest touch, which speeds up sandbox
+  creation.
 
 ### Removed
 
diff --git a/src/hyperlight_host/src/mem/shared_mem.rs b/src/hyperlight_host/src/mem/shared_mem.rs
index 42d01fac8..6baefda7f 100644
--- a/src/hyperlight_host/src/mem/shared_mem.rs
+++ b/src/hyperlight_host/src/mem/shared_mem.rs
@@ -1546,6 +1546,7 @@ unsafe impl Send for ReadonlySharedMemory {}
 unsafe impl Sync for ReadonlySharedMemory {}
 
 impl ReadonlySharedMemory {
+    #[cfg(any(test, hvf))]
     pub(crate) fn from_bytes(contents: &[u8], guest_mapped_size: usize) -> Result<Self> {
         if guest_mapped_size > contents.len() || !guest_mapped_size.is_multiple_of(page_size::get())
         {
diff --git a/src/hyperlight_host/src/sandbox/snapshot/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/mod.rs
index a8cfc388a..e01121d13 100644
--- a/src/hyperlight_host/src/sandbox/snapshot/mod.rs
+++ b/src/hyperlight_host/src/sandbox/snapshot/mod.rs
@@ -30,7 +30,7 @@ use crate::mem::exe::{ExeInfo, LoadInfo};
 use crate::mem::layout::SandboxMemoryLayout;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::mem::mgr::{BackingSource, GuestBacking, GuestPageTableBuffer, GuestPhysicalMemoryView};
-use crate::mem::shared_mem::{ExclusiveSharedMemory, HostSharedMemory, ReadonlySharedMemory};
+use crate::mem::shared_mem::{ExclusiveSharedMemory, HostSharedMemory};
 use crate::mem::virtq::VirtqSnapshot;
 use crate::sandbox::SandboxConfiguration;
 use crate::sandbox::uninitialized::{GuestBinary, GuestEnvironment};
@@ -550,18 +550,6 @@ impl Snapshot {
         layout.set_code_gva(code_gva)?;
         let regions = layout.get_memory_regions()?;
 
-        let mut memory = vec![0; layout.get_memory_size()?];
-
-        let load_info = exe_info.load(
-            layout.get_guest_code_gva() as u64,
-            &mut memory[layout.guest_code_offset()..],
-        )?;
-
-        layout.write_peb(&mut memory)?;
-
-        blob.map(|x| layout.write_init_data(&mut memory, x.data))
-            .transpose()?;
-
         // Set up page table entries for the snapshot
         let pt_buf = GuestPageTableBuffer::new(layout.get_pt_base_gpa() as usize);
 
@@ -615,13 +603,25 @@ impl Snapshot {
             .ok_or_else(|| crate::new_error!("ELF entrypoint GVA overflows"))?;
 
         // The page tables follow the data, which is padded to a host page.
-        let data_len = memory
-            .len()
-            .next_multiple_of(self::memory::host_page_size());
-        memory.resize(data_len, 0);
-        memory.extend_from_slice(&pt_bytes);
+        let host_page_size = self::memory::host_page_size();
+        let data_len = layout.get_memory_size()?.next_multiple_of(host_page_size);
+        let mut memory = ExclusiveSharedMemory::new(
+            (data_len + pt_bytes.len()).next_multiple_of(host_page_size),
+        )?;
+        let data = &mut memory.as_mut_slice()[..data_len];
+
+        let load_info = exe_info.load(
+            layout.get_guest_code_gva() as u64,
+            &mut data[layout.guest_code_offset()..],
+        )?;
+
+        layout.write_peb(data)?;
+
+        blob.map(|x| layout.write_init_data(data, x.data))
+            .transpose()?;
+        memory.copy_from_slice(&pt_bytes, data_len)?;
         let memory = Arc::new(SnapshotMemory::from_flat(
-            ReadonlySharedMemory::from_bytes(&memory, data_len)?,
+            memory.freeze(data_len)?,
             SandboxMemoryLayout::BASE_ADDRESS as u64,
             pt_bytes.len(),
             scratch_base_gpa(layout.get_scratch_size()),