rama_http/layer/decompression/
mod.rs

1//! Middleware that decompresses request and response bodies.
2//!
3//! # Examples
4//!
5//! #### Request
6//! ```rust
7//! use std::{error::Error, io::Write};
8//!
9//! use rama_core::bytes::{Bytes, BytesMut};
10//! use flate2::{write::GzEncoder, Compression};
11//!
12//! use rama_http::{Body, header, HeaderValue, Request, Response};
13//! use rama_core::service::service_fn;
14//! use rama_core::{Service, Layer};
15//! use rama_http::layer::decompression::{DecompressionBody, RequestDecompressionLayer};
16//! use rama_http::body::util::BodyExt;
17//! use rama_core::error::BoxError;
18//!
19//! # #[tokio::main]
20//! # async fn main() -> Result<(), BoxError> {
21//! // A request encoded with gzip coming from some HTTP client.
22//! let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
23//! encoder.write_all(b"Hello?")?;
24//! let request = Request::builder()
25//!     .header(header::CONTENT_ENCODING, "gzip")
26//!     .body(Body::from(encoder.finish()?))?;
27//!
28//! // Our HTTP server
29//! let mut server = (
30//!     // Automatically decompress request bodies.
31//!     RequestDecompressionLayer::new(),
32//! ).into_layer(service_fn(handler));
33//!
34//! // Send the request, with the gzip encoded body, to our server.
35//! let _response = server.serve(request).await?;
36//!
37//! // Handler receives request whose body is decoded when read
38//! async fn handler(mut req: Request<DecompressionBody<Body>>) -> Result<Response, BoxError>{
39//!     let data = req.into_body().collect().await?.to_bytes();
40//!     assert_eq!(&data[..], b"Hello?");
41//!     Ok(Response::new(Body::from("Hello, World!")))
42//! }
43//! # Ok(())
44//! # }
45//! ```
46//!
47//! #### Response
48//!
49//! ```rust
50//! use std::convert::Infallible;
51//!
52//! use rama_core::bytes::{Bytes, BytesMut};
53//!
54//! use rama_http::{Body, Request, Response};
55//! use rama_core::service::service_fn;
56//! use rama_core::{Service, Layer};
57//! use rama_http::layer::{compression::Compression, decompression::DecompressionLayer};
58//! use rama_http::body::util::BodyExt;
59//! use rama_core::error::BoxError;
60//!
61//! #
62//! # #[tokio::main]
63//! # async fn main() -> Result<(), BoxError> {
64//! # async fn handle(req: Request) -> Result<Response, Infallible> {
65//! #     let body = Body::from("Hello, World!");
66//! #     Ok(Response::new(body))
67//! # }
68//!
69//! // Some opaque service that applies compression.
70//! let service = Compression::new(service_fn(handle));
71//!
72//! // Our HTTP client.
73//! let mut client = (
74//!     // Automatically decompress response bodies.
75//!     DecompressionLayer::new(),
76//! ).into_layer(service);
77//!
78//! // Call the service.
79//! //
80//! // `DecompressionLayer` takes care of setting `Accept-Encoding`.
81//! let request = Request::new(Body::default());
82//!
83//! let response = client
84//!     .serve(request)
85//!     .await?;
86//!
87//! // Read the body
88//! let body = response.into_body();
89//! let bytes = body.collect().await?.to_bytes().to_vec();
90//! let body = String::from_utf8(bytes).map_err(Into::<BoxError>::into)?;
91//!
92//! assert_eq!(body, "Hello, World!");
93//! #
94//! # Ok(())
95//! # }
96//! ```
97
98mod request;
99
100use rama_core::extensions::Extension;
101
102pub(crate) mod body;
103mod layer;
104mod service;
105
106/// Marker extension inserted into a response when [`Decompression`] unwraps a
107/// compressed response body.
108#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Extension)]
109#[extension(tags(http))]
110pub enum DecompressedFrom {
111    Gzip,
112    Deflate,
113    Brotli,
114    Zstd,
115}
116
117#[doc(inline)]
118pub use self::{
119    body::DecompressionBody,
120    layer::DecompressionLayer,
121    service::{Decompression, DefaultDecompressionMatcher},
122};
123
124#[doc(inline)]
125pub use self::request::layer::RequestDecompressionLayer;
126#[doc(inline)]
127pub use self::request::service::RequestDecompression;
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    use std::convert::Infallible;
134    use std::io::Write;
135
136    use crate::layer::compression::Compression;
137    use crate::{Body, HeaderMap, HeaderName, Request, Response, body::util::BodyExt};
138    use rama_core::Service;
139    use rama_core::error::ErrorContext;
140    use rama_core::extensions::ExtensionsRef;
141    use rama_core::matcher::service::MatcherServicePair;
142    use rama_core::service::service_fn;
143
144    use rama_http_types::{BodyExtractExt, header};
145
146    #[tokio::test]
147    async fn works() {
148        let client = Decompression::new(Compression::new(service_fn(handle)));
149
150        let req = Request::builder()
151            .header("accept-encoding", "gzip")
152            .body(Body::empty())
153            .unwrap();
154        let res = client.serve(req).await.unwrap();
155
156        assert_eq!(
157            res.extensions().get_ref::<DecompressedFrom>(),
158            Some(&DecompressedFrom::Gzip)
159        );
160
161        // read the body, it will be decompressed automatically
162        let body = res.into_body();
163        let collected = body.collect().await.unwrap();
164        let decompressed_data = String::from_utf8(collected.to_bytes().to_vec()).unwrap();
165
166        assert_eq!(
167            decompressed_data,
168            "Hello, World! Hello, World! Hello, World!"
169        );
170    }
171
172    async fn handle(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
173        let mut trailers = HeaderMap::new();
174        trailers.insert(HeaderName::from_static("foo"), "bar".parse().unwrap());
175        let body = Body::from("Hello, World! Hello, World! Hello, World!");
176        Ok(Response::builder().body(body).unwrap())
177    }
178
179    #[tokio::test]
180    async fn decompress_multi_zstd() {
181        let client = Decompression::new(service_fn(handle_multi_zstd));
182
183        let req = Request::builder()
184            .header("accept-encoding", "zstd")
185            .body(Body::empty())
186            .unwrap();
187        let res = client.serve(req).await.unwrap();
188
189        // read the body, it will be decompressed automatically
190        let body = res.into_body();
191        let decompressed_data =
192            String::from_utf8(body.collect().await.unwrap().to_bytes().to_vec()).unwrap();
193
194        assert_eq!(decompressed_data, "Hello, World!");
195    }
196
197    async fn handle_multi_zstd(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
198        let mut buf = Vec::new();
199        let mut enc1 = zstd::Encoder::new(&mut buf, Default::default()).unwrap();
200        enc1.write_all(b"Hello, ").unwrap();
201        enc1.finish().unwrap();
202
203        let mut enc2 = zstd::Encoder::new(&mut buf, Default::default()).unwrap();
204        enc2.write_all(b"World!").unwrap();
205        enc2.finish().unwrap();
206
207        let mut res = Response::new(Body::from(buf));
208        res.headers_mut()
209            .insert("content-encoding", "zstd".parse().unwrap());
210        Ok(res)
211    }
212
213    #[tokio::test]
214    async fn decompress_empty() {
215        let client = Decompression::new(Compression::new(service_fn(handle_empty)));
216
217        let req = Request::builder()
218            .header("accept-encoding", "gzip")
219            .body(Body::empty())
220            .unwrap();
221        let res = client.serve(req).await.unwrap();
222
223        let decompressed_data = res.try_into_string().await.unwrap();
224
225        assert_eq!(decompressed_data, "");
226    }
227
228    async fn handle_empty(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
229        let mut res = Response::new(Body::empty());
230        res.headers_mut()
231            .insert("content-encoding", "gzip".parse().unwrap());
232        Ok(res)
233    }
234
235    #[tokio::test]
236    async fn decompress_empty_with_trailers() {
237        let client = Decompression::new(Compression::new(service_fn(handle_empty_with_trailers)));
238        let req = Request::builder()
239            .header("accept-encoding", "gzip")
240            .body(Body::empty())
241            .unwrap();
242        let res = client.serve(req).await.unwrap();
243        let body = res.into_body();
244        let collected = body.collect().await.unwrap();
245        let trailers = collected.trailers().cloned().unwrap(); // TODO
246        let decompressed_data = String::from_utf8(collected.to_bytes().to_vec()).unwrap();
247        assert_eq!(decompressed_data, "");
248        assert_eq!(trailers["foo"], "bar");
249    }
250
251    async fn handle_empty_with_trailers(_req: Request<Body>) -> Result<Response<Body>, Infallible> {
252        let mut trailers = HeaderMap::new();
253        trailers.insert(HeaderName::from_static("foo"), "bar".parse().unwrap());
254        let body = Body::empty().with_trailer_headers(trailers);
255        Ok(Response::builder()
256            .header("content-encoding", "gzip")
257            .body(body)
258            .unwrap())
259    }
260
261    #[tokio::test]
262    async fn does_not_insert_accept_encoding_when_disabled() {
263        let client = Decompression::new(service_fn(|req: Request<Body>| async move {
264            assert!(!req.headers().contains_key(header::ACCEPT_ENCODING));
265            Ok::<_, Infallible>(Response::new(Body::empty()))
266        }))
267        .with_insert_accept_encoding_header(false);
268
269        let req = Request::new(Body::empty());
270        _ = client.serve(req).await.unwrap();
271    }
272
273    // A long, mildly-varied plaintext so the brotli stream spans multiple chunks
274    // (guarantees the decoder yields some output before the truncation errors).
275    fn long_plaintext() -> String {
276        (0..4000)
277            .map(|i| format!("line {i}: the quick brown fox\n"))
278            .collect()
279    }
280
281    // Corrupt a run of bytes in the middle of a valid brotli stream: the decoder
282    // produces output for the valid prefix, then hits an invalid block and errors
283    // with InvalidData (the real "brotli error" case — a clean tail-truncation
284    // instead hits brotli's UnexpectedEof clean-EOF path and would NOT error).
285    fn corrupt_brotli(plaintext: &str) -> Vec<u8> {
286        let mut full = Vec::new();
287        {
288            let mut w = brotli::CompressorWriter::new(&mut full, 4096, 5, 22);
289            w.write_all(plaintext.as_bytes()).unwrap();
290        } // drop finishes the stream
291        let n = full.len();
292        let start = n / 3;
293        let end = (start + n / 4).min(n);
294        for b in &mut full[start..end] {
295            *b ^= 0xFF;
296        }
297        full
298    }
299
300    async fn collect_truncated_br(
301        tolerate: bool,
302    ) -> Result<rama_core::bytes::Bytes, rama_core::error::BoxError> {
303        let body = corrupt_brotli(&long_plaintext());
304        let client = Decompression::new(service_fn(move |_req: Request<Body>| {
305            let body = body.clone();
306            async move {
307                Ok::<_, Infallible>(
308                    Response::builder()
309                        .header(header::CONTENT_ENCODING, "br")
310                        .body(Body::from(body))
311                        .unwrap(),
312                )
313            }
314        }))
315        .with_tolerate_decode_errors(tolerate);
316        let req = Request::builder()
317            .header("accept-encoding", "br")
318            .body(Body::empty())
319            .unwrap();
320        let res = client.serve(req).await.unwrap();
321        res.into_body()
322            .collect()
323            .await
324            .map(|c| c.to_bytes())
325            .into_box_error()
326    }
327
328    #[tokio::test]
329    async fn truncated_brotli_with_tolerance_ends_cleanly() {
330        let collected = collect_truncated_br(true).await;
331        assert!(
332            collected.is_ok(),
333            "tolerant decode must end the stream cleanly, not error"
334        );
335        let bytes = collected.unwrap();
336        // The salvaged output must be a clean prefix of the original (no garbage
337        // from the corrupt block); it may be empty if the decoder had not flushed
338        // output before the corruption — the load-bearing property is that the
339        // stream ENDS rather than ABORTS.
340        assert!(
341            long_plaintext().as_bytes().starts_with(&bytes),
342            "decoded bytes must be a clean prefix of the original plaintext (len={})",
343            bytes.len()
344        );
345    }
346
347    #[tokio::test]
348    async fn truncated_brotli_without_tolerance_still_errors() {
349        // Regression guard: the default (strict) behavior must be unchanged.
350        assert!(
351            collect_truncated_br(false).await.is_err(),
352            "strict decode must surface the truncation as an error"
353        );
354    }
355
356    #[tokio::test]
357    async fn response_matcher_can_disable_decompression() {
358        let client = Decompression::new(Compression::new(service_fn(handle)))
359            .with_matcher(MatcherServicePair(true, MatcherServicePair(false, ())));
360
361        let req = Request::builder()
362            .header("accept-encoding", "gzip")
363            .body(Body::empty())
364            .unwrap();
365        let res = client.serve(req).await.unwrap();
366
367        assert!(res.extensions().get_ref::<DecompressedFrom>().is_none());
368        assert_eq!(res.headers().get(header::CONTENT_ENCODING).unwrap(), "gzip");
369
370        let compressed = res.into_body().collect().await.unwrap().to_bytes();
371        assert_ne!(
372            std::str::from_utf8(&compressed).unwrap_or_default(),
373            "Hello, World! Hello, World! Hello, World!"
374        );
375    }
376}
rama_http/layer/decompression/mod.rs

rama_http/layer/decompression/
mod.rs
