diff --git a/crates/wasmtime/src/func.rs b/crates/wasmtime/src/func.rs index 05b74cf91626..2a70df9e3cd2 100644 --- a/crates/wasmtime/src/func.rs +++ b/crates/wasmtime/src/func.rs @@ -1276,7 +1276,7 @@ pub unsafe trait WasmRet { // `invoke_wasm_and_catch_traps` is on the stack, and therefore this method // is unsafe. #[doc(hidden)] - unsafe fn into_abi_for_ret(self, store: &Store) -> Self::Abi; + unsafe fn into_abi_for_ret(self, store: &Store) -> Result; // Same as `WasmTy::push`. #[doc(hidden)] @@ -1303,7 +1303,9 @@ unsafe impl WasmRet for () { } #[inline] - unsafe fn into_abi_for_ret(self, _store: &Store) {} + unsafe fn into_abi_for_ret(self, _store: &Store) -> Result<(), Trap> { + Ok(()) + } #[inline] fn valtype() -> Option { @@ -1331,11 +1333,8 @@ unsafe impl WasmRet for Result<(), Trap> { } #[inline] - unsafe fn into_abi_for_ret(self, _store: &Store) { - match self { - Ok(()) => {} - Err(trap) => raise_user_trap(trap.into()), - } + unsafe fn into_abi_for_ret(self, _store: &Store) -> Result<(), Trap> { + self } #[inline] @@ -1365,8 +1364,8 @@ where ::compatible_with_store(self, store) } - unsafe fn into_abi_for_ret(self, store: &Store) -> Self::Abi { - ::into_abi(self, store) + unsafe fn into_abi_for_ret(self, store: &Store) -> Result { + Ok(::into_abi(self, store)) } fn valtype() -> Option { @@ -1396,15 +1395,8 @@ where } } - unsafe fn into_abi_for_ret(self, store: &Store) -> Self::Abi { - match self { - Ok(val) => return ::into_abi(val, store), - Err(trap) => handle_trap(trap), - } - - unsafe fn handle_trap(trap: Trap) -> ! { - raise_user_trap(trap.into()) - } + unsafe fn into_abi_for_ret(self, store: &Store) -> Result { + self.map(|val| ::into_abi(val, store)) } fn valtype() -> Option { @@ -1577,50 +1569,73 @@ macro_rules! impl_into_func { $( $args: WasmTy, )* R: WasmRet, { - let state = (*vmctx).host_state(); - // Double-check ourselves in debug mode, but we control - // the `Any` here so an unsafe downcast should also - // work. - debug_assert!(state.is::()); - let func = &*(state as *const _ as *const F); - - let store = wasmtime_runtime::with_last_info(|last| { - last.and_then(|s| s.downcast_ref::()) - .cloned() - .expect("function called without thread state") - }); - - let ret = { - panic::catch_unwind(AssertUnwindSafe(|| { - func( - Caller { store: &store, caller_vmctx }, - $( $args::from_abi($args, &store), )* - ) - })) - }; + enum CallResult { + Ok(T), + Trap(Trap), + Panic(Box), + } - // Note that we need to be careful when dealing with traps - // here. Traps are implemented with longjmp/setjmp meaning - // that it's not unwinding and consequently no Rust - // destructors are run. We need to be careful to ensure that - // nothing on the stack needs a destructor when we exit - // abnormally from this `match`, e.g. on `Err`, on - // cross-store-issues, or if `Ok(Err)` is raised. - match ret { - Err(panic) => wasmtime_runtime::resume_panic(panic), - Ok(ret) => { - // Because the wrapped function is not `unsafe`, we - // can't assume it returned a value that is - // compatible with this store. - if !ret.compatible_with_store(&store) { - // Explicitly drop all locals with destructors prior to raising the trap - drop(store); - drop(ret); - raise_cross_store_trap(); + // Note that this `result` is intentionally scoped into a + // separate block. Handling traps and panics will involve + // longjmp-ing from this function which means we won't run + // destructors. As a result anything requiring a destructor + // should be part of this block, and the long-jmp-ing + // happens after the block in handling `CallResult`. + let result = { + let state = (*vmctx).host_state(); + // Double-check ourselves in debug mode, but we control + // the `Any` here so an unsafe downcast should also + // work. + debug_assert!(state.is::()); + let func = &*(state as *const _ as *const F); + + let store = wasmtime_runtime::with_last_info(|last| { + last.and_then(|s| s.downcast_ref::()) + .cloned() + .expect("function called without thread state") + }); + + let ret = { + panic::catch_unwind(AssertUnwindSafe(|| { + func( + Caller { store: &store, caller_vmctx }, + $( $args::from_abi($args, &store), )* + ) + })) + }; + + // Note that we need to be careful when dealing with traps + // here. Traps are implemented with longjmp/setjmp meaning + // that it's not unwinding and consequently no Rust + // destructors are run. We need to be careful to ensure that + // nothing on the stack needs a destructor when we exit + // abnormally from this `match`, e.g. on `Err`, on + // cross-store-issues, or if `Ok(Err)` is raised. + match ret { + Err(panic) => CallResult::Panic(panic), + Ok(ret) => { + // Because the wrapped function is not `unsafe`, we + // can't assume it returned a value that is + // compatible with this store. + if !ret.compatible_with_store(&store) { + // Explicitly drop all locals with destructors prior to raising the trap + drop(store); + drop(ret); + raise_cross_store_trap(); + } + + match ret.into_abi_for_ret(&store) { + Ok(val) => CallResult::Ok(val), + Err(trap) => CallResult::Trap(trap), + } } - - ret.into_abi_for_ret(&store) } + }; + + match result { + CallResult::Ok(val) => val, + CallResult::Trap(trap) => raise_user_trap(trap.into()), + CallResult::Panic(panic) => wasmtime_runtime::resume_panic(panic), } } diff --git a/tests/all/func.rs b/tests/all/func.rs index b334056893ae..0e691993843f 100644 --- a/tests/all/func.rs +++ b/tests/all/func.rs @@ -1,4 +1,6 @@ use anyhow::Result; +use std::cell::Cell; +use std::rc::Rc; use std::sync::atomic::{AtomicUsize, Ordering::SeqCst}; use wasmtime::*; @@ -578,3 +580,42 @@ fn typed_multiple_results() -> anyhow::Result<()> { ); Ok(()) } + +#[test] +fn trap_doesnt_leak() -> anyhow::Result<()> { + struct Canary(Rc>); + + impl Drop for Canary { + fn drop(&mut self) { + self.0.set(true); + } + } + + let store = Store::default(); + + // test that `Func::wrap` is correct + let canary1 = Canary(Rc::new(Cell::new(false))); + let dtor1_run = canary1.0.clone(); + let f1 = Func::wrap(&store, move || -> Result<(), Trap> { + drop(&canary1); + Err(Trap::new("")) + }); + assert!(f1.typed::<(), ()>()?.call(()).is_err()); + assert!(f1.call(&[]).is_err()); + + // test that `Func::new` is correct + let canary2 = Canary(Rc::new(Cell::new(false))); + let dtor2_run = canary2.0.clone(); + let f2 = Func::new(&store, FuncType::new(None, None), move |_, _, _| { + drop(&canary2); + Err(Trap::new("")) + }); + assert!(f2.typed::<(), ()>()?.call(()).is_err()); + assert!(f2.call(&[]).is_err()); + + // drop everything and ensure dtors are run + drop((store, f1, f2)); + assert!(dtor1_run.get()); + assert!(dtor2_run.get()); + Ok(()) +}