account for call_host_function reading the return value

Signed-off-by: Jorge Prendes <jorge.prendes@gmail.com>

Author: jprendes

src/hyperlight_wasm_macro/src/wasmguest.rs

@@ -87,12 +87,11 @@ fn emit_import_extern_decl<'a, 'b, 'c>(
             };
             quote! {
                 #li.func_wrap::<_, (#(#pts,)*), #rt>(#edkn, |_, (#(#pds,)*)| {
-                    call_host_function(
+                    let #ret = call_host_function::<Vec<u8>>(
                         #fname,
                         ::core::option::Option::Some(vec![#(#pus,)*]),
                         ::hyperlight_common::flatbuffer_wrappers::function_types::ReturnType::VecBytes,
                     ).unwrap();
-                    let #ret = ::hyperlight_guest_bin::host_comm::get_host_return_value::<Vec<u8>>().unwrap();
                     ::core::result::Result::Ok(#ur)
                 });
             }

src/wasm_runtime/src/hostfuncs.rs

@@ -15,7 +15,6 @@ limitations under the License.
 */
 
 use alloc::collections::BTreeMap;
-use alloc::format;
 use alloc::string::{String, ToString};
 use alloc::vec::Vec;
 
@@ -88,21 +87,16 @@ pub(crate) fn hostfunc_type(d: &HostFunctionDefinition, e: &Engine) -> Result<Fu
         ReturnType::Void => {}
         ReturnType::Int | ReturnType::UInt => results.push(ValType::I32),
         ReturnType::Long | ReturnType::ULong => results.push(ValType::I64),
-        /* hyperlight_guest_bin::host_function_call::get_host_value_return_as_{bool,float,double,string} are missing */
-        // TODO: this comment is outdated, these are now implemented. Implement the other types
+        ReturnType::Bool => results.push(ValType::I32),
+        ReturnType::Float => results.push(ValType::F32),
+        ReturnType::Double => results.push(ValType::F64),
+        ReturnType::String => results.push(ValType::I32),
+        // TODO: this comment about using i64 for VecBytes doesn't seem to match with what
+        //       hl_return_to_val was doing, check if this is still correct.
         /* For compatibility with old host, we return
          * a packed i64 with a (wasm32) pointer in the lower half and
          * a length in the upper half. */
         ReturnType::VecBytes => results.push(ValType::I64),
-        _ => {
-            return Err(HyperlightGuestError::new(
-                ErrorCode::GuestError,
-                format!(
-                    "Host function return type {:?} must be (u)int or (u)long, if present",
-                    d.return_type
-                ),
-            ));
-        }
     }
     Ok(FuncType::new(e, params, results))
 }
@@ -121,7 +115,39 @@ pub(crate) fn call(
             marshal::val_to_hl_param(&mut c, |c, n| c.get_export(n), s, t)
         })
         .collect();
-    call_host_function::<ReturnValue>(&d.function_name, Some(params), d.return_type)
+
+    let rv = call_host_function::<ReturnValue>(&d.function_name, Some(params), d.return_type)
         .expect("Host function call failed");
-    marshal::hl_return_to_val(&mut c, |c, n| c.get_export(n), &d.return_type, rs)
+
+    assert!(
+        return_type_from_val(&rv) == d.return_type,
+        "Host function return type mismatch"
+    );
+
+    if rs.is_empty() {
+        assert!(
+            d.return_type == ReturnType::Void,
+            "Host function return type mismatch"
+        );
+        return Ok(());
+    }
+
+    rs[0] = marshal::hl_return_to_val(&mut c, |c, n| c.get_export(n), rv)?;
+
+    Ok(())
+}
+
+fn return_type_from_val(val: &ReturnValue) -> ReturnType {
+    match val {
+        ReturnValue::Int(_) => ReturnType::Int,
+        ReturnValue::UInt(_) => ReturnType::UInt,
+        ReturnValue::Long(_) => ReturnType::Long,
+        ReturnValue::ULong(_) => ReturnType::ULong,
+        ReturnValue::Float(_) => ReturnType::Float,
+        ReturnValue::Double(_) => ReturnType::Double,
+        ReturnValue::String(_) => ReturnType::String,
+        ReturnValue::VecBytes(_) => ReturnType::VecBytes,
+        ReturnValue::Bool(_) => ReturnType::Bool,
+        ReturnValue::Void(_) => ReturnType::Void,
+    }
 }

src/wasm_runtime/src/marshal.rs

@@ -22,12 +22,11 @@ use alloc::vec::Vec;
 use alloc::{format, vec};
 
 use hyperlight_common::flatbuffer_wrappers::function_types::{
-    ParameterType, ParameterValue, ReturnType,
+    ParameterType, ParameterValue, ReturnType, ReturnValue,
 };
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result;
 use hyperlight_guest::error::{HyperlightGuestError, Result};
-use hyperlight_guest_bin::host_comm::get_host_return_value;
 use wasmtime::{AsContextMut, Extern, Val};
 
 fn malloc<C: AsContextMut>(
@@ -252,34 +251,29 @@ pub fn val_to_hl_param<'a, C: AsContextMut>(
 pub fn hl_return_to_val<C: AsContextMut>(
     ctx: &mut C,
     get_export: impl Fn(&mut C, &str) -> Option<Extern>,
-    rt: &ReturnType,
-    ret: &mut [Val],
-) -> Result<()> {
-    match rt {
-        ReturnType::Void => get_host_return_value::<()>()?,
-        ReturnType::Int => {
-            ret[0] = Val::I32(get_host_return_value::<i32>()?);
-        }
-        ReturnType::Long => {
-            ret[0] = Val::I64(get_host_return_value::<i64>()?);
-        }
-        ReturnType::UInt => {
-            ret[0] = Val::I32(get_host_return_value::<u32>()? as i32);
-        }
-        ReturnType::ULong => {
-            ret[0] = Val::I64(get_host_return_value::<u64>()? as i64);
+    rv: ReturnValue,
+) -> Result<Val> {
+    match rv {
+        ReturnValue::Int(i) => Ok(Val::I32(i)),
+        ReturnValue::UInt(u) => Ok(Val::I32(u as i32)),
+        ReturnValue::Long(l) => Ok(Val::I64(l)),
+        ReturnValue::ULong(u) => Ok(Val::I64(u as i64)),
+        ReturnValue::Bool(b) => Ok(Val::I32(if b { 1 } else { 0 })),
+        ReturnValue::Float(f) => Ok(Val::F32(f.to_bits())),
+        ReturnValue::Double(f) => Ok(Val::F64(f.to_bits())),
+        ReturnValue::String(s) => {
+            let s = CString::new(s.as_str()).unwrap();
+            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())?;
+            Ok(Val::I32(addr))
         }
-        /* hyperlight_guest_bin::host_comm::get_host_value_return_as_{bool,float,double,string} are missing */
-        // TODO: this comment is outdated, implement these
-        ReturnType::VecBytes => {
-            let b = get_host_return_value::<Vec<u8>>()?;
+        ReturnValue::VecBytes(b) => {
             let addr = malloc(ctx, &get_export, b.len())?;
             write(ctx, &get_export, addr, b.as_ref())?;
-            ret[0] = Val::I32(addr);
-        }
-        _ => {
-            panic!("unimplemented");
+            Ok(Val::I32(addr))
+            // TODO: check that the next parameter is the correct length
         }
+        ReturnValue::Void(()) => Ok(Val::I32(0)),
     }
-    Ok(())
 }

src/wasm_runtime/src/wasip1.rs

@@ -22,7 +22,7 @@ use alloc::vec::Vec;
 
 use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType};
 use hyperlight_guest::error::Result;
-use hyperlight_guest_bin::host_comm::{call_host_function, get_host_return_value};
+use hyperlight_guest_bin::host_comm::call_host_function;
 use wasmtime::{Caller, Extern, Linker};
 
 pub(crate) fn register_handlers<T>(linker: &mut Linker<T>) -> Result<()> {
@@ -63,16 +63,13 @@ pub(crate) fn register_handlers<T>(linker: &mut Linker<T>) -> Result<()> {
                 let Ok(str) = core::str::from_utf8(&string_bytes) else {
                     return -2;
                 };
-                let Ok(()) = call_host_function(
+                let Ok(written) = call_host_function::<i32>(
                     "HostPrint",
                     Some(Vec::from(&[ParameterValue::String(str.to_string())])),
                     ReturnType::Int,
                 ) else {
                     return -3;
                 };
-                let Ok(written) = get_host_return_value::<i32>() else {
-                    return -4;
-                };
                 total_written += written;
             }
             memory