Fixes UB in stack::generator: Aliasing violations and references to uninit memory.

src/stack/generator.rs

@@ -1,25 +1,20 @@
+use std::{future::Future, mem::MaybeUninit, pin::Pin, ptr};
+
 use crate::{
     core::{advance, async_advance, Airlock as _, Next},
-    ext::MaybeUninitExt,
     ops::{Coroutine, GeneratorState},
     stack::engine::{Airlock, Co},
 };
-use std::{future::Future, mem, pin::Pin, ptr};
 
 /// This data structure holds the transient state of an executing generator.
 ///
 /// It's called "Shelf", rather than "State", to avoid confusion with the
 /// `GeneratorState` enum.
 ///
 /// [_See the module-level docs for examples._](.)
-// Safety: The lifetime of the data is controlled by a `Gen`, which constructs
-// it in place, and holds a mutable reference right up until dropping it in
-// place. Thus, the data inside is pinned and can never be moved.
-pub struct Shelf<Y, R, F: Future>(mem::MaybeUninit<State<Y, R, F>>);
-
-struct State<Y, R, F: Future> {
+pub struct Shelf<Y, R, F: Future> {
     airlock: Airlock<Y, R>,
-    future: F,
+    future: MaybeUninit<F>,
 }
 
 impl<Y, R, F: Future> Shelf<Y, R, F> {
@@ -28,7 +23,13 @@ impl<Y, R, F: Future> Shelf<Y, R, F> {
     /// [_See the module-level docs for examples._](.)
     #[must_use]
     pub fn new() -> Self {
-        Self(mem::MaybeUninit::uninit())
+        Self {
+            airlock: Airlock::default(),
+            // Safety: The lifetime of the data is controlled by a `Gen`, which constructs
+            // it in place, and holds a mutable reference right up until dropping it in
+            // place. Thus, the data inside is pinned and can never be moved.
+            future: MaybeUninit::uninit(),
+        }
     }
 }
 
@@ -43,7 +44,8 @@ impl<Y, R, F: Future> Default for Shelf<Y, R, F> {
 ///
 /// [_See the module-level docs for examples._](.)
 pub struct Gen<'s, Y, R, F: Future> {
-    state: Pin<&'s mut State<Y, R, F>>,
+    airlock: &'s Airlock<Y, R>,
+    future: Pin<&'s mut F>,
 }
 
 impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> {
@@ -89,19 +91,22 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> {
         shelf: &'s mut Shelf<Y, R, F>,
         producer: impl FnOnce(Co<'s, Y, R>) -> F,
     ) -> Self {
-        // Safety: Build the struct in place, by writing each field in place.
-        let p = &mut *shelf.0.as_mut_ptr() as *mut State<Y, R, F>;
-
-        let airlock = Airlock::default();
-        ptr::write(&mut (*p).airlock, airlock);
-
-        let future = producer(Co::new(&(*p).airlock));
-        ptr::write(&mut (*p).future, future);
-
-        // Safety: the state can never be moved again, because we store it inside a
-        // `Pin` until `Gen::drop`, where the contents are dropped in place.
-        let state = Pin::new_unchecked(shelf.0.assume_init_get_mut());
-        Self { state }
+        // By splitting the mutable `shelf` into a shared `airlock` and a unique
+        // pinned `future` reference we ensure the aliasing rules are not violated.
+        let airlock = &shelf.airlock;
+        // Safety: Initializes the future in-place using `ptr::write`, which is
+        // the correct way to initialize a `MaybeUninit`
+        shelf.future.as_mut_ptr().write(producer(Co::new(airlock)));
+        // Safety: The `MaybeUninit` is initialized by now, so its safe to create
+        // a reference to the future itself
+        // NB: can be replaced by `MaybeUninit::get_mut` once stabilized
+        let init = &mut *shelf.future.as_mut_ptr();
+        // Safety: The `shelf` remains borrowed during the entire lifetime of
+        // the `Gen`and is hence pinned.
+        Self {
+            airlock,
+            future: Pin::new_unchecked(init),
+        }
     }
 
     /// Resumes execution of the generator.
@@ -115,37 +120,22 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> {
     ///
     /// [_See the module-level docs for examples._](.)
     pub fn resume_with(&mut self, arg: R) -> GeneratorState<Y, F::Output> {
-        let (future, airlock) = self.project();
-        airlock.replace(Next::Resume(arg));
-        advance(future, &airlock)
-    }
-
-    fn project(&mut self) -> (Pin<&mut F>, &Airlock<Y, R>) {
-        unsafe {
-            // Safety: This is a pin projection. `future` is pinned, but never moved.
-            // `airlock` is never pinned.
-            let state = self.state.as_mut().get_unchecked_mut();
-
-            let future = Pin::new_unchecked(&mut state.future);
-            let airlock = &state.airlock;
-            (future, airlock)
-        }
+        self.airlock.replace(Next::Resume(arg));
+        advance(self.future.as_mut(), &self.airlock)
     }
 }
 
 impl<'s, Y, R, F: Future> Drop for Gen<'s, Y, R, F> {
     fn drop(&mut self) {
-        // Safety: `state` is a `MaybeUninit` which is guaranteed to be initialized,
-        // because the only way to construct a `Gen` is with `Gen::new`, which
-        // initializes it.
+        // Safety: `future` itself is a `MaybeUninit`, which is guaranteed to be
+        // initialized, because the only way to construct a `Gen` is with
+        // `Gen::new`, which initializes it.
         //
-        // Drop `state` in place, by dropping each field in place. Drop `future` first,
-        // since it likely contains a reference to `airlock` (through the `co` object).
-        // Since we drop everything in place, the `Pin` invariants are not violated.
+        // Drop `future` in place first (it likely contains a reference to airlock).
+        // Since we drop it in place, the `Pin` invariants are not violated.
+        // The airlock is regularly dropped when the `Shelf` goes out of scope.
         unsafe {
-            let state = self.state.as_mut().get_unchecked_mut();
-            ptr::drop_in_place(&mut state.future);
-            ptr::drop_in_place(&mut state.airlock);
+            ptr::drop_in_place(self.future.as_mut().get_unchecked_mut());
         }
     }
 }
@@ -171,9 +161,8 @@ impl<'s, Y, F: Future> Gen<'s, Y, (), F> {
     pub fn async_resume(
         &mut self,
     ) -> impl Future<Output = GeneratorState<Y, F::Output>> + '_ {
-        let (future, airlock) = self.project();
-        airlock.replace(Next::Resume(()));
-        async_advance(future, airlock)
+        self.airlock.replace(Next::Resume(()));
+        async_advance(self.future.as_mut(), self.airlock)
     }
 }
 

tests/stack.rs

@@ -117,6 +117,23 @@ fn stack_yield_match() {
     assert_eq!(vec![1, 3, 5, 7, 9], res)
 }
 
+#[test]
+fn stack_yield_closure_no_macro() {
+    let mut shelf = genawaiter::stack::Shelf::new();
+    let gen = unsafe {
+        Gen::new(&mut shelf, |co| async move {
+            let mut n = 1;
+            while n < 10 {
+                co.yield_(n).await;
+                n += 2;
+            }
+        })
+    };
+
+    let res = gen.into_iter().collect::<Vec<_>>();
+    assert_eq!(vec![1, 3, 5, 7, 9], res)
+}
+
 #[cfg(feature = "proc_macro")]
 #[test]
 fn stack_yield_closure() {