fix: improve WASM runtime memory management

- Add memory tracking for allocated parameters in guest function calls
- Implement automatic cleanup of allocated memory after function calls
- Component: Add post_return calls for proper WASM function cleanup
- Fix memory leaks in parameter marshaling

This ensures guest function parameters are properly freed after use.

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/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); // Track for later cleanup
             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); // Track for later cleanup
             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>,

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,71 @@ 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(
+    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,
+    );
+
+    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),
+        )
+    })?;
+    result
 }
 
 fn init_wasm_runtime() -> Result<Vec<u8>> {
@@ -124,8 +138,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 +166,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