Repository: bytecodealliance/wasm-tools
Pull request #1585 (2024-05-29 17:36:20)

`wasmparser`: Fix validation of the `return_call` family of instructions
We need to additionally check that the callee's results are an exact match of the caller's results. We were incorrectly allowing return calls that would push more values on the operand stack than would be returned. That is fine with a `call; return` sequence, where extra values on the stack are allowed to dangle, but not okay with a `return_call`. With a `return_call` it doesn't make sense because the callee might need a return pointer to put all its results into, but the caller can't supply one since its frame is going away, nor can the caller forward a return pointer that it received to the callee, since it might not return enough values to require a return pointer. This commit fixes the validation to match the spec and disallow `return_call`s that would leave dangling values on the operand stack.

cc https://github.com/bytecodealliance/wasmtime/issues/8704


Patch:
diff --git a/crates/wasmparser/src/validator/operators.rs b/crates/wasmparser/src/validator/operators.rs
index a3d1c313..396f6ac2 100644
--- a/crates/wasmparser/src/validator/operators.rs
+++ b/crates/wasmparser/src/validator/operators.rs
@@ -796,18 +796,22 @@ where
         }
     }
 
-    /// Validates a `call` instruction, ensuring that the function index is
-    /// in-bounds and the right types are on the stack to call the function.
-    fn check_call(&mut self, function_index: u32) -> Result<()> {
-        let ty = match self.resources.type_of_function(function_index) {
-            Some(i) => i,
+    fn type_of_function(&self, function_index: u32) -> Result<&'resources FuncType> {
+        match self.resources.type_of_function(function_index) {
+            Some(f) => Ok(f),
             None => {
                 bail!(
                     self.offset,
                     "unknown function {function_index}: function index out of bounds",
-                );
+                )
             }
-        };
+        }
+    }
+
+    /// Validates a `call` instruction, ensuring that the function index is
+    /// in-bounds and the right types are on the stack to call the function.
+    fn check_call(&mut self, function_index: u32) -> Result<()> {
+        let ty = self.type_of_function(function_index)?;
         self.check_call_ty(ty)
     }
 
@@ -864,6 +868,49 @@ where
         Ok(())
     }
 
+    /// Check that the function at the given index has the same result types as
+    /// the current function's results.
+    fn check_func_same_results(&self, function_index: u32) -> Result<()> {
+        let ty = self.type_of_function(function_index)?;
+        self.check_func_type_same_results(ty)
+    }
+
+    /// Check that the type at the given index has the same result types as the
+    /// current function's results.
+    fn check_func_type_index_same_results(&self, type_index: u32) -> Result<()> {
+        let ty = self.func_type_at(type_index)?;
+        self.check_func_type_same_results(ty)
+    }
+
+    /// Check that the given type has the same result types as the current
+    /// function's results.
+    fn check_func_type_same_results(&self, callee_ty: &FuncType) -> Result<()> {
+        let caller_rets = self.results(self.control[0].block_type)?;
+        if callee_ty.results().len() != caller_rets.len()
+            || !caller_rets
+                .zip(callee_ty.results())
+                .all(|(caller_ty, callee_ty)| self.resources.is_subtype(*callee_ty, caller_ty))
+        {
+            let caller_rets = self
+                .results(self.control[0].block_type)?
+                .map(|ty| format!("{ty}"))
+                .collect::<Vec<_>>()
+                .join(" ");
+            let callee_rets = callee_ty
+                .results()
+                .iter()
+                .map(|ty| format!("{ty}"))
+                .collect::<Vec<_>>()
+                .join(" ");
+            bail!(
+                self.offset,
+                "type mismatch: current function requires result type \
+                 [{caller_rets}] but callee returns [{callee_rets}]"
+            );
+        }
+        Ok(())
+    }
+
     /// Checks the validity of a common comparison operator.
     fn check_cmp_op(&mut self, ty: ValType) -> Result<()> {
         self.pop_operand(Some(ty))?;
@@ -1510,6 +1557,7 @@ where
     fn visit_return_call(&mut self, function_index: u32) -> Self::Output {
         self.check_call(function_index)?;
         self.check_return()?;
+        self.check_func_same_results(function_index)?;
         Ok(())
     }
     fn visit_call_ref(&mut self, type_index: u32) -> Self::Output {
@@ -1532,7 +1580,9 @@ where
     }
     fn visit_return_call_ref(&mut self, type_index: u32) -> Self::Output {
         self.visit_call_ref(type_index)?;
-        self.check_return()
+        self.check_return()?;
+        self.check_func_type_index_same_results(type_index)?;
+        Ok(())
     }
     fn visit_call_indirect(
         &mut self,
@@ -1549,9 +1599,10 @@ where
         self.check_call_indirect(index, table_index)?;
         Ok(())
     }
-    fn visit_return_call_indirect(&mut self, index: u32, table_index: u32) -> Self::Output {
-        self.check_call_indirect(index, table_index)?;
+    fn visit_return_call_indirect(&mut self, type_index: u32, table_index: u32) -> Self::Output {
+        self.check_call_indirect(type_index, table_index)?;
         self.check_return()?;
+        self.check_func_type_index_same_results(type_index)?;
         Ok(())
     }
     fn visit_drop(&mut self) -> Self::Output {


Test patch:
diff --git a/tests/local/function-references/return-call.wast b/tests/local/function-references/return-call.wast
new file mode 100644
index 00000000..c1c56079
--- /dev/null
+++ b/tests/local/function-references/return-call.wast
@@ -0,0 +1,37 @@
+;; Test that various return calls must exactly match the callee's returns, not
+;; simply leave the operand stack in a state where a `call; return` would
+;; otherwise be valid, but with some dangling stack values. Those dangling stack
+;; values are valid for regular calls, but not for return calls.
+
+(assert_invalid
+  (module
+    (func $f (result i32 i32) unreachable)
+    (func (result i32)
+      return_call $f
+    )
+  )
+  "type mismatch: current function requires result type [i32] but callee returns [i32 i32]"
+)
+
+(assert_invalid
+  (module
+    (type $ty (func (result i32 i32)))
+    (import "env" "table" (table $table 0 funcref))
+    (func (param i32) (result i32)
+      local.get 0
+      return_call_indirect $table (type $ty)
+    )
+  )
+  "type mismatch: current function requires result type [i32] but callee returns [i32 i32]"
+)
+
+(assert_invalid
+  (module
+    (type $ty (func (result i32 i32)))
+    (func (param funcref) (result i32)
+      local.get 0
+      return_call_ref $ty
+    )
+  )
+  "type mismatch: current function requires result type [i32] but callee returns [i32 i32]"
+)
diff --git a/tests/snapshots/local/function-references/return-call.wast.json b/tests/snapshots/local/function-references/return-call.wast.json
new file mode 100644
index 00000000..7ea31ec6
--- /dev/null
+++ b/tests/snapshots/local/function-references/return-call.wast.json
@@ -0,0 +1,26 @@
+{
+  "source_filename": "tests/local/function-references/return-call.wast",
+  "commands": [
+    {
+      "type": "assert_invalid",
+      "line": 7,
+      "filename": "return-call.0.wasm",
+      "text": "type mismatch: current function requires result type [i32] but callee returns [i32 i32]",
+      "module_type": "binary"
+    },
+    {
+      "type": "assert_invalid",
+      "line": 17,
+      "filename": "return-call.1.wasm",
+      "text": "type mismatch: current function requires result type [i32] but callee returns [i32 i32]",
+      "module_type": "binary"
+    },
+    {
+      "type": "assert_invalid",
+      "line": 29,
+      "filename": "return-call.2.wasm",
+      "text": "type mismatch: current function requires result type [i32] but callee returns [i32 i32]",
+      "module_type": "binary"
+    }
+  ]
+}
\ No newline at end of file