From 617be0198de30d0639d89f023a6fc66766c9d5ce Mon Sep 17 00:00:00 2001
From: Jef <jackefransham@gmail.com>
Date: Wed, 2 Jan 2019 12:18:24 +0100
Subject: [PATCH] Check type when resuming function (#152)

* Check type when resuming function

* Remove pub(crate)

* Update lib.rs
---
 src/func.rs       | 23 +++++++++++-------
 src/tests/host.rs | 60 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/src/func.rs b/src/func.rs
index b1d66db..bc650dd 100644
--- a/src/func.rs
+++ b/src/func.rs
@@ -293,16 +293,23 @@ impl<'args> FuncInvocation<'args> {
         return_val: Option<RuntimeValue>,
         externals: &'externals mut E,
     ) -> Result<Option<RuntimeValue>, ResumableError> {
-        match self.kind {
-            FuncInvocationKind::Internal(ref mut interpreter) => {
-                if !interpreter.state().is_resumable() {
-                    return Err(ResumableError::AlreadyStarted);
+        use crate::TrapKind;
+
+        if return_val.map(|v| v.value_type()) != self.resumable_value_type() {
+            return Err(ResumableError::Trap(Trap::new(
+                TrapKind::UnexpectedSignature,
+            )));
+        }
+
+        match &mut self.kind {
+            FuncInvocationKind::Internal(interpreter) => {
+                if interpreter.state().is_resumable() {
+                    Ok(interpreter.resume_execution(return_val, externals)?)
+                } else {
+                    Err(ResumableError::AlreadyStarted)
                 }
-                Ok(interpreter.resume_execution(return_val, externals)?)
-            }
-            FuncInvocationKind::Host { .. } => {
-                return Err(ResumableError::NotResumable);
             }
+            FuncInvocationKind::Host { .. } => Err(ResumableError::NotResumable),
         }
     }
 }
diff --git a/src/tests/host.rs b/src/tests/host.rs
index 81b2100..99b7416 100644
--- a/src/tests/host.rs
+++ b/src/tests/host.rs
@@ -302,6 +302,66 @@ fn resume_call_host_func() {
     );
 }
 
+#[test]
+fn resume_call_host_func_type_mismatch() {
+    fn resume_with_val(val: Option<RuntimeValue>) {
+        let module = parse_wat(
+            r#"
+            (module
+                (import "env" "trap_sub" (func $trap_sub (param i32 i32) (result i32)))
+
+                (func (export "test") (result i32)
+                    (call $trap_sub
+                        (i32.const 5)
+                        (i32.const 7)
+                    )
+                )
+            )
+            "#,
+        );
+
+        let mut env = TestHost::new();
+
+        let instance =
+            ModuleInstance::new(&module, &ImportsBuilder::new().with_resolver("env", &env))
+                .expect("Failed to instantiate module")
+                .assert_no_start();
+
+        let export = instance.export_by_name("test").unwrap();
+        let func_instance = export.as_func().unwrap();
+
+        let mut invocation = FuncInstance::invoke_resumable(&func_instance, &[]).unwrap();
+        let result = invocation.start_execution(&mut env);
+        match result {
+            Err(ResumableError::Trap(_)) => {}
+            _ => panic!(),
+        }
+
+        assert!(invocation.is_resumable());
+        let err = invocation.resume_execution(val, &mut env).unwrap_err();
+
+        match &err {
+            ResumableError::Trap(trap) => {
+                if let TrapKind::UnexpectedSignature = trap.kind() {
+                    return;
+                }
+            }
+            _ => {}
+        }
+
+        // If didn't return in the previous `match`...
+
+        panic!(
+            "Expected `ResumableError::Trap(Trap {{ kind: \
+             TrapKind::UnexpectedSignature, }})`, got `{:?}`",
+            err
+        )
+    }
+
+    resume_with_val(None);
+    resume_with_val(Some((-1i64).into()));
+}
+
 #[test]
 fn host_err() {
     let module = parse_wat(