feat: close ImapStream after an error

This prevents calls to poll_next()
from reaching the underlying stream
after an error is returned once.

Previously calls to poll_next()
to ImapStream built on top
of a network connection
could have returned infinite
stream of Some(Err(_))
values.
This is dangerous
for the code that looks for the end
of stream without processing
errors such as
`while stream.next().await.is_some() {}`
because it may result in
infinite loop.

It is still better to process
errors by writing
`while stream.try_next().await?.is_some() {}`,
but the change protects
in case of incorrect library user code.

Author: link2xt

src/imap_stream.rs

@@ -25,6 +25,15 @@ pub struct ImapStream<R: Read + Write> {
     decode_needs: usize,
     /// The buffer.
     buffer: Buffer,
+
+    /// True if the stream should not return any more items.
+    ///
+    /// This is set when reading from a stream
+    /// returns an error.
+    /// Afterwards the stream returns only `None`
+    /// and `poll_next()` does not read
+    /// from the underlying stream.
+    read_closed: bool,
 }
 
 impl<R: Read + Write + Unpin> ImapStream<R> {
@@ -34,6 +43,7 @@ impl<R: Read + Write + Unpin> ImapStream<R> {
             inner,
             buffer: Buffer::new(),
             decode_needs: 0,
+            read_closed: false,
         }
     }
 
@@ -132,6 +142,52 @@ impl<R: Read + Write + Unpin> ImapStream<R> {
             }
         }
     }
+
+    fn do_poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<io::Result<ResponseData>>> {
+        let this = &mut *self;
+        if let Some(response) = this.decode()? {
+            return Poll::Ready(Some(Ok(response)));
+        }
+        loop {
+            this.buffer.ensure_capacity(this.decode_needs)?;
+            let buf = this.buffer.free_as_mut_slice();
+
+            // The buffer should have at least one byte free
+            // before we try reading into it
+            // so we can treat 0 bytes read as EOF.
+            // This is guaranteed by `ensure_capacity()` above
+            // even if it is called with 0 as an argument.
+            debug_assert!(!buf.is_empty());
+
+            #[cfg(feature = "runtime-async-std")]
+            let num_bytes_read = ready!(Pin::new(&mut this.inner).poll_read(cx, buf))?;
+
+            #[cfg(feature = "runtime-tokio")]
+            let num_bytes_read = {
+                let buf = &mut tokio::io::ReadBuf::new(buf);
+                let start = buf.filled().len();
+                ready!(Pin::new(&mut this.inner).poll_read(cx, buf))?;
+                buf.filled().len() - start
+            };
+
+            if num_bytes_read == 0 {
+                if this.buffer.used() > 0 {
+                    return Poll::Ready(Some(Err(io::Error::new(
+                        io::ErrorKind::UnexpectedEof,
+                        "bytes remaining in stream",
+                    ))));
+                }
+                return Poll::Ready(None);
+            }
+            this.buffer.extend_used(num_bytes_read);
+            if let Some(response) = this.decode()? {
+                return Poll::Ready(Some(Ok(response)));
+            }
+        }
+    }
 }
 
 /// Abstraction around needed buffer management.
@@ -273,46 +329,18 @@ impl<R: Read + Write + Unpin> Stream for ImapStream<R> {
     type Item = io::Result<ResponseData>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let this = &mut *self;
-        if let Some(response) = this.decode()? {
-            return Poll::Ready(Some(Ok(response)));
+        if self.read_closed {
+            return Poll::Ready(None);
         }
-        loop {
-            this.buffer.ensure_capacity(this.decode_needs)?;
-            let buf = this.buffer.free_as_mut_slice();
-
-            // The buffer should have at least one byte free
-            // before we try reading into it
-            // so we can treat 0 bytes read as EOF.
-            // This is guaranteed by `ensure_capacity()` above
-            // even if it is called with 0 as an argument.
-            debug_assert!(!buf.is_empty());
-
-            #[cfg(feature = "runtime-async-std")]
-            let num_bytes_read = ready!(Pin::new(&mut this.inner).poll_read(cx, buf))?;
-
-            #[cfg(feature = "runtime-tokio")]
-            let num_bytes_read = {
-                let buf = &mut tokio::io::ReadBuf::new(buf);
-                let start = buf.filled().len();
-                ready!(Pin::new(&mut this.inner).poll_read(cx, buf))?;
-                buf.filled().len() - start
-            };
-
-            if num_bytes_read == 0 {
-                if this.buffer.used() > 0 {
-                    return Poll::Ready(Some(Err(io::Error::new(
-                        io::ErrorKind::UnexpectedEof,
-                        "bytes remaining in stream",
-                    ))));
-                }
-                return Poll::Ready(None);
-            }
-            this.buffer.extend_used(num_bytes_read);
-            if let Some(response) = this.decode()? {
-                return Poll::Ready(Some(Ok(response)));
+        let res = match ready!(self.as_mut().do_poll_next(cx)) {
+            None => None,
+            Some(Err(err)) => {
+                self.read_closed = true;
+                Some(Err(err))
             }
-        }
+            Some(Ok(item)) => Some(Ok(item)),
+        };
+        Poll::Ready(res)
     }
 }