[wasmtime-api] Record original Trap from API callback. (#657)

* Record original Trap from API callback.

Fixes #645

* use TrapRegistry

* comment about magic number

crates/api/src/callable.rs

@@ -1,7 +1,7 @@
 use crate::data_structures::ir;
 use crate::r#ref::HostRef;
 use crate::runtime::Store;
-use crate::trampoline::generate_func_export;
+use crate::trampoline::{generate_func_export, take_api_trap};
 use crate::trap::Trap;
 use crate::types::FuncType;
 use crate::values::Val;
@@ -90,7 +90,8 @@ impl WrappedCallable for WasmtimeFn {
                 values_vec.as_mut_ptr() as *mut u8,
             )
         } {
-            return Err(HostRef::new(Trap::new(message)));
+            let trap = take_api_trap().unwrap_or_else(|| HostRef::new(Trap::new(message)));
+            return Err(trap);
         }
 
         // Load the return values out of `values_vec`.

crates/api/src/instance.rs

@@ -3,6 +3,7 @@ use crate::externals::Extern;
 use crate::module::Module;
 use crate::r#ref::HostRef;
 use crate::runtime::Store;
+use crate::trampoline::take_api_trap;
 use crate::types::{ExportType, ExternType, Name};
 use anyhow::Result;
 use std::cell::RefCell;
@@ -40,7 +41,12 @@ pub fn instantiate_in_context(
         &mut resolver,
         exports,
         debug_info,
-    )?;
+    )
+    .map_err(|e| {
+        // TODO wrap HostRef<Trap> into Error
+        drop(take_api_trap());
+        e
+    })?;
     contexts.insert(context);
     Ok((instance, contexts))
 }

crates/api/src/trampoline/func.rs

@@ -1,28 +1,62 @@
 //! Support for a calling of an imported function.
 
 use super::create_handle::create_handle;
-use super::ir::{
-    ExternalName, Function, InstBuilder, MemFlags, StackSlotData, StackSlotKind, TrapCode,
-};
+use super::ir::{ExternalName, Function, InstBuilder, MemFlags, StackSlotData, StackSlotKind};
+use super::trap::{record_api_trap, TrapSink, API_TRAP_CODE};
 use super::{binemit, pretty_error, TargetIsa};
 use super::{Context, FunctionBuilder, FunctionBuilderContext};
 use crate::data_structures::ir::{self, types};
 use crate::data_structures::wasm::{DefinedFuncIndex, FuncIndex};
 use crate::data_structures::{native_isa_builder, settings, EntityRef, PrimaryMap};
 use crate::r#ref::HostRef;
-use crate::{Callable, FuncType, Store, Trap, Val};
+use crate::{Callable, FuncType, Store, Val};
 use anyhow::Result;
 use std::cmp;
+use std::convert::TryFrom;
 use std::rc::Rc;
-use wasmtime_environ::{CompiledFunction, Export, Module};
+use wasmtime_environ::{CompiledFunction, Export, Module, TrapInformation};
 use wasmtime_jit::CodeMemory;
-use wasmtime_runtime::{InstanceHandle, VMContext, VMFunctionBody};
+use wasmtime_runtime::{
+    get_mut_trap_registry, InstanceHandle, TrapRegistrationGuard, VMContext, VMFunctionBody,
+};
 
 struct TrampolineState {
     func: Rc<dyn Callable + 'static>,
-    trap: Option<HostRef<Trap>>,
     #[allow(dead_code)]
     code_memory: CodeMemory,
+    trap_registration_guards: Vec<TrapRegistrationGuard>,
+}
+
+impl TrampolineState {
+    fn new(
+        func: Rc<dyn Callable + 'static>,
+        code_memory: CodeMemory,
+        func_addr: *const VMFunctionBody,
+        func_traps: &[TrapInformation],
+    ) -> Self {
+        let mut trap_registry = get_mut_trap_registry();
+        let mut trap_registration_guards = Vec::new();
+        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);
+        }
+        TrampolineState {
+            func,
+            code_memory,
+            trap_registration_guards,
+        }
+    }
+}
+
+impl Drop for TrampolineState {
+    fn drop(&mut self) {
+        // We must deregister traps before freeing the code memory.
+        self.trap_registration_guards.clear();
+    }
 }
 
 unsafe extern "C" fn stub_fn(vmctx: *mut VMContext, call_id: u32, values_vec: *mut i64) -> u32 {
@@ -58,12 +92,7 @@ unsafe extern "C" fn stub_fn(vmctx: *mut VMContext, call_id: u32, values_vec: *m
             0
         }
         Err(trap) => {
-            // TODO read custom exception
-            InstanceHandle::from_vmctx(vmctx)
-                .host_state()
-                .downcast_mut::<TrampolineState>()
-                .expect("state")
-                .trap = Some(trap);
+            record_api_trap(trap);
             1
         }
     }
@@ -76,7 +105,7 @@ fn make_trampoline(
     fn_builder_ctx: &mut FunctionBuilderContext,
     call_id: u32,
     signature: &ir::Signature,
-) -> *const VMFunctionBody {
+) -> (*const VMFunctionBody, Vec<TrapInformation>) {
     // Mostly reverse copy of the similar method from wasmtime's
     // wasmtime-jit/src/compiler.rs.
     let pointer_type = isa.pointer_type();
@@ -147,7 +176,7 @@ fn make_trampoline(
             .call_indirect(new_sig, callee_value, &callee_args);
 
         let call_result = builder.func.dfg.inst_results(call)[0];
-        builder.ins().trapnz(call_result, TrapCode::User(0));
+        builder.ins().trapnz(call_result, API_TRAP_CODE);
 
         let mflags = MemFlags::trusted();
         let mut results = Vec::new();
@@ -166,7 +195,7 @@ fn make_trampoline(
 
     let mut code_buf: Vec<u8> = Vec::new();
     let mut reloc_sink = binemit::TrampolineRelocSink {};
-    let mut trap_sink = binemit::NullTrapSink {};
+    let mut trap_sink = TrapSink::new();
     let mut stackmap_sink = binemit::NullStackmapSink {};
     context
         .compile_and_emit(
@@ -182,14 +211,17 @@ fn make_trampoline(
     let mut unwind_info = Vec::new();
     context.emit_unwind_info(isa, &mut unwind_info);
 
-    code_memory
+    let traps = trap_sink.traps;
+
+    let addr = code_memory
         .allocate_for_function(&CompiledFunction {
             body: code_buf,
             jt_offsets: context.func.jt_offsets,
             unwind_info,
         })
         .expect("allocate_for_function")
-        .as_ptr()
+        .as_ptr();
+    (addr, traps)
 }
 
 pub fn create_handle_with_function(
@@ -216,7 +248,7 @@ pub fn create_handle_with_function(
     module
         .exports
         .insert("trampoline".to_string(), Export::Function(func_id));
-    let trampoline = make_trampoline(
+    let (trampoline, traps) = make_trampoline(
         isa.as_ref(),
         &mut code_memory,
         &mut fn_builder_ctx,
@@ -227,11 +259,7 @@ pub fn create_handle_with_function(
 
     finished_functions.push(trampoline);
 
-    let trampoline_state = TrampolineState {
-        func: func.clone(),
-        trap: None,
-        code_memory,
-    };
+    let trampoline_state = TrampolineState::new(func.clone(), code_memory, trampoline, &traps);
 
     create_handle(
         module,

crates/api/src/trampoline/mod.rs

@@ -5,6 +5,7 @@ mod func;
 mod global;
 mod memory;
 mod table;
+mod trap;
 
 use self::func::create_handle_with_function;
 use self::global::create_global;
@@ -16,6 +17,7 @@ use anyhow::Result;
 use std::rc::Rc;
 
 pub use self::global::GlobalState;
+pub use self::trap::take_api_trap;
 
 pub fn generate_func_export(
     ft: &FuncType,
@@ -53,7 +55,7 @@ pub fn generate_table_export(
 pub(crate) use cranelift_codegen::print_errors::pretty_error;
 
 pub(crate) mod binemit {
-    pub(crate) use cranelift_codegen::binemit::{NullStackmapSink, NullTrapSink};
+    pub(crate) use cranelift_codegen::binemit::{CodeOffset, NullStackmapSink, TrapSink};
 
     pub use cranelift_codegen::{binemit, ir};
 
@@ -100,7 +102,8 @@ pub(crate) mod binemit {
 
 pub(crate) mod ir {
     pub(crate) use cranelift_codegen::ir::{
-        ExternalName, Function, InstBuilder, MemFlags, StackSlotData, StackSlotKind, TrapCode,
+        ExternalName, Function, InstBuilder, MemFlags, SourceLoc, StackSlotData, StackSlotKind,
+        TrapCode,
     };
 }
 pub(crate) use cranelift_codegen::isa::TargetIsa;

crates/api/src/trampoline/trap.rs

@@ -0,0 +1,54 @@
+use std::cell::Cell;
+
+use super::binemit;
+use super::ir::{SourceLoc, TrapCode};
+use crate::r#ref::HostRef;
+use crate::Trap;
+use wasmtime_environ::TrapInformation;
+
+// Randomly selected user TrapCode magic number 13.
+pub const API_TRAP_CODE: TrapCode = TrapCode::User(13);
+
+thread_local! {
+    static RECORDED_API_TRAP: Cell<Option<HostRef<Trap>>> = Cell::new(None);
+}
+
+pub fn record_api_trap(trap: HostRef<Trap>) {
+    RECORDED_API_TRAP.with(|data| {
+        let trap = Cell::new(Some(trap));
+        data.swap(&trap);
+        assert!(
+            trap.take().is_none(),
+            "Only one API trap per thread can be recorded at a moment!"
+        );
+    });
+}
+
+pub fn take_api_trap() -> Option<HostRef<Trap>> {
+    RECORDED_API_TRAP.with(|data| data.take())
+}
+
+pub(crate) struct TrapSink {
+    pub traps: Vec<TrapInformation>,
+}
+
+impl TrapSink {
+    pub fn new() -> Self {
+        Self { traps: Vec::new() }
+    }
+}
+
+impl binemit::TrapSink for TrapSink {
+    fn trap(
+        &mut self,
+        code_offset: binemit::CodeOffset,
+        source_loc: SourceLoc,
+        trap_code: TrapCode,
+    ) {
+        self.traps.push(TrapInformation {
+            code_offset,
+            source_loc,
+            trap_code,
+        });
+    }
+}

crates/api/tests/traps.rs

@@ -0,0 +1,49 @@
+use std::rc::Rc;
+use wasmtime::*;
+use wat::parse_str;
+
+#[test]
+fn test_trap_return() -> Result<(), String> {
+    struct HelloCallback;
+
+    impl Callable for HelloCallback {
+        fn call(&self, _params: &[Val], _results: &mut [Val]) -> Result<(), HostRef<Trap>> {
+            Err(HostRef::new(Trap::new("test 123".into())))
+        }
+    }
+
+    let engine = HostRef::new(Engine::default());
+    let store = HostRef::new(Store::new(&engine));
+    let binary = parse_str(
+        r#"
+                (module
+                (func $hello (import "" "hello"))
+                (func (export "run") (call $hello))
+                )
+            "#,
+    )
+    .map_err(|e| format!("failed to parse WebAssembly text source: {}", e))?;
+
+    let module = HostRef::new(
+        Module::new(&store, &binary).map_err(|e| format!("failed to compile module: {}", e))?,
+    );
+    let hello_type = FuncType::new(Box::new([]), Box::new([]));
+    let hello_func = HostRef::new(Func::new(&store, hello_type, Rc::new(HelloCallback)));
+
+    let imports = vec![hello_func.into()];
+    let instance = Instance::new(&store, &module, imports.as_slice())
+        .map_err(|e| format!("failed to instantiate module: {}", e))?;
+    let run_func = instance.exports()[0]
+        .func()
+        .expect("expected function export");
+
+    let e = run_func
+        .borrow()
+        .call(&[])
+        .err()
+        .expect("error calling function");
+
+    assert_eq!(e.borrow().message(), "test 123");
+
+    Ok(())
+}