Unsoundness due to thread locals

[hanna-kruppe]

Hi, thanks for publishing this crate! I found it while looking for proc-macro-free async-stream alternatives, that would be a great thing to have. Unfortunately the current implementation is unsound, because thread locals can be used to smuggle YieldFutures out of the async scope, breaking the stack discipline w.r.t. STATE that the library tries to enforce. Proof of concept:

#![forbid(unsafe_code)]
use std::{
    cell::Cell,
    pin::pin,
    sync::Arc,
    task::{Context, Poll, Wake, Waker},
};

// async-stream-lite = "=0.1.0"
use async_stream_lite::{async_stream, YieldFut};
// futures-core = "0.3.31"
use futures_core::Stream;

thread_local! {
    static YIELD_FUT_USIZE: Cell<Option<YieldFut<usize>>> = const { Cell::new(None) };
}

fn main() {
    // Step 0: get a Context to run some futures by hand
    let waker = Waker::from(Arc::new(NoopWaker));
    let mut cx = Context::from_waker(&waker);

    // Step 1: get our hands on a YieldFut that's ready to write Some(0usize) through the pointer
    // stored in `async_stream_lite::STORE`.
    let stream1 = async_stream(|r#yield| async move {
        YIELD_FUT_USIZE.with(|cell| cell.set(Some(r#yield(0usize))));
    });
    let _ = pin!(stream1).poll_next(&mut cx);

    // Step 2: create a stream that's supposed to yield Box<i32>...
    let stream2 = async_stream(|r#yield| async move {
        drop(r#yield(Box::new(0i32)));
        // ... but actually writes Some(0usize) in the `Option<Box<i23>>` destination, using the
        // YieldFuture from step 1. Incidentially, this also scribbles over adjacent stack memory
        // (`Option<usize>` is twice as big as `Option<Box<i32>>`) so we already have UB here.
        let yield_0_fut: YieldFut<usize> = YIELD_FUT_USIZE.with(|cell| cell.take().unwrap());
        yield_0_fut.await;
    });

    // Step 3: poll the stream and get the bogus `Box` out of it.
    let bogus_box: Box<i32> = match pin!(stream2).poll_next(&mut cx) {
        Poll::Ready(elem) => elem.expect("stream should yield an item"),
        Poll::Pending => unreachable!(),
    };
    // Step 4: exhibit UB as segmentation fault (release mode) or failing UB-check (debug mode).
    let number = *bogus_box;
    println!("{number}");
}

struct NoopWaker;

impl Wake for NoopWaker {
    fn wake(self: Arc<Self>) {}
}

[decahedron1]

I believe https://github.com/pykeio/async-stream-lite/commit/d382efa5a2f3c43c04e1dcc48bbc354f94e67e8e fixes this? (Not a huge fan of the added verbosity but it's better than allowing UB - would love to hear if there's a better way!)

[hanna-kruppe]

Thanks for the super quick response! Adding a lifetime somehow might prevent this kind of escaping. Maybe some GhostCell-style "branded types" can be adapted. The linked commit is not sufficient because the new Yielder type doesn't have a lifetime so we can just smuggle that one out instead of the actual future. Diff:

--- src/main_old.rs     2024-11-19 21:38:23.777954798 +0100
+++ src/main.rs 2024-11-19 21:37:35.257954174 +0100
@@ -7,12 +7,12 @@
 };

 // async-stream-lite = "=0.1.0"
-use async_stream_lite::{async_stream, YieldFut};
+use async_stream_lite::{async_stream, Yielder};
 // futures-core = "0.3.31"
 use futures_core::Stream;

 thread_local! {
-    static YIELD_FUT_USIZE: Cell<Option<YieldFut<usize>>> = const { Cell::new(None) };
+    static YIELD_FUT_USIZE: Cell<Option<Yielder<usize>>> = const { Cell::new(None) };
 }

 fn main() {
@@ -22,19 +22,18 @@

     // Step 1: get our hands on a YieldFut that's ready to write Some(0usize) through the pointer
     // stored in `async_stream_lite::STORE`.
-    let stream1 = async_stream(|r#yield| async move {
-        YIELD_FUT_USIZE.with(|cell| cell.set(Some(r#yield(0usize))));
+    let stream1 = async_stream(|yielder| async move {
+        YIELD_FUT_USIZE.with(|cell| cell.set(Some(yielder)));
     });
     let _ = pin!(stream1).poll_next(&mut cx);

     // Step 2: create a stream that's supposed to yield Box<i32>...
-    let stream2 = async_stream(|r#yield| async move {
-        drop(r#yield(Box::new(0i32)));
+    let stream2 = async_stream(|_: Yielder<Box<i32>>| async move {
         // ... but actually writes Some(0usize) in the `Option<Box<i23>>` destination, using the
         // YieldFuture from step 1. Incidentially, this also scribbles over adjacent stack memory
         // (`Option<usize>` is twice as big as `Option<Box<i32>>`) so we already have UB here.
-        let yield_0_fut: YieldFut<usize> = YIELD_FUT_USIZE.with(|cell| cell.take().unwrap());
-        yield_0_fut.await;
+        let yielder: Yielder<usize> = YIELD_FUT_USIZE.with(|cell| cell.take().unwrap());
+        yielder.r#yield(0usize).await;
     });

     // Step 3: poll the stream and get the bogus `Box` out of it.