From 7ab27c083dc00e952df2248faab6dd01fdfc589d Mon Sep 17 00:00:00 2001
From: Kevin Mehall <km@kevinmehall.net>
Date: Sat, 16 Nov 2024 13:03:42 -0700
Subject: [PATCH] Fix race condition that could lose a transfer completion
 wakup

If `poll_completion_generic` runs between `notify_completion`'s
`waker.take()` and `state.swap(STATE_COMPLETED, ...)`, it could see the
transfer as pending and register its waker, but that waker would never
be called.

This adds an Arc because once the transfer is set to `STATE_COMPLETED`,
the other thread might free it, so we need to ensure the AtomicWaker
stays alive while we use it. That's what the `take()` was intended
to work around, but reading the waker before writing the status
introduces the race.
---
 src/transfer/internal.rs | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/src/transfer/internal.rs b/src/transfer/internal.rs
index dd778f7..44cfe1d 100644
--- a/src/transfer/internal.rs
+++ b/src/transfer/internal.rs
@@ -2,7 +2,10 @@ use std::{
     cell::UnsafeCell,
     ffi::c_void,
     ptr::NonNull,
-    sync::atomic::{AtomicU8, Ordering},
+    sync::{
+        atomic::{AtomicU8, Ordering},
+        Arc,
+    },
     task::{Context, Poll},
 };
 
@@ -44,7 +47,7 @@ struct TransferInner<P: PlatformTransfer> {
     state: AtomicU8,
 
     /// Waker that is notified when transfer completes.
-    waker: AtomicWaker,
+    waker: Arc<AtomicWaker>,
 }
 
 /// Handle to a transfer.
@@ -81,7 +84,7 @@ impl<P: PlatformTransfer> TransferHandle<P> {
         let b = Box::new(TransferInner {
             platform_data: UnsafeCell::new(inner),
             state: AtomicU8::new(STATE_IDLE),
-            waker: AtomicWaker::new(),
+            waker: Arc::new(AtomicWaker::new()),
         });
 
         TransferHandle {
@@ -180,13 +183,9 @@ impl<P: PlatformTransfer> Drop for TransferHandle<P> {
 pub(crate) unsafe fn notify_completion<P: PlatformTransfer>(transfer: *mut c_void) {
     unsafe {
         let transfer = transfer as *mut TransferInner<P>;
-        let waker = (*transfer).waker.take();
+        let waker = (*transfer).waker.clone();
         match (*transfer).state.swap(STATE_COMPLETED, Ordering::Release) {
-            STATE_PENDING => {
-                if let Some(waker) = waker {
-                    waker.wake()
-                }
-            }
+            STATE_PENDING => waker.wake(),
             STATE_ABANDONED => {
                 drop(Box::from_raw(transfer));
             }