feat(decompression): support HTTP responses containing multiple ZSTD frames (#548)

0x676e67 · web-flow · commit 711a86d6cbe2 · 2025-03-12T16:03:06.000-04:00
diff --git a/tower-http/src/decompression/body.rs b/tower-http/src/decompression/body.rs
@@ -397,7 +397,9 @@ where
     type Output = ZstdDecoder<Self::Input>;
 
     fn apply(input: Self::Input, _quality: CompressionLevel) -> Self::Output {
-        ZstdDecoder::new(input)
+        let mut decoder = ZstdDecoder::new(input);
+        decoder.multiple_members(true);
+        decoder
     }
 
     fn get_pin_mut(pinned: Pin<&mut Self::Output>) -> Pin<&mut Self::Input> {
diff --git a/tower-http/src/decompression/mod.rs b/tower-http/src/decompression/mod.rs
@@ -168,6 +168,24 @@ mod tests {
         assert_eq!(decompressed_data, "Hello, World!");
     }
 
+    #[tokio::test]
+    async fn decompress_multi_zstd() {
+        let mut client = Decompression::new(service_fn(handle_multi_zstd));
+
+        let req = Request::builder()
+            .header("accept-encoding", "zstd")
+            .body(Body::empty())
+            .unwrap();
+        let res = client.ready().await.unwrap().call(req).await.unwrap();
+
+        // read the body, it will be decompressed automatically
+        let body = res.into_body();
+        let decompressed_data =
+            String::from_utf8(body.collect().await.unwrap().to_bytes().to_vec()).unwrap();
+
+        assert_eq!(decompressed_data, "Hello, World!");
+    }
+
     async fn handle_multi_gz(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
         let mut buf = Vec::new();
         let mut enc1 = GzEncoder::new(&mut buf, Default::default());
@@ -184,6 +202,22 @@ mod tests {
         Ok(res)
     }
 
+    async fn handle_multi_zstd(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
+        let mut buf = Vec::new();
+        let mut enc1 = zstd::Encoder::new(&mut buf, Default::default()).unwrap();
+        enc1.write_all(b"Hello, ").unwrap();
+        enc1.finish().unwrap();
+
+        let mut enc2 = zstd::Encoder::new(&mut buf, Default::default()).unwrap();
+        enc2.write_all(b"World!").unwrap();
+        enc2.finish().unwrap();
+
+        let mut res = Response::new(Body::from(buf));
+        res.headers_mut()
+            .insert("content-encoding", "zstd".parse().unwrap());
+        Ok(res)
+    }
+
     #[allow(dead_code)]
     async fn is_compatible_with_hyper() {
         let client =

Original file line number	Diff line number	Diff line change
`@@ -397,7 +397,9 @@ where`
`397`	`397`	`type Output = ZstdDecoder<Self::Input>;`
`398`	`398`
`399`	`399`	`fn apply(input: Self::Input, _quality: CompressionLevel) -> Self::Output {`
`400`		`- ZstdDecoder::new(input)`
	`400`	`+ let mut decoder = ZstdDecoder::new(input);`
	`401`	`+ decoder.multiple_members(true);`
	`402`	`+ decoder`
`401`	`403`	`}`
`402`	`404`
`403`	`405`	`fn get_pin_mut(pinned: Pin<&mut Self::Output>) -> Pin<&mut Self::Input> {`