diff --git a/CHANGELOG.md b/CHANGELOG.md index f32204c..420d949 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0. ## [Unreleased] +### Fixed +- Fix race condition that could lead to use-after-free if the `Receiver` was polled asynchronously, + but then dropped before completion. https://github.com/faern/oneshot/pull/74 ## [0.1.11] - 2025-02-22 diff --git a/src/lib.rs b/src/lib.rs index 18846e5..c8d0c1c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1114,6 +1114,21 @@ impl Drop for Receiver { // that signal to the sender to deallocate the channel. So the channel pointer is valid let channel = unsafe { self.channel_ptr.as_ref() }; + // If this receiver was previously polled, but was not polled to completion, then the + // channel is in the RECEIVING state and with a waker written. We must tell the sender + // we are no longer receiving, and then drop the waker. We must first move away from + // the RECEIVING state in order to not race with the sender around taking the waker. + #[cfg(feature = "async")] + if channel.state.load(Relaxed) == RECEIVING + && channel + .state + .compare_exchange(RECEIVING, EMPTY, Relaxed, Relaxed) + .is_ok() + { + // SAFETY: The RECEIVING state guarantees we have written a waker. + unsafe { channel.drop_waker() }; + } + // Set the channel state to disconnected and read what state the channel was in // ORDERING: Release is required so that in the states where the sender becomes responsible // for deallocating the channel, they can synchronize with this final state swap store from @@ -1133,15 +1148,6 @@ impl Drop for Receiver { // final write of the state. So we can safely deallocate it. unsafe { dealloc(self.channel_ptr) }; } - // This receiver was previously polled. But now we are dropped instead. The sender - // has never observed our RECEIVING state and we are resposnible for dropping the - // waker we previously wrote. The sender becomes responsible for deallocating the - // channel. - #[cfg(feature = "async")] - RECEIVING => { - // SAFETY: The RECEIVING state guarantees we have written a waker. - unsafe { channel.drop_waker() }; - } // The sender was already dropped. We are responsible for freeing the channel. DISCONNECTED => { // SAFETY: The acquire ordering of the swap above synchronize with the sender's