Update fibers to avoid no-return functions (#12928)

* Update fibers to avoid no-return functions

This commit is aimed at fixing the ASAN false positives in #12899.
Initially the fix there was to invoke some `__asan_*` intrinsics, and
I ended up finding a sort of smaller set of `__asan_*` intrinsics to
call as well. In the end what's happening though is that fibers, upon
terminating, have a few frames of Rust code on the stack before
switching off. To ASAN these frames never returned so when a stack is
subsequently reused ASAN is tricked into thinking this is buffer
overflow or use-after-free since it's stomping on frames that haven't
returned.

The fix in this commit is to avoid this style of function which doesn't
returns. Functions which don't return in Rust are easy to leak memory
from and are a hazard from a safety perspective as well (e.g. it's
unsafe to skip running destructors of stack variables). I feel we've had
better success over time with "all Rust functions always return" and so
what's what was applied here. Unlike #12899 or my thoughts on that PR
this does not have any new `__asan_*` intrinsic calls. Instead what this
does is it shuffles around responsibility for what exact piece of the
infrastructure is responsible for what. Specifically `fiber_start`
functions now actually return, meaning the `wasmtime_fiber_start` naked
function actually resumes execution, unlike before. The
`wasmtime_fiber_start` then delegates to `wasmtime_fiber_switch`
immediately to perform the final switch.

Effectively there's now only two function frames that never return, and
both of these frames are handwritten inline assembly. This means that
ASAN gets to see that all normal functions return and updates all of its
metadata accordingly. The end result is that the original issue from #12899
is fixed and this I feel is in general more robust as well.

One caveat is that the handwritten `wasmtime_fiber_start` assembly needs
to invoke a sibling `wasmtime_fiber_switch_` function. In lieu of trying
to figure out how to get PIC-vs-not calls working (e.g. static calls)
I've opted to use indirect function calls and pointers instead. This
mirrors historical changes in our fiber implementation too.

* Fix CI builds

* Fix miri

Author: alexcrichton

crates/fiber/src/lib.rs

@@ -255,7 +255,7 @@ impl<Resume, Yield, Return> Suspend<Resume, Yield, Return> {
         inner: imp::Suspend,
         initial: Resume,
         func: impl FnOnce(Resume, &mut Suspend<Resume, Yield, Return>) -> Return,
-    ) {
+    ) -> imp::Suspend {
         let mut suspend = Suspend {
             inner,
             _phantom: PhantomData,
@@ -278,7 +278,8 @@ impl<Resume, Yield, Return> Suspend<Resume, Yield, Return> {
         #[cfg(not(feature = "std"))]
         let result = RunResult::Returned((func)(initial, &mut suspend));
 
-        suspend.inner.exit::<Resume, Yield, Return>(result);
+        suspend.inner.start_exit::<Resume, Yield, Return>(result);
+        suspend.inner
     }
 }
 
@@ -431,7 +432,7 @@ mod tests {
 
     #[test]
     fn fiber_stack_max_size() {
-        if cfg!(windows) {
+        if cfg!(windows) || cfg!(miri) {
             return;
         }
         assert!(FiberStack::new(usize::MAX, true).is_err());

crates/fiber/src/miri.rs

@@ -116,13 +116,17 @@ where
     // Execute this fiber through `Suspend::execute` and once that's done
     // deallocate the `state` that we have.
     let state = Arc::into_raw(state);
-    super::Suspend::<A, B, C>::execute(
+    let mut suspend = super::Suspend::<A, B, C>::execute(
         Suspend {
             state: state.cast(),
         },
         init,
         func.0,
     );
+    match suspend.block_until_notified::<A, B, C>() {
+        State::Exiting => {}
+        _ => unreachable!(),
+    }
     unsafe {
         drop(Arc::from_raw(state));
     }
@@ -150,8 +154,7 @@ impl Fiber {
                     let state = state.clone();
                     let func = IgnoreSendSync(func);
                     move || run(state, func)
-                })
-                .unwrap()
+                })?
         };
 
         // Cast the fiber back into a raw pointer to lose the type parameters
@@ -210,7 +213,7 @@ impl Fiber {
 }
 
 impl Suspend {
-    fn suspend<A, B, C>(&mut self, result: RunResult<A, B, C>) -> State<A, B, C> {
+    fn set_result<A, B, C>(&mut self, result: RunResult<A, B, C>) {
         let state = unsafe { self.state() };
         let mut lock = state.state.lock().unwrap();
 
@@ -219,9 +222,11 @@ impl Suspend {
         assert!(matches!(*lock, State::None));
         *lock = State::SuspendWith(result);
         state.cond.notify_one();
+    }
 
-        // Wait for the resumption to come back, which is returned from this
-        // method.
+    fn block_until_notified<A, B, C>(&mut self) -> State<A, B, C> {
+        let state = unsafe { self.state() };
+        let mut lock = state.state.lock().unwrap();
         lock = state
             .cond
             .wait_while(lock, |s| {
@@ -232,17 +237,18 @@ impl Suspend {
     }
 
     pub(crate) fn switch<A, B, C>(&mut self, result: RunResult<A, B, C>) -> A {
-        match self.suspend(result) {
+        self.set_result(result);
+
+        // Wait for the resumption to come back, which is returned from this
+        // method.
+        match self.block_until_notified::<A, B, C>() {
             State::ResumeWith(RunResult::Resuming(a)) => a,
             _ => unreachable!(),
         }
     }
 
-    pub(crate) fn exit<A, B, C>(&mut self, result: RunResult<A, B, C>) {
-        match self.suspend(result) {
-            State::Exiting => {}
-            _ => unreachable!(),
-        }
+    pub(crate) fn start_exit<A, B, C>(&mut self, result: RunResult<A, B, C>) {
+        self.set_result(result);
     }
 
     unsafe fn state<A, B, C>(&self) -> &SharedFiberState<A, B, C> {

crates/fiber/src/nostd.rs

@@ -115,14 +115,16 @@ pub struct Suspend {
     top_of_stack: *mut u8,
 }
 
-extern "C" fn fiber_start<F, A, B, C>(arg0: *mut u8, top_of_stack: *mut u8)
+extern "C" fn fiber_start<F, A, B, C>(arg0: *mut u8, top_of_stack: *mut u8) -> *mut u8
 where
     F: FnOnce(A, &mut super::Suspend<A, B, C>) -> C,
 {
     unsafe {
         let inner = Suspend { top_of_stack };
         let initial = inner.take_resume::<A, B, C>();
-        super::Suspend::<A, B, C>::execute(inner, initial, Box::from_raw(arg0.cast::<F>()))
+        let inner =
+            super::Suspend::<A, B, C>::execute(inner, initial, Box::from_raw(arg0.cast::<F>()));
+        inner.top_of_stack
     }
 }
 
@@ -176,9 +178,10 @@ impl Suspend {
         }
     }
 
-    pub(crate) fn exit<A, B, C>(&mut self, result: RunResult<A, B, C>) {
-        self.switch(result);
-        unreachable!();
+    pub(crate) fn start_exit<A, B, C>(&mut self, result: RunResult<A, B, C>) {
+        unsafe {
+            (*self.result_location::<A, B, C>()).set(result);
+        }
     }
 
     unsafe fn take_resume<A, B, C>(&self) -> A {

crates/fiber/src/stackswitch.rs

@@ -65,7 +65,7 @@ mod unsupported {
 
     pub(crate) unsafe fn wasmtime_fiber_init(
         _top_of_stack: *mut u8,
-        _entry: extern "C" fn(*mut u8, *mut u8),
+        _entry: extern "C" fn(*mut u8, *mut u8) -> *mut u8,
         _entry_arg0: *mut u8,
     ) {
         unreachable!();

crates/fiber/src/stackswitch/aarch64.rs

@@ -92,7 +92,7 @@ unsafe extern "C" fn wasmtime_fiber_switch_(top_of_stack: *mut u8 /* x0 */) {
 
 pub(crate) unsafe fn wasmtime_fiber_init(
     top_of_stack: *mut u8,
-    entry_point: extern "C" fn(*mut u8, *mut u8),
+    entry_point: extern "C" fn(*mut u8, *mut u8) -> *mut u8,
     entry_arg0: *mut u8, // x2
 ) {
     #[repr(C)]
@@ -132,6 +132,7 @@ pub(crate) unsafe fn wasmtime_fiber_init(
             x19: top_of_stack,
             x20: entry_point as *mut u8,
             x21: entry_arg0,
+            x22: wasmtime_fiber_switch_ as *mut u8,
 
             // We set up the newly initialized fiber, so that it resumes
             // execution from wasmtime_fiber_start(). As a result, we need a
@@ -208,6 +209,11 @@ unsafe extern "C" fn wasmtime_fiber_start() -> ! {
         // ... and then we call the function! Note that this is a function call
         // so our frame stays on the stack to backtrace through.
         blr x20
+
+        // The entry function returns where to switch to as the final switch, so
+        // that's performed here in inline assembly.
+        blr x22
+
         // Unreachable, here for safety. This should help catch unexpected
         // behaviors.  Use a noticeable payload so one can grep for it in the
         // codebase.

crates/fiber/src/stackswitch/arm.rs

@@ -36,7 +36,7 @@ unsafe extern "C" fn wasmtime_fiber_switch_(top_of_stack: *mut u8 /* r0 */) {
 
 pub(crate) unsafe fn wasmtime_fiber_init(
     top_of_stack: *mut u8,
-    entry_point: extern "C" fn(*mut u8, *mut u8),
+    entry_point: extern "C" fn(*mut u8, *mut u8) -> *mut u8,
     entry_arg0: *mut u8,
 ) {
     #[repr(C)]
@@ -60,6 +60,7 @@ pub(crate) unsafe fn wasmtime_fiber_init(
     unsafe {
         let initial_stack = top_of_stack.cast::<InitialStack>().sub(1);
         initial_stack.write(InitialStack {
+            r8: wasmtime_fiber_switch_ as *mut u8,
             r9: entry_arg0,
             r10: entry_point as *mut u8,
             r11: top_of_stack,
@@ -103,6 +104,7 @@ unsafe extern "C" fn wasmtime_fiber_start() -> ! {
         mov r1, r11
         mov r0, r9
         blx r10
+        blx r8
         .cfi_endproc
         ",
     );

crates/fiber/src/stackswitch/riscv32imac.rs

@@ -72,7 +72,7 @@ unsafe extern "C" fn wasmtime_fiber_switch_(top_of_stack: *mut u8 /* a0 */) {
 
 pub(crate) unsafe fn wasmtime_fiber_init(
     top_of_stack: *mut u8,
-    entry_point: extern "C" fn(*mut u8, *mut u8),
+    entry_point: extern "C" fn(*mut u8, *mut u8) -> *mut u8,
     entry_arg0: *mut u8,
 ) {
     #[repr(C)]
@@ -106,6 +106,7 @@ pub(crate) unsafe fn wasmtime_fiber_init(
         initial_stack.write(InitialStack {
             s1: entry_point as *mut u8,
             s2: entry_arg0,
+            s3: wasmtime_fiber_switch_ as *mut u8,
             fp: top_of_stack,
             ra: wasmtime_fiber_start as *mut u8,
             last_sp: initial_stack.cast(),
@@ -147,6 +148,7 @@ unsafe extern "C" fn wasmtime_fiber_start() -> ! {
       mv a0, s2
       mv a1, fp
       jalr s1
+      jalr s3
       // .4byte 0 will cause panic.
       // for safety just like x86_64.rs and riscv64.rs.
       .4byte 0

crates/fiber/src/stackswitch/riscv64.rs

@@ -91,7 +91,7 @@ unsafe extern "C" fn wasmtime_fiber_switch_(top_of_stack: *mut u8 /* a0 */) {
 
 pub(crate) unsafe fn wasmtime_fiber_init(
     top_of_stack: *mut u8,
-    entry_point: extern "C" fn(*mut u8, *mut u8),
+    entry_point: extern "C" fn(*mut u8, *mut u8) -> *mut u8,
     entry_arg0: *mut u8,
 ) {
     #[repr(C)]
@@ -126,6 +126,7 @@ pub(crate) unsafe fn wasmtime_fiber_init(
         initial_stack.write(InitialStack {
             s1: entry_point as *mut u8,
             s2: entry_arg0,
+            s3: wasmtime_fiber_switch_ as *mut u8,
             fp: top_of_stack,
             ra: wasmtime_fiber_start as *mut u8,
             last_sp: initial_stack.cast(),
@@ -143,10 +144,10 @@ unsafe extern "C" fn wasmtime_fiber_start() -> ! {
 
 
     .cfi_escape 0x0f, /* DW_CFA_def_cfa_expression */ \
-      5,             /* the byte length of this expression */ \
-      0x52,          /* DW_OP_reg2 (sp) */ \
-      0x06,          /* DW_OP_deref */ \
-      0x08, 0xd0,    /* DW_OP_const1u 0xc8 */ \
+        5,             /* the byte length of this expression */ \
+        0x52,          /* DW_OP_reg2 (sp) */ \
+        0x06,          /* DW_OP_deref */ \
+        0x08, 0xd0,    /* DW_OP_const1u 0xc8 */ \
       0x22           /* DW_OP_plus */
 
 
@@ -178,7 +179,9 @@ unsafe extern "C" fn wasmtime_fiber_start() -> ! {
 
       mv a0, s2
       mv a1, fp
-      jalr s1
+      jalr s1 // entry_point
+      jalr s3 // wasmtime_fiber_switch_
+
       // .4byte 0 will cause panic.
       // for safety just like x86_64.rs.
       .4byte 0

crates/fiber/src/stackswitch/s390x.rs

@@ -52,7 +52,7 @@ unsafe extern "C" fn wasmtime_fiber_switch_(top_of_stack: *mut u8 /* x0 */) {
 
 pub(crate) unsafe fn wasmtime_fiber_init(
     top_of_stack: *mut u8,
-    entry_point: extern "C" fn(*mut u8, *mut u8),
+    entry_point: extern "C" fn(*mut u8, *mut u8) -> *mut u8,
     entry_arg0: *mut u8,
 ) {
     #[repr(C)]
@@ -104,6 +104,7 @@ pub(crate) unsafe fn wasmtime_fiber_init(
             r6: top_of_stack,
             r7: entry_point as *mut u8,
             r8: entry_arg0,
+            r9: wasmtime_fiber_switch_ as *mut u8,
 
             last_sp: initial_stack.cast(),
             ..InitialStack::default()
@@ -145,6 +146,10 @@ unsafe extern "C" fn wasmtime_fiber_start() -> ! {
         // ... and then we call the function! Note that this is a function call so
         // our frame stays on the stack to backtrace through.
         basr %r14, %r7  // entry_point
+
+        // Perform the final switch.
+        basr %r14, %r9  // wasmtime_fiber_switch_
+
         // .. technically we shouldn't get here, so just trap.
         .word 0x0000
         .cfi_endproc

crates/fiber/src/stackswitch/x86.rs

@@ -47,7 +47,7 @@ unsafe extern "C" fn wasmtime_fiber_switch_(top_of_stack: *mut u8) {
 
 pub(crate) unsafe fn wasmtime_fiber_init(
     top_of_stack: *mut u8,
-    entry_point: extern "C" fn(*mut u8, *mut u8),
+    entry_point: extern "C" fn(*mut u8, *mut u8) -> *mut u8,
     entry_arg0: *mut u8,
 ) {
     // Our stack from top-to-bottom looks like:
@@ -84,6 +84,7 @@ pub(crate) unsafe fn wasmtime_fiber_init(
         let initial_stack = top_of_stack.cast::<InitialStack>().sub(1);
         initial_stack.write(InitialStack {
             ebp: entry_point as *mut u8,
+            esi: wasmtime_fiber_switch_ as *mut u8,
             return_address: wasmtime_fiber_start as *mut u8,
             arg1: entry_arg0,
             arg2: top_of_stack,
@@ -112,8 +113,15 @@ unsafe extern "C" fn wasmtime_fiber_start() -> ! {
         .cfi_rel_offset edi, -20
 
         // Our arguments and stack alignment are all prepped by
-        // `wasmtime_fiber_init`.
+        // `wasmtime_fiber_init`. After calling `entry_point` clean up the
+        // stack arguments.
         call ebp
+        pop edi
+        pop edi
+
+        // Call `wasmtime_fiber_switch_`, pushing its argument onto the stack.
+        push eax
+        call esi
         ud2
         .cfi_endproc
         ",