fix: add returned type(s) check for functions created with Extern::func (untyped one)

Author: WhaleKit

crates/tinywasm/src/error.rs

@@ -29,6 +29,9 @@ pub enum Error {
     /// The store is not the one that the module instance was instantiated in
     InvalidStore,
 
+    /// A host function did not return results expected of it's type
+    InvalidHostFnReturn,
+
     #[cfg(feature = "std")]
     /// An I/O error occurred
     Io(crate::std::io::Error),
@@ -185,6 +188,7 @@ impl Display for Error {
             Self::UnsupportedFeature(feature) => write!(f, "unsupported feature: {feature}"),
             Self::FuncDidNotReturn => write!(f, "function did not return"),
             Self::InvalidStore => write!(f, "invalid store"),
+            Self::InvalidHostFnReturn=>write!(f, "Host function returned wrong type(s)")
         }
     }
 }

crates/tinywasm/src/func.rs

@@ -51,9 +51,9 @@ impl FuncHandle {
         let func_inst = store.get_func(self.addr);
         let wasm_func = match &func_inst.func {
             Function::Host(host_func) => {
-                let func = &host_func.clone().func;
+                let host_func = host_func.clone();
                 let ctx = FuncContext { store, module_addr: self.module_addr };
-                return (func)(ctx, params);
+                return host_func.call(ctx, params);
             }
             Function::Wasm(wasm_func) => wasm_func,
         };

crates/tinywasm/src/imports.rs

@@ -6,7 +6,7 @@ use alloc::vec::Vec;
 use core::fmt::Debug;
 
 use crate::func::{FromWasmValueTuple, IntoWasmValueTuple, ValTypesFromTuple};
-use crate::{log, LinkingError, MemoryRef, MemoryRefMut, Result};
+use crate::{cold, log, Error, LinkingError, MemoryRef, MemoryRefMut, Result};
 use tinywasm_types::*;
 
 /// The internal representation of a function
@@ -42,11 +42,12 @@ impl HostFunction {
 
     /// Call the function
     pub fn call(&self, ctx: FuncContext<'_>, args: &[WasmValue]) -> Result<Vec<WasmValue>> {
-        (self.func)(ctx, args)
+        (self.func)(ctx, args, &self.ty.results)
     }
 }
 
-pub(crate) type HostFuncInner = Box<dyn Fn(FuncContext<'_>, &[WasmValue]) -> Result<Vec<WasmValue>>>;
+// third argument - expected return types, for validation
+pub(crate) type HostFuncInner = Box<dyn Fn(FuncContext<'_>, &[WasmValue], &[ValType]) -> Result<Vec<WasmValue>>>;
 
 /// The context of a host-function call
 #[derive(Debug)]
@@ -139,7 +140,19 @@ impl Extern {
         ty: &tinywasm_types::FuncType,
         func: impl Fn(FuncContext<'_>, &[WasmValue]) -> Result<Vec<WasmValue>> + 'static,
     ) -> Self {
-        Self::Function(Function::Host(Rc::new(HostFunction { func: Box::new(func), ty: ty.clone() })))
+        let wrapper = move |ctx: FuncContext<'_>, params: &[WasmValue], res_ty: &[ValType]| {
+            let ret = func(ctx, params)?;
+            let types_match = res_ty.iter().cloned().eq(ret.iter().map(WasmValue::val_type));
+            if types_match {
+                Ok(ret)
+            } else {
+                cold();
+                let got_types = ret.iter().map(WasmValue::val_type).collect::<Box<[_]>>();
+                log::error!("Return from host function type mismatch: expected {:?}, got {:?}", res_ty, got_types);
+                Err(Error::InvalidHostFnReturn)
+            }
+        };
+        Self::Function(Function::Host(Rc::new(HostFunction { func: Box::new(wrapper), ty: ty.clone() })))
     }
 
     /// Create a new typed function import
@@ -148,7 +161,7 @@ impl Extern {
         P: FromWasmValueTuple + ValTypesFromTuple,
         R: IntoWasmValueTuple + ValTypesFromTuple + Debug,
     {
-        let inner_func = move |ctx: FuncContext<'_>, args: &[WasmValue]| -> Result<Vec<WasmValue>> {
+        let inner_func = move |ctx: FuncContext<'_>, args: &[WasmValue], _: &[ValType]| -> Result<Vec<WasmValue>> {
             let args = P::from_wasm_value_tuple(args)?;
             let result = func(ctx, args)?;
             Ok(result.into_wasm_value_tuple().to_vec())

crates/tinywasm/src/interpreter/executor.rs

@@ -331,7 +331,7 @@ impl<'store, 'stack> Executor<'store, 'stack> {
                 let func = &host_func.clone();
                 let params = self.stack.values.pop_params(&host_func.ty.params);
                 let res =
-                    (func.func)(FuncContext { store: self.store, module_addr: self.module.id() }, &params).to_cf()?;
+                    func.call(FuncContext { store: self.store, module_addr: self.module.id() }, &params).to_cf()?;
                 self.stack.values.extend_from_wasmvalues(&res);
                 self.cf.incr_instr_ptr();
                 return ControlFlow::Continue(());
@@ -370,7 +370,7 @@ impl<'store, 'stack> Executor<'store, 'stack> {
                 let host_func = host_func.clone();
                 let params = self.stack.values.pop_params(&host_func.ty.params);
                 let res =
-                    match (host_func.func)(FuncContext { store: self.store, module_addr: self.module.id() }, &params) {
+                    match host_func.call(FuncContext { store: self.store, module_addr: self.module.id() }, &params) {
                         Ok(res) => res,
                         Err(e) => return ControlFlow::Break(Some(e)),
                     };