tokio-util: stop polling StreamReader after EOF (#8332)

tokio-util/src/io/stream_reader.rs

@@ -158,6 +158,7 @@ pub struct StreamReader<S, B> {
     inner: S,
     // This field is not pinned.
     chunk: Option<B>,
+    eof: bool,
 }
 
 impl<S, B, E> StreamReader<S, B>
@@ -179,6 +180,7 @@ where
         Self {
             inner: stream,
             chunk: None,
+            eof: false,
         }
     }
 
@@ -277,6 +279,8 @@ where
                 // This unwrap is very sad, but it can't be avoided.
                 let buf = self.project().chunk.as_ref().unwrap().chunk();
                 return Poll::Ready(Ok(buf));
+            } else if *self.as_mut().project().eof {
+                return Poll::Ready(Ok(&[]));
             } else {
                 match self.as_mut().project().inner.poll_next(cx) {
                     Poll::Ready(Some(Ok(chunk))) => {
@@ -284,7 +288,10 @@ where
                         *self.as_mut().project().chunk = Some(chunk);
                     }
                     Poll::Ready(Some(Err(err))) => return Poll::Ready(Err(err.into())),
-                    Poll::Ready(None) => return Poll::Ready(Ok(&[])),
+                    Poll::Ready(None) => {
+                        *self.as_mut().project().eof = true;
+                        return Poll::Ready(Ok(&[]));
+                    }
                     Poll::Pending => return Poll::Pending,
                 }
             }
@@ -311,6 +318,7 @@ impl<S: Unpin, B> Unpin for StreamReader<S, B> {}
 struct StreamReaderProject<'a, S, B> {
     inner: Pin<&'a mut S>,
     chunk: &'a mut Option<B>,
+    eof: &'a mut bool,
 }
 
 impl<S, B> StreamReader<S, B> {
@@ -322,6 +330,7 @@ impl<S, B> StreamReader<S, B> {
         StreamReaderProject {
             inner: unsafe { Pin::new_unchecked(&mut me.inner) },
             chunk: &mut me.chunk,
+            eof: &mut me.eof,
         }
     }
 }

tokio-util/tests/io_stream_reader.rs

@@ -33,3 +33,22 @@ async fn test_stream_reader() -> std::io::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_stream_reader_does_not_poll_after_eof() -> std::io::Result<()> {
+    // the first poll of this stream will return `Poll::Ready(None)`,
+    // and the second poll will panic
+    let stream = futures::stream::unfold((), |_| async { None::<(std::io::Result<Bytes>, ())> });
+    let read = StreamReader::new(stream);
+    tokio::pin!(read);
+    let mut buf = [0; 1];
+
+    // the first poll hits the inner stream,
+    // and the inner stream returns `Poll::Ready(None)`.
+    assert_eq!(read.read(&mut buf).await?, 0);
+    // the second poll doesn't hit the inner stream,
+    // so this `.read()` doesn't panic.
+    assert_eq!(read.read(&mut buf).await?, 0);
+
+    Ok(())
+}