Repository: bytecodealliance/wasmtime (rust)

Issue:
Trap in host function leaks memory, 0.25 regression 
### Test Case

```rust
use wasmtime::*;

fn main() {
    let wat = r#"
        (module
            (import "" "" (func $host_hello))
            (memory (export "memory") 0)
            (func (export "hello") call $host_hello)
        )
        "#;
    let big = vec![1; u32::MAX as usize / 8];

    loop {
        let store = Store::default();
        let module = Module::new(store.engine(), wat).unwrap();
        let host_hello = Func::wrap(&store, || -> Result<(), Trap> {
            // If `Ok(())` is returned instead, memory will not leak.
            Err(Trap::new("it's a trap"))
        });
        let instance = Instance::new(&store, &module, &[host_hello.into()]).unwrap();

        // Do a large allocation to speed up memory leak.
        let mem = instance.get_memory("memory").unwrap();
        mem.grow(10_000).unwrap();
        mem.write(0, &big).unwrap();

        let hello = instance.get_typed_func::<(), ()>("hello").unwrap();
        let _ = hello.call(());

        // In wasmtime 0.24, this does not reproduce:
        // let hello = instance.get_func("hello").unwrap().get0::<()>().unwrap();
        // let _ = hello();
    }
}
```

### Steps to Reproduce

Run the above program on Linux x86_64.

### Expected Results

Program can run forever without exhausting system memory.

### Actual Results

Program leaks memory until the system memory is exhausted.

### Versions and Environment

Wasmtime version or commit: 0.25
Operating system: Linux
Architecture:  x86_64

### Extra Info

It does not reproduce on version 0.24, indicating a 0.25 regression. It seems the `Store` is being leaked. This issue was observed in our application as the error `Insufficient resources: mmap failed: Cannot allocate memory (os error 12)`, because the virtual memory space would be exhausted.

Pull request:
Move some scopes around to fix a leak on raising a trap
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

<!--

Please ensure that the following steps are all taken care of before submitting
the PR.

- [ ] This has been discussed in issue #..., or if not, please tell us why
  here.
- [ ] A short description of what this does, why it is needed; if the
  description becomes long, the matter should probably be discussed in an issue
  first.
- [ ] This PR contains test cases, if meaningful.
- [ ] A reviewer from the core maintainer team has been assigned for this PR.
  If you don't know who could review this, please indicate so. The list of
  suggested reviewers on the right can help you.

Please ensure all communication adheres to the [code of
conduct](https://github.com/bytecodealliance/wasmtime/blob/master/CODE_OF_CONDUCT.md).
-->


Patch:
diff --git a/crates/wasmtime/src/func.rs b/crates/wasmtime/src/func.rs
index 05b74cf91..2a70df9e3 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<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) {
-        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
         <Self as WasmTy>::compatible_with_store(self, store)
     }
 
-    unsafe fn into_abi_for_ret(self, store: &Store) -> Self::Abi {
-        <Self as WasmTy>::into_abi(self, store)
+    unsafe fn into_abi_for_ret(self, store: &Store) -> Result<Self::Abi, Trap> {
+        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 {
-        match self {
-            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())
-        }
+    unsafe fn into_abi_for_ret(self, store: &Store) -> Result<Self::Abi, Trap> {
+        self.map(|val| <T as WasmTy>::into_abi(val, store))
     }
 
     fn valtype() -> Option<ValType> {
@@ -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::<F>());
-                    let func = &*(state as *const _ as *const F);
-
-                    let store = wasmtime_runtime::with_last_info(|last| {
-                        last.and_then(|s| s.downcast_ref::<Store>())
-                            .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<T> {
+                        Ok(T),
+                        Trap(Trap),
+                        Panic(Box<dyn std::any::Any + Send>),
+                    }
 
-                    // 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::<F>());
+                        let func = &*(state as *const _ as *const F);
+
+                        let store = wasmtime_runtime::with_last_info(|last| {
+                            last.and_then(|s| s.downcast_ref::<Store>())
+                                .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),
                     }
                 }
 


Test patch:
diff --git a/tests/all/func.rs b/tests/all/func.rs
index b33405689..0e6919938 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<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(())
+}