Move some scopes around to fix a leak on raising a trap (#2803)

Some recent refactorings accidentally had a local `Store` on the stack
when a longjmp was initiated, bypassing its destructor and causing
`Store` to leak.

Closes #2802

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<Self::Abi, Trap>;
 
     // 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<ValType> {
@@ -1331,11 +1333,8 @@ unsafe impl WasmRet for Result<(), Trap> {
     }
 
     #[inline]
-    unsafe fn into_abi_for_ret(self, _store: &Store) {
+    unsafe fn into_abi_for_ret(self, _store: &Store) -> Result<(), Trap> {
-        match self {
+        self
-            Ok(()) => {}
-            Err(trap) => raise_user_trap(trap.into()),
-        }
     }
 
     #[inline]
@@ -1365,8 +1364,8 @@ where
         <Self as WasmTy>::compatible_with_store(self, store)
     }
 
-    unsafe fn into_abi_for_ret(self, store: &Store) -> Self::Abi {
+    unsafe fn into_abi_for_ret(self, store: &Store) -> Result<Self::Abi, Trap> {
-        <Self as WasmTy>::into_abi(self, store)
+        Ok(<Self as WasmTy>::into_abi(self, store))
     }
 
     fn valtype() -> Option<ValType> {
@@ -1396,15 +1395,8 @@ where
         }
     }
 
-    unsafe fn into_abi_for_ret(self, store: &Store) -> Self::Abi {
+    unsafe fn into_abi_for_ret(self, store: &Store) -> Result<Self::Abi, Trap> {
-        match self {
+        self.map(|val| <T as WasmTy>::into_abi(val, store))
-            Ok(val) => return <T as WasmTy>::into_abi(val, store),
-            Err(trap) => handle_trap(trap),
-        }
-
-        unsafe fn handle_trap(trap: Trap) -> ! {
-            raise_user_trap(trap.into())
-        }
     }
 
     fn valtype() -> Option<ValType> {
@@ -1577,6 +1569,19 @@ macro_rules! impl_into_func {
                     $( $args: WasmTy, )*
                     R: WasmRet,
                 {
+                    enum CallResult<T> {
+                        Ok(T),
+                        Trap(Trap),
+                        Panic(Box<dyn std::any::Any + Send>),
+                    }
+
+                    // 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
@@ -1607,7 +1612,7 @@ macro_rules! impl_into_func {
                         // 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),
+                            Err(panic) => CallResult::Panic(panic),
                             Ok(ret) => {
                                 // Because the wrapped function is not `unsafe`, we
                                 // can't assume it returned a value that is
@@ -1619,10 +1624,20 @@ macro_rules! impl_into_func {
                                     raise_cross_store_trap();
                                 }
 
-                            ret.into_abi_for_ret(&store)
+                                match ret.into_abi_for_ret(&store) {
+                                    Ok(val) => CallResult::Ok(val),
+                                    Err(trap) => CallResult::Trap(trap),
                                 }
                             }
                         }
+                    };
+
+                    match result {
+                        CallResult::Ok(val) => val,
+                        CallResult::Trap(trap) => raise_user_trap(trap.into()),
+                        CallResult::Panic(panic) => wasmtime_runtime::resume_panic(panic),
+                    }
+                }
 
                 /// This trampoline allows host code to indirectly call the
                 /// wrapped function (e.g. via `Func::call` on a `funcref` that

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<Cell<bool>>);
+
+    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(())
+}