diff --git a/src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs b/src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs index b468bd0..bb99473 100644 --- a/src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs +++ b/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>, + 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, ) -> Result { - 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) -> Result { + 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) -> Result { + 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 { + self.restore_if_needed()?; let inner = self .inner .as_mut() diff --git a/src/hyperlight_wasm_runtime/src/platform.rs b/src/hyperlight_wasm_runtime/src/platform.rs index ff6b660..2ef7f60 100644 --- a/src/hyperlight_wasm_runtime/src/platform.rs +++ b/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"); }