diff --git a/tokio-util/src/io/stream_reader.rs b/tokio-util/src/io/stream_reader.rs index 849538a50..b954d73f5 100644 --- a/tokio-util/src/io/stream_reader.rs +++ b/tokio-util/src/io/stream_reader.rs @@ -158,6 +158,7 @@ pub struct StreamReader { inner: S, // This field is not pinned. chunk: Option, + eof: bool, } impl StreamReader @@ -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 Unpin for StreamReader {} struct StreamReaderProject<'a, S, B> { inner: Pin<&'a mut S>, chunk: &'a mut Option, + eof: &'a mut bool, } impl StreamReader { @@ -322,6 +330,7 @@ impl StreamReader { StreamReaderProject { inner: unsafe { Pin::new_unchecked(&mut me.inner) }, chunk: &mut me.chunk, + eof: &mut me.eof, } } } diff --git a/tokio-util/tests/io_stream_reader.rs b/tokio-util/tests/io_stream_reader.rs index 59759941c..36f2c204d 100644 --- a/tokio-util/tests/io_stream_reader.rs +++ b/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, ())> }); + 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(()) +}