From 421a4538a96577450dd1b31dd65f2d280717c95c 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 c7dfde709..ab1c24875 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,6 +23,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   regions that the guest has not mapped into its page tables are lost.
 * `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()),