From 5bc82b431dc14412a1a03c157f9fcded42fd84c4 Mon Sep 17 00:00:00 2001
From: Jamey Sharp <jsharp@fastly.com>
Date: Tue, 18 Jun 2024 16:19:43 -0700
Subject: [PATCH] wasmtime: Consume fuel on all return paths

As far as I can tell, when functions use a `return` instruction rather
than falling off the end, fuel was not being tracked correctly. The
`fuel_function_exit` method was only called from
`after_translate_function`, which is only called once all the
instructions in the function have been translated, not at each return.

In this commit I switched to calling `fuel_function_exit` from
`handle_before_return` instead, alongside some of the `wmemcheck` hooks.
That should ensure that it happens on every function exit, regardless of
what form that exit takes.

I think `after_translate_function` is fundamentally difficult to use
correctly, and it wasn't used for anything else, so I've also removed it
in this commit. And I did a minor cleanup at the site of one of the
calls to `handle_before_return` while I was looking at it.
---
 cranelift/wasm/src/code_translator.rs |  7 ++----
 cranelift/wasm/src/environ/spec.rs    | 10 --------
 cranelift/wasm/src/func_translator.rs |  1 -
 crates/cranelift/src/func_environ.rs  | 34 ++++++++++++---------------
 4 files changed, 17 insertions(+), 35 deletions(-)

diff --git a/cranelift/wasm/src/code_translator.rs b/cranelift/wasm/src/code_translator.rs
index 365a4877149e..b535e2fd83d7 100644
--- a/cranelift/wasm/src/code_translator.rs
+++ b/cranelift/wasm/src/code_translator.rs
@@ -577,13 +577,10 @@ pub fn translate_operator<FE: FuncEnvironment + ?Sized>(
             state.reachable = false;
         }
         Operator::Return => {
-            let return_count = {
-                let frame = &mut state.control_stack[0];
-                frame.num_return_values()
-            };
+            let return_count = state.control_stack[0].num_return_values();
             {
                 let return_args = state.peekn_mut(return_count);
-                environ.handle_before_return(&return_args, builder);
+                environ.handle_before_return(return_args, builder);
                 bitcast_wasm_returns(environ, return_args, builder);
                 builder.ins().return_(return_args);
             }
diff --git a/cranelift/wasm/src/environ/spec.rs b/cranelift/wasm/src/environ/spec.rs
index 7fea12696f85..d2cf84d79dcf 100644
--- a/cranelift/wasm/src/environ/spec.rs
+++ b/cranelift/wasm/src/environ/spec.rs
@@ -584,16 +584,6 @@ pub trait FuncEnvironment: TargetEnvironment {
         Ok(())
     }
 
-    /// Optional callback for the `FunctionEnvironment` performing this translation to perform work
-    /// after the function body is translated.
-    fn after_translate_function(
-        &mut self,
-        _builder: &mut FunctionBuilder,
-        _state: &FuncTranslationState,
-    ) -> WasmResult<()> {
-        Ok(())
-    }
-
     /// Whether or not to force relaxed simd instructions to have deterministic
     /// lowerings meaning they will produce the same results across all hosts,
     /// regardless of the cost to performance.
diff --git a/cranelift/wasm/src/func_translator.rs b/cranelift/wasm/src/func_translator.rs
index 15bc960e72f7..b7572783b30e 100644
--- a/cranelift/wasm/src/func_translator.rs
+++ b/cranelift/wasm/src/func_translator.rs
@@ -264,7 +264,6 @@ fn parse_function_body<FE: FuncEnvironment + ?Sized>(
         translate_operator(validator, &op, builder, state, environ)?;
         environ.after_translate_operator(&op, builder, state)?;
     }
-    environ.after_translate_function(builder, state)?;
     let pos = reader.original_position();
     validator.finish(pos)?;
 
diff --git a/crates/cranelift/src/func_environ.rs b/crates/cranelift/src/func_environ.rs
index 0dbffa4b4e11..4c927f4cdc1e 100644
--- a/crates/cranelift/src/func_environ.rs
+++ b/crates/cranelift/src/func_environ.rs
@@ -2810,15 +2810,23 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         Ok(())
     }
 
-    fn after_translate_function(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        state: &FuncTranslationState,
-    ) -> WasmResult<()> {
-        if self.tunables.consume_fuel && state.reachable() {
+    fn handle_before_return(&mut self, retvals: &[ir::Value], builder: &mut FunctionBuilder) {
+        // Ignore unused-argument warnings
+        let _ = retvals;
+
+        if self.tunables.consume_fuel {
             self.fuel_function_exit(builder);
         }
-        Ok(())
+
+        #[cfg(feature = "wmemcheck")]
+        if self.wmemcheck {
+            let func_name = self.current_func_name(builder);
+            if func_name == Some("malloc") {
+                self.hook_malloc_exit(builder, retvals);
+            } else if func_name == Some("free") {
+                self.hook_free_exit(builder);
+            }
+        }
     }
 
     fn relaxed_simd_deterministic(&self) -> bool {
@@ -2849,18 +2857,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         self.isa.has_x86_pmaddubsw_lowering()
     }
 
-    #[cfg(feature = "wmemcheck")]
-    fn handle_before_return(&mut self, retvals: &[ir::Value], builder: &mut FunctionBuilder) {
-        if self.wmemcheck {
-            let func_name = self.current_func_name(builder);
-            if func_name == Some("malloc") {
-                self.hook_malloc_exit(builder, retvals);
-            } else if func_name == Some("free") {
-                self.hook_free_exit(builder);
-            }
-        }
-    }
-
     #[cfg(feature = "wmemcheck")]
     fn before_load(
         &mut self,