fix: prevent UB in mapped RwLock downgrades and oneshot reclamation

CHANGELOG.md

@@ -23,12 +23,17 @@ All notable changes to this project will be documented in this file.
 
 ## Unreleased
 
+### Breaking changes
+
+* Remove `MappedRwLockWriteGuard::downgrade` and `OwnedMappedRwLockWriteGuard::downgrade`, which could cause undefined behavior with projections into interior-mutable data. Keep the original write guard and downgrade it before projecting with `RwLockReadGuard::map`/`filter_map` or their owned counterparts; projections that require exclusive access must retain the mapped write guard.
+
 ### New features
 
 * Add bounded MPMC `reserve` and `try_reserve` methods returning a borrowed `Permit`, so callers can wait for capacity before constructing a value; sends and reservations receive capacity in wait-queue order, and unused permits release it.
 
 ### Bug fixes
 
+* Prevent undefined behavior when a oneshot receiver frees the channel while the sender is finishing a send or disconnection notification.
 * Prevent deadlocks in `singleflight::Group::work` and `try_work` when a duplicate key's destructor calls back into the same group.
 
 ### Notable changes

asyncband/src/oneshot/mod.rs

@@ -305,6 +305,8 @@ impl<T> Channel<T> {
     ///
     /// # Safety
     ///
+    /// * `channel_ptr` must point to the live channel owned by this sender. The caller must not
+    ///   access the channel after this call unless it returns `false` for receiver ownership.
     /// * `final_state` must be `MESSAGE` or `DISCONNECTED`.
     /// * The caller must have just observed `RECEIVING` with an atomic read-modify-write that
     ///   changed the state to `AWAKING`. This gives the caller exclusive ownership of the
@@ -313,21 +315,29 @@ impl<T> Channel<T> {
     /// * When publishing `MESSAGE`, the caller must own an initialized message that precedes the
     ///   release operation in this method.
     #[inline(always)]
-    unsafe fn finish_sender_awakening(&self, final_state: u8) -> (Waker, bool) {
+    unsafe fn finish_sender_awakening(
+        channel_ptr: NonNull<Self>,
+        final_state: u8,
+    ) -> (Waker, bool) {
         debug_assert!(matches!(final_state, MESSAGE | DISCONNECTED));
 
+        // An &self parameter would keep its borrow protected until this function returns, but
+        // the receiver may reclaim the channel as soon as the terminal state is published below.
+        // SAFETY: The caller still owns the live channel's AWAKING state at this point.
+        let channel = unsafe { channel_ptr.as_ref() };
+
         // ORDERING: The caller's Release RMW read RECEIVING with a Relaxed load. Acquire
         // synchronizes that read with the receiver's Release publication before taking the waker.
         fence(Ordering::Acquire);
 
         // SAFETY: The caller's RECEIVING-to-AWAKING transition transferred exclusive ownership of
         // the initialized waker to the sender.
-        let waker = unsafe { self.take_waker() };
+        let waker = unsafe { channel.take_waker() };
 
         // ORDERING: Release publishes the message or disconnected state when this replaces
         // AWAKING. The RMW's load half is Relaxed; if it reads a receiver-written DISCONNECTED, the
         // conditional Acquire below completes the reverse allocation-ownership handoff.
-        let previous_state = self.state.swap(final_state, Ordering::Release);
+        let previous_state = channel.state.swap(final_state, Ordering::Release);
         if matches!(previous_state, AWAKING) {
             (waker, true)
         } else {

asyncband/src/oneshot/sender.rs

@@ -61,9 +61,8 @@ impl<T> Sender<T> {
         let sender = ManuallyDrop::new(self);
         let channel_ptr = sender.channel_ptr;
 
-        // SAFETY: The channel exists on the heap for the entire duration of this method, and we
-        // only ever acquire shared references to it. Dropping the receiver does not immediately
-        // free the channel.
+        // SAFETY: This sender keeps the channel alive until it publishes a terminal state. No
+        // access through this reference follows a transfer of allocation ownership to the receiver.
         let channel = unsafe { channel_ptr.as_ref() };
 
         // Write the message into the channel on the heap.
@@ -91,7 +90,7 @@ impl<T> Sender<T> {
                 // exclusive ownership of the published waker to this sender. The message was
                 // initialized before the RMW and is ready to publish.
                 let (waker, receiver_owns_allocation) =
-                    unsafe { channel.finish_sender_awakening(MESSAGE) };
+                    unsafe { Channel::finish_sender_awakening(channel_ptr, MESSAGE) };
                 if receiver_owns_allocation {
                     waker.wake();
                 } else {
@@ -180,7 +179,7 @@ impl<T> Drop for Sender<T> {
                 // exclusive ownership of the published waker to this sender. Sender::drop has no
                 // message to publish.
                 let (waker, receiver_owns_allocation) =
-                    unsafe { channel.finish_sender_awakening(DISCONNECTED) };
+                    unsafe { Channel::finish_sender_awakening(self.channel_ptr, DISCONNECTED) };
                 if receiver_owns_allocation {
                     waker.wake();
                 } else {

asyncband/src/oneshot/tests.rs

@@ -78,7 +78,7 @@ fn poll_returns_while_sender_owns_waker() {
     // SAFETY: fetch_add observed RECEIVING and changed it to AWAKING after the message and waker
     // were initialized.
     let (sender_waker, receiver_owns_allocation) =
-        unsafe { channel.finish_sender_awakening(super::MESSAGE) };
+        unsafe { super::Channel::finish_sender_awakening(channel_ptr, super::MESSAGE) };
     assert!(receiver_owns_allocation);
     sender_waker.wake();
     assert_eq!(stored_probe.wake_count(), 1);
@@ -144,7 +144,7 @@ fn drop_transfers_cleanup_while_sender_owns_waker() {
     // SAFETY: fetch_add observed RECEIVING and changed it to AWAKING after the message and waker
     // were initialized.
     let (sender_waker, receiver_owns_allocation) =
-        unsafe { channel.finish_sender_awakening(super::MESSAGE) };
+        unsafe { super::Channel::finish_sender_awakening(channel_ptr, super::MESSAGE) };
     if receiver_owns_allocation {
         sender_waker.wake();
     } else {

asyncband/src/rwlock/mapped_write_guard.rs

@@ -29,13 +29,17 @@ use std::ops::DerefMut;
 use std::ptr::NonNull;
 
 use crate::internal::semaphore::Semaphore;
-use crate::rwlock::MappedRwLockReadGuard;
 use crate::rwlock::access::WriteAccess;
 
 /// Exclusive access to a projection of a locked value borrowed for the guard lifetime.
 ///
 /// Use [`RwLockWriteGuard::map`](crate::rwlock::RwLockWriteGuard::map) to select a component.
 /// Dropping the guard releases its access.
+///
+/// This guard cannot be downgraded: a mutable projection may expose data that other readers
+/// could mutate through interior mutability. To retain a shared projection, first downgrade the
+/// original [`RwLockWriteGuard`](crate::rwlock::RwLockWriteGuard), then use
+/// [`RwLockReadGuard::map`](crate::rwlock::RwLockReadGuard::map) or `filter_map`.
 #[must_use = "dropping the guard releases its write access immediately"]
 pub struct MappedRwLockWriteGuard<'a, T: ?Sized> {
     d: NonNull<T>,
@@ -123,12 +127,4 @@ impl<'a, T: ?Sized> MappedRwLockWriteGuard<'a, T> {
             None => Err(orig),
         }
     }
-
-    /// Retains shared access to the same projection while releasing exclusive access.
-    ///
-    /// There is no unlocked interval in which another writer can modify the value. Queued
-    /// requests retain their order, so a waiting writer can prevent later readers from joining.
-    pub fn downgrade(self) -> MappedRwLockReadGuard<'a, T> {
-        MappedRwLockReadGuard::new(self.d, self.access.downgrade())
-    }
 }

asyncband/src/rwlock/mod.rs

@@ -71,11 +71,31 @@
 //! let slot = RwLockWriteGuard::map(queue.write().await, |items| &mut items[0]);
 //! let mut message = MappedRwLockWriteGuard::filter_map(slot, Option::as_mut).unwrap();
 //! message.push_str(" review");
-//! let message = message.downgrade();
 //! assert_eq!(&*message, "pending review");
 //! # }
 //! ```
 //!
+//! Mapped write guards cannot be downgraded: mutable projections may reach data that other
+//! readers could change through interior mutability. For a shared projection, update through
+//! the original write guard, downgrade it, then map the resulting read guard. This retains
+//! access without an unlocked interval and derives the projection from a shared reference.
+//!
+//! ```
+//! use asyncband::rwlock::RwLock;
+//! use asyncband::rwlock::RwLockReadGuard;
+//!
+//! # #[tokio::main]
+//! # async fn main() {
+//! let queue = RwLock::new(vec![String::from("pending")]);
+//! let mut edit = queue.write().await;
+//! edit[0].push_str(" review");
+//! let message = RwLockReadGuard::map(edit.downgrade(), |items| &items[0]);
+//! let other_reader = queue.try_read().unwrap();
+//! assert_eq!(&*message, &other_reader[0]);
+//! assert!(queue.try_write().is_none());
+//! # }
+//! ```
+//!
 //! # Keeping the lock alive
 //!
 //! Owned guards retain the `Arc` passed to acquisition, allowing the guard to outlive that call's
@@ -195,6 +215,41 @@ impl<T: ?Sized> RwLock<T> {
 
 #[cfg(doctest)]
 mod compile_fail_tests {
+    /// A mutable projection into a mutex must not coexist with readers that can lock it.
+    ///
+    /// ```compile_fail,E0599
+    /// use asyncband::rwlock::{RwLock, RwLockWriteGuard};
+    ///
+    /// let lock = RwLock::new(std::sync::Mutex::new(0u8));
+    /// let write = RwLockWriteGuard::map(lock.try_write().unwrap(), |inner| {
+    ///     inner.get_mut().unwrap()
+    /// });
+    /// let read = write.downgrade();
+    /// let shared = &*read;
+    /// let outer = lock.try_read().unwrap();
+    /// *outer.lock().unwrap() = 1;
+    /// println!("{shared}");
+    /// ```
+    struct MappedWriteGuardCannotDowngrade;
+
+    /// Keeping the lock alive with an `Arc` does not make a mutable projection shareable.
+    ///
+    /// ```compile_fail,E0599
+    /// use std::sync::{Arc, Mutex};
+    /// use asyncband::rwlock::{OwnedRwLockWriteGuard, RwLock};
+    ///
+    /// let lock = Arc::new(RwLock::new(Mutex::new(0u8)));
+    /// let write = OwnedRwLockWriteGuard::map(lock.clone().try_write_owned().unwrap(), |inner| {
+    ///     inner.get_mut().unwrap()
+    /// });
+    /// let read = write.downgrade();
+    /// let shared = &*read;
+    /// let outer = lock.try_read().unwrap();
+    /// *outer.lock().unwrap() = 1;
+    /// println!("{shared}");
+    /// ```
+    struct OwnedMappedWriteGuardCannotDowngrade;
+
     /// ```compile_fail
     /// use asyncband::rwlock::RwLockWriteGuard;
     ///

asyncband/src/rwlock/owned_mapped_write_guard.rs

@@ -29,14 +29,18 @@ use std::ops::DerefMut;
 use std::ptr::NonNull;
 use std::sync::Arc;
 
-use crate::rwlock::OwnedMappedRwLockReadGuard;
 use crate::rwlock::RwLock;
 use crate::rwlock::access::WriteAccess;
 
 /// Exclusive access to a projection of a locked value kept alive by an `Arc`.
 ///
 /// Use [`OwnedRwLockWriteGuard::map`](crate::rwlock::OwnedRwLockWriteGuard::map) to select a
 /// component. Dropping the guard releases its access.
+///
+/// This guard cannot be downgraded: a mutable projection may expose data that other readers
+/// could mutate through interior mutability. To retain a shared projection, first downgrade the
+/// original [`OwnedRwLockWriteGuard`](crate::rwlock::OwnedRwLockWriteGuard), then use
+/// [`OwnedRwLockReadGuard::map`](crate::rwlock::OwnedRwLockReadGuard::map) or `filter_map`.
 #[must_use = "dropping the guard releases its write access immediately"]
 pub struct OwnedMappedRwLockWriteGuard<T: ?Sized, U: ?Sized> {
     d: NonNull<U>,
@@ -125,12 +129,4 @@ impl<T: ?Sized, U: ?Sized> OwnedMappedRwLockWriteGuard<T, U> {
             None => Err(orig),
         }
     }
-
-    /// Retains shared access to the same projection while releasing exclusive access.
-    ///
-    /// There is no unlocked interval in which another writer can modify the value. Queued
-    /// requests retain their order, so a waiting writer can prevent later readers from joining.
-    pub fn downgrade(self) -> OwnedMappedRwLockReadGuard<T, U> {
-        OwnedMappedRwLockReadGuard::new(self.d, self.access.downgrade())
-    }
 }

asyncband/src/rwlock/owned_write_guard.rs

@@ -141,6 +141,9 @@ impl<T: ?Sized> OwnedRwLockWriteGuard<T> {
     ///
     /// There is no unlocked interval in which another writer can modify the value. Queued
     /// requests retain their order, so a waiting writer can prevent later readers from joining.
+    ///
+    /// To retain a shared projection, call this before [`OwnedRwLockReadGuard::map`] or
+    /// [`OwnedRwLockReadGuard::filter_map`]. Mapped write guards cannot be downgraded.
     pub fn downgrade(self) -> OwnedRwLockReadGuard<T> {
         OwnedRwLockReadGuard {
             access: self.access.downgrade(),

asyncband/src/rwlock/write_guard.rs

@@ -136,6 +136,9 @@ impl<'a, T: ?Sized> RwLockWriteGuard<'a, T> {
     ///
     /// There is no unlocked interval in which another writer can modify the value. Queued
     /// requests retain their order, so a waiting writer can prevent later readers from joining.
+    ///
+    /// To retain a shared projection, call this before [`RwLockReadGuard::map`] or
+    /// [`RwLockReadGuard::filter_map`]. Mapped write guards cannot be downgraded.
     pub fn downgrade(self) -> RwLockReadGuard<'a, T> {
         RwLockReadGuard {
             access: self.access.downgrade(),

tests-integration/tests/rwlock_test.rs

@@ -141,7 +141,6 @@ async fn owned_write_mappings_preserve_lock_ownership() {
     let mut value =
         OwnedMappedRwLockWriteGuard::filter_map(value, |value| value.first_mut()).unwrap();
     *value = 100;
-    let value = value.downgrade();
 
     drop(rwlock);
     assert!(weak.upgrade().is_some());
@@ -181,7 +180,7 @@ async fn owned_read_mappings_preserve_lock_ownership() {
 }
 
 #[tokio::test]
-async fn every_write_guard_type_can_downgrade() {
+async fn write_guards_downgrade_before_mapping() {
     let rwlock = Arc::new(RwLock::new((42, "test".to_string())));
 
     {
@@ -210,24 +209,25 @@ async fn every_write_guard_type_can_downgrade() {
     }
 
     {
-        let write_guard = rwlock.write().await;
-        let mut mapped_write = RwLockWriteGuard::map(write_guard, |data| &mut data.0);
-        *mapped_write = 200;
+        let mut write_guard = rwlock.write().await;
+        write_guard.0 = 200;
 
-        let mapped_read = mapped_write.downgrade();
+        let mapped_read = RwLockReadGuard::map(write_guard.downgrade(), |data| &data.0);
         assert_eq!(*mapped_read, 200);
+        assert_eq!(rwlock.try_read().unwrap().0, 200);
 
         assert!(rwlock.try_write().is_none());
         drop(mapped_read);
     }
 
     {
-        let owned_write = rwlock.clone().write_owned().await;
-        let mut owned_mapped = OwnedRwLockWriteGuard::map(owned_write, |data| &mut data.1);
-        *owned_mapped = "final".to_string();
+        let mut owned_write = rwlock.clone().write_owned().await;
+        owned_write.1 = "final".to_string();
 
-        let owned_mapped_read = owned_mapped.downgrade();
+        let owned_mapped_read = OwnedRwLockReadGuard::map(owned_write.downgrade(), |data| &data.1);
         assert_eq!(*owned_mapped_read, "final");
+        assert_eq!(rwlock.try_read().unwrap().1, "final");
+        assert!(rwlock.try_write().is_none());
 
         let moved_task = tokio::spawn(async move {
             assert_eq!(*owned_mapped_read, "final");

tests-integration/tests/unsafe_paths_test.rs

@@ -91,8 +91,8 @@ fn mapped_rwlock_guards_preserve_lock_ownership() {
     let mapped = RwLockWriteGuard::map(write, |data| &mut data.values);
     let mut mapped = MappedRwLockWriteGuard::map(mapped, |values| &mut values[1]);
     *mapped = 3;
-    let mapped = mapped.downgrade();
     assert_eq!(*mapped, 3);
+    assert!(lock.try_read().is_none());
     assert!(lock.try_write().is_none());
     drop(mapped);
 
@@ -113,8 +113,8 @@ fn mapped_rwlock_guards_preserve_lock_ownership() {
     let mut mapped =
         OwnedMappedRwLockWriteGuard::filter_map(mapped, |values| values.first_mut()).unwrap();
     *mapped = 6;
-    let mapped = mapped.downgrade();
     assert_eq!(*mapped, 6);
+    assert!(lock.try_read().is_none());
     drop(mapped);
 
     let read = lock.clone().try_read_owned().unwrap();
@@ -128,6 +128,46 @@ fn mapped_rwlock_guards_preserve_lock_ownership() {
     assert!(weak.upgrade().is_none());
 }
 
+#[test]
+fn downgraded_shared_projection_preserves_interior_mutability() {
+    let lock = RwLock::new(Some(std::sync::Mutex::new(0u8)));
+    let mut write = lock.try_write().unwrap();
+    *write.as_mut().unwrap().get_mut().unwrap() = 1;
+    let read = RwLockReadGuard::map(write.downgrade(), |value| value.as_ref().unwrap());
+
+    // The shared projection retains the mutex, so access to its contents remains synchronized.
+    let shared = &*read;
+    let outer = lock.try_read().unwrap();
+    *outer.as_ref().unwrap().lock().unwrap() = 2;
+    assert_eq!(*shared.lock().unwrap(), 2);
+    assert!(lock.try_write().is_none());
+
+    drop(outer);
+    drop(read);
+    assert!(lock.try_write().is_some());
+}
+
+#[test]
+fn owned_downgraded_shared_projection_preserves_interior_mutability() {
+    let lock = Arc::new(RwLock::new(Some(std::sync::Mutex::new(0u8))));
+    let weak = Arc::downgrade(&lock);
+    let mut write = lock.clone().try_write_owned().unwrap();
+    *write.as_mut().unwrap().get_mut().unwrap() = 1;
+    let read = OwnedRwLockReadGuard::filter_map(write.downgrade(), Option::as_ref).unwrap();
+
+    let shared = &*read;
+    let outer = lock.try_read().unwrap();
+    *outer.as_ref().unwrap().lock().unwrap() = 2;
+    assert_eq!(*shared.lock().unwrap(), 2);
+    assert!(lock.try_write().is_none());
+
+    drop(outer);
+    drop(lock);
+    assert_eq!(*shared.lock().unwrap(), 2);
+    drop(read);
+    assert!(weak.upgrade().is_none());
+}
+
 struct AddressSensitiveFuture {
     address: Cell<*const Self>,
     _pin: PhantomPinned,