fix: improve WASM runtime memory management

- Add memory tracking for allocated parameters in guest function calls
- Add memory tracking for allocated parameters in host function calls
- Free allocated parameters after each guest function call (implies that
  host function return values are returned by ref, so caller must copy
values if they want to keep them around as they will get freed after
returning).
- Component: Add post_return calls for proper WASM function cleanup
- Fix ABI mismatch in parameter of guest_dispatch_function

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_wasm_macro/src/wasmguest.rs

@@ -212,15 +212,17 @@ fn emit_wasm_function_call(
     let rwt = match result {
         None => {
             quote! {
-                instance.get_typed_func::<(#(#pwts,)*), ()>(&mut *store, func_idx)?
-                    .call(&mut *store, (#(#pus,)*))?;
+                let func = instance.get_typed_func::<(#(#pwts,)*), ()>(&mut *store, func_idx)?;
+                func.call(&mut *store, (#(#pus,)*))?;
+                func.post_return(&mut *store)?;
             }
         }
         _ => {
             let r = rtypes::emit_func_result(s, result);
             quote! {
-                let #ret = instance.get_typed_func::<(#(#pwts,)*), ((#r,))>(&mut *store, func_idx)?
-                    .call(&mut *store, (#(#pus,)*))?.0;
+                let func = instance.get_typed_func::<(#(#pwts,)*), ((#r,))>(&mut *store, func_idx)?;
+                let #ret = func.call(&mut *store, (#(#pus,)*))?.0;
+                func.post_return(&mut *store)?;
             }
         }
     };

src/wasm_runtime/src/hostfuncs.rs

@@ -15,6 +15,7 @@ limitations under the License.
 */
 
 use alloc::string::ToString;
+use alloc::vec;
 use alloc::vec::Vec;
 
 use hyperlight_common::flatbuffer_wrappers::function_types::{
@@ -23,10 +24,23 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 use hyperlight_guest_bin::host_comm::call_host_function;
+use spin::Mutex;
 use wasmtime::{Caller, Engine, FuncType, Val, ValType};
 
 use crate::marshal;
 
+// Track memory addresses allocated by host function return values
+static HOSTFUNC_ALLOCATED_ADDRS: Mutex<Vec<i32>> = Mutex::new(Vec::new());
+
+/// Retrieve and clear all addresses allocated by host function return values.
+/// Caller is responsible for making sure the addresses are freed
+pub(crate) fn take_hostfunc_allocated_addrs() -> Vec<i32> {
+    let mut addrs = HOSTFUNC_ALLOCATED_ADDRS.lock();
+    let mut result = Vec::new();
+    core::mem::swap(&mut *addrs, &mut result);
+    result
+}
+
 pub(crate) type HostFunctionDefinition =
     hyperlight_common::flatbuffer_wrappers::host_function_definition::HostFunctionDefinition;
 pub(crate) type HostFunctionDetails =
@@ -79,9 +93,9 @@ pub(crate) fn hostfunc_type(d: &HostFunctionDefinition, e: &Engine) -> Result<Fu
     Ok(FuncType::new(e, params, results))
 }
 
-pub(crate) fn call(
+pub(crate) fn call<T>(
     d: &HostFunctionDefinition,
-    mut c: Caller<'_, ()>,
+    mut c: Caller<'_, T>,
     ps: &[Val],
     rs: &mut [Val],
 ) -> Result<()> {
@@ -110,7 +124,19 @@ pub(crate) fn call(
         return Ok(());
     }
 
-    rs[0] = marshal::hl_return_to_val(&mut c, |c, n| c.get_export(n), rv)?;
+    let mut allocated_addrs = vec![];
+    rs[0] = marshal::hl_return_to_val_with_tracking(
+        &mut c,
+        |c, n| c.get_export(n),
+        rv,
+        &mut allocated_addrs,
+    )?;
+
+    // Track any allocations for later cleanup
+    if !allocated_addrs.is_empty() {
+        let mut global_addrs = HOSTFUNC_ALLOCATED_ADDRS.lock();
+        global_addrs.extend(allocated_addrs);
+    }
 
     Ok(())
 }

src/wasm_runtime/src/marshal.rs

@@ -46,6 +46,21 @@ fn malloc<C: AsContextMut>(
     Ok(addr)
 }
 
+fn free<C: AsContextMut>(
+    ctx: &mut C,
+    get_export: &impl Fn(&mut C, &str) -> Option<Extern>,
+    addr: i32,
+) -> Result<()> {
+    let free = get_export(&mut *ctx, "free")
+        .and_then(Extern::into_func)
+        .ok_or(HyperlightGuestError::new(
+            ErrorCode::GuestError,
+            "free function not exported".to_string(),
+        ))?;
+    free.typed::<i32, ()>(&mut *ctx)?.call(&mut *ctx, addr)?;
+    Ok(())
+}
+
 fn write<C: AsContextMut>(
     ctx: &mut C,
     get_export: &impl Fn(&mut C, &str) -> Option<Extern>,
@@ -126,10 +141,11 @@ fn read_cstr<C: AsContextMut>(
     })
 }
 
-pub fn hl_param_to_val<C: AsContextMut>(
+pub fn hl_param_to_val_with_tracking<C: AsContextMut>(
     mut ctx: C,
     get_export: impl Fn(&mut C, &str) -> Option<Extern>,
     param: &ParameterValue,
+    allocated_addrs: &mut Vec<i32>,
 ) -> Result<Val> {
     match param {
         ParameterValue::Int(i) => Ok(Val::I32(*i)),
@@ -144,17 +160,31 @@ pub fn hl_param_to_val<C: AsContextMut>(
             let nbytes = s.count_bytes() + 1; // include the NUL terminator
             let addr = malloc(&mut ctx, &get_export, nbytes)?;
             write(&mut ctx, &get_export, addr, s.as_bytes_with_nul())?;
+            allocated_addrs.push(addr);
             Ok(Val::I32(addr))
         }
         ParameterValue::VecBytes(b) => {
             let addr = malloc(&mut ctx, &get_export, b.len())?;
             write(&mut ctx, &get_export, addr, b)?;
+            allocated_addrs.push(addr);
             Ok(Val::I32(addr))
             // TODO: check that the next parameter is the correct length
         }
     }
 }
 
+/// Helper function to free all tracked allocated addresses
+pub fn free_allocated_addrs<C: AsContextMut>(
+    mut ctx: C,
+    get_export: impl Fn(&mut C, &str) -> Option<Extern>,
+    allocated_addrs: &[i32],
+) -> Result<()> {
+    for &addr in allocated_addrs {
+        free(&mut ctx, &get_export, addr)?;
+    }
+    Ok(())
+}
+
 pub fn val_to_hl_result<C: AsContextMut>(
     mut ctx: C,
     get_export: impl Fn(&mut C, &str) -> Option<Extern>,
@@ -248,10 +278,11 @@ pub fn val_to_hl_param<'a, C: AsContextMut>(
     }
 }
 
-pub fn hl_return_to_val<C: AsContextMut>(
+pub fn hl_return_to_val_with_tracking<C: AsContextMut>(
     ctx: &mut C,
     get_export: impl Fn(&mut C, &str) -> Option<Extern>,
     rv: ReturnValue,
+    allocated_addrs: &mut Vec<i32>,
 ) -> Result<Val> {
     match rv {
         ReturnValue::Int(i) => Ok(Val::I32(i)),
@@ -266,11 +297,13 @@ pub fn hl_return_to_val<C: AsContextMut>(
             let nbytes = s.count_bytes() + 1; // include the NUL terminator
             let addr = malloc(ctx, &get_export, nbytes)?;
             write(ctx, &get_export, addr, s.as_bytes_with_nul())?;
+            allocated_addrs.push(addr);
             Ok(Val::I32(addr))
         }
         ReturnValue::VecBytes(b) => {
             let addr = malloc(ctx, &get_export, b.len())?;
             write(ctx, &get_export, addr, b.as_ref())?;
+            allocated_addrs.push(addr);
             Ok(Val::I32(addr))
             // TODO: check that the next parameter is the correct length
         }

src/wasm_runtime/src/module.rs

@@ -17,7 +17,7 @@ limitations under the License.
 use alloc::string::ToString;
 use alloc::vec::Vec;
 use alloc::{format, vec};
-use core::ops::Deref;
+use core::ops::{Deref, DerefMut};
 
 use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall;
 use hyperlight_common::flatbuffer_wrappers::function_types::{
@@ -32,59 +32,95 @@ use hyperlight_guest_bin::host_comm::print_output_with_host_print;
 use spin::Mutex;
 use wasmtime::{Config, Engine, Linker, Module, Store, Val};
 
+use crate::hostfuncs::take_hostfunc_allocated_addrs;
 use crate::{hostfuncs, marshal, platform, wasip1};
 
+// Set by transition to WasmSandbox (by init_wasm_runtime)
 static CUR_ENGINE: Mutex<Option<Engine>> = Mutex::new(None);
 static CUR_LINKER: Mutex<Option<Linker<()>>> = Mutex::new(None);
+// Set by transition to LoadedWasmSandbox (by load_wasm_module/load_wasm_module_phys)
 static CUR_MODULE: Mutex<Option<Module>> = Mutex::new(None);
+static CUR_STORE: Mutex<Option<Store<()>>> = Mutex::new(None);
+static CUR_INSTANCE: Mutex<Option<wasmtime::Instance>> = Mutex::new(None);
 
 #[no_mangle]
-pub fn guest_dispatch_function(function_call: &FunctionCall) -> Result<Vec<u8>> {
-    let engine = CUR_ENGINE.lock();
-    let engine = engine.deref().as_ref().ok_or(HyperlightGuestError::new(
+pub fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>> {
+    let mut store = CUR_STORE.lock();
+    let store = store.deref_mut().as_mut().ok_or(HyperlightGuestError::new(
         ErrorCode::GuestError,
-        "Wasm runtime is not initialized".to_string(),
+        "No wasm store available".to_string(),
     ))?;
-    let linker = CUR_LINKER.lock();
-    let linker = linker.deref().as_ref().ok_or(HyperlightGuestError::new(
+    let instance = CUR_INSTANCE.lock();
+    let instance = instance.deref().as_ref().ok_or(HyperlightGuestError::new(
         ErrorCode::GuestError,
-        "impossible: wasm runtime has no valid linker".to_string(),
+        "No wasm instance available".to_string(),
     ))?;
-    let module = CUR_MODULE.lock();
-    let module = module.deref().as_ref().ok_or(HyperlightGuestError::new(
-        ErrorCode::GuestError,
-        "No wasm module loaded".to_string(),
-    ))?;
-    let mut store = Store::new(engine, ());
-    let instance = linker.instantiate(&mut store, module)?;
+
     let func = instance
-        .get_func(&mut store, &function_call.function_name)
+        .get_func(&mut *store, &function_call.function_name)
         .ok_or(HyperlightGuestError::new(
             ErrorCode::GuestError,
             "Function not found".to_string(),
         ))?;
+
     let mut w_params = vec![];
+    let mut allocated_addrs = vec![];
     for f_param in (function_call.parameters)
         .as_ref()
         .unwrap_or(&vec![])
         .iter()
     {
-        w_params.push(marshal::hl_param_to_val(
-            &mut store,
+        w_params.push(marshal::hl_param_to_val_with_tracking(
+            &mut *store,
             |ctx, name| instance.get_export(ctx, name),
             f_param,
+            &mut allocated_addrs,
         )?);
     }
     let is_void = ReturnType::Void == function_call.expected_return_type;
     let n_results = if is_void { 0 } else { 1 };
     let mut results = vec![Val::I32(0); n_results];
-    func.call(&mut store, &w_params, &mut results)?;
-    marshal::val_to_hl_result(
-        &mut store,
+    func.call(&mut *store, &w_params, &mut results)?;
+    let result = marshal::val_to_hl_result(
+        &mut *store,
         |ctx, name| instance.get_export(ctx, name),
         function_call.expected_return_type,
         &results,
+    );
+
+    // Free memory allocated during marshalling of hyperlight parameters into wasm parameters
+    marshal::free_allocated_addrs(
+        &mut *store,
+        |ctx, name| instance.get_export(ctx, name),
+        &allocated_addrs,
     )
+    .map_err(|e| {
+        HyperlightGuestError::new(
+            ErrorCode::GuestError,
+            format!("Failed to free memory allocated for params: {:?}", e),
+        )
+    })?;
+
+    // Free memory allocated by marshalling host function return values into wasm values
+    let hostfunc_addrs = take_hostfunc_allocated_addrs();
+    if !hostfunc_addrs.is_empty() {
+        marshal::free_allocated_addrs(
+            &mut *store,
+            |ctx, name| instance.get_export(ctx, name),
+            &hostfunc_addrs,
+        )
+        .map_err(|e| {
+            HyperlightGuestError::new(
+                ErrorCode::GuestError,
+                format!(
+                    "Failed to free memory allocated for host function returns: {:?}",
+                    e
+                ),
+            )
+        })?;
+    }
+
+    result
 }
 
 fn init_wasm_runtime() -> Result<Vec<u8>> {
@@ -124,8 +160,19 @@ fn load_wasm_module(function_call: &FunctionCall) -> Result<Vec<u8>> {
         &function_call.parameters.as_ref().unwrap()[1],
         &*CUR_ENGINE.lock(),
     ) {
+        let linker = CUR_LINKER.lock();
+        let linker = linker.deref().as_ref().ok_or(HyperlightGuestError::new(
+            ErrorCode::GuestError,
+            "impossible: wasm runtime has no valid linker".to_string(),
+        ))?;
+
         let module = unsafe { Module::deserialize(engine, wasm_bytes)? };
+        let mut store = Store::new(engine, ());
+        let instance = linker.instantiate(&mut store, &module)?;
+
         *CUR_MODULE.lock() = Some(module);
+        *CUR_STORE.lock() = Some(store);
+        *CUR_INSTANCE.lock() = Some(instance);
         Ok(get_flatbuffer_result::<i32>(0))
     } else {
         Err(HyperlightGuestError::new(
@@ -141,8 +188,19 @@ fn load_wasm_module_phys(function_call: &FunctionCall) -> Result<Vec<u8>> {
         &function_call.parameters.as_ref().unwrap()[1],
         &*CUR_ENGINE.lock(),
     ) {
+        let linker = CUR_LINKER.lock();
+        let linker = linker.deref().as_ref().ok_or(HyperlightGuestError::new(
+            ErrorCode::GuestError,
+            "impossible: wasm runtime has no valid linker".to_string(),
+        ))?;
+
         let module = unsafe { Module::deserialize_raw(engine, platform::map_buffer(*phys, *len))? };
+        let mut store = Store::new(engine, ());
+        let instance = linker.instantiate(&mut store, &module)?;
+
         *CUR_MODULE.lock() = Some(module);
+        *CUR_STORE.lock() = Some(store);
+        *CUR_INSTANCE.lock() = Some(instance);
         Ok(get_flatbuffer_result::<()>(()))
     } else {
         Err(HyperlightGuestError::new(

src/wasm_runtime/src/platform.rs

@@ -59,7 +59,7 @@ pub(crate) fn register_page_fault_handler() {
     // See AMD64 Architecture Programmer's Manual, Volume 2
     //    §8.2 Vectors, p. 245
     //      Table 8-1: Interrupt Vector Source and Cause
-    handler::handlers[14].store(page_fault_handler as usize as u64, Ordering::Release);
+    handler::HANDLERS[14].store(page_fault_handler as usize as u64, Ordering::Release);
 }
 
 // Wasmtime Embedding Interface
@@ -155,7 +155,7 @@ pub extern "C" fn wasmtime_init_traps(handler: wasmtime_trap_handler_t) -> i32 {
     // See AMD64 Architecture Programmer's Manual, Volume 2
     //    §8.2 Vectors, p. 245
     //      Table 8-1: Interrupt Vector Source and Cause
-    handler::handlers[6].store(wasmtime_trap_handler as usize as u64, Ordering::Release);
+    handler::HANDLERS[6].store(wasmtime_trap_handler as usize as u64, Ordering::Release);
     // TODO: Add handlers for any other traps that wasmtime needs,
     //       probably including at least some floating-point
     //       exceptions