Update may_enter flag handling in components (#4354)

This commit updates the management of the `may_enter` flag in line with
WebAssembly/component-model#57. Namely the `may_enter` flag is now
exclusively managed in the `canon lift` function (which is
`TypedFunc::call`) and is only unset after post-return completes
successfully. This implements semantics where if any trap happens for
any reason (lifting, lowering, execution, imports, etc) then the
instance is considered permanently poisoned and can no longer be
entered.

Tests needed many updates to create new instances where previously the
same instance was reused after it had an erroneous state.

crates/wasmtime/src/component/func/host.rs

@@ -154,10 +154,6 @@ where
         bail!("cannot leave component instance");
     }
 
-    // While we're lifting and lowering this instance cannot be reentered, so
-    // unset the flag here. This is also reset back to `true` on exit.
-    let _reset_may_enter = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_enter);
-
     // There's a 2x2 matrix of whether parameters and results are stored on the
     // stack or on the heap. Each of the 4 branches here have a different
     // representation of the storage of arguments/returns which is represented
@@ -168,13 +164,12 @@ where
     // trivially DCE'd by LLVM. Perhaps one day with enough const programming in
     // Rust we can make monomorphizations of this function codegen only one
     // branch, but today is not that day.
-    let reset_may_leave;
     if Params::flatten_count() <= MAX_STACK_PARAMS {
         if Return::flatten_count() <= MAX_STACK_RESULTS {
             let storage = cast_storage::<ReturnStack<Params::Lower, Return::Lower>>(storage);
             let params = Params::lift(cx.0, &options, &storage.assume_init_ref().args)?;
             let ret = closure(cx.as_context_mut(), params)?;
-            reset_may_leave = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_leave);
+            (*flags).set_may_leave(false);
             ret.lower(&mut cx, &options, map_maybe_uninit!(storage.ret))?;
         } else {
             let storage = cast_storage::<ReturnPointer<Params::Lower>>(storage).assume_init_ref();
@@ -182,7 +177,7 @@ where
             let ret = closure(cx.as_context_mut(), params)?;
             let mut memory = MemoryMut::new(cx.as_context_mut(), &options);
             let ptr = validate_inbounds::<Return>(memory.as_slice_mut(), &storage.retptr)?;
-            reset_may_leave = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_leave);
+            (*flags).set_may_leave(false);
             ret.store(&mut memory, ptr)?;
         }
     } else {
@@ -193,7 +188,7 @@ where
                 validate_inbounds::<Params>(memory.as_slice(), &storage.assume_init_ref().args)?;
             let params = Params::load(&memory, &memory.as_slice()[ptr..][..Params::SIZE32])?;
             let ret = closure(cx.as_context_mut(), params)?;
-            reset_may_leave = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_leave);
+            (*flags).set_may_leave(false);
             ret.lower(&mut cx, &options, map_maybe_uninit!(storage.ret))?;
         } else {
             let storage = cast_storage::<ReturnPointer<ValRaw>>(storage).assume_init_ref();
@@ -202,32 +197,14 @@ where
             let ret = closure(cx.as_context_mut(), params)?;
             let mut memory = MemoryMut::new(cx.as_context_mut(), &options);
             let ptr = validate_inbounds::<Return>(memory.as_slice_mut(), &storage.retptr)?;
-            reset_may_leave = unset_and_reset_on_drop(flags, VMComponentFlags::set_may_leave);
+            (*flags).set_may_leave(false);
             ret.store(&mut memory, ptr)?;
         }
     }
 
-    drop(reset_may_leave);
+    (*flags).set_may_leave(true);
 
     return Ok(());
-
-    unsafe fn unset_and_reset_on_drop(
-        slot: *mut VMComponentFlags,
-        set: fn(&mut VMComponentFlags, bool),
-    ) -> impl Drop {
-        set(&mut *slot, false);
-        return Reset(slot, set);
-
-        struct Reset(*mut VMComponentFlags, fn(&mut VMComponentFlags, bool));
-
-        impl Drop for Reset {
-            fn drop(&mut self) {
-                unsafe {
-                    (self.1)(&mut *self.0, true);
-                }
-            }
-        }
-    }
 }
 
 fn validate_inbounds<T: ComponentType>(memory: &[u8], ptr: &ValRaw) -> Result<usize> {

crates/wasmtime/src/component/func/typed.rs

@@ -333,11 +333,19 @@ where
         let flags = instance.flags(component_instance);
 
         unsafe {
+            // Test the "may enter" flag which is a "lock" on this instance.
+            // This is immediately set to `false` afterwards and note that
+            // there's no on-cleanup setting this flag back to true. That's an
+            // intentional design aspect where if anything goes wrong internally
+            // from this point on the instance is considered "poisoned" and can
+            // never be entered again. The only time this flag is set to `true`
+            // again is after post-return logic has completed successfully.
             if !(*flags).may_enter() {
                 bail!("cannot reenter component instance");
             }
-            debug_assert!((*flags).may_leave());
+            (*flags).set_may_enter(false);
 
+            debug_assert!((*flags).may_leave());
             (*flags).set_may_leave(false);
             let result = lower(store, &options, params, map_maybe_uninit!(space.params));
             (*flags).set_may_leave(true);
@@ -370,41 +378,21 @@ where
             // Lift the result into the host while managing post-return state
             // here as well.
             //
-            // Initially the `may_enter` flag is set to `false` for this
-            // component instance and additionally we set a flag indicating that
-            // a post-return is required. This isn't specified by the component
-            // model itself but is used for our implementation of the API of
-            // `post_return` as a separate function call.
-            //
-            // FIXME(WebAssembly/component-model#55) it's not really clear what
-            // the semantics should be in the face of a lift error/trap. For now
-            // the flags are reset so the instance can continue to be reused in
-            // tests but that probably isn't what's desired.
-            //
-            // Otherwise though after a successful lift the return value of the
-            // function, which is currently required to be 0 or 1 values
-            // according to the canonical ABI, is saved within the `Store`'s
-            // `FuncData`. This'll later get used in post-return.
-            (*flags).set_may_enter(false);
+            // After a successful lift the return value of the function, which
+            // is currently required to be 0 or 1 values according to the
+            // canonical ABI, is saved within the `Store`'s `FuncData`. This'll
+            // later get used in post-return.
             (*flags).set_needs_post_return(true);
-            match lift(store.0, &options, ret) {
-                Ok(val) => {
-                    let ret_slice = cast_storage(ret);
-                    let data = &mut store.0[self.func.0];
-                    assert!(data.post_return_arg.is_none());
-                    match ret_slice.len() {
-                        0 => data.post_return_arg = Some(ValRaw::i32(0)),
-                        1 => data.post_return_arg = Some(ret_slice[0]),
-                        _ => unreachable!(),
-                    }
-                    return Ok(val);
-                }
-                Err(err) => {
-                    (*flags).set_may_enter(true);
-                    (*flags).set_needs_post_return(false);
-                    return Err(err);
-                }
+            let val = lift(store.0, &options, ret)?;
+            let ret_slice = cast_storage(ret);
+            let data = &mut store.0[self.func.0];
+            assert!(data.post_return_arg.is_none());
+            match ret_slice.len() {
+                0 => data.post_return_arg = Some(ValRaw::i32(0)),
+                1 => data.post_return_arg = Some(ret_slice[0]),
+                _ => unreachable!(),
             }
+            return Ok(val);
         }
 
         unsafe fn cast_storage<T>(storage: &T) -> &[ValRaw] {
@@ -480,23 +468,30 @@ where
             // This is a sanity-check assert which shouldn't ever trip.
             assert!(!(*flags).may_enter());
 
-            // With the state of the world validated these flags are updated to
-            // their component-model-defined states.
-            (*flags).set_may_enter(true);
+            // Unset the "needs post return" flag now that post-return is being
+            // processed. This will cause future invocations of this method to
+            // panic, even if the function call below traps.
             (*flags).set_needs_post_return(false);
 
-            // And finally if the function actually had a `post-return`
-            // configured in its canonical options that's executed here.
-            let (func, trampoline) = match post_return {
-                Some(pair) => pair,
-                None => return Ok(()),
-            };
-            crate::Func::call_unchecked_raw(
-                &mut store,
-                func.anyfunc,
-                trampoline,
-                &post_return_arg as *const ValRaw as *mut ValRaw,
-            )?;
+            // If the function actually had a `post-return` configured in its
+            // canonical options that's executed here.
+            //
+            // Note that if this traps (returns an error) this function
+            // intentionally leaves the instance in a "poisoned" state where it
+            // can no longer be entered because `may_enter` is `false`.
+            if let Some((func, trampoline)) = post_return {
+                crate::Func::call_unchecked_raw(
+                    &mut store,
+                    func.anyfunc,
+                    trampoline,
+                    &post_return_arg as *const ValRaw as *mut ValRaw,
+                )?;
+            }
+
+            // And finally if everything completed successfully then the "may
+            // enter" flag is set to `true` again here which enables further use
+            // of the component.
+            (*flags).set_may_enter(true);
         }
         Ok(())
     }

tests/all/component_model/func.rs

@@ -188,7 +188,8 @@ fn integers() -> Result<()> {
     let engine = super::engine();
     let component = Component::new(&engine, component)?;
     let mut store = Store::new(&engine, ());
-    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let new_instance = |store: &mut Store<()>| Linker::new(&engine).instantiate(store, &component);
+    let instance = new_instance(&mut store)?;
 
     // Passing in 100 is valid for all primitives
     instance
@@ -217,42 +218,42 @@ fn integers() -> Result<()> {
         .call_and_post_return(&mut store, (100,))?;
 
     // This specific wasm instance traps if any value other than 100 is passed
-    instance
+    new_instance(&mut store)?
         .get_typed_func::<(u8,), (), _>(&mut store, "take-u8")?
         .call(&mut store, (101,))
         .unwrap_err()
         .downcast::<Trap>()?;
-    instance
+    new_instance(&mut store)?
         .get_typed_func::<(i8,), (), _>(&mut store, "take-s8")?
         .call(&mut store, (101,))
         .unwrap_err()
         .downcast::<Trap>()?;
-    instance
+    new_instance(&mut store)?
         .get_typed_func::<(u16,), (), _>(&mut store, "take-u16")?
         .call(&mut store, (101,))
         .unwrap_err()
         .downcast::<Trap>()?;
-    instance
+    new_instance(&mut store)?
         .get_typed_func::<(i16,), (), _>(&mut store, "take-s16")?
         .call(&mut store, (101,))
         .unwrap_err()
         .downcast::<Trap>()?;
-    instance
+    new_instance(&mut store)?
         .get_typed_func::<(u32,), (), _>(&mut store, "take-u32")?
         .call(&mut store, (101,))
         .unwrap_err()
         .downcast::<Trap>()?;
-    instance
+    new_instance(&mut store)?
         .get_typed_func::<(i32,), (), _>(&mut store, "take-s32")?
         .call(&mut store, (101,))
         .unwrap_err()
         .downcast::<Trap>()?;
-    instance
+    new_instance(&mut store)?
         .get_typed_func::<(u64,), (), _>(&mut store, "take-u64")?
         .call(&mut store, (101,))
         .unwrap_err()
         .downcast::<Trap>()?;
-    instance
+    new_instance(&mut store)?
         .get_typed_func::<(i64,), (), _>(&mut store, "take-s64")?
         .call(&mut store, (101,))
         .unwrap_err()
@@ -606,13 +607,26 @@ fn chars() -> Result<()> {
     roundtrip('\n')?;
     roundtrip('💝')?;
 
-    let err = u32_to_char.call(&mut store, (0xd800,)).unwrap_err();
+    let u32_to_char = |store: &mut Store<()>| {
+        Linker::new(&engine)
+            .instantiate(&mut *store, &component)?
+            .get_typed_func::<(u32,), char, _>(&mut *store, "u32-to-char")
+    };
+    let err = u32_to_char(&mut store)?
+        .call(&mut store, (0xd800,))
+        .unwrap_err();
     assert!(err.to_string().contains("integer out of range"), "{}", err);
-    let err = u32_to_char.call(&mut store, (0xdfff,)).unwrap_err();
+    let err = u32_to_char(&mut store)?
+        .call(&mut store, (0xdfff,))
+        .unwrap_err();
     assert!(err.to_string().contains("integer out of range"), "{}", err);
-    let err = u32_to_char.call(&mut store, (0x110000,)).unwrap_err();
+    let err = u32_to_char(&mut store)?
+        .call(&mut store, (0x110000,))
+        .unwrap_err();
     assert!(err.to_string().contains("integer out of range"), "{}", err);
-    let err = u32_to_char.call(&mut store, (u32::MAX,)).unwrap_err();
+    let err = u32_to_char(&mut store)?
+        .call(&mut store, (u32::MAX,))
+        .unwrap_err();
     assert!(err.to_string().contains("integer out of range"), "{}", err);
 
     Ok(())
@@ -1068,10 +1082,10 @@ fn some_traps() -> Result<()> {
     let engine = super::engine();
     let component = Component::new(&engine, component)?;
     let mut store = Store::new(&engine, ());
-    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let instance = |store: &mut Store<()>| Linker::new(&engine).instantiate(store, &component);
 
     // This should fail when calling the allocator function for the argument
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&[u8],), (), _>(&mut store, "take-list-unreachable")?
         .call(&mut store, (&[],))
         .unwrap_err()
@@ -1079,7 +1093,7 @@ fn some_traps() -> Result<()> {
     assert_eq!(err.trap_code(), Some(TrapCode::UnreachableCodeReached));
 
     // This should fail when calling the allocator function for the argument
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&str,), (), _>(&mut store, "take-string-unreachable")?
         .call(&mut store, ("",))
         .unwrap_err()
@@ -1088,7 +1102,7 @@ fn some_traps() -> Result<()> {
 
     // This should fail when calling the allocator function for the space
     // to store the arguments (before arguments are even lowered)
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&str, &str, &str, &str, &str, &str, &str, &str, &str, &str), (), _>(
             &mut store,
             "take-many-unreachable",
@@ -1113,27 +1127,27 @@ fn some_traps() -> Result<()> {
             err,
         );
     }
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&[u8],), (), _>(&mut store, "take-list-base-oob")?
         .call(&mut store, (&[],))
         .unwrap_err();
     assert_oob(&err);
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&[u8],), (), _>(&mut store, "take-list-base-oob")?
         .call(&mut store, (&[1],))
         .unwrap_err();
     assert_oob(&err);
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&str,), (), _>(&mut store, "take-string-base-oob")?
         .call(&mut store, ("",))
         .unwrap_err();
     assert_oob(&err);
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&str,), (), _>(&mut store, "take-string-base-oob")?
         .call(&mut store, ("x",))
         .unwrap_err();
     assert_oob(&err);
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&str, &str, &str, &str, &str, &str, &str, &str, &str, &str), (), _>(
             &mut store,
             "take-many-base-oob",
@@ -1145,29 +1159,29 @@ fn some_traps() -> Result<()> {
     // Test here that when the returned pointer from malloc is one byte from the
     // end of memory that empty things are fine, but larger things are not.
 
-    instance
+    instance(&mut store)?
         .get_typed_func::<(&[u8],), (), _>(&mut store, "take-list-end-oob")?
         .call_and_post_return(&mut store, (&[],))?;
-    instance
+    instance(&mut store)?
         .get_typed_func::<(&[u8],), (), _>(&mut store, "take-list-end-oob")?
         .call_and_post_return(&mut store, (&[1, 2, 3, 4],))?;
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&[u8],), (), _>(&mut store, "take-list-end-oob")?
         .call(&mut store, (&[1, 2, 3, 4, 5],))
         .unwrap_err();
     assert_oob(&err);
-    instance
+    instance(&mut store)?
         .get_typed_func::<(&str,), (), _>(&mut store, "take-string-end-oob")?
         .call_and_post_return(&mut store, ("",))?;
-    instance
+    instance(&mut store)?
         .get_typed_func::<(&str,), (), _>(&mut store, "take-string-end-oob")?
         .call_and_post_return(&mut store, ("abcd",))?;
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&str,), (), _>(&mut store, "take-string-end-oob")?
         .call(&mut store, ("abcde",))
         .unwrap_err();
     assert_oob(&err);
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&str, &str, &str, &str, &str, &str, &str, &str, &str, &str), (), _>(
             &mut store,
             "take-many-end-oob",
@@ -1179,7 +1193,7 @@ fn some_traps() -> Result<()> {
     // For this function the first allocation, the space to store all the
     // arguments, is in-bounds but then all further allocations, such as for
     // each individual string, are all out of bounds.
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(&str, &str, &str, &str, &str, &str, &str, &str, &str, &str), (), _>(
             &mut store,
             "take-many-second-oob",
@@ -1304,9 +1318,12 @@ fn string_list_oob() -> Result<()> {
     let engine = super::engine();
     let component = Component::new(&engine, component)?;
     let mut store = Store::new(&engine, ());
-    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
-    let ret_list_u8 = instance.get_typed_func::<(), WasmList<u8>, _>(&mut store, "ret-list-u8")?;
-    let ret_string = instance.get_typed_func::<(), WasmStr, _>(&mut store, "ret-string")?;
+    let ret_list_u8 = Linker::new(&engine)
+        .instantiate(&mut store, &component)?
+        .get_typed_func::<(), WasmList<u8>, _>(&mut store, "ret-list-u8")?;
+    let ret_string = Linker::new(&engine)
+        .instantiate(&mut store, &component)?
+        .get_typed_func::<(), WasmStr, _>(&mut store, "ret-string")?;
 
     let err = ret_list_u8.call(&mut store, ()).err().unwrap();
     assert!(err.to_string().contains("out of bounds"), "{}", err);
@@ -1460,7 +1477,8 @@ fn option() -> Result<()> {
     let engine = super::engine();
     let component = Component::new(&engine, component)?;
     let mut store = Store::new(&engine, ());
-    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let linker = Linker::new(&engine);
+    let instance = linker.instantiate(&mut store, &component)?;
     let option_unit_to_u32 =
         instance.get_typed_func::<(Option<()>,), u32, _>(&mut store, "option-unit-to-u32")?;
     assert_eq!(option_unit_to_u32.call(&mut store, (None,))?, 0);
@@ -1506,6 +1524,7 @@ fn option() -> Result<()> {
     assert_eq!(b.to_str(&store)?, "hello");
     option_string_to_tuple.post_return(&mut store)?;
 
+    let instance = linker.instantiate(&mut store, &component)?;
     let to_option_unit =
         instance.get_typed_func::<(u32,), Option<()>, _>(&mut store, "to-option-unit")?;
     assert_eq!(to_option_unit.call(&mut store, (0,))?, None);
@@ -1515,6 +1534,7 @@ fn option() -> Result<()> {
     let err = to_option_unit.call(&mut store, (2,)).unwrap_err();
     assert!(err.to_string().contains("invalid option"), "{}", err);
 
+    let instance = linker.instantiate(&mut store, &component)?;
     let to_option_u8 =
         instance.get_typed_func::<(u32, u32), Option<u8>, _>(&mut store, "to-option-u8")?;
     assert_eq!(to_option_u8.call(&mut store, (0x00_00, 0))?, None);
@@ -1525,6 +1545,7 @@ fn option() -> Result<()> {
     to_option_u8.post_return(&mut store)?;
     assert!(to_option_u8.call(&mut store, (0x00_02, 0)).is_err());
 
+    let instance = linker.instantiate(&mut store, &component)?;
     let to_option_u32 =
         instance.get_typed_func::<(u32, u32), Option<u32>, _>(&mut store, "to-option-u32")?;
     assert_eq!(to_option_u32.call(&mut store, (0, 0))?, None);
@@ -1538,6 +1559,7 @@ fn option() -> Result<()> {
     to_option_u32.post_return(&mut store)?;
     assert!(to_option_u32.call(&mut store, (2, 0)).is_err());
 
+    let instance = linker.instantiate(&mut store, &component)?;
     let to_option_string = instance
         .get_typed_func::<(u32, &str), Option<WasmStr>, _>(&mut store, "to-option-string")?;
     let ret = to_option_string.call(&mut store, (0, ""))?;
@@ -1637,7 +1659,8 @@ fn expected() -> Result<()> {
     let engine = super::engine();
     let component = Component::new(&engine, component)?;
     let mut store = Store::new(&engine, ());
-    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let linker = Linker::new(&engine);
+    let instance = linker.instantiate(&mut store, &component)?;
     let take_expected_unit =
         instance.get_typed_func::<(Result<(), ()>,), u32, _>(&mut store, "take-expected-unit")?;
     assert_eq!(take_expected_unit.call(&mut store, (Ok(()),))?, 0);
@@ -1669,6 +1692,7 @@ fn expected() -> Result<()> {
     assert_eq!(b.to_str(&store)?, "goodbye");
     take_expected_string.post_return(&mut store)?;
 
+    let instance = linker.instantiate(&mut store, &component)?;
     let to_expected_unit =
         instance.get_typed_func::<(u32,), Result<(), ()>, _>(&mut store, "to-expected-unit")?;
     assert_eq!(to_expected_unit.call(&mut store, (0,))?, Ok(()));
@@ -1678,6 +1702,7 @@ fn expected() -> Result<()> {
     let err = to_expected_unit.call(&mut store, (2,)).unwrap_err();
     assert!(err.to_string().contains("invalid expected"), "{}", err);
 
+    let instance = linker.instantiate(&mut store, &component)?;
     let to_expected_s16_f32 = instance
         .get_typed_func::<(u32, u32), Result<i16, f32>, _>(&mut store, "to-expected-s16-f32")?;
     assert_eq!(to_expected_s16_f32.call(&mut store, (0, 0))?, Ok(0));
@@ -1869,9 +1894,9 @@ fn invalid_alignment() -> Result<()> {
     let engine = super::engine();
     let component = Component::new(&engine, component)?;
     let mut store = Store::new(&engine, ());
-    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let instance = |store: &mut Store<()>| Linker::new(&engine).instantiate(store, &component);
 
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(
             &str,
             &str,
@@ -1895,7 +1920,7 @@ fn invalid_alignment() -> Result<()> {
         err
     );
 
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(), WasmStr, _>(&mut store, "string-ret")?
         .call(&mut store, ())
         .err()
@@ -1906,7 +1931,7 @@ fn invalid_alignment() -> Result<()> {
         err
     );
 
-    let err = instance
+    let err = instance(&mut store)?
         .get_typed_func::<(), WasmList<u32>, _>(&mut store, "list-u32-ret")?
         .call(&mut store, ())
         .err()
@@ -2253,3 +2278,88 @@ fn lower_then_lift() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn errors_that_poison_instance() -> Result<()> {
+    let component = format!(
+        r#"
+(component $c
+  (core module $m1
+    (func (export "f1") unreachable)
+    (func (export "f2"))
+  )
+  (core instance $m1 (instantiate $m1))
+  (func (export "f1") (canon lift (core func $m1 "f1")))
+  (func (export "f2") (canon lift (core func $m1 "f2")))
+
+  (core module $m2
+    (func (export "f") (param i32 i32))
+    (func (export "r") (param i32 i32 i32 i32) (result i32) unreachable)
+    (memory (export "m") 1)
+  )
+  (core instance $m2 (instantiate $m2))
+  (func (export "f3") (param string)
+    (canon lift (core func $m2 "f") (realloc (func $m2 "r")) (memory $m2 "m"))
+  )
+
+  (core module $m3
+    (func (export "f") (result i32) i32.const 1)
+    (memory (export "m") 1)
+  )
+  (core instance $m3 (instantiate $m3))
+  (func (export "f4") (result string)
+    (canon lift (core func $m3 "f") (memory $m3 "m"))
+  )
+)
+    "#
+    );
+
+    let engine = super::engine();
+    let component = Component::new(&engine, component)?;
+    let mut store = Store::new(&engine, ());
+    let linker = Linker::new(&engine);
+    let instance = linker.instantiate(&mut store, &component)?;
+    let f1 = instance.get_typed_func::<(), (), _>(&mut store, "f1")?;
+    let f2 = instance.get_typed_func::<(), (), _>(&mut store, "f2")?;
+    assert_unreachable(f1.call(&mut store, ()));
+    assert_poisoned(f1.call(&mut store, ()));
+    assert_poisoned(f2.call(&mut store, ()));
+
+    let instance = linker.instantiate(&mut store, &component)?;
+    let f3 = instance.get_typed_func::<(&str,), (), _>(&mut store, "f3")?;
+    assert_unreachable(f3.call(&mut store, ("x",)));
+    assert_poisoned(f3.call(&mut store, ("x",)));
+
+    let instance = linker.instantiate(&mut store, &component)?;
+    let f4 = instance.get_typed_func::<(), WasmStr, _>(&mut store, "f4")?;
+    assert!(f4.call(&mut store, ()).is_err());
+    assert_poisoned(f4.call(&mut store, ()));
+
+    return Ok(());
+
+    #[track_caller]
+    fn assert_unreachable<T>(err: Result<T>) {
+        let err = match err {
+            Ok(_) => panic!("expected an error"),
+            Err(e) => e,
+        };
+        assert_eq!(
+            err.downcast::<Trap>().unwrap().trap_code(),
+            Some(TrapCode::UnreachableCodeReached)
+        );
+    }
+
+    #[track_caller]
+    fn assert_poisoned<T>(err: Result<T>) {
+        let err = match err {
+            Ok(_) => panic!("expected an error"),
+            Err(e) => e,
+        };
+        assert!(
+            err.to_string()
+                .contains("cannot reenter component instance"),
+            "{}",
+            err,
+        );
+    }
+}

tests/all/component_model/import.rs

@@ -221,11 +221,11 @@ fn attempt_to_leave_during_malloc() -> Result<()> {
         })?;
     let component = Component::new(&engine, component)?;
     let mut store = Store::new(&engine, ());
-    let instance = linker.instantiate(&mut store, &component)?;
 
     // Assert that during a host import if we return values to wasm that a trap
     // happens if we try to leave the instance.
-    let trap = instance
+    let trap = linker
+        .instantiate(&mut store, &component)?
         .get_typed_func::<(), (), _>(&mut store, "run")?
         .call(&mut store, ())
         .unwrap_err()
@@ -261,7 +261,8 @@ fn attempt_to_leave_during_malloc() -> Result<()> {
     // In addition to the above trap also ensure that when we enter a wasm
     // component if we try to leave while lowering then that's also a dynamic
     // trap.
-    let trap = instance
+    let trap = linker
+        .instantiate(&mut store, &component)?
         .get_typed_func::<(&str,), (), _>(&mut store, "take-string")?
         .call(&mut store, ("x",))
         .unwrap_err()
@@ -599,14 +600,15 @@ fn bad_import_alignment() -> Result<()> {
     )?;
     let component = Component::new(&engine, component)?;
     let mut store = Store::new(&engine, ());
-    let instance = linker.instantiate(&mut store, &component)?;
-    let trap = instance
+    let trap = linker
+        .instantiate(&mut store, &component)?
         .get_typed_func::<(), (), _>(&mut store, "unaligned-retptr")?
         .call(&mut store, ())
         .unwrap_err()
         .downcast::<Trap>()?;
     assert!(trap.to_string().contains("pointer not aligned"), "{}", trap);
-    let trap = instance
+    let trap = linker
+        .instantiate(&mut store, &component)?
         .get_typed_func::<(), (), _>(&mut store, "unaligned-argptr")?
         .call(&mut store, ())
         .unwrap_err()

tests/all/component_model/post_return.rs

@@ -1,6 +1,6 @@
 use anyhow::Result;
 use wasmtime::component::*;
-use wasmtime::{Store, StoreContextMut};
+use wasmtime::{Store, StoreContextMut, Trap, TrapCode};
 
 #[test]
 fn invalid_api() -> Result<()> {
@@ -257,3 +257,44 @@ fn post_return_string() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn trap_in_post_return_poisons_instance() -> Result<()> {
+    let component = r#"
+        (component
+            (core module $m
+                (func (export "f"))
+                (func (export "post") unreachable)
+            )
+            (core instance $i (instantiate $m))
+            (func (export "f")
+                (canon lift
+                    (core func $i "f")
+                    (post-return (func $i "post"))
+                )
+            )
+        )
+    "#;
+
+    let engine = super::engine();
+    let component = Component::new(&engine, component)?;
+    let mut store = Store::new(&engine, ());
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let f = instance.get_typed_func::<(), (), _>(&mut store, "f")?;
+    f.call(&mut store, ())?;
+    let trap = f.post_return(&mut store).unwrap_err().downcast::<Trap>()?;
+    assert_eq!(trap.trap_code(), Some(TrapCode::UnreachableCodeReached));
+    let err = f.call(&mut store, ()).unwrap_err();
+    assert!(
+        err.to_string()
+            .contains("cannot reenter component instance"),
+        "{}",
+        err
+    );
+    assert_panics(
+        || drop(f.post_return(&mut store)),
+        "can only be called after",
+    );
+
+    Ok(())
+}