Fix severe compression bug introduced by #97

Cargo.toml

@@ -18,13 +18,18 @@ edition = "2021"
 clap = { version = "3", default-features = false, features = ["std", "cargo"] }
 # Server
 tokio = { version = "1", features = ["rt-multi-thread", "macros"] }
+tokio-util = { version = "0.7", features = ["io"] }
 hyper = { version = "0.14.20", features = ["http1", "server", "tcp", "stream"] }
 headers = "0.3"
 mime_guess = "2.0"
 percent-encoding = "2.1"
 # Compression
-brotli = "3"
-flate2 = "1"
+async-compression = { version = "0.3.7", features = [
+  "brotli",
+  "deflate",
+  "gzip",
+  "tokio",
+] }
 # Rendering
 tera = "1"
 serde = { version = "1.0", features = [

src/http/content_encoding.rs

@@ -7,11 +7,17 @@
 // except according to those terms.
 
 use std::cmp::Ordering;
-use std::io::{self, BufReader};
+use std::io;
 
-use flate2::read::{DeflateEncoder, GzEncoder};
-use flate2::Compression;
+use async_compression::{
+    tokio::bufread::{BrotliEncoder, DeflateEncoder, GzipEncoder},
+    Level,
+};
+use bytes::Bytes;
+use futures::Stream;
 use hyper::header::HeaderValue;
+use hyper::Body;
+use tokio_util::io::{ReaderStream, StreamReader};
 
 pub const IDENTITY: &str = "identity";
 pub const DEFLATE: &str = "deflate";
@@ -122,30 +128,22 @@ pub fn get_prior_encoding<'a>(accept_encoding: &'a HeaderValue) -> &'static str
         .unwrap_or(IDENTITY)
 }
 
-/// Compress data.
-///
-/// # Parameters
-///
-/// * `data` - Data to be compressed.
-/// * `encoding` - Only support `br`, `gzip`, `deflate` and `identity`.
-pub fn compress(data: &[u8], encoding: &str) -> io::Result<Vec<u8>> {
-    use std::io::prelude::*;
-    let mut buf = Vec::new();
+pub fn compress_stream(
+    input: impl Stream<Item = io::Result<Bytes>> + std::marker::Send + 'static,
+    encoding: &str,
+) -> io::Result<hyper::Body> {
     match encoding {
-        BR => {
-            BufReader::new(brotli::CompressorReader::new(data, 4096, 6, 20))
-                .read_to_end(&mut buf)?;
-        }
-        GZIP => {
-            BufReader::new(GzEncoder::new(data, Compression::default())).read_to_end(&mut buf)?;
-        }
-        DEFLATE => {
-            BufReader::new(DeflateEncoder::new(data, Compression::default()))
-                .read_to_end(&mut buf)?;
-        }
-        _ => return Err(io::Error::new(io::ErrorKind::Other, "Unsupported Encoding")),
-    };
-    Ok(buf)
+        BR => Ok(Body::wrap_stream(ReaderStream::new(
+            BrotliEncoder::with_quality(StreamReader::new(input), Level::Fastest),
+        ))),
+        DEFLATE => Ok(Body::wrap_stream(ReaderStream::new(DeflateEncoder::new(
+            StreamReader::new(input),
+        )))),
+        GZIP => Ok(Body::wrap_stream(ReaderStream::new(GzipEncoder::new(
+            StreamReader::new(input),
+        )))),
+        _ => Err(io::Error::new(io::ErrorKind::Other, "Unsupported Encoding")),
+    }
 }
 
 pub fn should_compress(enc: &str) -> bool {
@@ -248,20 +246,27 @@ mod t_prior {
 #[cfg(test)]
 mod t_compress {
     use super::*;
+    use bytes::Bytes;
 
     #[test]
     fn failed() {
-        let error = compress(b"hello", "unrecognized").unwrap_err();
+        let s = futures::stream::iter(vec![Ok::<_, io::Error>(Bytes::from_static(b"hello"))]);
+        let error = compress_stream(s, "unrecognized").unwrap_err();
         assert_eq!(error.kind(), io::ErrorKind::Other);
     }
 
-    #[test]
-    fn compressed() {
-        let buf = compress(b"xxxxx", DEFLATE);
-        assert!(!buf.unwrap().is_empty());
-        let buf = compress(b"xxxxx", GZIP);
-        assert!(!buf.unwrap().is_empty());
-        let buf = compress(b"xxxxx", BR);
-        assert!(!buf.unwrap().is_empty());
+    #[tokio::test]
+    async fn compressed() {
+        let s = futures::stream::iter(vec![Ok::<_, io::Error>(Bytes::from_static(b"xxxxx"))]);
+        let body = compress_stream(s, BR).unwrap();
+        assert_eq!(hyper::body::to_bytes(body).await.unwrap().len(), 9);
+
+        let s = futures::stream::iter(vec![Ok::<_, io::Error>(Bytes::from_static(b"xxxxx"))]);
+        let body = compress_stream(s, DEFLATE).unwrap();
+        assert_eq!(hyper::body::to_bytes(body).await.unwrap().len(), 5);
+
+        let s = futures::stream::iter(vec![Ok::<_, io::Error>(Bytes::from_static(b"xxxxx"))]);
+        let body = compress_stream(s, GZIP).unwrap();
+        assert_eq!(hyper::body::to_bytes(body).await.unwrap().len(), 23);
     }
 }

src/server/serve.rs

@@ -14,7 +14,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use chrono::Local;
-use futures::{StreamExt as _, TryStreamExt as _};
+use futures::TryStreamExt as _;
 use headers::{
     AcceptRanges, AccessControlAllowHeaders, AccessControlAllowOrigin, CacheControl, ContentLength,
     ContentType, ETag, HeaderMapExt, LastModified, Range, Server,
@@ -32,7 +32,7 @@ use serde::Serialize;
 use crate::cli::Args;
 use crate::extensions::{MimeExt, PathExt, SystemTimeExt};
 use crate::http::conditional_requests::{is_fresh, is_precondition_failed};
-use crate::http::content_encoding::{compress, get_prior_encoding, should_compress};
+use crate::http::content_encoding::{compress_stream, get_prior_encoding, should_compress};
 use crate::http::range_requests::{is_range_fresh, is_satisfiable_range};
 
 use crate::server::send::{send_dir, send_dir_as_zip, send_file, send_file_with_range};
@@ -411,13 +411,10 @@ impl InnerService {
             if let Some(encoding) = req.headers().get(hyper::header::ACCEPT_ENCODING) {
                 let content_encoding = get_prior_encoding(encoding);
                 if should_compress(content_encoding) {
-                    let stream = body
-                        .map_err(|_e| io::Error::from(io::ErrorKind::InvalidData))
-                        .map(|b| match b {
-                            Ok(b) => compress(&b, content_encoding),
-                            Err(e) => Err(e),
-                        });
-                    let body = Body::wrap_stream(stream);
+                    let b = compress_stream(
+                        body.map_err(|e| io::Error::new(io::ErrorKind::Other, e)),
+                        content_encoding,
+                    )?;
                     res.headers_mut().insert(
                         hyper::header::CONTENT_ENCODING,
                         hyper::header::HeaderValue::from_static(content_encoding),
@@ -427,7 +424,7 @@ impl InnerService {
                         hyper::header::VARY,
                         hyper::header::HeaderValue::from_name(hyper::header::ACCEPT_ENCODING),
                     );
-                    body
+                    b
                 } else {
                     body
                 }