Trap registry

src/bin/wasm2obj.rs

@@ -233,7 +233,8 @@ fn handle_module(
         )
     };
 
-    let (compilation, relocations, address_transform, value_ranges, stack_slots) =
+    // TODO: use the traps information
+    let (compilation, relocations, address_transform, value_ranges, stack_slots, _traps) =
         Cranelift::compile_module(
             &module,
             lazy_function_body_inputs,

wasmtime-environ/src/cache.rs

@@ -1,5 +1,5 @@
 use crate::address_map::{ModuleAddressMap, ValueLabelsRanges};
-use crate::compilation::{Compilation, Relocations};
+use crate::compilation::{Compilation, Relocations, Traps};
 use crate::module::Module;
 use crate::module_environ::FunctionBodyData;
 use core::hash::Hasher;
@@ -64,6 +64,7 @@ pub struct ModuleCacheData {
     address_transforms: ModuleAddressMap,
     value_ranges: ValueLabelsRanges,
     stack_slots: PrimaryMap<DefinedFuncIndex, ir::StackSlots>,
+    traps: Traps,
 }
 
 type ModuleCacheDataTupleType = (
@@ -72,6 +73,7 @@ type ModuleCacheDataTupleType = (
     ModuleAddressMap,
     ValueLabelsRanges,
     PrimaryMap<DefinedFuncIndex, ir::StackSlots>,
+    Traps,
 );
 
 struct Sha256Hasher(Sha256);
@@ -231,6 +233,7 @@ impl ModuleCacheData {
             address_transforms: data.2,
             value_ranges: data.3,
             stack_slots: data.4,
+            traps: data.5,
         }
     }
 
@@ -241,6 +244,7 @@ impl ModuleCacheData {
             self.address_transforms,
             self.value_ranges,
             self.stack_slots,
+            self.traps,
         )
     }
 }

wasmtime-environ/src/cache/tests.rs

@@ -1,7 +1,7 @@
 use super::config::tests::test_prolog;
 use super::*;
 use crate::address_map::{FunctionAddressMap, InstructionAddressMap};
-use crate::compilation::{CodeAndJTOffsets, Relocation, RelocationTarget};
+use crate::compilation::{CodeAndJTOffsets, Relocation, RelocationTarget, TrapInformation};
 use crate::module::{MemoryPlan, MemoryStyle, Module};
 use cranelift_codegen::{binemit, ir, isa, settings, ValueLocRange};
 use cranelift_entity::EntityRef;
@@ -330,11 +330,24 @@ fn new_module_cache_data(rng: &mut impl Rng) -> ModuleCacheData {
         })
         .collect();
 
+    let traps = (0..rng.gen_range(0, 0xd))
+        .map(|i| {
+            ((i..i + rng.gen_range(0, 4))
+                .map(|_| TrapInformation {
+                    code_offset: rng.gen(),
+                    source_loc: ir::SourceLoc::new(rng.gen()),
+                    trap_code: ir::TrapCode::StackOverflow,
+                })
+                .collect())
+        })
+        .collect();
+
     ModuleCacheData::from_tuple((
         Compilation::new(funcs),
         relocs,
         trans,
         value_ranges,
         stack_slots,
+        traps,
     ))
 }

wasmtime-environ/src/compilation.rs

@@ -129,6 +129,20 @@ pub enum RelocationTarget {
 /// Relocations to apply to function bodies.
 pub type Relocations = PrimaryMap<DefinedFuncIndex, Vec<Relocation>>;
 
+/// Information about trap.
+#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
+pub struct TrapInformation {
+    /// The offset of the trapping instruction in native code. It is relative to the beginning of the function.
+    pub code_offset: binemit::CodeOffset,
+    /// Location of trapping instruction in WebAssembly binary module.
+    pub source_loc: ir::SourceLoc,
+    /// Code of the trap.
+    pub trap_code: ir::TrapCode,
+}
+
+/// Information about traps associated with the functions where the traps are placed.
+pub type Traps = PrimaryMap<DefinedFuncIndex, Vec<TrapInformation>>;
+
 /// An error while compiling WebAssembly to machine code.
 #[derive(Fail, Debug)]
 pub enum CompileError {
@@ -156,6 +170,7 @@ pub trait Compiler {
             ModuleAddressMap,
             ValueLabelsRanges,
             PrimaryMap<DefinedFuncIndex, ir::StackSlots>,
+            Traps,
         ),
         CompileError,
     >;

wasmtime-environ/src/cranelift.rs

@@ -6,6 +6,7 @@ use crate::address_map::{
 use crate::cache::{ModuleCacheData, ModuleCacheEntry};
 use crate::compilation::{
     CodeAndJTOffsets, Compilation, CompileError, Relocation, RelocationTarget, Relocations,
+    TrapInformation, Traps,
 };
 use crate::func_environ::{
     get_func_name, get_imported_memory32_grow_name, get_imported_memory32_size_name,
@@ -103,6 +104,31 @@ impl RelocSink {
     }
 }
 
+struct TrapSink {
+    pub traps: Vec<TrapInformation>,
+}
+
+impl TrapSink {
+    fn new() -> Self {
+        Self { traps: Vec::new() }
+    }
+}
+
+impl binemit::TrapSink for TrapSink {
+    fn trap(
+        &mut self,
+        code_offset: binemit::CodeOffset,
+        source_loc: ir::SourceLoc,
+        trap_code: ir::TrapCode,
+    ) {
+        self.traps.push(TrapInformation {
+            code_offset,
+            source_loc,
+            trap_code,
+        });
+    }
+}
+
 fn get_function_address_map<'data>(
     context: &Context,
     data: &FunctionBodyData<'data>,
@@ -161,6 +187,7 @@ impl crate::compilation::Compiler for Cranelift {
             ModuleAddressMap,
             ValueLabelsRanges,
             PrimaryMap<DefinedFuncIndex, ir::StackSlots>,
+            Traps,
         ),
         CompileError,
     > {
@@ -180,6 +207,7 @@ impl crate::compilation::Compiler for Cranelift {
                 let mut address_transforms = PrimaryMap::with_capacity(function_body_inputs.len());
                 let mut value_ranges = PrimaryMap::with_capacity(function_body_inputs.len());
                 let mut stack_slots = PrimaryMap::with_capacity(function_body_inputs.len());
+                let mut traps = PrimaryMap::with_capacity(function_body_inputs.len());
 
                 function_body_inputs
                     .into_iter()
@@ -207,7 +235,7 @@ impl crate::compilation::Compiler for Cranelift {
 
                         let mut code_buf: Vec<u8> = Vec::new();
                         let mut reloc_sink = RelocSink::new(func_index);
-                        let mut trap_sink = binemit::NullTrapSink {};
+                        let mut trap_sink = TrapSink::new();
                         let mut stackmap_sink = binemit::NullStackmapSink {};
                         context
                             .compile_and_emit(
@@ -247,12 +275,21 @@ impl crate::compilation::Compiler for Cranelift {
                             address_transform,
                             ranges,
                             stack_slots,
+                            trap_sink.traps,
                         ))
                     })
                     .collect::<Result<Vec<_>, CompileError>>()?
                     .into_iter()
                     .for_each(
-                        |(function, func_jt_offsets, relocs, address_transform, ranges, sss)| {
+                        |(
+                            function,
+                            func_jt_offsets,
+                            relocs,
+                            address_transform,
+                            ranges,
+                            sss,
+                            function_traps,
+                        )| {
                             functions.push(CodeAndJTOffsets {
                                 body: function,
                                 jt_offsets: func_jt_offsets,
@@ -263,6 +300,7 @@ impl crate::compilation::Compiler for Cranelift {
                             }
                             value_ranges.push(ranges.unwrap_or(std::collections::HashMap::new()));
                             stack_slots.push(sss);
+                            traps.push(function_traps);
                         },
                     );
 
@@ -274,6 +312,7 @@ impl crate::compilation::Compiler for Cranelift {
                     address_transforms,
                     value_ranges,
                     stack_slots,
+                    traps,
                 ));
                 cache_entry.update_data(&data);
                 data

wasmtime-environ/src/lib.rs

@@ -57,6 +57,7 @@ pub use crate::address_map::{
 pub use crate::cache::{create_new_config as cache_create_new_config, init as cache_init};
 pub use crate::compilation::{
     Compilation, CompileError, Compiler, Relocation, RelocationTarget, Relocations,
+    TrapInformation, Traps,
 };
 pub use crate::cranelift::Cranelift;
 pub use crate::func_environ::BuiltinFunctionIndex;

wasmtime-jit/src/compiler.rs

@@ -4,6 +4,7 @@ use super::HashMap;
 use crate::code_memory::CodeMemory;
 use crate::instantiate::SetupError;
 use crate::target_tunables::target_tunables;
+use core::convert::TryFrom;
 use cranelift_codegen::ir::InstBuilder;
 use cranelift_codegen::isa::{TargetFrontendConfig, TargetIsa};
 use cranelift_codegen::Context;
@@ -17,9 +18,12 @@ use std::vec::Vec;
 use wasmtime_debug::{emit_debugsections_image, DebugInfoData};
 use wasmtime_environ::{
     Compilation, CompileError, Compiler as _C, FunctionBodyData, Module, ModuleVmctxInfo,
-    Relocations, Tunables, VMOffsets,
+    Relocations, Traps, Tunables, VMOffsets,
+};
+use wasmtime_runtime::{
+    get_mut_trap_registry, InstantiationError, SignatureRegistry, TrapRegistrationGuard,
+    VMFunctionBody,
 };
-use wasmtime_runtime::{InstantiationError, SignatureRegistry, VMFunctionBody};
 
 /// A WebAssembly code JIT compiler.
 ///
@@ -33,6 +37,7 @@ pub struct Compiler {
     isa: Box<dyn TargetIsa>,
 
     code_memory: CodeMemory,
+    trap_registration_guards: Vec<TrapRegistrationGuard>,
     trampoline_park: HashMap<*const VMFunctionBody, *const VMFunctionBody>,
     signatures: SignatureRegistry,
 
@@ -46,6 +51,7 @@ impl Compiler {
         Self {
             isa,
             code_memory: CodeMemory::new(),
+            trap_registration_guards: Vec::new(),
             trampoline_park: HashMap::new(),
             signatures: SignatureRegistry::new(),
             fn_builder_ctx: FunctionBuilderContext::new(),
@@ -53,6 +59,20 @@ impl Compiler {
     }
 }
 
+impl Drop for Compiler {
+    fn drop(&mut self) {
+        // We must deregister traps before freeing the code memory.
+        // Otherwise, we have a race:
+        // - Compiler #1 dropped code memory, but hasn't deregistered the trap yet
+        // - Compiler #2 allocated code memory and tries to register a trap,
+        //   but the trap at certain address happens to be already registered,
+        //   since Compiler #1 hasn't deregistered it yet => assertion in trap registry fails.
+        // Having a custom drop implementation we are independent from the field order
+        // in the struct what reduces potential human error.
+        self.trap_registration_guards.clear();
+    }
+}
+
 #[cfg(feature = "lightbeam")]
 type DefaultCompiler = wasmtime_environ::lightbeam::Lightbeam;
 #[cfg(not(feature = "lightbeam"))]
@@ -84,7 +104,7 @@ impl Compiler {
         ),
         SetupError,
     > {
-        let (compilation, relocations, address_transform, value_ranges, stack_slots) =
+        let (compilation, relocations, address_transform, value_ranges, stack_slots, traps) =
             DefaultCompiler::compile_module(
                 module,
                 function_body_inputs,
@@ -101,6 +121,12 @@ impl Compiler {
                 )))
             })?;
 
+        register_traps(
+            &allocated_functions,
+            &traps,
+            &mut self.trap_registration_guards,
+        );
+
         let dbg = if let Some(debug_data) = debug_data {
             let target_config = self.isa.frontend_config();
             let triple = self.isa.triple().clone();
@@ -309,6 +335,24 @@ fn allocate_functions(
     Ok(result)
 }
 
+fn register_traps(
+    allocated_functions: &PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
+    traps: &Traps,
+    trap_registration_guards: &mut Vec<TrapRegistrationGuard>,
+) {
+    let mut trap_registry = get_mut_trap_registry();
+    for (func_addr, func_traps) in allocated_functions.values().zip(traps.values()) {
+        for trap_desc in func_traps.iter() {
+            let func_addr = *func_addr as *const u8 as usize;
+            let offset = usize::try_from(trap_desc.code_offset).unwrap();
+            let trap_addr = func_addr + offset;
+            let guard =
+                trap_registry.register_trap(trap_addr, trap_desc.source_loc, trap_desc.trap_code);
+            trap_registration_guards.push(guard);
+        }
+    }
+}
+
 /// We don't expect trampoline compilation to produce any relocations, so
 /// this `RelocSink` just asserts that it doesn't recieve any.
 struct RelocSink {}

wasmtime-runtime/signalhandlers/SignalHandlers.cpp

@@ -408,6 +408,10 @@ HandleTrap(CONTEXT* context)
 {
     assert(sAlreadyHandlingTrap);
 
+    if (!CheckIfTrapAtAddress(ContextToPC(context))) {
+        return false;
+    }
+
     RecordTrap(ContextToPC(context));
 
     // Unwind calls longjmp, so it doesn't run the automatic

wasmtime-runtime/signalhandlers/SignalHandlers.hpp

@@ -11,6 +11,7 @@
 extern "C" {
 #endif
 
+int8_t CheckIfTrapAtAddress(const uint8_t* pc);
 // Record the Trap code and wasm bytecode offset in TLS somewhere
 void RecordTrap(const uint8_t* pc);
 

wasmtime-runtime/src/lib.rs

@@ -38,6 +38,7 @@ mod mmap;
 mod sig_registry;
 mod signalhandlers;
 mod table;
+mod trap_registry;
 mod traphandlers;
 mod vmcontext;
 
@@ -50,6 +51,7 @@ pub use crate::jit_int::GdbJitImageRegistration;
 pub use crate::mmap::Mmap;
 pub use crate::sig_registry::SignatureRegistry;
 pub use crate::signalhandlers::{wasmtime_init_eager, wasmtime_init_finish};
+pub use crate::trap_registry::{get_mut_trap_registry, get_trap_registry, TrapRegistrationGuard};
 pub use crate::traphandlers::{wasmtime_call, wasmtime_call_trampoline};
 pub use crate::vmcontext::{
     VMCallerCheckedAnyfunc, VMContext, VMFunctionBody, VMFunctionImport, VMGlobalDefinition,

wasmtime-runtime/src/trap_registry.rs

@@ -0,0 +1,71 @@
+use cranelift_codegen::ir;
+use lazy_static::lazy_static;
+use std::collections::HashMap;
+use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+lazy_static! {
+    static ref REGISTRY: RwLock<TrapRegistry> = RwLock::new(TrapRegistry::default());
+}
+
+/// The registry maintains descriptions of traps in currently allocated functions.
+#[derive(Default)]
+pub struct TrapRegistry {
+    traps: HashMap<usize, TrapDescription>,
+}
+
+/// Description of a trap.
+#[derive(Clone, Copy, PartialEq, Debug)]
+pub struct TrapDescription {
+    /// Location of the trap in source binary module.
+    pub source_loc: ir::SourceLoc,
+    /// Code of the trap.
+    pub trap_code: ir::TrapCode,
+}
+
+/// RAII guard for deregistering traps
+pub struct TrapRegistrationGuard(usize);
+
+impl TrapRegistry {
+    /// Registers a new trap.
+    /// Returns a RAII guard that deregisters the trap when dropped.
+    pub fn register_trap(
+        &mut self,
+        address: usize,
+        source_loc: ir::SourceLoc,
+        trap_code: ir::TrapCode,
+    ) -> TrapRegistrationGuard {
+        let entry = TrapDescription {
+            source_loc,
+            trap_code,
+        };
+        let previous_trap = self.traps.insert(address, entry);
+        assert!(previous_trap.is_none());
+        TrapRegistrationGuard(address)
+    }
+
+    fn deregister_trap(&mut self, address: usize) {
+        assert!(self.traps.remove(&address).is_some());
+    }
+
+    /// Gets a trap description at given address.
+    pub fn get_trap(&self, address: usize) -> Option<TrapDescription> {
+        self.traps.get(&address).copied()
+    }
+}
+
+impl Drop for TrapRegistrationGuard {
+    fn drop(&mut self) {
+        let mut registry = get_mut_trap_registry();
+        registry.deregister_trap(self.0);
+    }
+}
+
+/// Gets guarded writable reference to traps registry
+pub fn get_mut_trap_registry() -> RwLockWriteGuard<'static, TrapRegistry> {
+    REGISTRY.write().expect("trap registry lock got poisoned")
+}
+
+/// Gets guarded readable reference to traps registry
+pub fn get_trap_registry() -> RwLockReadGuard<'static, TrapRegistry> {
+    REGISTRY.read().expect("trap registry lock got poisoned")
+}

wasmtime-runtime/src/traphandlers.rs

@@ -1,9 +1,12 @@
 //! WebAssembly trap handling, which is built on top of the lower-level
 //! signalhandling mechanisms.
 
+use crate::trap_registry::get_trap_registry;
+use crate::trap_registry::TrapDescription;
 use crate::vmcontext::{VMContext, VMFunctionBody};
 use core::cell::Cell;
 use core::ptr;
+use cranelift_codegen::ir;
 use std::string::String;
 
 extern "C" {
@@ -16,17 +19,44 @@ extern "C" {
 }
 
 thread_local! {
-    static TRAP_PC: Cell<*const u8> = Cell::new(ptr::null());
+    static RECORDED_TRAP: Cell<Option<TrapDescription>> = Cell::new(None);
     static JMP_BUF: Cell<*const u8> = Cell::new(ptr::null());
 }
 
+/// Check if there is a trap at given PC
+#[doc(hidden)]
+#[allow(non_snake_case)]
+#[no_mangle]
+pub extern "C" fn CheckIfTrapAtAddress(_pc: *const u8) -> i8 {
+    // TODO: stack overflow can happen at any random time (i.e. in malloc() in memory.grow)
+    // and it's really hard to determine if the cause was stack overflow and if it happened
+    // in WebAssembly module.
+    // So, let's assume that any untrusted code called from WebAssembly doesn't trap.
+    // Then, if we have called some WebAssembly code, it means the trap is stack overflow.
+    JMP_BUF.with(|ptr| !ptr.get().is_null()) as i8
+}
+
 /// Record the Trap code and wasm bytecode offset in TLS somewhere
 #[doc(hidden)]
 #[allow(non_snake_case)]
 #[no_mangle]
 pub extern "C" fn RecordTrap(pc: *const u8) {
-    // TODO: Look up the wasm bytecode offset and trap code and record them instead.
+    // TODO: please see explanation in CheckIfTrapAtAddress.
-    TRAP_PC.with(|data| data.set(pc));
+    let registry = get_trap_registry();
+    let trap_desc = registry
+        .get_trap(pc as usize)
+        .unwrap_or_else(|| TrapDescription {
+            source_loc: ir::SourceLoc::default(),
+            trap_code: ir::TrapCode::StackOverflow,
+        });
+    RECORDED_TRAP.with(|data| {
+        assert_eq!(
+            data.get(),
+            None,
+            "Only one trap per thread can be recorded at a moment!"
+        );
+        data.set(Some(trap_desc))
+    });
 }
 
 #[doc(hidden)]
@@ -50,13 +80,17 @@ pub extern "C" fn LeaveScope(ptr: *const u8) {
     JMP_BUF.with(|buf| buf.set(ptr))
 }
 
-fn trap_message(_vmctx: *mut VMContext) -> String {
+fn trap_message() -> String {
-    let pc = TRAP_PC.with(|data| data.replace(ptr::null()));
+    let trap_desc = RECORDED_TRAP
+        .with(|data| data.replace(None))
+        .expect("trap_message must be called after trap occurred");
 
-    // TODO: Record trap metadata in the VMContext, and look up the
+    format!(
-    // pc to obtain the TrapCode and SourceLoc.
+        "wasm trap: code {:?}, source location: {}",
-
+        // todo print the error message from wast tests
-    format!("wasm trap at {:?}", pc)
+        trap_desc.trap_code,
+        trap_desc.source_loc,
+    )
 }
 
 /// Call the wasm function pointed to by `callee`. `values_vec` points to
@@ -69,7 +103,7 @@ pub unsafe extern "C" fn wasmtime_call_trampoline(
     values_vec: *mut u8,
 ) -> Result<(), String> {
     if WasmtimeCallTrampoline(vmctx as *mut u8, callee, values_vec) == 0 {
-        Err(trap_message(vmctx))
+        Err(trap_message())
     } else {
         Ok(())
     }
@@ -83,7 +117,7 @@ pub unsafe extern "C" fn wasmtime_call(
     callee: *const VMFunctionBody,
 ) -> Result<(), String> {
     if WasmtimeCall(vmctx as *mut u8, callee) == 0 {
-        Err(trap_message(vmctx))
+        Err(trap_message())
     } else {
         Ok(())
     }