1 files changed, 16 insertions, 9 deletions
diff --git a/src/sync/notify.rs b/src/sync/notify.rs
index 74b97cc..c93ce3b 100644
--- a/src/sync/notify.rs
+++ b/src/sync/notify.rs
@@ -20,7 +20,7 @@ use std::task::{Context, Poll, Waker};
 
 type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>;
 
-/// Notify a single task to wake up.
+/// Notifies a single task to wake up.
 ///
 /// `Notify` provides a basic mechanism to notify a single task of an event.
 /// `Notify` itself does not carry any data. Instead, it is to be used to signal
@@ -57,13 +57,16 @@ type WaitList = LinkedList<Waiter, <Waiter as linked_list::Link>::Target>;
 ///     let notify = Arc::new(Notify::new());
 ///     let notify2 = notify.clone();
 ///
-///     tokio::spawn(async move {
+///     let handle = tokio::spawn(async move {
 ///         notify2.notified().await;
 ///         println!("received notification");
 ///     });
 ///
 ///     println!("sending notification");
 ///     notify.notify_one();
+///
+///     // Wait for task to receive notification.
+///     handle.await.unwrap();
 /// }
 /// ```
 ///
@@ -128,10 +131,10 @@ enum NotificationType {
 
 #[derive(Debug)]
 struct Waiter {
-    /// Intrusive linked-list pointers
+    /// Intrusive linked-list pointers.
     pointers: linked_list::Pointers<Waiter>,
 
-    /// Waiting task's waker
+    /// Waiting task's waker.
     waker: Option<Waker>,
 
     /// `true` if the notification has been assigned to this waiter.
@@ -168,13 +171,13 @@ const NOTIFY_WAITERS_SHIFT: usize = 2;
 const STATE_MASK: usize = (1 << NOTIFY_WAITERS_SHIFT) - 1;
 const NOTIFY_WAITERS_CALLS_MASK: usize = !STATE_MASK;
 
-/// Initial "idle" state
+/// Initial "idle" state.
 const EMPTY: usize = 0;
 
 /// One or more threads are currently waiting to be notified.
 const WAITING: usize = 1;
 
-/// Pending notification
+/// Pending notification.
 const NOTIFIED: usize = 2;
 
 fn set_state(data: usize, state: usize) -> usize {
@@ -289,7 +292,7 @@ impl Notify {
         }
     }
 
-    /// Notifies a waiting task
+    /// Notifies a waiting task.
     ///
     /// If a task is currently waiting, that task is notified. Otherwise, a
     /// permit is stored in this `Notify` value and the **next** call to
@@ -359,7 +362,7 @@ impl Notify {
         }
     }
 
-    /// Notifies all waiting tasks
+    /// Notifies all waiting tasks.
     ///
     /// If a task is currently waiting, that task is notified. Unlike with
     /// `notify_one()`, no permit is stored to be used by the next call to
@@ -551,6 +554,10 @@ impl Future for Notified<'_> {
                         return Poll::Ready(());
                     }
 
+                    // Clone the waker before locking, a waker clone can be
+                    // triggering arbitrary code.
+                    let waker = cx.waker().clone();
+
                     // Acquire the lock and attempt to transition to the waiting
                     // state.
                     let mut waiters = notify.waiters.lock();
@@ -612,7 +619,7 @@ impl Future for Notified<'_> {
 
                     // Safety: called while locked.
                     unsafe {
-                        (*waiter.get()).waker = Some(cx.waker().clone());
+                        (*waiter.get()).waker = Some(waker);
                     }
 
                     // Insert the waiter into the linked list