Take better advantage of CoW/snapshots

src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs

@@ -41,6 +41,7 @@ pub struct WasmSandbox {
     // Snapshot of state of an initial WasmSandbox (runtime loaded, but no guest module code loaded).
     // Used for LoadedWasmSandbox to be able restore state back to WasmSandbox
     snapshot: Option<Arc<Snapshot>>,
+    needs_restore: bool,
 }
 
 const MAPPED_BINARY_VA: u64 = 0x1_0000_0000u64;
@@ -56,6 +57,7 @@ impl WasmSandbox {
         Ok(WasmSandbox {
             inner: Some(inner),
             snapshot: Some(snapshot),
+            needs_restore: false,
         })
     }
 
@@ -64,24 +66,33 @@ impl WasmSandbox {
     /// the snapshot has already been created in that case.
     /// Expects a snapshot of the state where wasm runtime is loaded, but no guest module code is loaded.
     pub(super) fn new_from_loaded(
-        mut loaded: MultiUseSandbox,
+        loaded: MultiUseSandbox,
         snapshot: Arc<Snapshot>,
     ) -> Result<Self> {
-        loaded.restore(snapshot.clone())?;
         metrics::gauge!(METRIC_ACTIVE_WASM_SANDBOXES).increment(1);
         metrics::counter!(METRIC_TOTAL_WASM_SANDBOXES).increment(1);
         Ok(WasmSandbox {
             inner: Some(loaded),
             snapshot: Some(snapshot),
+            needs_restore: true,
         })
     }
 
+    fn restore_if_needed(&mut self) -> Result<()> {
+        if self.needs_restore {
+            self.inner.as_mut().ok_or(new_error!("WasmSandbox is none"))?.restore(self.snapshot.as_ref().ok_or(new_error!("Snapshot is none"))?.clone())?;
+            self.needs_restore = false;
+        }
+        Ok(())
+    }
+
     /// Load a Wasm module at the given path into the sandbox and return a `LoadedWasmSandbox`
     /// able to execute code in the loaded Wasm Module.
     ///
     /// Before you can call guest functions in the sandbox, you must call
     /// this function and use the returned value to call guest functions.
     pub fn load_module(mut self, file: impl AsRef<Path>) -> Result<LoadedWasmSandbox> {
+        self.restore_if_needed()?;
         let inner = self
             .inner
             .as_mut()
@@ -97,6 +108,17 @@ impl WasmSandbox {
         self.finalize_module_load()
     }
 
+    /// Load a Wasm module by restoring a Hyperlight snapshot taken
+    /// from a `LoadedWasmSandbox`.
+    pub fn load_from_snapshot(mut self, snapshot: Arc<Snapshot>) -> Result<LoadedWasmSandbox> {
+        let mut sb = self.inner.take().unwrap();
+        sb.restore(snapshot)?;
+        LoadedWasmSandbox::new(
+            sb,
+            self.snapshot.take().unwrap(),
+        )
+    }
+
     /// Load a Wasm module that is currently present in a buffer in
     /// host memory, by mapping the host memory directly into the
     /// sandbox.
@@ -114,6 +136,7 @@ impl WasmSandbox {
         base: *mut libc::c_void,
         len: usize,
     ) -> Result<LoadedWasmSandbox> {
+        self.restore_if_needed()?;
         let inner = self
             .inner
             .as_mut()

src/hyperlight_wasm_runtime/src/platform.rs

@@ -172,32 +172,95 @@ pub extern "C" fn wasmtime_init_traps(handler: wasmtime_trap_handler_t) -> i32 {
     0
 }
 
-// The wasmtime_memory_image APIs are not yet supported.
+// Copy a VA range to a new VA. Old and new VA, and len, must be
+// page-aligned.
+fn copy_va_mapping(base: *const u8, len: usize, to_va: *mut u8, remap_original: bool) {
+    // TODO: make this more efficient by directly exposing the ability
+    // to traverse an entire VA range in
+    // hyperlight_guest_bin::paging::virt_to_phys, and coalescing
+    // continuous ranges there.
+    let base_u = base as u64;
+    let va_page_bases = (base_u..(base_u + len as u64)).step_by(vmem::PAGE_SIZE);
+    let mappings = va_page_bases.flat_map(paging::virt_to_phys);
+    for mapping in mappings {
+        // TODO: Deduplicate with identical logic in hyperlight_host snapshot.
+        let (new_kind, was_writable) = match mapping.kind {
+            // Skip unmapped pages, since they will be unmapped in
+            // both the original and the new copy
+            vmem::MappingKind::Unmapped => continue,
+            vmem::MappingKind::Basic(bm) if bm.writable => (vmem::MappingKind::Cow(vmem::CowMapping {
+                readable: bm.readable,
+                executable: bm.executable,
+            }), true),
+            vmem::MappingKind::Basic(bm) => (vmem::MappingKind::Basic(vmem::BasicMapping {
+                readable: bm.readable,
+                writable: false,
+                executable: bm.executable,
+            }), false),
+            vmem::MappingKind::Cow(cm) => (vmem::MappingKind::Cow(cm), false),
+        };
+        if remap_original && was_writable {
+            // If necessary, remap the original page as Cow, instead
+            // of whatever it is now, to ensure that any more writes to
+            // that region do not change the image base.
+            //
+            // TODO: could the table traversal needed for this be fused
+            // with the table traversal that got the original mapping,
+            // above?
+            unsafe {
+                paging::map_region(
+                    mapping.phys_base,
+                    mapping.virt_base as *mut u8,
+                    vmem::PAGE_SIZE as u64,
+                    new_kind,
+                );
+            }
+        }
+        // map the same pages to the new VA
+        unsafe {
+            paging::map_region(
+                mapping.phys_base,
+                to_va.wrapping_add((mapping.virt_base - base_u) as usize),
+                vmem::PAGE_SIZE as u64,
+                new_kind,
+            );
+        }
+    }
+}
+
+// Create a copy-on-write memory image from some existing VA range.
+// `ptr` and `len` must be page-aligned (which is guaranteed by the
+// wasmtime-platform.h interface).
 #[no_mangle]
 pub extern "C" fn wasmtime_memory_image_new(
-    _ptr: *const u8,
-    _len: usize,
+    ptr: *const u8,
+    len: usize,
     ret: &mut *mut c_void,
 ) -> i32 {
-    *ret = core::ptr::null_mut();
+    // Choose an arbitrary VA, which we will use as the memory image
+    // identifier. We will construct the image by mapping a copy of
+    // the original VA range here, making the original copy CoW as we
+    // go.
+    let new_virt = FIRST_VADDR.fetch_add(0x100_0000_0000, Ordering::Relaxed) as *mut u8;
+    copy_va_mapping(ptr, len, new_virt, true);
+    *ret = new_virt as *mut c_void;
     0
 }
 
 #[no_mangle]
 pub extern "C" fn wasmtime_memory_image_map_at(
-    _image: *mut c_void,
-    _addr: *mut u8,
-    _len: usize,
+    image: *mut c_void,
+    addr: *mut u8,
+    len: usize,
 ) -> i32 {
-    /* This should never be called because wasmtime_memory_image_new
-     * returns NULL */
-    panic!("wasmtime_memory_image_map_at");
+    copy_va_mapping(image as *mut u8, len, addr, false);
+    0
 }
 
 #[no_mangle]
 pub extern "C" fn wasmtime_memory_image_free(_image: *mut c_void) {
-    /* This should never be called because wasmtime_memory_image_new
-     * returns NULL */
+    /* This should never be called in practice, because we simple
+     * restore the snapshot rather than actually unload/destroy instances */
     panic!("wasmtime_memory_image_free");
 }