Fix unsoundness issues tokio-rs#106 and tokio-rs#107

v3xro · Jun 26, 2024 · da7886f · da7886f
1 parent e8abfaa
commit da7886f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 8 deletions.
diff --git a/async-stream-impl/src/lib.rs b/async-stream-impl/src/lib.rs
@@ -120,10 +120,19 @@ impl VisitMut for Scrub<'_> {
 
                 // let ident = &self.yielder;
 
-                *i = if self.is_try {
-                    syn::parse_quote! { __yield_tx.send(::core::result::Result::Ok(#value_expr)).await }
+                let yield_expr = if self.is_try {
+                    quote! { __yield_tx.send(::core::result::Result::Ok(#value_expr)).await }
                 } else {
-                    syn::parse_quote! { __yield_tx.send(#value_expr).await }
+                    quote! { __yield_tx.send(#value_expr).await }
+                };
+                *i = syn::parse_quote! {
+                    {
+                        #[allow(unreachable_code)]
+                        if false {
+                            break 'check_stream_scope (loop {});
+                        }
+                        #yield_expr
+                    }
                 };
             }
             syn::Expr::Try(try_expr) => {
@@ -225,8 +234,10 @@ pub fn stream_inner(input: TokenStream) -> TokenStream {
     quote!({
         let (mut __yield_tx, __yield_rx) = unsafe { #crate_path::__private::yielder::pair() };
         #crate_path::__private::AsyncStream::new(__yield_rx, async move {
-            #dummy_yield
-            #(#stmts)*
+            'check_stream_scope: {
+                #dummy_yield
+                #(#stmts)*
+            }
         })
     })
     .into()
@@ -259,8 +270,10 @@ pub fn try_stream_inner(input: TokenStream) -> TokenStream {
     quote!({
         let (mut __yield_tx, __yield_rx) = unsafe { #crate_path::__private::yielder::pair() };
         #crate_path::__private::AsyncStream::new(__yield_rx, async move {
-            #dummy_yield
-            #(#stmts)*
+            'check_stream_scope: {
+                #dummy_yield
+                #(#stmts)*
+            }
         })
     })
     .into()

diff --git a/async-stream/src/yielder.rs b/async-stream/src/yielder.rs
@@ -7,7 +7,7 @@ use std::task::{Context, Poll};
 
 #[derive(Debug)]
 pub struct Sender<T> {
-    _p: PhantomData<T>,
+    _p: PhantomData<fn(T) -> T>,
 }
 
 #[derive(Debug)]