fix(restore): replace scratch zero with fresh allocation on Windows

src/hyperlight_host/src/mem/mgr.rs

@@ -490,8 +490,9 @@ impl SandboxMemoryManager<HostSharedMemory> {
         };
         let new_scratch_size = snapshot.layout().get_scratch_size();
         let gscratch = if new_scratch_size == self.scratch_mem.mem_size() {
-            self.scratch_mem.zero()?;
-            None
+            // zero_or_replace picks the fastest zeroing strategy for
+            // the current platform (see SharedMemory::zero_or_replace).
+            self.scratch_mem.zero_or_replace()?
         } else {
             let new_scratch_mem = ExclusiveSharedMemory::new(new_scratch_size)?;
             let (hscratch, gscratch) = new_scratch_mem.build();
@@ -501,7 +502,6 @@ impl SandboxMemoryManager<HostSharedMemory> {
             // mapping, so it won't actually be deallocated until it
             // has been unmapped from the VM.
             self.scratch_mem = hscratch;
-
             Some(gscratch)
         };
         self.layout = *snapshot.layout();

src/hyperlight_host/src/mem/shared_mem.rs

@@ -590,31 +590,6 @@ pub trait SharedMemory {
     fn with_contents<T, F: FnOnce(&[u8]) -> T>(&mut self, f: F) -> Result<T> {
         self.with_exclusivity(|m| f(m.as_slice()))
     }
-
-    /// Zero a shared memory region
-    fn zero(&mut self) -> Result<()> {
-        self.with_exclusivity(|e| {
-            #[allow(unused_mut)] // unused on some platforms, although not others
-            let mut do_copy = true;
-            // TODO: Compare & add heuristic thresholds: mmap, MADV_DONTNEED, MADV_REMOVE, MADV_FREE (?)
-            // TODO: Find a similar lazy zeroing approach that works on MSHV.
-            //       (See Note [Keeping mappings in sync between userspace and the guest])
-            #[cfg(all(target_os = "linux", 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;
-                }
-            }
-            if do_copy {
-                e.as_mut_slice().fill(0);
-            }
-        })
-    }
 }
 
 fn mapping_at(
@@ -1503,6 +1478,58 @@ impl HostSharedMemory {
     }
 }
 
+impl HostSharedMemory {
+    /// Reset this memory region to all-zeros, choosing the fastest
+    /// strategy for the current platform and hypervisor configuration.
+    ///
+    /// On Linux/KVM (without mshv3), uses `MADV_DONTNEED` for lazy
+    /// zeroing.  On Linux/mshv3, falls through to `fill(0)`.
+    ///
+    /// 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.
+    ///
+    // 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<Option<GuestSharedMemory>> {
+        #[cfg(target_os = "windows")]
+        {
+            let new_mem = ExclusiveSharedMemory::new(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;
+                    }
+                }
+                if do_copy {
+                    e.as_mut_slice().fill(0);
+                }
+            })?;
+            Ok(None)
+        }
+    }
+}
+
 impl SharedMemory for HostSharedMemory {
     fn region(&self) -> &HostMapping {
         &self.region

src/hyperlight_host/src/sandbox/initialized_multi_use.rs

@@ -1747,7 +1747,12 @@ mod tests {
         sandbox.restore(snapshot).unwrap();
 
         assert_eq!(sandbox.status(), SandboxStatus::Ready);
-        assert_eq!(sandbox.vm.base_mapping_state(), mappings);
+        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.
+        assert_eq!(new_mappings.1.map(|m| m.1), mappings.1.map(|m| m.1));
         assert!(!fault_plan.is_consumed());
         assert_eq!(sandbox.call::<i32>("GetStatic", ()).unwrap(), 0);
     }