Skip to main content

rama_http/layer/compression/
service.rs

1#![expect(
2    clippy::allow_attributes,
3    reason = "macro-generated `#[allow]` attributes whose underlying lints fire only for some expansions"
4)]
5
6use super::CompressionBody;
7use super::CompressionLevel;
8use super::body::BodyInner;
9use super::predicate::{DefaultPredicate, Predicate, PreferredEncoding};
10use crate::headers::encoding::{
11    AcceptEncoding, Encoding, maybe_preferred_encoding_with_wildcard,
12    parse_accept_encoding_headers, parse_accept_encoding_wildcard_quality,
13};
14use crate::layer::remove_header::remove_payload_metadata_headers;
15use crate::layer::util::compression::WrapBody;
16use crate::{Request, Response, StatusCode, header};
17use rama_core::Service;
18use rama_core::extensions::ExtensionsRef;
19use rama_http_headers::specifier::{Quality, QualityValue};
20use rama_http_types::HeaderValue;
21use rama_http_types::Method;
22use rama_http_types::StreamingBody;
23use rama_utils::collections::smallvec::SmallVec;
24use rama_utils::macros::define_inner_service_accessors;
25use rama_utils::str::submatch_ignore_ascii_case;
26
27/// Compress response bodies of the underlying service.
28///
29/// This uses the `Accept-Encoding` header to pick an appropriate encoding and adds the
30/// `Content-Encoding` header to responses.
31///
32/// See the [module docs](crate::layer::compression) for more details.
33#[derive(Debug, Clone)]
34pub struct Compression<S, P = DefaultPredicate> {
35    pub(crate) inner: S,
36    pub(crate) accept: AcceptEncoding,
37    pub(crate) predicate: P,
38    pub(crate) respect_content_encoding_if_possible: bool,
39    pub(crate) quality: CompressionLevel,
40    pub(crate) enforce_not_acceptable: bool,
41}
42
43impl<S> Compression<S, DefaultPredicate> {
44    /// Creates a new `Compression` wrapping the `service`.
45    pub fn new(service: S) -> Self {
46        Self {
47            inner: service,
48            accept: AcceptEncoding::default(),
49            predicate: DefaultPredicate::default(),
50            respect_content_encoding_if_possible: false,
51            quality: CompressionLevel::default(),
52            enforce_not_acceptable: true,
53        }
54    }
55}
56
57impl<S, P> Compression<S, P> {
58    define_inner_service_accessors!();
59
60    rama_utils::macros::generate_set_and_with! {
61        /// Sets whether to enable the gzip encoding.
62        pub fn gzip(mut self, enable: bool) -> Self {
63            self.accept.set_gzip(enable);
64            self
65        }
66    }
67
68    rama_utils::macros::generate_set_and_with! {
69        /// Sets whether to enable the Deflate encoding.
70        pub fn deflate(mut self, enable: bool) -> Self {
71            self.accept.set_deflate(enable);
72            self
73        }
74    }
75
76    rama_utils::macros::generate_set_and_with! {
77        /// Sets whether to enable the Brotli encoding.
78        pub fn br(mut self, enable: bool) -> Self {
79            self.accept.set_br(enable);
80            self
81        }
82    }
83
84    rama_utils::macros::generate_set_and_with! {
85        /// Sets whether to enable the Zstd encoding.
86        pub fn zstd(mut self, enable: bool) -> Self {
87            self.accept.set_zstd(enable);
88            self
89        }
90    }
91
92    rama_utils::macros::generate_set_and_with! {
93        /// Sets the compression quality.
94        pub fn quality(mut self, quality: CompressionLevel) -> Self {
95            self.quality = quality;
96            self
97        }
98    }
99
100    rama_utils::macros::generate_set_and_with! {
101        /// Allow responses with content-encoding.
102        ///
103        /// Useful in case your stack uses that response header as preference.
104        /// Not something you want for regular servers or proxies however,
105        /// or most use cases for that matter.
106        pub fn respect_content_encoding_if_possible(mut self) -> Self {
107            self.respect_content_encoding_if_possible = true;
108            self
109        }
110    }
111
112    rama_utils::macros::generate_set_and_with! {
113        /// Sets whether to respond with `406 Not Acceptable` when the client's
114        /// `Accept-Encoding` header rejects every available representation
115        /// (e.g. `*;q=0` or a lone `identity;q=0`), as recommended by RFC 9110 §12.5.3.
116        ///
117        /// Enabled by default. Disable to opt out and instead fall back to sending an
118        /// uncompressed (identity) response regardless of the client's stated preference.
119        pub fn enforce_not_acceptable(mut self, enable: bool) -> Self {
120            self.enforce_not_acceptable = enable;
121            self
122        }
123    }
124
125    /// Replace the current compression predicate.
126    ///
127    /// Predicates are used to determine whether a response should be compressed or not.
128    ///
129    /// The default predicate is [`DefaultPredicate`]. See its documentation for more
130    /// details on which responses it won't compress.
131    ///
132    /// # Changing the compression predicate
133    ///
134    /// ```
135    /// use rama_utils::str::arcstr::arcstr;
136    /// use rama_http::layer::compression::{
137    ///     Compression,
138    ///     predicate::{Predicate, NotForContentType, DefaultPredicate},
139    /// };
140    /// use rama_core::service::service_fn;
141    ///
142    /// // Placeholder service_fn
143    /// let service = service_fn(async |_: ()| {
144    ///     Ok::<_, std::io::Error>(rama_http::Response::new(()))
145    /// });
146    ///
147    /// // build our custom compression predicate
148    /// // its recommended to still include `DefaultPredicate` as part of
149    /// // custom predicates
150    /// let predicate = DefaultPredicate::new()
151    ///     // don't compress responses who's `content-type` starts with `application/json`
152    ///     .and(NotForContentType::new(arcstr!("application/json")));
153    ///
154    /// let service = Compression::new(service).with_compress_predicate(predicate);
155    /// ```
156    ///
157    /// See [`predicate`](super::predicate) for more utilities for building compression predicates.
158    ///
159    /// Responses that are already compressed (ie have a `content-encoding` header) will _never_ be
160    /// recompressed, regardless what they predicate says.
161    #[must_use]
162    pub fn with_compress_predicate<C>(self, predicate: C) -> Compression<S, C>
163    where
164        C: Predicate,
165    {
166        Compression {
167            inner: self.inner,
168            accept: self.accept,
169            predicate,
170            respect_content_encoding_if_possible: self.respect_content_encoding_if_possible,
171            quality: self.quality,
172            enforce_not_acceptable: self.enforce_not_acceptable,
173        }
174    }
175}
176
177impl<ReqBody, ResBody, S, P> Service<Request<ReqBody>> for Compression<S, P>
178where
179    S: Service<Request<ReqBody>, Output = Response<ResBody>>,
180    ResBody: StreamingBody<Data: Send + 'static, Error: Send + 'static> + Send + 'static,
181    P: Predicate + Send + Sync + 'static,
182    ReqBody: Send + 'static,
183{
184    type Output = Response<CompressionBody<ResBody>>;
185    type Error = S::Error;
186
187    #[allow(unreachable_code, unused_mut, unused_variables, unreachable_patterns)]
188    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
189        let accepted_encodings: SmallVec<[QualityValue<Encoding>; 4]> =
190            parse_accept_encoding_headers(req.headers(), self.accept).collect();
191        let wildcard_quality = parse_accept_encoding_wildcard_quality(req.headers());
192        let req_method = req.method().clone();
193
194        let mut res = self.inner.serve(req).await?;
195        let mut respected_encoding = None;
196
197        // RFC 9110 §9.3.2 (HEAD) / §9.3.6 (CONNECT) and the body-prohibiting status codes
198        // (1xx Informational, 204 No Content, 205 Reset Content, 304 Not Modified) mean there
199        // is no representation body to encode (or to reject for the client).
200        let body_allowed = !matches!(req_method, Method::HEAD | Method::CONNECT)
201            && !matches!(res.status().as_u16(), 100..=199 | 204 | 205 | 304);
202
203        let should_compress = body_allowed &&
204            //never compress responses that are ranges
205            !res.headers().contains_key(header::CONTENT_RANGE) &&
206            self.predicate.should_compress(&mut res) &&
207            if self.respect_content_encoding_if_possible {
208                respected_encoding = Encoding::maybe_from_content_encoding_header(res.headers(), self.accept);
209                true
210            } else {
211                // unless requested do not recompress responses that are already compressed
212                !res.headers().contains_key(header::CONTENT_ENCODING)
213            };
214
215        let negotiated = negotiate_response_encoding(
216            &accepted_encodings,
217            wildcard_quality,
218            self.accept,
219            respected_encoding,
220            res.extensions().get_ref::<PreferredEncoding>().copied(),
221        );
222
223        // RFC 9110 §12.5.3: when the client's `Accept-Encoding` rejects every available
224        // representation (e.g. `*;q=0` or a lone `identity;q=0`), respond 406 Not Acceptable.
225        // This can be opted out of, in which case we fall back to an identity response.
226        let selected_encoding = match negotiated {
227            Some(encoding) => encoding,
228            None if self.enforce_not_acceptable && body_allowed => {
229                let (mut parts, body) = res.into_parts();
230                parts.status = StatusCode::NOT_ACCEPTABLE;
231                ensure_vary_accept_encoding(&mut parts.headers);
232                return Ok(Response::from_parts(
233                    parts,
234                    CompressionBody::new(BodyInner::identity(body)),
235                ));
236            }
237            None => Encoding::Identity,
238        };
239
240        let (mut parts, body) = res.into_parts();
241
242        if should_compress {
243            ensure_vary_accept_encoding(&mut parts.headers);
244        }
245
246        let body = match (should_compress, selected_encoding) {
247            // if compression is _not_ supported or the client doesn't accept it
248            (false, _) | (_, Encoding::Identity) => {
249                return Ok(Response::from_parts(
250                    parts,
251                    CompressionBody::new(BodyInner::identity(body)),
252                ));
253            }
254
255            (_, Encoding::Gzip) => {
256                CompressionBody::new(BodyInner::gzip(WrapBody::new(body, self.quality)))
257            }
258            (_, Encoding::Deflate) => {
259                CompressionBody::new(BodyInner::deflate(WrapBody::new(body, self.quality)))
260            }
261            (_, Encoding::Brotli) => {
262                CompressionBody::new(BodyInner::brotli(WrapBody::new(body, self.quality)))
263            }
264            (_, Encoding::Zstd) => {
265                CompressionBody::new(BodyInner::zstd(WrapBody::new(body, self.quality)))
266            }
267            #[allow(unreachable_patterns)]
268            (true, _) => {
269                // This should never happen because the `AcceptEncoding` struct which is used to determine
270                // `self.encoding` will only enable the different compression algorithms if the
271                // corresponding crate feature has been enabled. This means
272                // Encoding::[Gzip|Brotli|Deflate] should be impossible at this point without the
273                // features enabled.
274                //
275                // The match arm is still required though because the `fs` feature uses the
276                // Encoding struct independently and requires no compression logic to be enabled.
277                // This means a combination of an individual compression feature and `fs` will fail
278                // to compile without this branch even though it will never be reached.
279                //
280                // To safeguard against refactors that changes this relationship or other bugs the
281                // server will return an uncompressed response instead of panicking since that could
282                // become a ddos attack vector.
283                return Ok(Response::from_parts(
284                    parts,
285                    CompressionBody::new(BodyInner::identity(body)),
286                ));
287            }
288        };
289
290        remove_payload_metadata_headers(&mut parts.headers);
291
292        parts.headers.insert(
293            header::CONTENT_ENCODING,
294            HeaderValue::from(selected_encoding),
295        );
296
297        let res = Response::from_parts(parts, body);
298        Ok(res)
299    }
300}
301
302/// Picks the response encoding, returning `None` when the client's `Accept-Encoding` rejects
303/// every available representation (406 Not Acceptable per RFC 9110 §12.5.3).
304fn negotiate_response_encoding(
305    accepted_encodings: &[QualityValue<Encoding>],
306    wildcard_quality: Option<Quality>,
307    supported: AcceptEncoding,
308    respected: Option<Encoding>,
309    preferred: Option<PreferredEncoding>,
310) -> Option<Encoding> {
311    if let Some(respected) = respected
312        && accepted_encodings
313            .iter()
314            .any(|qval| qval.value == respected && qval.quality.as_u16() > 0)
315    {
316        return Some(respected);
317    }
318
319    if let Some(preferred) = preferred.map(PreferredEncoding::as_encoding)
320        && accepted_encodings
321            .iter()
322            .any(|qval| qval.value == preferred && qval.quality.as_u16() > 0)
323    {
324        return Some(preferred);
325    }
326
327    maybe_preferred_encoding_with_wildcard(accepted_encodings, wildcard_quality, supported)
328}
329
330/// Appends `Vary: accept-encoding` unless an equivalent value is already present.
331fn ensure_vary_accept_encoding(headers: &mut rama_http_types::HeaderMap) {
332    if !headers.get_all(header::VARY).iter().any(|value| {
333        submatch_ignore_ascii_case(
334            value.as_bytes(),
335            header::ACCEPT_ENCODING.as_str().as_bytes(),
336        )
337    }) {
338        headers.append(header::VARY, header::ACCEPT_ENCODING.into());
339    }
340}