Undefined behavior from poll_fn + futures::join! / std::future::join!

[jplatte]

Recently, when working on a library for observables, I found new miri test failures (I had been getting them on a specific test before, but didn't investigate after ensuring it wasn't my own UB). This lead to a bug report to tokio where Alice Ryhl showed that it didn't actually require tokio::spawn (which I had failed to get rid of in minimizing the issue), and the same issue was reproducible with join! + futures_executor.

I have now reduced the example much further, to this (works with the nightly std::future::join!, or futures::join!; playground):

#![feature(future_join)]

use std::{
    future::{join, poll_fn, Future},
    pin::pin,
    sync::Arc,
    task::{Context, Poll, Wake},
};

fn main() {
    let mut fut = pin!(async {
        let mut thing = "hello".to_owned();
        let fut1 = async move {
            poll_fn(|_| {
                thing += ", world";
                Poll::<()>::Pending
            })
            .await;
        };

        let fut2 = async {};

        join!(fut1, fut2).await;
    });

    let waker = Arc::new(DummyWaker).into();
    let mut ctx = Context::from_waker(&waker);
    _ = fut.as_mut().poll(&mut ctx);
    _ = fut.as_mut().poll(&mut ctx);
}

struct DummyWaker;

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

Alice suggested that this may be an instance of rust-lang/rust#63818, is that correct? I was under the impression that the compiler for now doesn't implement alias restrictions for !Unpin types, but maybe this only applies to the LLVM backend, not miri?

[RalfJung]

Miri honors !Unpin, so there must be something more subtle going on here.

Please include the error when reporting a bug. Here it is:

error: Undefined Behavior: trying to retag from <3188> for SharedReadWrite permission at alloc1123[0x30], but that tag does not exist in the borrow stack for this location
  --> src/main.rs:15:17
   |
15 |                 thing += ", world";
   |                 ^^^^^
   |                 |
   |                 trying to retag from <3188> for SharedReadWrite permission at alloc1123[0x30], but that tag does not exist in the borrow stack for this location
   |                 this error occurs as part of two-phase retag at alloc1123[0x30..0x48]
   |
   = help: this indicates a potential bug in the program: it performed an invalid operation, but the Stacked Borrows rules it violated are still experimental
   = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/stacked-borrows.md for further information
help: <3188> was created by a Unique retag at offsets [0x30..0x48]
  --> src/main.rs:14:13
   |
14 | /             poll_fn(|_| {
15 | |                 thing += ", world";
16 | |                 Poll::<()>::Pending
17 | |             })
18 | |             .await;
   | |__________________^
help: <3188> was later invalidated at offsets [0x30..0x38] by a read access
  --> src/main.rs:23:9
   |
23 |         join!(fut1, fut2).await;
   |         ^^^^^^^^^^^^^^^^^
   = note: BACKTRACE (of the first span):
   = note: inside closure at src/main.rs:15:17: 15:22
   = note: inside `<std::future::PollFn<{closure@src/main.rs:14:21: 14:24}> as std::future::Future>::poll` at /playground/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/future/poll_fn.rs:151:9: 151:57
note: inside closure
  --> src/main.rs:18:14
   |
18 |             .await;
   |              ^^^^^
   = note: inside `<std::future::join::MaybeDone<{async block@src/main.rs:13:20: 13:30}> as std::future::Future>::poll` at /playground/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/future/join.rs:183:38: 183:68
note: inside closure
  --> src/main.rs:23:9
   |
23 |         join!(fut1, fut2).await;
   |         ^^^^^^^^^^^^^^^^^
   = note: inside `<std::future::PollFn<{closure@/playground/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/future/join.rs:110:21: 110:30}> as std::future::Future>::poll` at /playground/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/future/poll_fn.rs:151:9: 151:57
note: inside closure
  --> src/main.rs:23:9
   |
23 |         join!(fut1, fut2).await;
   |         ^^^^^^^^^^^^^^^^^
note: inside closure
  --> src/main.rs:23:27
   |
23 |         join!(fut1, fut2).await;
   |                           ^^^^^
note: inside `main`
  --> src/main.rs:29:9
   |
29 |     _ = fut.as_mut().poll(&mut ctx);
   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
   = note: this error originates in the macro `join_internal` which comes from the expansion of the macro `join` (in Nightly builds, run with -Z macro-backtrace for more info)

[Darksonn]

In the original bug report to Tokio, the UB was happening even though tokio::spawn was used instead of join!. So it may not be the join! macro's fault.

[taiki-e]

It seems adding move to poll_fn silences SB violation error: playground

[taiki-e]

adding move to poll_fn silences SB violation error

Or making thing &mut _: playground
(This reminds me of https://internals.rust-lang.org/t/surprising-soundness-trouble-around-pollfn/17484)

Btw, it is also possible to reproduce the error with a function, such as futures::future::join(), rather than a macro: playground

[oli-obk]

Minimized:

use std::{
    future::Future,
    pin::Pin,
    sync::Arc,
    task::{Context, Poll, Wake},
};

struct ThingAdder<'a> {
    thing: &'a mut String,
}

impl Future for ThingAdder<'_> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
        unsafe {
            *self.get_unchecked_mut().thing += ", world";
        }
        Poll::Pending
    }
}

fn main() {
    let mut thing = "hello".to_owned();
    let fut = async move { ThingAdder { thing: &mut thing }.await };

    let mut fut = MaybeDone::Future(fut);
    let mut fut = unsafe { Pin::new_unchecked(&mut fut) };

    let waker = Arc::new(DummyWaker).into();
    let mut ctx = Context::from_waker(&waker);
    assert_eq!(fut.as_mut().poll(&mut ctx), Poll::Pending);
    assert_eq!(fut.as_mut().poll(&mut ctx), Poll::Pending);
}

struct DummyWaker;

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

pub enum MaybeDone<F: Future> {
    Future(F),
    Done,
}
impl<F: Future<Output = ()>> Future for MaybeDone<F> {
    type Output = ();

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        unsafe {
            match *self.as_mut().get_unchecked_mut() {
                MaybeDone::Future(ref mut f) => Pin::new_unchecked(f).poll(cx),
                MaybeDone::Done => unreachable!(),
            }
        }
    }
}

Removing the MaybeDone::Done variant will also hide the UB

edit: was able to get rid of poll_fn, too

[RalfJung]

The memory that the conflict occurs on is the fut itself. Do I understand correctly that async move means we are moving thing into the future, and then creating a reference to it? So the future has (at least) two fields, a String and a ThingAdder pointing to that string.

The match *self.as_mut().get_unchecked_mut() will read the discriminant of the MaybeDone, which has been niche-optimized to be stored inside the String. This read overlaps with where the ThingAdder points to, thus invalidating the mutable reference. Finally, the mutable reference is used again, causing UB.

In other words, this is the "mutable reference aliasing" version of why UnsafeCell blocks niches: we also need to block niches for fields that have self-referential references pointing to them. So the fix here is to change async lowering to use UnsafePinned and then have UnsafePinned block niches the same way UnsafeCell blocks niches.

[RalfJung]

Since Miri is behaving as intended here, I'm going to close the issue in favor of rust-lang/rust#63818 and rust-lang/rust#125735. I don't see a way that we can work around this akin to the !Unpin hack; this fundamentally needs the layout of these futures to be different from what it is today.

Thanks for bringing this up, and for the very helpful minimization! In retrospect it is somewhat surprising that it took us so long to discover this, given that we already knew about the exact same issue with shared references and UnsafeCell...

Cc @rust-lang/opsem, just an FYI :)

[jplatte]

In retrospect it is somewhat surprising that it took us so long to discover this

I'm just as surprised! 😄

Given that the &mut method I'm trying to call in the real code that made me look into this is on my own type (not String), is there something I can do hide its niches without using unsafe, as a temporary measure to unblock me from using miri to detect other potential unsoundness?

[RalfJung]

If you only need &mut access, you could wrap the thing the mutable reference points to in an UnsafeCell. 😂 UnsafeCell::get_mut is safe, after all.

[Darksonn]

Is there anything Tokio should be doing to avoid UB with niche-optimized enums? We have the MaybeDone enum:

pin_project! {
    /// A future that may have completed.
    #[derive(Debug)]
    #[project = MaybeDoneProj]
    #[project_replace = MaybeDoneProjReplace]
    pub enum MaybeDone<Fut: Future> {
        /// A not-yet-completed future.
        Future { #[pin] future: Fut },
        /// The output of the completed future.
        Done { output: Fut::Output },
        /// The empty variant after the result of a [`MaybeDone`] has been
        /// taken using the [`take_output`](MaybeDone::take_output) method.
        Gone,
    }
}

[RalfJung]

I don't know if you should be doing anything, since it seems unlikely that this will lead to actual miscompilations. Though maybe @comex is able to trick LLVM into generating bad code here again. ;)

What you could do it mark that enum with a repr, which AFAIK prevents it from being niche-optimized.

The best thing to do would be to find someone who can implement UnsafePinned. :)

[taiki-e]

What you could do it mark that enum with a repr, which AFAIK prevents it from being niche-optimized.

While this is adequate as a workaround for the specific case here, it does not appear to be a reasonable workaround for this problem as a whole. This is because the same issue can be caused by using Option instead of a custom enum like MaybeDone.

https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=63ab1f219beadfb593aac4ec5e0952d7

[RalfJung]

Yeah I did not mean to imply that this is a reasonable work-around for this case, I just pointed out the possibility.

I don't think there is a good work-around, this fundamentally needs rustc layout computation to change -- either by special-casing generators, or by implementing rust-lang/rust#125735.

[taiki-e]

I don't think there is a good work-around, this fundamentally needs rustc layout computation to change -- either by special-casing generators, or by implementing rust-lang/rust#125735.

I thought wrapping a coroutine with a MaybeUninit inside a future returned by an async block/async function (like crossbeam-rs/crossbeam#834) as a potential workaround.

In previous async lowing I think it was enough to change the GenFuture field from T to MaybeUninit<T>, but I'm not sure what should be done since such a type is not used in recent async lowing.

[RalfJung]

Here's a proper work-around: rust-lang/rust#129313. It increases the size of some futures, though...