[guest] improve WASM runtime memory management

- Reuse wasmtime Store and Instance across guest function calls instead
  of creating new one per call.
- Establish memory contract between host and guest.
 - Guest functions takes ownership of input parameters
 - Guest transfer ownership of return values
 - Host functions parameters are borrowed from guest
 - Host function return values are owned by guest and guest must free them.
- Component: Add post_return calls for proper WASM function cleanup
- Fix ABI mismatch in parameter of guest_dispatch_function

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

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

@@ -79,9 +79,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<()> {

src/wasm_runtime/src/marshal.rs

@@ -14,6 +14,34 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+//! Parameter and return value marshalling for WASM guest function calls.
+//!
+//! # Memory Management Contract
+//!
+//! This module implements a clear memory ownership model for both guest function calls and host function calls:
+//!
+//! ## Guest Function Parameters (Host → Guest)
+//! - When calling guest functions with String or VecBytes parameters, the host allocates memory
+//!   in the guest's memory space and passes pointers to the guest.
+//! - **The guest owns these allocations and must free them** when no longer needed using the
+//!   `free` function exported from the guest module.
+//!
+//! ## Guest Function Return Values (Guest → Host)  
+//! - When guest functions return String or VecBytes values, the guest allocates memory in its
+//!   own memory space and returns pointers to the host.
+//! - **The host takes ownership of these allocations and will free them** on the next VM entry
+//!   to prevent memory leaks.
+//!
+//! ## Host Function Parameters (Guest → Host)
+//! - When guest code calls host functions with String or VecBytes parameters, the guest passes
+//!   pointers to data in its own memory space.
+//! - **The guest retains ownership** of these allocations and remains responsible for freeing them.
+//!
+//! ## Host Function Return Values (Host → Guest)
+//! - When host functions return String or VecBytes values to the guest, the host allocates memory
+//!   in the guest's memory space and returns pointers.
+//! - **The guest owns these allocations and must free them** when no longer needed.
+
 extern crate alloc;
 
 use alloc::ffi::CString;
@@ -29,6 +57,28 @@ use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
 use hyperlight_guest::error::{HyperlightGuestError, Result};
 use wasmtime::{AsContextMut, Extern, Val};
 
+use spin::Mutex;
+
+// Global tracking for return value allocations that need to be freed on next VM entry
+static RETURN_VALUE_ALLOCATIONS: Mutex<Vec<i32>> = Mutex::new(Vec::new());
+
+/// Track a return value allocation that should be freed on the next VM entry
+fn track_return_value_allocation(addr: i32) {
+    RETURN_VALUE_ALLOCATIONS.lock().push(addr);
+}
+
+/// Free all tracked return value allocations from previous VM calls
+pub fn free_return_value_allocations<C: AsContextMut>(
+    ctx: &mut C,
+    get_export: &impl Fn(&mut C, &str) -> Option<Extern>,
+) -> Result<()> {
+    let mut allocations = RETURN_VALUE_ALLOCATIONS.lock();
+    for addr in allocations.drain(..) {
+        free(ctx, get_export, addr)?;
+    }
+    Ok(())
+}
+
 fn malloc<C: AsContextMut>(
     ctx: &mut C,
     get_export: &impl Fn(&mut C, &str) -> Option<Extern>,
@@ -46,6 +96,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,6 +191,11 @@ fn read_cstr<C: AsContextMut>(
     })
 }
 
+/// Convert a hyperlight parameter value to a wasmtime value.
+///
+/// For String and VecBytes parameter types, this allocates memory in the guest's memory space
+/// and returns a pointer. The guest function is responsible for freeing this memory when it is no
+/// longer needed using the `free` function exported from the guest module.
 pub fn hl_param_to_val<C: AsContextMut>(
     mut ctx: C,
     get_export: impl Fn(&mut C, &str) -> Option<Extern>,
@@ -155,6 +225,11 @@ pub fn hl_param_to_val<C: AsContextMut>(
     }
 }
 
+/// Convert guest function return values to hyperlight return value.
+///
+/// For String and VecBytes return types, the guest has allocated memory in its own memory space
+/// and returned pointers. The host takes ownership of these allocations and tracks them for
+/// automatic cleanup on the next VM entry to prevent memory leaks.
 pub fn val_to_hl_result<C: AsContextMut>(
     mut ctx: C,
     get_export: impl Fn(&mut C, &str) -> Option<Extern>,
@@ -172,15 +247,21 @@ pub fn val_to_hl_result<C: AsContextMut>(
         /* todo: get_flatbuffer_result_from_bool is missing */
         (ReturnType::Float, Val::F32(f)) => Ok(get_flatbuffer_result::<f32>(f32::from_bits(f))),
         (ReturnType::Double, Val::F64(f)) => Ok(get_flatbuffer_result::<f64>(f64::from_bits(f))),
-        (ReturnType::String, Val::I32(p)) => Ok(get_flatbuffer_result::<&str>(
-            read_cstr(&mut ctx, &get_export, p)?.to_str().map_err(|e| {
-                HyperlightGuestError::new(
-                    ErrorCode::GuestError,
-                    format!("non-UTF-8 c string in guest function return: {}", e),
-                )
-            })?,
-        )),
+        (ReturnType::String, Val::I32(p)) => {
+            // Track this allocation so it can be freed on next VM entry
+            track_return_value_allocation(p);
+            Ok(get_flatbuffer_result::<&str>(
+                read_cstr(&mut ctx, &get_export, p)?.to_str().map_err(|e| {
+                    HyperlightGuestError::new(
+                        ErrorCode::GuestError,
+                        format!("non-UTF-8 c string in guest function return: {}", e),
+                    )
+                })?,
+            ))
+        }
         (ReturnType::VecBytes, Val::I32(ret)) => {
+            // Track this allocation so it can be freed on next VM entry
+            track_return_value_allocation(ret);
             let mut size_bytes = [0; 4];
             read(&mut ctx, &get_export, ret, &mut size_bytes)?;
             let size = i32::from_le_bytes(size_bytes);
@@ -198,6 +279,11 @@ pub fn val_to_hl_result<C: AsContextMut>(
     }
 }
 
+/// Convert guest-provided WASM values to hyperlight parameters for host function calls.
+///
+/// For String and VecBytes parameter types, the guest passes pointers to data in its own
+/// memory space. The guest retains ownership of these allocations and remains responsible
+/// for freeing them. This function only reads the data without taking ownership.
 pub fn val_to_hl_param<'a, C: AsContextMut>(
     ctx: &mut C,
     get_export: impl Fn(&mut C, &str) -> Option<Extern>,
@@ -248,6 +334,11 @@ pub fn val_to_hl_param<'a, C: AsContextMut>(
     }
 }
 
+/// Convert a hyperlight return value to a wasmtime value for host function returns.
+///
+/// For String and VecBytes return types, this allocates memory in the guest's memory space
+/// and returns a pointer. The guest owns these allocations and must free them when no longer needed
+/// using the `free` function exported from the guest module.
 pub fn hl_return_to_val<C: AsContextMut>(
     ctx: &mut C,
     get_export: impl Fn(&mut C, &str) -> Option<Extern>,

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::{
@@ -34,57 +34,65 @@ use wasmtime::{Config, Engine, Linker, Module, Store, Val};
 
 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)?;
+
+    // Free any return value allocations from the previous VM call
+    // This implements the memory model where hyperlight owns return values
+    // and frees them on the next VM entry
+    marshal::free_return_value_allocations(&mut *store, &|ctx, name| {
+        instance.get_export(ctx, name)
+    })?;
+
     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![];
     for f_param in (function_call.parameters)
         .as_ref()
         .unwrap_or(&vec![])
         .iter()
     {
         w_params.push(marshal::hl_param_to_val(
-            &mut store,
+            &mut *store,
             |ctx, name| instance.get_export(ctx, name),
             f_param,
         )?);
     }
     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,
-    )
+    );
+
+    result
 }
 
 fn init_wasm_runtime() -> Result<Vec<u8>> {
@@ -124,8 +132,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 +160,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