Respect signals-based-traps in component trampolines (#12626)

This commit fixes a regression from #12547 where traps happening in
component trampolines weren't respecting the signals-based-traps
configuration of the `Engine` meaning that native CLIF traps could be
executed in environments which weren't configured to catch it. This
refactors the implementation internally to share more code with
`FuncEnvironment` through some traits to ensure that the same knobs and
plumbing are used to translate traps.

Author: alexcrichton

crates/cranelift/src/bounds_checks.rs

@@ -23,6 +23,7 @@ use crate::{
     Reachability,
     func_environ::FuncEnvironment,
     translate::{HeapData, TargetEnvironment},
+    trap::TranslateTrap,
 };
 use Reachability::*;
 use cranelift_codegen::{

crates/cranelift/src/compiler/component.rs

@@ -1,14 +1,18 @@
 //! Compilation support for the component model.
 
+use crate::func_environ::BuiltinFunctions;
+use crate::trap::TranslateTrap;
 use crate::{TRAP_CANNOT_LEAVE_COMPONENT, TRAP_INTERNAL_ASSERT, compiler::Compiler};
+use cranelift_codegen::cursor::FuncCursor;
 use cranelift_codegen::ir::condcodes::IntCC;
 use cranelift_codegen::ir::{self, InstBuilder, MemFlags, Value};
 use cranelift_codegen::isa::{CallConv, TargetIsa};
 use cranelift_frontend::FunctionBuilder;
 use wasmtime_environ::error::{Result, bail};
 use wasmtime_environ::{
-    Abi, CompiledFunctionBody, EntityRef, FuncKey, HostCall, PtrSize, TrapSentinel, Tunables,
-    WasmFuncType, WasmValType, component::*, fact::PREPARE_CALL_FIXED_PARAMS,
+    Abi, BuiltinFunctionIndex, CompiledFunctionBody, EntityRef, FuncKey, HostCall, PtrSize,
+    TrapSentinel, Tunables, WasmFuncType, WasmValType, component::*,
+    fact::PREPARE_CALL_FIXED_PARAMS,
 };
 
 struct TrampolineCompiler<'a> {
@@ -20,6 +24,7 @@ struct TrampolineCompiler<'a> {
     offsets: VMComponentOffsets<u8>,
     block0: ir::Block,
     signature: &'a WasmFuncType,
+    builtins: BuiltinFunctions,
 }
 
 /// What host functions can be called, used in `translate_hostcall` below.
@@ -112,6 +117,7 @@ impl<'a> TrampolineCompiler<'a> {
             offsets: VMComponentOffsets::new(isa.pointer_bytes(), component),
             block0,
             signature,
+            builtins: BuiltinFunctions::new(compiler),
         }
     }
 
@@ -1412,8 +1418,12 @@ impl<'a> TrampolineCompiler<'a> {
         self.builder.ins().return_(results);
     }
 
+    fn caller_vmctx(&self) -> ir::Value {
+        self.builder.func.dfg.block_params(self.block0)[1]
+    }
+
     fn raise_if_host_trapped(&mut self, succeeded: ir::Value) {
-        let caller_vmctx = self.builder.func.dfg.block_params(self.block0)[1];
+        let caller_vmctx = self.caller_vmctx();
         self.compiler
             .raise_if_host_trapped(&mut self.builder, caller_vmctx, succeeded);
     }
@@ -1533,9 +1543,45 @@ impl<'a> TrampolineCompiler<'a> {
             .builder
             .ins()
             .band_imm(flags, i64::from(FLAG_MAY_LEAVE));
-        self.builder
-            .ins()
-            .trapz(may_leave_bit, TRAP_CANNOT_LEAVE_COMPONENT);
+        let (mut traps, builder) = self.traps();
+        traps.trapz(builder, may_leave_bit, TRAP_CANNOT_LEAVE_COMPONENT);
+    }
+
+    fn traps(&mut self) -> (TrapTranslator<'_>, &mut FunctionBuilder<'a>) {
+        (
+            TrapTranslator {
+                compiler: self.compiler,
+                vmctx: self.builder.func.dfg.block_params(self.block0)[0],
+                builtins: &mut self.builtins,
+            },
+            &mut self.builder,
+        )
+    }
+}
+
+// Helper structure to implement `TranslateTrap`. This isn't possible to do
+// natively for `TrampolineCompiler` because it stores `FunctionBuilder`
+// internally. This differs from `FuncEnvironment` for core wasm which stores it
+// externally, hence the slightly different idioms to bridge here.
+struct TrapTranslator<'a> {
+    compiler: &'a Compiler,
+    vmctx: ir::Value,
+    builtins: &'a mut BuiltinFunctions,
+}
+
+impl TranslateTrap for TrapTranslator<'_> {
+    fn compiler(&self) -> &Compiler {
+        self.compiler
+    }
+    fn vmctx_val(&mut self, _: &mut FuncCursor<'_>) -> ir::Value {
+        self.vmctx
+    }
+    fn builtin_funcref(
+        &mut self,
+        builder: &mut FunctionBuilder<'_>,
+        index: BuiltinFunctionIndex,
+    ) -> ir::FuncRef {
+        self.builtins.load_builtin(builder.func, index)
     }
 }
 

crates/cranelift/src/func_environ.rs

@@ -1,12 +1,13 @@
 mod gc;
 pub(crate) mod stack_switching;
 
+use crate::BuiltinFunctionSignatures;
 use crate::compiler::Compiler;
 use crate::translate::{
     FuncTranslationStacks, GlobalVariable, Heap, HeapData, StructFieldsVec, TableData, TableSize,
     TargetEnvironment,
 };
-use crate::{BuiltinFunctionSignatures, TRAP_INTERNAL_ASSERT};
+use crate::trap::TranslateTrap;
 use cranelift_codegen::cursor::FuncCursor;
 use cranelift_codegen::ir::condcodes::{FloatCC, IntCC};
 use cranelift_codegen::ir::immediates::{Imm64, Offset32, V128Imm};
@@ -29,9 +30,9 @@ use wasmtime_environ::{
     BuiltinFunctionIndex, DataIndex, DefinedFuncIndex, ElemIndex, EngineOrModuleTypeIndex,
     FrameStateSlotBuilder, FrameValType, FuncIndex, FuncKey, GlobalConstValue, GlobalIndex,
     IndexType, Memory, MemoryIndex, Module, ModuleInternedTypeIndex, ModuleTranslation,
-    ModuleTypesBuilder, PtrSize, Table, TableIndex, TagIndex, TripleExt, Tunables, TypeConvert,
-    TypeIndex, VMOffsets, WasmCompositeInnerType, WasmFuncType, WasmHeapTopType, WasmHeapType,
-    WasmRefType, WasmResult, WasmValType,
+    ModuleTypesBuilder, PtrSize, Table, TableIndex, TagIndex, Tunables, TypeConvert, TypeIndex,
+    VMOffsets, WasmCompositeInnerType, WasmFuncType, WasmHeapTopType, WasmHeapType, WasmRefType,
+    WasmResult, WasmValType,
 };
 use wasmtime_environ::{FUNCREF_INIT_BIT, FUNCREF_MASK};
 
@@ -51,15 +52,19 @@ pub(crate) struct BuiltinFunctions {
 }
 
 impl BuiltinFunctions {
-    fn new(compiler: &Compiler) -> Self {
+    pub(crate) fn new(compiler: &Compiler) -> Self {
         Self {
             types: BuiltinFunctionSignatures::new(compiler),
             builtins: [None; BuiltinFunctionIndex::len() as usize],
             breakpoint_trampoline: None,
         }
     }
 
-    fn load_builtin(&mut self, func: &mut Function, builtin: BuiltinFunctionIndex) -> ir::FuncRef {
+    pub(crate) fn load_builtin(
+        &mut self,
+        func: &mut Function,
+        builtin: BuiltinFunctionIndex,
+    ) -> ir::FuncRef {
         let cache = &mut self.builtins[builtin.index() as usize];
         if let Some(f) = cache {
             return *f;
@@ -312,12 +317,6 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         })
     }
 
-    pub(crate) fn vmctx_val(&mut self, pos: &mut FuncCursor<'_>) -> ir::Value {
-        let pointer_type = self.pointer_type();
-        let vmctx = self.vmctx(&mut pos.func);
-        pos.ins().global_value(pointer_type, vmctx)
-    }
-
     fn get_table_copy_func(
         &mut self,
         func: &mut Function,
@@ -1034,34 +1033,6 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         self.global_load_with_memory_type(func, vmctx, offset, flags, self.pcc_vmctx_memtype)
     }
 
-    /// Helper to emit a conditional trap based on `trap_cond`.
-    ///
-    /// This should only be used if `self.clif_instruction_traps_enabled()` is
-    /// false, otherwise native CLIF instructions should be used instead.
-    pub fn conditionally_trap(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        trap_cond: ir::Value,
-        trap: ir::TrapCode,
-    ) {
-        assert!(!self.clif_instruction_traps_enabled());
-
-        let trap_block = builder.create_block();
-        builder.set_cold_block(trap_block);
-        let continuation_block = builder.create_block();
-
-        builder
-            .ins()
-            .brif(trap_cond, trap_block, &[], continuation_block, &[]);
-
-        builder.seal_block(trap_block);
-        builder.seal_block(continuation_block);
-
-        builder.switch_to_block(trap_block);
-        self.trap(builder, trap);
-        builder.switch_to_block(continuation_block);
-    }
-
     /// Helper used when `!self.clif_instruction_traps_enabled()` is enabled to
     /// test whether the divisor is zero.
     fn guard_zero_divisor(&mut self, builder: &mut FunctionBuilder, rhs: ir::Value) {
@@ -1395,6 +1366,26 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
     }
 }
 
+impl TranslateTrap for FuncEnvironment<'_> {
+    fn compiler(&self) -> &Compiler {
+        &self.compiler
+    }
+
+    fn vmctx_val(&mut self, pos: &mut FuncCursor<'_>) -> ir::Value {
+        let pointer_type = self.pointer_type();
+        let vmctx = self.vmctx(&mut pos.func);
+        pos.ins().global_value(pointer_type, vmctx)
+    }
+
+    fn builtin_funcref(
+        &mut self,
+        builder: &mut FunctionBuilder<'_>,
+        index: BuiltinFunctionIndex,
+    ) -> ir::FuncRef {
+        self.builtin_functions.load_builtin(builder.func, index)
+    }
+}
+
 #[derive(Default)]
 pub(crate) struct WasmEntities {
     /// Map from a Wasm global index from this module to its implementation in
@@ -4431,70 +4422,6 @@ impl FuncEnvironment<'_> {
         &*self.isa
     }
 
-    pub fn trap(&mut self, builder: &mut FunctionBuilder, trap: ir::TrapCode) {
-        match (
-            self.clif_instruction_traps_enabled(),
-            crate::clif_trap_to_env_trap(trap),
-        ) {
-            // If libcall traps are disabled or there's no wasmtime-defined trap
-            // code for this, then emit a native trap instruction.
-            (true, _) | (_, None) => {
-                builder.ins().trap(trap);
-            }
-            // ... otherwise with libcall traps explicitly enabled and a
-            // wasmtime-based trap code invoke the libcall to raise a trap and
-            // pass in our trap code. Leave a debug `unreachable` in place
-            // afterwards as a defense-in-depth measure.
-            (false, Some(trap)) => {
-                let libcall = self.builtin_functions.trap(&mut builder.func);
-                let vmctx = self.vmctx_val(&mut builder.cursor());
-                let trap_code = builder.ins().iconst(I8, i64::from(trap as u8));
-                builder.ins().call(libcall, &[vmctx, trap_code]);
-                let raise = self.builtin_functions.raise(&mut builder.func);
-                builder.ins().call(raise, &[vmctx]);
-                builder.ins().trap(TRAP_INTERNAL_ASSERT);
-            }
-        }
-    }
-
-    pub fn trapz(&mut self, builder: &mut FunctionBuilder, value: ir::Value, trap: ir::TrapCode) {
-        if self.clif_instruction_traps_enabled() {
-            builder.ins().trapz(value, trap);
-        } else {
-            let ty = builder.func.dfg.value_type(value);
-            let zero = builder.ins().iconst(ty, 0);
-            let cmp = builder.ins().icmp(IntCC::Equal, value, zero);
-            self.conditionally_trap(builder, cmp, trap);
-        }
-    }
-
-    pub fn trapnz(&mut self, builder: &mut FunctionBuilder, value: ir::Value, trap: ir::TrapCode) {
-        if self.clif_instruction_traps_enabled() {
-            builder.ins().trapnz(value, trap);
-        } else {
-            let ty = builder.func.dfg.value_type(value);
-            let zero = builder.ins().iconst(ty, 0);
-            let cmp = builder.ins().icmp(IntCC::NotEqual, value, zero);
-            self.conditionally_trap(builder, cmp, trap);
-        }
-    }
-
-    pub fn uadd_overflow_trap(
-        &mut self,
-        builder: &mut FunctionBuilder,
-        lhs: ir::Value,
-        rhs: ir::Value,
-        trap: ir::TrapCode,
-    ) -> ir::Value {
-        if self.clif_instruction_traps_enabled() {
-            builder.ins().uadd_overflow_trap(lhs, rhs, trap)
-        } else {
-            let (ret, overflow) = builder.ins().uadd_overflow(lhs, rhs);
-            self.conditionally_trap(builder, overflow, trap);
-            ret
-        }
-    }
-
     pub fn translate_sdiv(
         &mut self,
         builder: &mut FunctionBuilder,
@@ -4572,19 +4499,6 @@ impl FuncEnvironment<'_> {
         self.tunables.signals_based_traps && !self.is_pulley()
     }
 
-    /// Returns whether it's acceptable to have CLIF instructions natively trap,
-    /// such as division-by-zero.
-    ///
-    /// This is enabled if `signals_based_traps` is `true` or on
-    /// Pulley unconditionally since Pulley doesn't use hardware-based
-    /// traps in its runtime. However, if guest debugging is enabled,
-    /// then we cannot rely on Pulley traps and still need a libcall
-    /// to gain proper ownership of the store in the runtime's
-    /// debugger hooks.
-    pub fn clif_instruction_traps_enabled(&self) -> bool {
-        self.tunables.signals_based_traps || (self.is_pulley() && !self.tunables.debug_guest)
-    }
-
     /// Returns whether loads from the null address are allowed as signals of
     /// whether to trap or not.
     pub fn load_from_zero_allowed(&self) -> bool {
@@ -4594,11 +4508,6 @@ impl FuncEnvironment<'_> {
             || (self.clif_memory_traps_enabled() && self.heap_access_spectre_mitigation())
     }
 
-    /// Returns whether translation is happening for Pulley bytecode.
-    pub fn is_pulley(&self) -> bool {
-        self.isa.triple().is_pulley()
-    }
-
     /// Returns whether the current location is reachable.
     pub fn is_reachable(&self) -> bool {
         self.stacks.reachable()

crates/cranelift/src/func_environ/gc/enabled.rs

@@ -2,6 +2,7 @@ use super::{ArrayInit, GcCompiler};
 use crate::bounds_checks::BoundsCheck;
 use crate::func_environ::{Extension, FuncEnvironment};
 use crate::translate::{Heap, HeapData, StructFieldsVec, TargetEnvironment};
+use crate::trap::TranslateTrap;
 use crate::{Reachability, TRAP_INTERNAL_ASSERT};
 use cranelift_codegen::ir::immediates::Offset32;
 use cranelift_codegen::ir::{BlockArg, ExceptionTableData, ExceptionTableItem};

crates/cranelift/src/func_environ/gc/enabled/drc.rs

@@ -3,6 +3,7 @@
 
 use super::*;
 use crate::translate::TargetEnvironment;
+use crate::trap::TranslateTrap;
 use crate::{TRAP_INTERNAL_ASSERT, func_environ::FuncEnvironment};
 use cranelift_codegen::ir::condcodes::IntCC;
 use cranelift_codegen::ir::{self, InstBuilder};

crates/cranelift/src/func_environ/gc/enabled/null.rs

@@ -6,6 +6,7 @@
 
 use super::*;
 use crate::func_environ::FuncEnvironment;
+use crate::trap::TranslateTrap;
 use cranelift_codegen::ir::{self, InstBuilder};
 use cranelift_frontend::FunctionBuilder;
 use wasmtime_environ::VMSharedTypeIndex;

crates/cranelift/src/func_environ/stack_switching/instructions.rs

@@ -1,6 +1,7 @@
 use cranelift_codegen::ir::BlockArg;
 use itertools::{Either, Itertools};
 
+use crate::trap::TranslateTrap;
 use cranelift_codegen::ir::condcodes::*;
 use cranelift_codegen::ir::types::*;
 use cranelift_codegen::ir::{self, MemFlags};

crates/cranelift/src/lib.rs

@@ -43,6 +43,7 @@ mod compiler;
 mod debug;
 mod func_environ;
 mod translate;
+mod trap;
 
 use self::compiler::Compiler;
 

crates/cranelift/src/translate/code_translator.rs

@@ -80,6 +80,7 @@ use crate::translate::stack::{ControlStackFrame, ElseData};
 use crate::translate::translation_utils::{
     block_with_params, blocktype_params_results, f32_translation, f64_translation,
 };
+use crate::trap::TranslateTrap;
 use cranelift_codegen::ir::condcodes::{FloatCC, IntCC};
 use cranelift_codegen::ir::immediates::Offset32;
 use cranelift_codegen::ir::{

crates/cranelift/src/translate/table.rs

@@ -1,4 +1,5 @@
 use crate::func_environ::FuncEnvironment;
+use crate::trap::TranslateTrap;
 use cranelift_codegen::cursor::FuncCursor;
 use cranelift_codegen::ir::{self, InstBuilder, condcodes::IntCC, immediates::Imm64};
 use cranelift_codegen::isa::TargetIsa;