Support loading a wasm module/component via direct mapping

This adds support for directly mapping a host buffer containing a wasm
module/component into the guest, enabling the use of mmap() on the host
to share a single module/component across multiple sandboxes.

Signed-off-by: Lucy Menon <[email protected]>

Author: syntactically

src/hyperlight_wasm/src/sandbox/wasm_sandbox.rs

@@ -17,6 +17,7 @@ limitations under the License.
 use std::path::Path;
 
 use hyperlight_host::func::call_ctx::MultiUseGuestCallContext;
+use hyperlight_host::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
 use hyperlight_host::sandbox::Callable;
 use hyperlight_host::sandbox_state::sandbox::{EvolvableSandbox, Sandbox};
 use hyperlight_host::sandbox_state::transition::MultiUseContextCallback;
@@ -59,34 +60,87 @@ impl WasmSandbox {
     /// 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(self, file: impl AsRef<Path>) -> Result<LoadedWasmSandbox> {
-        let wasm_bytes = std::fs::read(file)?;
-        self.load_module_inner(wasm_bytes)
+        let func = Box::new(move |call_ctx: &mut MultiUseGuestCallContext| {
+            if let Ok(len) = call_ctx.map_file_cow(file.as_ref(), 0x1_0000_0000) {
+                call_ctx.call("LoadWasmModulePhys", (0x1_0000_0000u64, len))
+            } else {
+                let wasm_bytes = std::fs::read(file)?;
+                Self::load_module_from_buffer_transition_func(wasm_bytes)(call_ctx)
+            }
+        });
+        self.load_module_inner(func)
     }
 
-    /// Load a Wasm module from a buffer of bytes into the sandbox and return a `LoadedWasmSandbox`
-    /// able to execute code in the loaded Wasm Module.
+    /// Load a Wasm module that is currently present in a buffer in
+    /// host memory, by mapping the host memory directly into the
+    /// sandbox.
     ///
-    /// 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_from_buffer(self, buffer: &[u8]) -> Result<LoadedWasmSandbox> {
-        self.load_module_inner(buffer.to_vec())
+    /// Depending on the host platform, there are likely alignment
+    /// requirements of at least one page for base and len
+    ///
+    /// # Safety
+    /// It is the caller's responsibility to ensure that the host side
+    /// of the region remains intact and is not written to until the
+    /// produced LoadedWasmSandbox is discarded or devolved.
+    pub unsafe fn load_module_by_mapping(
+        self,
+        base: *mut libc::c_void,
+        len: usize,
+    ) -> Result<LoadedWasmSandbox> {
+        let func = Box::new(move |call_ctx: &mut MultiUseGuestCallContext| {
+            let guest_base: usize = 0x1_0000_0000;
+            let rgn = MemoryRegion {
+                host_region: base as usize..base.wrapping_add(len) as usize,
+                guest_region: guest_base..guest_base + len,
+                flags: MemoryRegionFlags::READ | MemoryRegionFlags::EXECUTE,
+                region_type: MemoryRegionType::Heap,
+            };
+            if let Ok(()) = unsafe { call_ctx.map_region(&rgn) } {
+                call_ctx.call("LoadWasmModulePhys", (0x1_0000_0000u64, len as u64))
+            } else {
+                let wasm_bytes =
+                    unsafe { std::slice::from_raw_parts(base as *const u8, len).to_vec() };
+                Self::load_module_from_buffer_transition_func(wasm_bytes)(call_ctx)
+            }
+        });
+        self.load_module_inner(func)
     }
 
-    fn load_module_inner(mut self, wasm_bytes: Vec<u8>) -> Result<LoadedWasmSandbox> {
-        let func = Box::new(move |call_ctx: &mut MultiUseGuestCallContext| {
-            let len = wasm_bytes.len() as i32;
-            let res: i32 = call_ctx.call("LoadWasmModule", (wasm_bytes, len))?;
+    // todo: take a slice rather than a vec (requires somewhat
+    // refactoring the flatbuffers stuff maybe)
+    fn load_module_from_buffer_transition_func(
+        buffer: Vec<u8>,
+    ) -> impl FnOnce(&mut MultiUseGuestCallContext) -> Result<()> {
+        move |call_ctx: &mut MultiUseGuestCallContext| {
+            let len = buffer.len() as i32;
+            let res: i32 = call_ctx.call("LoadWasmModule", (buffer, len))?;
             if res != 0 {
                 return Err(new_error!(
                     "LoadWasmModule Failed with error code {:?}",
                     res
                 ));
             }
             Ok(())
-        });
+        }
+    }
 
-        let transition_func = MultiUseContextCallback::from(func);
+    /// Load a Wasm module from a buffer of bytes 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_from_buffer(self, buffer: &[u8]) -> Result<LoadedWasmSandbox> {
+        // TODO: get rid of this clone
+        let func = Self::load_module_from_buffer_transition_func(buffer.to_vec());
 
+        self.load_module_inner(func)
+    }
+
+    fn load_module_inner<F: FnOnce(&mut MultiUseGuestCallContext) -> Result<()>>(
+        mut self,
+        func: F,
+    ) -> Result<LoadedWasmSandbox> {
+        let transition_func = MultiUseContextCallback::from(func);
         match self.inner.take() {
             Some(sbox) => {
                 let new_sbox: MultiUseSandbox = sbox.evolve(transition_func)?;

src/wasm_runtime/src/component.rs

@@ -17,6 +17,7 @@ limitations under the License.
 use alloc::string::ToString;
 use alloc::vec;
 use alloc::vec::Vec;
+use core::ptr::NonNull;
 use core::result::Result::*;
 
 use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall;
@@ -72,13 +73,46 @@ fn load_wasm_module(function_call: &FunctionCall) -> Result<Vec<u8>> {
     }
 }
 
+fn load_wasm_module_phys(function_call: &FunctionCall) -> Result<Vec<u8>> {
+    if let (ParameterValue::ULong(ref phys), ParameterValue::ULong(ref len), Some(ref engine)) = (
+        &function_call.parameters.as_ref().unwrap()[0],
+        &function_call.parameters.as_ref().unwrap()[1],
+        &*CUR_ENGINE.lock(),
+    ) {
+        use hyperlight_guest_bin::paging;
+        // TODO: make sure this address is sensible
+        let virt = *phys as *mut u8;
+        let component = unsafe {
+            paging::map_region(*phys, virt, *len + 4096);
+            Component::deserialize_raw(
+                engine,
+                NonNull::new_unchecked(core::ptr::slice_from_raw_parts_mut(virt, *len as usize)),
+            )?
+        };
+        let mut store = Store::new(engine, ());
+        let instance = (*CUR_LINKER.lock())
+            .as_ref()
+            .unwrap()
+            .instantiate(&mut store, &component)?;
+        *CUR_STORE.lock() = Some(store);
+        *CUR_INSTANCE.lock() = Some(instance);
+        Ok(get_flatbuffer_result::<()>(()))
+    } else {
+        Err(HyperlightGuestError::new(
+            ErrorCode::GuestFunctionParameterTypeMismatch,
+            "Invalid parameters passed to LoadWasmModulePhys".to_string(),
+        ))
+    }
+}
+
 #[no_mangle]
 pub extern "C" fn hyperlight_main() {
     let mut config = Config::new();
     config.memory_reservation(0);
     config.memory_guard_size(0);
     config.memory_reservation_for_growth(0);
     config.guard_before_linear_memory(false);
+    config.with_custom_code_memory(Some(alloc::sync::Arc::new(crate::platform::WasmtimeCodeMemory {})));
     let engine = Engine::new(&config).unwrap();
     let linker = Linker::new(&engine);
     *CUR_ENGINE.lock() = Some(engine);
@@ -98,6 +132,12 @@ pub extern "C" fn hyperlight_main() {
         ReturnType::Int,
         load_wasm_module as usize,
     ));
+    register_function(GuestFunctionDefinition::new(
+        "LoadWasmModulePhys".to_string(),
+        vec![ParameterType::ULong, ParameterType::ULong],
+        ReturnType::Void,
+        load_wasm_module_phys as usize,
+    ));
 }
 
 #[no_mangle]

src/wasm_runtime/src/module.rs

@@ -18,6 +18,7 @@ use alloc::string::ToString;
 use alloc::vec::Vec;
 use alloc::{format, vec};
 use core::ops::Deref;
+use core::ptr::NonNull;
 
 use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall;
 use hyperlight_common::flatbuffer_wrappers::function_types::{
@@ -93,6 +94,7 @@ fn init_wasm_runtime() -> Result<Vec<u8>> {
     config.memory_guard_size(0);
     config.memory_reservation_for_growth(0);
     config.guard_before_linear_memory(false);
+    config.with_custom_code_memory(Some(alloc::sync::Arc::new(crate::platform::WasmtimeCodeMemory {})));
     let engine = Engine::new(&config)?;
     let mut linker = Linker::new(&engine);
     wasip1::register_handlers(&mut linker)?;
@@ -138,6 +140,32 @@ fn load_wasm_module(function_call: &FunctionCall) -> Result<Vec<u8>> {
     }
 }
 
+fn load_wasm_module_phys(function_call: &FunctionCall) -> Result<Vec<u8>> {
+    if let (ParameterValue::ULong(ref phys), ParameterValue::ULong(ref len), Some(ref engine)) = (
+        &function_call.parameters.as_ref().unwrap()[0],
+        &function_call.parameters.as_ref().unwrap()[1],
+        &*CUR_ENGINE.lock(),
+    ) {
+        use hyperlight_guest_bin::paging;
+        // TODO: make sure this address is sensible
+        let virt = *phys as *mut u8;
+        let module = unsafe {
+            paging::map_region(*phys, virt, *len + 4096);
+            Module::deserialize_raw(
+                engine,
+                NonNull::new_unchecked(core::ptr::slice_from_raw_parts_mut(virt, *len as usize)),
+            )?
+        };
+        *CUR_MODULE.lock() = Some(module);
+        Ok(get_flatbuffer_result::<()>(()))
+    } else {
+        Err(HyperlightGuestError::new(
+            ErrorCode::GuestFunctionParameterTypeMismatch,
+            "Invalid parameters passed to LoadWasmModulePhys".to_string(),
+        ))
+    }
+}
+
 #[no_mangle]
 #[allow(clippy::fn_to_numeric_cast)] // GuestFunctionDefinition expects a function pointer as i64
 pub extern "C" fn hyperlight_main() {
@@ -161,4 +189,10 @@ pub extern "C" fn hyperlight_main() {
         ReturnType::Int,
         load_wasm_module as usize,
     ));
+    register_function(GuestFunctionDefinition::new(
+        "LoadWasmModulePhys".to_string(),
+        vec![ParameterType::ULong, ParameterType::ULong],
+        ReturnType::Void,
+        load_wasm_module_phys as usize,
+    ));
 }

src/wasm_runtime/src/platform.rs

@@ -145,3 +145,25 @@ pub extern "C" fn wasmtime_tls_get() -> *mut u8 {
 pub extern "C" fn wasmtime_tls_set(ptr: *mut u8) {
     FAKE_TLS.store(ptr, Ordering::Release)
 }
+
+pub struct WasmtimeCodeMemory {}
+// TODO: Actually change the page tables for W^X
+impl wasmtime::CustomCodeMemory for WasmtimeCodeMemory {
+    fn required_alignment(&self) -> usize {
+        unsafe { hyperlight_guest_bin::OS_PAGE_SIZE as usize }
+    }
+    fn publish_executable(
+        &self,
+        _ptr: *const u8,
+        _len: usize,
+    ) -> core::result::Result<(), wasmtime::Error> {
+        Ok(())
+    }
+    fn unpublish_executable(
+        &self,
+        _ptr: *const u8,
+        _len: usize,
+    ) -> core::result::Result<(), wasmtime::Error> {
+        Ok(())
+    }
+}