From 6599ea656de5a35be7c866a45c1d9f6ab5392df5 Mon Sep 17 00:00:00 2001
From: Floris Bruynooghe <flub@n0.computer>
Date: Fri, 3 Jan 2025 16:35:31 +0100
Subject: [PATCH] fix-or-feat(iroh): Set MaybeFuture to None on Poll::Ready
 (#3090)

## Description

This removes a footgun where after polling the future it is easy to
forget to call MaybeFuture::set_none, which can result in a panic if
the inner future is not fused.

## Breaking Changes

<!-- Optional, if there are any breaking changes document them,
including how to migrate older code. -->

## Notes & open questions

<!-- Any notes, remarks or open questions you have to make about the PR.
-->

## Change checklist

- [x] Self-review.
- [x] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.
- [x] Tests if relevant.
- [x] All breaking changes documented.
---
 iroh/src/util.rs | 37 +++++++++++++++++++++++++++++++++----
 1 file changed, 33 insertions(+), 4 deletions(-)

diff --git a/iroh/src/util.rs b/iroh/src/util.rs
index 4c80f4c551..a545156ef7 100644
--- a/iroh/src/util.rs
+++ b/iroh/src/util.rs
@@ -53,11 +53,18 @@ impl<T> MaybeFuture<T> {
 impl<T: Future> Future for MaybeFuture<T> {
     type Output = T::Output;
 
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let this = self.project();
-        match this {
-            MaybeFutureProj::Some(t) => t.poll(cx),
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let mut this = self.as_mut().project();
+        let poll_res = match this {
+            MaybeFutureProj::Some(ref mut t) => t.as_mut().poll(cx),
             MaybeFutureProj::None => Poll::Pending,
+        };
+        match poll_res {
+            Poll::Ready(val) => {
+                self.as_mut().project_replace(Self::None);
+                Poll::Ready(val)
+            }
+            Poll::Pending => Poll::Pending,
         }
     }
 }
@@ -70,3 +77,25 @@ impl<T: Future> Future for MaybeFuture<T> {
 pub(crate) fn relay_only_mode() -> bool {
     std::option_env!("DEV_RELAY_ONLY").is_some()
 }
+
+#[cfg(test)]
+mod tests {
+    use std::pin::pin;
+
+    use tokio::time::Duration;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_maybefuture_poll_after_use() {
+        let fut = async move { "hello" };
+        let mut maybe_fut = pin!(MaybeFuture::Some(fut));
+        let res = (&mut maybe_fut).await;
+
+        assert_eq!(res, "hello");
+
+        // Now poll again
+        let res = tokio::time::timeout(Duration::from_millis(10), maybe_fut).await;
+        assert!(res.is_err());
+    }
+}