macros: don't take ownership of futures in macros

tokio/src/future/poll_fn.rs

@@ -4,22 +4,25 @@
 
 use std::fmt;
 use std::future::Future;
+use std::marker::PhantomPinned;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
 /// Future for the [`poll_fn`] function.
 pub struct PollFn<F> {
     f: F,
+    _pinned: PhantomPinned,
 }
 
-impl<F> Unpin for PollFn<F> {}
-
 /// Creates a new future wrapping around a function returning [`Poll`].
 pub fn poll_fn<T, F>(f: F) -> PollFn<F>
 where
     F: FnMut(&mut Context<'_>) -> Poll<T>,
 {
-    PollFn { f }
+    PollFn {
+        f,
+        _pinned: PhantomPinned,
+    }
 }
 
 impl<F> fmt::Debug for PollFn<F> {
@@ -34,7 +37,15 @@ where
 {
     type Output = T;
 
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
-        (self.f)(cx)
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
+        // SAFETY: We never construct a `Pin<&mut F>` anywhere, so accessing `f`
+        // mutably in an unpinned way is sound.
+        //
+        // This is, strictly speaking, not necessary. We could make `PollFn`
+        // unconditionally `Unpin` and avoid this unsafe. However, making this
+        // struct `!Unpin` mitigates the issues described here:
+        // <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484>
+        let me = unsafe { Pin::into_inner_unchecked(self) };
+        (me.f)(cx)
     }
 }

tokio/src/macros/join.rs

@@ -74,6 +74,12 @@ macro_rules! join {
         // the requirement of `Pin::new_unchecked` called below.
         let mut futures = ( $( maybe_done($e), )* );
 
+        // This assignment makes sure that the `poll_fn` closure only has a
+        // reference to the futures, instead of taking ownership of them. This
+        // mitigates the issue described in
+        // <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484>
+        let mut futures = &mut futures;
+
         // Each time the future created by poll_fn is polled, a different future will be polled first
         // to ensure every future passed to join! gets a chance to make progress even if
         // one of the futures consumes the whole budget.
@@ -106,7 +112,7 @@ macro_rules! join {
                     to_run -= 1;
 
                     // Extract the future for this branch from the tuple.
-                    let ( $($skip,)* fut, .. ) = &mut futures;
+                    let ( $($skip,)* fut, .. ) = &mut *futures;
 
                     // Safety: future is stored on the stack above
                     // and never moved.

tokio/src/macros/select.rs

@@ -462,6 +462,12 @@ macro_rules! select {
             // satisfy the requirement of `Pin::new_unchecked` called below.
             let mut futures = ( $( $fut , )+ );
 
+            // This assignment makes sure that the `poll_fn` closure only has a
+            // reference to the futures, instead of taking ownership of them.
+            // This mitigates the issue described in
+            // <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484>
+            let mut futures = &mut futures;
+
             $crate::macros::support::poll_fn(|cx| {
                 // Track if any branch returns pending. If no branch completes
                 // **or** returns pending, this implies that all branches are
@@ -497,7 +503,7 @@ macro_rules! select {
 
                                 // Extract the future for this branch from the
                                 // tuple
-                                let ( $($skip,)* fut, .. ) = &mut futures;
+                                let ( $($skip,)* fut, .. ) = &mut *futures;
 
                                 // Safety: future is stored on the stack above
                                 // and never moved.

tokio/src/macros/try_join.rs

@@ -120,6 +120,12 @@ macro_rules! try_join {
         // the requirement of `Pin::new_unchecked` called below.
         let mut futures = ( $( maybe_done($e), )* );
 
+        // This assignment makes sure that the `poll_fn` closure only has a
+        // reference to the futures, instead of taking ownership of them. This
+        // mitigates the issue described in
+        // <https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484>
+        let mut futures = &mut futures;
+
         // Each time the future created by poll_fn is polled, a different future will be polled first
         // to ensure every future passed to join! gets a chance to make progress even if
         // one of the futures consumes the whole budget.
@@ -152,7 +158,7 @@ macro_rules! try_join {
                     to_run -= 1;
 
                     // Extract the future for this branch from the tuple.
-                    let ( $($skip,)* fut, .. ) = &mut futures;
+                    let ( $($skip,)* fut, .. ) = &mut *futures;
 
                     // Safety: future is stored on the stack above
                     // and never moved.

tokio/tests/macros_join.rs

@@ -3,6 +3,7 @@
 use std::sync::Arc;
 
 #[cfg(tokio_wasm_not_wasi)]
+#[cfg(target_pointer_width = "64")]
 use wasm_bindgen_test::wasm_bindgen_test as test;
 #[cfg(tokio_wasm_not_wasi)]
 use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test;
@@ -64,6 +65,7 @@ async fn two_await() {
 }
 
 #[test]
+#[cfg(target_pointer_width = "64")]
 fn join_size() {
     use futures::future;
     use std::mem;
@@ -72,14 +74,14 @@ fn join_size() {
         let ready = future::ready(0i32);
         tokio::join!(ready)
     };
-    assert_eq!(mem::size_of_val(&fut), 20);
+    assert_eq!(mem::size_of_val(&fut), 32);
 
     let fut = async {
         let ready1 = future::ready(0i32);
         let ready2 = future::ready(0i32);
         tokio::join!(ready1, ready2)
     };
-    assert_eq!(mem::size_of_val(&fut), 32);
+    assert_eq!(mem::size_of_val(&fut), 48);
 }
 
 async fn non_cooperative_task(permits: Arc<Semaphore>) -> usize {

tokio/tests/macros_select.rs

@@ -207,6 +207,7 @@ async fn nested() {
 }
 
 #[maybe_tokio_test]
+#[cfg(target_pointer_width = "64")]
 async fn struct_size() {
     use futures::future;
     use std::mem;
@@ -219,7 +220,7 @@ async fn struct_size() {
         }
     };
 
-    assert!(mem::size_of_val(&fut) <= 32);
+    assert_eq!(mem::size_of_val(&fut), 40);
 
     let fut = async {
         let ready1 = future::ready(0i32);
@@ -231,7 +232,7 @@ async fn struct_size() {
         }
     };
 
-    assert!(mem::size_of_val(&fut) <= 40);
+    assert_eq!(mem::size_of_val(&fut), 48);
 
     let fut = async {
         let ready1 = future::ready(0i32);
@@ -245,7 +246,7 @@ async fn struct_size() {
         }
     };
 
-    assert!(mem::size_of_val(&fut) <= 48);
+    assert_eq!(mem::size_of_val(&fut), 56);
 }
 
 #[maybe_tokio_test]

tokio/tests/macros_try_join.rs

@@ -88,6 +88,7 @@ async fn err_abort_early() {
 }
 
 #[test]
+#[cfg(target_pointer_width = "64")]
 fn join_size() {
     use futures::future;
     use std::mem;
@@ -96,14 +97,14 @@ fn join_size() {
         let ready = future::ready(ok(0i32));
         tokio::try_join!(ready)
     };
-    assert_eq!(mem::size_of_val(&fut), 20);
+    assert_eq!(mem::size_of_val(&fut), 32);
 
     let fut = async {
         let ready1 = future::ready(ok(0i32));
         let ready2 = future::ready(ok(0i32));
         tokio::try_join!(ready1, ready2)
     };
-    assert_eq!(mem::size_of_val(&fut), 32);
+    assert_eq!(mem::size_of_val(&fut), 48);
 }
 
 fn ok<T>(val: T) -> Result<T, ()> {