From efba8ccafb56fdd7b5577069eae65a785ccf2e1a Mon Sep 17 00:00:00 2001 From: oliver-giersch Date: Wed, 25 Mar 2020 08:17:45 +0100 Subject: [PATCH 1/3] Replaces some instances of unsound code (UB) around uninitialized memory and aliasing of mutable references. --- src/stack/generator.rs | 65 +++++++++++++++++++++--------------------- tests/stack.rs | 17 +++++++++++ 2 files changed, 49 insertions(+), 33 deletions(-) diff --git a/src/stack/generator.rs b/src/stack/generator.rs index 46aa6e5..9eb0520 100644 --- a/src/stack/generator.rs +++ b/src/stack/generator.rs @@ -1,10 +1,11 @@ +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}; +use std::cell::UnsafeCell; /// This data structure holds the transient state of an executing generator. /// @@ -12,14 +13,14 @@ use std::{future::Future, mem, pin::Pin, ptr}; /// `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(mem::MaybeUninit>); +pub struct Shelf(UnsafeCell>); -struct State { +struct State { airlock: Airlock, - future: F, + // 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, } impl Shelf { @@ -28,7 +29,10 @@ impl Shelf { /// [_See the module-level docs for examples._](.) #[must_use] pub fn new() -> Self { - Self(mem::MaybeUninit::uninit()) + Self(UnsafeCell::new(State { + airlock: Airlock::default(), + future: MaybeUninit::uninit(), + })) } } @@ -43,7 +47,7 @@ impl Default for Shelf { /// /// [_See the module-level docs for examples._](.) pub struct Gen<'s, Y, R, F: Future> { - state: Pin<&'s mut State>, + state: Pin<&'s mut Shelf>, } impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { @@ -89,19 +93,16 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { shelf: &'s mut Shelf, 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; - - 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 } + let state = shelf.0.get(); + let future = producer(Co::new(&(*state).airlock)); + // initializes the future in-place + (*state).future.as_mut_ptr().write(future); + + Self { + // Safety: The shelf is borrowed by the resulting `Gen` is hence + // pinned + state: Pin::new_unchecked(shelf), + } } /// Resumes execution of the generator. @@ -124,10 +125,10 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { 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 state = self.state.0.get(); - let future = Pin::new_unchecked(&mut state.future); - let airlock = &state.airlock; + let future = Pin::new_unchecked(&mut *(*state).future.as_mut_ptr()); + let airlock = &(*state).airlock; (future, airlock) } } @@ -135,17 +136,15 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { 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, + // Safety: `future` 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 (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.state.0.get()).future.as_mut_ptr()); } } } diff --git a/tests/stack.rs b/tests/stack.rs index 1cccc03..4883d76 100644 --- a/tests/stack.rs +++ b/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::>(); + assert_eq!(vec![1, 3, 5, 7, 9], res) +} + #[cfg(feature = "proc_macro")] #[test] fn stack_yield_closure() { From 189d2833bb4c8b8bd39db3522a08158aed85b898 Mon Sep 17 00:00:00 2001 From: oliver-giersch Date: Wed, 25 Mar 2020 09:27:01 +0100 Subject: [PATCH 2/3] fixes aliasing violation by splitting the mutable `shelf` borrow into a shared and a unique part --- src/stack/generator.rs | 88 +++++++++++++++++++++--------------------- 1 file changed, 44 insertions(+), 44 deletions(-) diff --git a/src/stack/generator.rs b/src/stack/generator.rs index 9eb0520..71ed01b 100644 --- a/src/stack/generator.rs +++ b/src/stack/generator.rs @@ -1,11 +1,15 @@ -use std::{future::Future, mem::MaybeUninit, pin::Pin, ptr}; +use std::{ + future::Future, + mem::{ManuallyDrop, MaybeUninit}, + pin::Pin, + ptr, +}; use crate::{ core::{advance, async_advance, Airlock as _, Next}, ops::{Coroutine, GeneratorState}, stack::engine::{Airlock, Co}, }; -use std::cell::UnsafeCell; /// This data structure holds the transient state of an executing generator. /// @@ -13,13 +17,8 @@ use std::cell::UnsafeCell; /// `GeneratorState` enum. /// /// [_See the module-level docs for examples._](.) -pub struct Shelf(UnsafeCell>); - -struct State { +pub struct Shelf { airlock: Airlock, - // 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, } @@ -29,10 +28,13 @@ impl Shelf { /// [_See the module-level docs for examples._](.) #[must_use] pub fn new() -> Self { - Self(UnsafeCell::new(State { + 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(), - })) + } } } @@ -47,7 +49,8 @@ impl Default for Shelf { /// /// [_See the module-level docs for examples._](.) pub struct Gen<'s, Y, R, F: Future> { - state: Pin<&'s mut Shelf>, + airlock: &'s Airlock, + future: ManuallyDrop>, } impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { @@ -93,16 +96,21 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { shelf: &'s mut Shelf, producer: impl FnOnce(Co<'s, Y, R>) -> F, ) -> Self { - let state = shelf.0.get(); - let future = producer(Co::new(&(*state).airlock)); - // initializes the future in-place - (*state).future.as_mut_ptr().write(future); - - Self { - // Safety: The shelf is borrowed by the resulting `Gen` is hence - // pinned - state: Pin::new_unchecked(shelf), - } + // 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 + // todo: 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. + let future = ManuallyDrop::new(Pin::new_unchecked(init)); + + Self { airlock, future } } /// Resumes execution of the generator. @@ -116,35 +124,28 @@ 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 { - let (future, airlock) = self.project(); - airlock.replace(Next::Resume(arg)); - advance(future, &airlock) - } - - fn project(&mut self) -> (Pin<&mut F>, &Airlock) { - unsafe { - // Safety: This is a pin projection. `future` is pinned, but never moved. - // `airlock` is never pinned. - let state = self.state.0.get(); - - let future = Pin::new_unchecked(&mut *(*state).future.as_mut_ptr()); - 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: `future` 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. + // + // The pinned reference to the initialized future is wrapped in a + // `ManuallyDrop` because `Pin::get_unchecked_mut` consumes the `Pin`, + // which would require moving it (the pin) out of the `Gen` first. // - // Drop `future` in place first (likely contains a reference to airlock), + // 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 { - ptr::drop_in_place((*self.state.0.get()).future.as_mut_ptr()); + // todo: can be replaced by `ManuallyDrop::take`, stabilized in 1.42 + let future = ptr::read(&*self.future); + ptr::drop_in_place(future.get_unchecked_mut()); } } } @@ -170,9 +171,8 @@ impl<'s, Y, F: Future> Gen<'s, Y, (), F> { pub fn async_resume( &mut self, ) -> impl Future> + '_ { - 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) } } From 5185280d7032890666f44ff8534bb5a782badee6 Mon Sep 17 00:00:00 2001 From: oliver-giersch Date: Wed, 25 Mar 2020 11:00:07 +0100 Subject: [PATCH 3/3] replaces usage of ManuallyDrop in favor of `self.future.as_mut().get_unchecked_mut()` --- src/stack/generator.rs | 26 ++++++++------------------ 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/src/stack/generator.rs b/src/stack/generator.rs index 71ed01b..05e0802 100644 --- a/src/stack/generator.rs +++ b/src/stack/generator.rs @@ -1,9 +1,4 @@ -use std::{ - future::Future, - mem::{ManuallyDrop, MaybeUninit}, - pin::Pin, - ptr, -}; +use std::{future::Future, mem::MaybeUninit, pin::Pin, ptr}; use crate::{ core::{advance, async_advance, Airlock as _, Next}, @@ -50,7 +45,7 @@ impl Default for Shelf { /// [_See the module-level docs for examples._](.) pub struct Gen<'s, Y, R, F: Future> { airlock: &'s Airlock, - future: ManuallyDrop>, + future: Pin<&'s mut F>, } impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { @@ -104,13 +99,14 @@ impl<'s, Y, R, F: Future> Gen<'s, Y, R, F> { 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 - // todo: can be replaced by `MaybeUninit::get_mut` once stabilized + // 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. - let future = ManuallyDrop::new(Pin::new_unchecked(init)); - - Self { airlock, future } + Self { + airlock, + future: Pin::new_unchecked(init), + } } /// Resumes execution of the generator. @@ -135,17 +131,11 @@ impl<'s, Y, R, F: Future> Drop for Gen<'s, Y, R, F> { // initialized, because the only way to construct a `Gen` is with // `Gen::new`, which initializes it. // - // The pinned reference to the initialized future is wrapped in a - // `ManuallyDrop` because `Pin::get_unchecked_mut` consumes the `Pin`, - // which would require moving it (the pin) out of the `Gen` first. - // // 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 { - // todo: can be replaced by `ManuallyDrop::take`, stabilized in 1.42 - let future = ptr::read(&*self.future); - ptr::drop_in_place(future.get_unchecked_mut()); + ptr::drop_in_place(self.future.as_mut().get_unchecked_mut()); } } }