Poll inner connection until EOF after gzip decoder is complete

src/async_impl/decoder.rs

@@ -9,6 +9,14 @@ use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 
+#[cfg(any(
+    feature = "gzip",
+    feature = "zstd",
+    feature = "brotli",
+    feature = "deflate"
+))]
+use futures_util::stream::Fuse;
+
 #[cfg(feature = "gzip")]
 use async_compression::tokio::bufread::GzipDecoder;
 
@@ -108,19 +116,19 @@ enum Inner {
 
     /// A `Gzip` decoder will uncompress the gzipped response content before returning it.
     #[cfg(feature = "gzip")]
-    Gzip(Pin<Box<FramedRead<GzipDecoder<PeekableIoStreamReader>, BytesCodec>>>),
+    Gzip(Pin<Box<Fuse<FramedRead<GzipDecoder<PeekableIoStreamReader>, BytesCodec>>>>),
 
     /// A `Brotli` decoder will uncompress the brotlied response content before returning it.
     #[cfg(feature = "brotli")]
-    Brotli(Pin<Box<FramedRead<BrotliDecoder<PeekableIoStreamReader>, BytesCodec>>>),
+    Brotli(Pin<Box<Fuse<FramedRead<BrotliDecoder<PeekableIoStreamReader>, BytesCodec>>>>),
 
     /// A `Zstd` decoder will uncompress the zstd compressed response content before returning it.
     #[cfg(feature = "zstd")]
-    Zstd(Pin<Box<FramedRead<ZstdDecoder<PeekableIoStreamReader>, BytesCodec>>>),
+    Zstd(Pin<Box<Fuse<FramedRead<ZstdDecoder<PeekableIoStreamReader>, BytesCodec>>>>),
 
     /// A `Deflate` decoder will uncompress the deflated response content before returning it.
     #[cfg(feature = "deflate")]
-    Deflate(Pin<Box<FramedRead<ZlibDecoder<PeekableIoStreamReader>, BytesCodec>>>),
+    Deflate(Pin<Box<Fuse<FramedRead<ZlibDecoder<PeekableIoStreamReader>, BytesCodec>>>>),
 
     /// A decoder that doesn't have a value yet.
     #[cfg(any(
@@ -365,34 +373,74 @@ impl HttpBody for Decoder {
             }
             #[cfg(feature = "gzip")]
             Inner::Gzip(ref mut decoder) => {
-                match futures_core::ready!(Pin::new(decoder).poll_next(cx)) {
+                match futures_core::ready!(Pin::new(&mut *decoder).poll_next(cx)) {
                     Some(Ok(bytes)) => Poll::Ready(Some(Ok(Frame::data(bytes.freeze())))),
                     Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
-                    None => Poll::Ready(None),
+                    None => {
+                        // poll inner connection until EOF after gzip stream is finished
+                        let inner_stream = decoder.get_mut().get_mut().get_mut().get_mut();
+                        match futures_core::ready!(Pin::new(inner_stream).poll_next(cx)) {
+                            Some(Ok(_)) => Poll::Ready(Some(Err(crate::error::decode(
+                                "there are extra bytes after body has been decompressed",
+                            )))),
+                            Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
+                            None => Poll::Ready(None),
+                        }
+                    }
                 }
             }
             #[cfg(feature = "brotli")]
             Inner::Brotli(ref mut decoder) => {
-                match futures_core::ready!(Pin::new(decoder).poll_next(cx)) {
+                match futures_core::ready!(Pin::new(&mut *decoder).poll_next(cx)) {
                     Some(Ok(bytes)) => Poll::Ready(Some(Ok(Frame::data(bytes.freeze())))),
                     Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
-                    None => Poll::Ready(None),
+                    None => {
+                        // poll inner connection until EOF after brotli stream is finished
+                        let inner_stream = decoder.get_mut().get_mut().get_mut().get_mut();
+                        match futures_core::ready!(Pin::new(inner_stream).poll_next(cx)) {
+                            Some(Ok(_)) => Poll::Ready(Some(Err(crate::error::decode(
+                                "there are extra bytes after body has been decompressed",
+                            )))),
+                            Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
+                            None => Poll::Ready(None),
+                        }
+                    }
                 }
             }
             #[cfg(feature = "zstd")]
             Inner::Zstd(ref mut decoder) => {
-                match futures_core::ready!(Pin::new(decoder).poll_next(cx)) {
+                match futures_core::ready!(Pin::new(&mut *decoder).poll_next(cx)) {
                     Some(Ok(bytes)) => Poll::Ready(Some(Ok(Frame::data(bytes.freeze())))),
                     Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
-                    None => Poll::Ready(None),
+                    None => {
+                        // poll inner connection until EOF after zstd stream is finished
+                        let inner_stream = decoder.get_mut().get_mut().get_mut().get_mut();
+                        match futures_core::ready!(Pin::new(inner_stream).poll_next(cx)) {
+                            Some(Ok(_)) => Poll::Ready(Some(Err(crate::error::decode(
+                                "there are extra bytes after body has been decompressed",
+                            )))),
+                            Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
+                            None => Poll::Ready(None),
+                        }
+                    }
                 }
             }
             #[cfg(feature = "deflate")]
             Inner::Deflate(ref mut decoder) => {
-                match futures_core::ready!(Pin::new(decoder).poll_next(cx)) {
+                match futures_core::ready!(Pin::new(&mut *decoder).poll_next(cx)) {
                     Some(Ok(bytes)) => Poll::Ready(Some(Ok(Frame::data(bytes.freeze())))),
                     Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
-                    None => Poll::Ready(None),
+                    None => {
+                        // poll inner connection until EOF after deflate stream is finished
+                        let inner_stream = decoder.get_mut().get_mut().get_mut().get_mut();
+                        match futures_core::ready!(Pin::new(inner_stream).poll_next(cx)) {
+                            Some(Ok(_)) => Poll::Ready(Some(Err(crate::error::decode(
+                                "there are extra bytes after body has been decompressed",
+                            )))),
+                            Some(Err(err)) => Poll::Ready(Some(Err(crate::error::decode_io(err)))),
+                            None => Poll::Ready(None),
+                        }
+                    }
                 }
             }
         }
@@ -456,25 +504,37 @@ impl Future for Pending {
 
         match self.1 {
             #[cfg(feature = "brotli")]
-            DecoderType::Brotli => Poll::Ready(Ok(Inner::Brotli(Box::pin(FramedRead::new(
-                BrotliDecoder::new(StreamReader::new(_body)),
-                BytesCodec::new(),
-            ))))),
+            DecoderType::Brotli => Poll::Ready(Ok(Inner::Brotli(Box::pin(
+                FramedRead::new(
+                    BrotliDecoder::new(StreamReader::new(_body)),
+                    BytesCodec::new(),
+                )
+                .fuse(),
+            )))),
             #[cfg(feature = "zstd")]
-            DecoderType::Zstd => Poll::Ready(Ok(Inner::Zstd(Box::pin(FramedRead::new(
-                ZstdDecoder::new(StreamReader::new(_body)),
-                BytesCodec::new(),
-            ))))),
+            DecoderType::Zstd => Poll::Ready(Ok(Inner::Zstd(Box::pin(
+                FramedRead::new(
+                    ZstdDecoder::new(StreamReader::new(_body)),
+                    BytesCodec::new(),
+                )
+                .fuse(),
+            )))),
             #[cfg(feature = "gzip")]
-            DecoderType::Gzip => Poll::Ready(Ok(Inner::Gzip(Box::pin(FramedRead::new(
-                GzipDecoder::new(StreamReader::new(_body)),
-                BytesCodec::new(),
-            ))))),
+            DecoderType::Gzip => Poll::Ready(Ok(Inner::Gzip(Box::pin(
+                FramedRead::new(
+                    GzipDecoder::new(StreamReader::new(_body)),
+                    BytesCodec::new(),
+                )
+                .fuse(),
+            )))),
             #[cfg(feature = "deflate")]
-            DecoderType::Deflate => Poll::Ready(Ok(Inner::Deflate(Box::pin(FramedRead::new(
-                ZlibDecoder::new(StreamReader::new(_body)),
-                BytesCodec::new(),
-            ))))),
+            DecoderType::Deflate => Poll::Ready(Ok(Inner::Deflate(Box::pin(
+                FramedRead::new(
+                    ZlibDecoder::new(StreamReader::new(_body)),
+                    BytesCodec::new(),
+                )
+                .fuse(),
+            )))),
         }
     }
 }

tests/brotli.rs

@@ -1,6 +1,7 @@
 mod support;
 use std::io::Read;
 use support::server;
+use tokio::io::AsyncWriteExt;
 
 #[tokio::test]
 async fn brotli_response() {
@@ -145,3 +146,212 @@ async fn brotli_case(response_size: usize, chunk_size: usize) {
     let body = res.text().await.expect("text");
     assert_eq!(body, content);
 }
+
+const COMPRESSED_RESPONSE_HEADERS: &[u8] = b"HTTP/1.1 200 OK\x0d\x0a\
+            Content-Type: text/plain\x0d\x0a\
+            Connection: keep-alive\x0d\x0a\
+            Content-Encoding: br\x0d\x0a";
+
+const RESPONSE_CONTENT: &str = "some message here";
+
+fn brotli_compress(input: &[u8]) -> Vec<u8> {
+    let mut encoder = brotli_crate::CompressorReader::new(input, 4096, 5, 20);
+    let mut brotlied_content = Vec::new();
+    encoder.read_to_end(&mut brotlied_content).unwrap();
+    brotlied_content
+}
+
+#[tokio::test]
+async fn test_non_chunked_non_fragmented_response() {
+    let server = server::low_level_with_response(|_raw_request, client_socket| {
+        Box::new(async move {
+            let brotlied_content = brotli_compress(RESPONSE_CONTENT.as_bytes());
+            let content_length_header =
+                format!("Content-Length: {}\r\n\r\n", brotlied_content.len()).into_bytes();
+            let response = [
+                COMPRESSED_RESPONSE_HEADERS,
+                &content_length_header,
+                &brotlied_content,
+            ]
+            .concat();
+
+            client_socket
+                .write_all(response.as_slice())
+                .await
+                .expect("response write_all failed");
+            client_socket.flush().await.expect("response flush failed");
+        })
+    });
+
+    let res = reqwest::Client::new()
+        .get(&format!("http://{}/", server.addr()))
+        .send()
+        .await
+        .expect("response");
+
+    assert_eq!(res.text().await.expect("text"), RESPONSE_CONTENT);
+}
+
+#[tokio::test]
+async fn test_chunked_fragmented_response_1() {
+    const DELAY_BETWEEN_RESPONSE_PARTS: tokio::time::Duration =
+        tokio::time::Duration::from_millis(1000);
+    const DELAY_MARGIN: tokio::time::Duration = tokio::time::Duration::from_millis(50);
+
+    let server = server::low_level_with_response(|_raw_request, client_socket| {
+        Box::new(async move {
+            let brotlied_content = brotli_compress(RESPONSE_CONTENT.as_bytes());
+            let response_first_part = [
+                COMPRESSED_RESPONSE_HEADERS,
+                format!(
+                    "Transfer-Encoding: chunked\r\n\r\n{:x}\r\n",
+                    brotlied_content.len()
+                )
+                .as_bytes(),
+                &brotlied_content,
+            ]
+            .concat();
+            let response_second_part = b"\r\n0\r\n\r\n";
+
+            client_socket
+                .write_all(response_first_part.as_slice())
+                .await
+                .expect("response_first_part write_all failed");
+            client_socket
+                .flush()
+                .await
+                .expect("response_first_part flush failed");
+
+            tokio::time::sleep(DELAY_BETWEEN_RESPONSE_PARTS).await;
+
+            client_socket
+                .write_all(response_second_part)
+                .await
+                .expect("response_second_part write_all failed");
+            client_socket
+                .flush()
+                .await
+                .expect("response_second_part flush failed");
+        })
+    });
+
+    let start = tokio::time::Instant::now();
+    let res = reqwest::Client::new()
+        .get(&format!("http://{}/", server.addr()))
+        .send()
+        .await
+        .expect("response");
+
+    assert_eq!(res.text().await.expect("text"), RESPONSE_CONTENT);
+    assert!(start.elapsed() >= DELAY_BETWEEN_RESPONSE_PARTS - DELAY_MARGIN);
+}
+
+#[tokio::test]
+async fn test_chunked_fragmented_response_2() {
+    const DELAY_BETWEEN_RESPONSE_PARTS: tokio::time::Duration =
+        tokio::time::Duration::from_millis(1000);
+    const DELAY_MARGIN: tokio::time::Duration = tokio::time::Duration::from_millis(50);
+
+    let server = server::low_level_with_response(|_raw_request, client_socket| {
+        Box::new(async move {
+            let brotlied_content = brotli_compress(RESPONSE_CONTENT.as_bytes());
+            let response_first_part = [
+                COMPRESSED_RESPONSE_HEADERS,
+                format!(
+                    "Transfer-Encoding: chunked\r\n\r\n{:x}\r\n",
+                    brotlied_content.len()
+                )
+                .as_bytes(),
+                &brotlied_content,
+                b"\r\n",
+            ]
+            .concat();
+            let response_second_part = b"0\r\n\r\n";
+
+            client_socket
+                .write_all(response_first_part.as_slice())
+                .await
+                .expect("response_first_part write_all failed");
+            client_socket
+                .flush()
+                .await
+                .expect("response_first_part flush failed");
+
+            tokio::time::sleep(DELAY_BETWEEN_RESPONSE_PARTS).await;
+
+            client_socket
+                .write_all(response_second_part)
+                .await
+                .expect("response_second_part write_all failed");
+            client_socket
+                .flush()
+                .await
+                .expect("response_second_part flush failed");
+        })
+    });
+
+    let start = tokio::time::Instant::now();
+    let res = reqwest::Client::new()
+        .get(&format!("http://{}/", server.addr()))
+        .send()
+        .await
+        .expect("response");
+
+    assert_eq!(res.text().await.expect("text"), RESPONSE_CONTENT);
+    assert!(start.elapsed() >= DELAY_BETWEEN_RESPONSE_PARTS - DELAY_MARGIN);
+}
+
+#[tokio::test]
+async fn test_chunked_fragmented_response_with_extra_bytes() {
+    const DELAY_BETWEEN_RESPONSE_PARTS: tokio::time::Duration =
+        tokio::time::Duration::from_millis(1000);
+    const DELAY_MARGIN: tokio::time::Duration = tokio::time::Duration::from_millis(50);
+
+    let server = server::low_level_with_response(|_raw_request, client_socket| {
+        Box::new(async move {
+            let brotlied_content = brotli_compress(RESPONSE_CONTENT.as_bytes());
+            let response_first_part = [
+                COMPRESSED_RESPONSE_HEADERS,
+                format!(
+                    "Transfer-Encoding: chunked\r\n\r\n{:x}\r\n",
+                    brotlied_content.len()
+                )
+                .as_bytes(),
+                &brotlied_content,
+            ]
+            .concat();
+            let response_second_part = b"\r\n2ab\r\n0\r\n\r\n";
+
+            client_socket
+                .write_all(response_first_part.as_slice())
+                .await
+                .expect("response_first_part write_all failed");
+            client_socket
+                .flush()
+                .await
+                .expect("response_first_part flush failed");
+
+            tokio::time::sleep(DELAY_BETWEEN_RESPONSE_PARTS).await;
+
+            client_socket
+                .write_all(response_second_part)
+                .await
+                .expect("response_second_part write_all failed");
+            client_socket
+                .flush()
+                .await
+                .expect("response_second_part flush failed");
+        })
+    });
+
+    let start = tokio::time::Instant::now();
+    let res = reqwest::Client::new()
+        .get(&format!("http://{}/", server.addr()))
+        .send()
+        .await
+        .expect("response");
+
+    let err = res.text().await.expect_err("there must be an error");
+    assert!(err.is_decode());
+    assert!(start.elapsed() >= DELAY_BETWEEN_RESPONSE_PARTS - DELAY_MARGIN);
+}