fix(cursors/bytes): protect against empty chunks

Author: loyd

src/cursors/bytes.rs

@@ -1,5 +1,5 @@
-use crate::{bytes_ext::BytesExt, cursors::RawCursor, error::Result, response::Response};
-use bytes::{Bytes, BytesMut};
+use crate::{cursors::RawCursor, error::Result, response::Response};
+use bytes::{Buf, Bytes, BytesMut};
 use std::{
     io::Result as IoResult,
     pin::Pin,
@@ -33,7 +33,7 @@ use tokio::io::{AsyncBufRead, AsyncRead, ReadBuf};
 /// [`Query::fetch_bytes`]: crate::query::Query::fetch_bytes
 pub struct BytesCursor {
     raw: RawCursor,
-    bytes: BytesExt,
+    bytes: Bytes,
 }
 
 // TODO: what if any next/poll_* called AFTER error returned?
@@ -42,7 +42,7 @@ impl BytesCursor {
     pub(crate) fn new(response: Response) -> Self {
         Self {
             raw: RawCursor::new(response),
-            bytes: BytesExt::default(),
+            bytes: Bytes::default(),
         }
     }
 
@@ -82,17 +82,19 @@ impl BytesCursor {
 
     #[cold]
     fn poll_refill(&mut self, cx: &mut Context<'_>) -> Poll<IoResult<bool>> {
-        debug_assert_eq!(self.bytes.remaining(), 0);
-
-        // TODO: should we repeat if `poll_next` returns an empty buffer?
-
-        match ready!(self.raw.poll_next(cx)?) {
-            Some(chunk) => {
-                self.bytes.extend(chunk);
-                Poll::Ready(Ok(true))
+        debug_assert_eq!(self.bytes.len(), 0);
+
+        // Theoretically, `self.raw.poll_next(cx)` can return empty chunks.
+        // In this case, we should continue polling until we get a non-empty chunk or
+        // end of stream in order to avoid false positive `Ok(0)` in I/O traits.
+        while self.bytes.is_empty() {
+            match ready!(self.raw.poll_next(cx)?) {
+                Some(chunk) => self.bytes = chunk,
+                None => return Poll::Ready(Ok(false)),
             }
-            None => Poll::Ready(Ok(false)),
         }
+
+        Poll::Ready(Ok(true))
     }
 
     /// Returns the total size in bytes received from the CH server since
@@ -125,9 +127,9 @@ impl AsyncRead for BytesCursor {
                 break;
             }
 
-            let bytes = self.bytes.slice();
-            let len = bytes.len().min(buf.remaining());
-            buf.put_slice(&bytes[..len]);
+            let len = self.bytes.len().min(buf.remaining());
+            let bytes = self.bytes.slice(..len);
+            buf.put_slice(&bytes[0..len]);
             self.bytes.advance(len);
         }
 
@@ -138,17 +140,17 @@ impl AsyncRead for BytesCursor {
 impl AsyncBufRead for BytesCursor {
     #[inline]
     fn poll_fill_buf(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<&[u8]>> {
-        if self.bytes.is_empty() && !ready!(self.poll_refill(cx)?) {
-            return Poll::Ready(Ok(&[]));
+        if self.bytes.is_empty() {
+            ready!(self.poll_refill(cx)?);
         }
 
-        Poll::Ready(Ok(self.get_mut().bytes.slice()))
+        Poll::Ready(Ok(&self.get_mut().bytes))
     }
 
     #[inline]
     fn consume(mut self: Pin<&mut Self>, amt: usize) {
         assert!(
-            amt <= self.bytes.remaining(),
+            amt <= self.bytes.len(),
             "invalid `AsyncBufRead::consume` usage"
         );
         self.bytes.advance(amt);

src/response.rs

@@ -147,7 +147,7 @@ impl Chunks {
     fn new(stream: Incoming, compression: Compression) -> Self {
         let stream = IncomingStream(stream);
         let stream = Decompress::new(stream, compression);
-        let stream = DetectDbException::new(stream);
+        let stream = DetectDbException(stream);
         Self(Some(Box::new(stream)))
     }
 }
@@ -156,14 +156,12 @@ impl Stream for Chunks {
     type Item = Result<Chunk>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        // `take()` prevents from use after caught panic.
+        // We use `take()` to make the stream fused, including the case of panics.
         if let Some(mut stream) = self.0.take() {
             let res = Pin::new(&mut stream).poll_next(cx);
 
             if matches!(res, Poll::Pending | Poll::Ready(Some(Ok(_)))) {
                 self.0 = Some(stream);
-            } else {
-                assert!(self.0.is_none());
             }
 
             res
@@ -244,16 +242,7 @@ where
 
 // === DetectDbException ===
 
-enum DetectDbException<S> {
-    Stream(S),
-    Exception(Option<Error>),
-}
-
-impl<S> DetectDbException<S> {
-    fn new(stream: S) -> Self {
-        Self::Stream(stream)
-    }
-}
+struct DetectDbException<S>(S);
 
 impl<S> Stream for DetectDbException<S>
 where
@@ -262,30 +251,23 @@ where
     type Item = Result<Chunk>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        match &mut *self {
-            Self::Stream(stream) => {
-                let mut res = Pin::new(stream).poll_next(cx);
-
-                if let Poll::Ready(Some(Ok(chunk))) = &mut res {
-                    if let Some(err) = extract_exception(&mut chunk.data) {
-                        *self = Self::Exception(Some(err));
+        let res = Pin::new(&mut self.0).poll_next(cx);
 
-                        // NOTE: `chunk` can be empty, but it's ok for callers.
-                    }
-                }
-
-                res
+        if let Poll::Ready(Some(Ok(chunk))) = &res {
+            if let Some(err) = extract_exception(&chunk.data) {
+                return Poll::Ready(Some(Err(err)));
             }
-            Self::Exception(err) => Poll::Ready(err.take().map(Err)),
         }
+
+        res
     }
 }
 
 // Format:
 // ```
 //   <data>Code: <code>. DB::Exception: <desc> (version <version> (official build))\n
 // ```
-fn extract_exception(chunk: &mut Bytes) -> Option<Error> {
+fn extract_exception(chunk: &[u8]) -> Option<Error> {
     // `))\n` is very rare in real data, so it's fast dirty check.
     // In random data, it occurs with a probability of ~6*10^-8 only.
     if chunk.ends_with(b"))\n") {
@@ -297,14 +279,13 @@ fn extract_exception(chunk: &mut Bytes) -> Option<Error> {
 
 #[cold]
 #[inline(never)]
-fn extract_exception_slow(chunk: &mut Bytes) -> Option<Error> {
+fn extract_exception_slow(chunk: &[u8]) -> Option<Error> {
     let index = chunk.rfind(b"Code:")?;
 
     if !chunk[index..].contains_str(b"DB::Exception:") {
         return None;
     }
 
-    let exception = chunk.split_off(index);
-    let exception = String::from_utf8_lossy(&exception[..exception.len() - 1]);
+    let exception = String::from_utf8_lossy(&chunk[index..chunk.len() - 1]);
     Some(Error::BadResponse(exception.into()))
 }