rama_http/layer/compression/stream/
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::CompressionLevel;
7use super::body::StreamCompressionBody;
8use crate::HeaderValue;
9use crate::StreamingBody;
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::compression::predicate::{DefaultStreamPredicate, Predicate, PreferredEncoding};
15use crate::layer::remove_header::remove_payload_metadata_headers;
16use crate::{Request, Response, StatusCode, header};
17use rama_core::Service;
18use rama_core::extensions::ExtensionsRef;
19use rama_core::telemetry::tracing;
20use rama_http_headers::HeaderMapExt;
21use rama_http_headers::TransferEncoding;
22use rama_http_headers::specifier::{Quality, QualityValue};
23use rama_http_types::Method;
24use rama_utils::collections::smallvec::SmallVec;
25use rama_utils::macros::define_inner_service_accessors;
26use rama_utils::str::submatch_ignore_ascii_case;
27
28/// Compress response bodies of the underlying service.
29///
30/// This uses the `Accept-Encoding` header to pick an appropriate encoding and adds the
31/// `Content-Encoding` header to responses.
32///
33/// See the [module docs](super::StreamCompression) for more details.
34#[derive(Debug, Clone)]
35pub struct StreamCompression<S, P = DefaultStreamPredicate> {
36    pub(crate) inner: S,
37    pub(crate) accept: AcceptEncoding,
38    pub(crate) predicate: P,
39    pub(crate) quality: CompressionLevel,
40    pub(crate) enforce_not_acceptable: bool,
41}
42
43impl<S> StreamCompression<S, DefaultStreamPredicate> {
44    /// Creates a new `StreamCompression` wrapping the `service`.
45    pub fn new(service: S) -> Self {
46        Self {
47            inner: service,
48            accept: AcceptEncoding::default(),
49            predicate: DefaultStreamPredicate::default(),
50            quality: CompressionLevel::default(),
51            enforce_not_acceptable: true,
52        }
53    }
54}
55
56impl<S, P> StreamCompression<S, P> {
57    define_inner_service_accessors!();
58
59    rama_utils::macros::generate_set_and_with! {
60        /// Sets whether to enable the gzip encoding.
61        pub fn gzip(mut self, enable: bool) -> Self {
62            self.accept.set_gzip(enable);
63            self
64        }
65    }
66
67    rama_utils::macros::generate_set_and_with! {
68        /// Sets whether to enable the Deflate encoding.
69        pub fn deflate(mut self, enable: bool) -> Self {
70            self.accept.set_deflate(enable);
71            self
72        }
73    }
74
75    rama_utils::macros::generate_set_and_with! {
76        /// Sets whether to enable the Brotli encoding.
77        pub fn br(mut self, enable: bool) -> Self {
78            self.accept.set_br(enable);
79            self
80        }
81    }
82
83    rama_utils::macros::generate_set_and_with! {
84        /// Sets whether to enable the Zstd encoding.
85        pub fn zstd(mut self, enable: bool) -> Self {
86            self.accept.set_zstd(enable);
87            self
88        }
89    }
90
91    rama_utils::macros::generate_set_and_with! {
92        /// Sets the StreamCompression quality.
93        pub fn quality(mut self, quality: CompressionLevel) -> Self {
94            self.quality = quality;
95            self
96        }
97    }
98
99    rama_utils::macros::generate_set_and_with! {
100        /// Sets whether to respond with `406 Not Acceptable` when the client's
101        /// `Accept-Encoding` header rejects every available representation
102        /// (e.g. `*;q=0` or a lone `identity;q=0`), as recommended by RFC 9110 §12.5.3.
103        ///
104        /// Enabled by default. Disable to opt out and instead fall back to sending an
105        /// uncompressed (identity) response regardless of the client's stated preference.
106        pub fn enforce_not_acceptable(mut self, enable: bool) -> Self {
107            self.enforce_not_acceptable = enable;
108            self
109        }
110    }
111
112    /// Replace the current StreamCompression predicate.
113    ///
114    /// Predicates are used to determine whether a response should be compressed or not.
115    ///
116    /// The default predicate is [`DefaultPredicate`]. See its documentation for more
117    /// details on which responses it won't compress.
118    ///
119    /// [`DefaultPredicate`]: crate::layer::compression::DefaultPredicate
120    #[must_use]
121    pub fn with_compress_predicate<C>(self, predicate: C) -> StreamCompression<S, C>
122    where
123        C: Predicate,
124    {
125        StreamCompression {
126            inner: self.inner,
127            accept: self.accept,
128            predicate,
129            quality: self.quality,
130            enforce_not_acceptable: self.enforce_not_acceptable,
131        }
132    }
133}
134
135impl<ReqBody, ResBody, S, P> Service<Request<ReqBody>> for StreamCompression<S, P>
136where
137    S: Service<Request<ReqBody>, Output = Response<ResBody>>,
138    ResBody: StreamingBody<Data: Send + 'static, Error: Send + 'static> + Send + 'static,
139    P: Predicate + Send + Sync + 'static,
140    ReqBody: Send + 'static,
141{
142    type Output = Response<StreamCompressionBody<ResBody>>;
143    type Error = S::Error;
144
145    #[allow(unreachable_code, unused_mut, unused_variables, unreachable_patterns)]
146    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
147        let accepted_encodings: SmallVec<[QualityValue<Encoding>; 4]> =
148            parse_accept_encoding_headers(req.headers(), self.accept).collect();
149        let wildcard_quality = parse_accept_encoding_wildcard_quality(req.headers());
150        let req_method = req.method().clone();
151
152        let mut res = self.inner.serve(req).await?;
153
154        // RFC 9110 §9.3.2 (HEAD) / §9.3.6 (CONNECT) and the body-prohibiting status codes
155        // (1xx Informational, 204 No Content, 205 Reset Content, 304 Not Modified) mean there
156        // is no representation body to encode (or to reject for the client).
157        let body_allowed = !matches!(req_method, Method::HEAD | Method::CONNECT)
158            && !matches!(res.status().as_u16(), 100..=199 | 204 | 205 | 304);
159
160        let should_compress = body_allowed
161            // never recompress responses that are already compressed
162            && !res.headers().contains_key(header::CONTENT_ENCODING)
163            // never compress responses that are ranges
164            && !res.headers().contains_key(header::CONTENT_RANGE)
165            && self.predicate.should_compress(&mut res);
166
167        let negotiated = negotiate_response_encoding(
168            &accepted_encodings,
169            wildcard_quality,
170            self.accept,
171            res.extensions().get_ref::<PreferredEncoding>().copied(),
172        );
173
174        // RFC 9110 §12.5.3: when the client's `Accept-Encoding` rejects every available
175        // representation (e.g. `*;q=0` or a lone `identity;q=0`), respond 406 Not Acceptable.
176        // This can be opted out of, in which case we fall back to an identity response.
177        let encoding = match negotiated {
178            Some(encoding) => encoding,
179            None if self.enforce_not_acceptable && body_allowed => {
180                let (mut parts, body) = res.into_parts();
181                parts.status = StatusCode::NOT_ACCEPTABLE;
182                ensure_vary_accept_encoding(&mut parts.headers);
183                return Ok(Response::from_parts(
184                    parts,
185                    StreamCompressionBody::identity(body),
186                ));
187            }
188            None => Encoding::Identity,
189        };
190
191        let (mut parts, body) = res.into_parts();
192
193        if should_compress {
194            ensure_vary_accept_encoding(&mut parts.headers);
195        }
196
197        let always_flush = parts
198            .headers
199            .get("x-accel-buffering")
200            .and_then(|v| v.to_str().ok())
201            .is_some_and(|v| v.eq_ignore_ascii_case("no"))
202            || is_streaming_content_type(&parts.headers)
203            || is_chunked_encoding(&parts.headers);
204
205        tracing::trace!(
206            "should_compress = {should_compress}; always_flush = {always_flush}; encoding = {encoding}"
207        );
208
209        let body = match (should_compress, encoding) {
210            // if StreamCompression is _not_ supported or the client doesn't accept it
211            (false, _) | (_, Encoding::Identity) => {
212                return Ok(Response::from_parts(
213                    parts,
214                    StreamCompressionBody::identity(body),
215                ));
216            }
217
218            (true, Encoding::Gzip) => StreamCompressionBody::gzip(body, self.quality, always_flush),
219            (true, Encoding::Deflate) => {
220                StreamCompressionBody::deflate(body, self.quality, always_flush)
221            }
222            (true, Encoding::Brotli) => {
223                StreamCompressionBody::brotli(body, self.quality, always_flush)
224            }
225            (true, Encoding::Zstd) => StreamCompressionBody::zstd(body, self.quality, always_flush),
226        };
227
228        remove_payload_metadata_headers(&mut parts.headers);
229
230        parts
231            .headers
232            .insert(header::CONTENT_ENCODING, HeaderValue::from(encoding));
233
234        let res = Response::from_parts(parts, body);
235        Ok(res)
236    }
237}
238
239/// Picks the response encoding, returning `None` when the client's `Accept-Encoding` rejects
240/// every available representation (406 Not Acceptable per RFC 9110 §12.5.3).
241fn negotiate_response_encoding(
242    accepted_encodings: &[QualityValue<Encoding>],
243    wildcard_quality: Option<Quality>,
244    supported: AcceptEncoding,
245    preferred: Option<PreferredEncoding>,
246) -> Option<Encoding> {
247    if let Some(preferred) = preferred.map(PreferredEncoding::as_encoding)
248        && accepted_encodings
249            .iter()
250            .any(|qval| qval.value == preferred && qval.quality.as_u16() > 0)
251    {
252        return Some(preferred);
253    }
254
255    maybe_preferred_encoding_with_wildcard(accepted_encodings, wildcard_quality, supported)
256}
257
258/// Appends `Vary: accept-encoding` unless an equivalent value is already present.
259fn ensure_vary_accept_encoding(headers: &mut header::HeaderMap) {
260    if !headers.get_all(header::VARY).iter().any(|value| {
261        submatch_ignore_ascii_case(
262            value.as_bytes(),
263            header::ACCEPT_ENCODING.as_str().as_bytes(),
264        )
265    }) {
266        headers.append(header::VARY, header::ACCEPT_ENCODING.into());
267    }
268}
269
270fn is_streaming_content_type(headers: &header::HeaderMap) -> bool {
271    headers
272        .get(header::CONTENT_TYPE)
273        .and_then(|v| v.to_str().ok())
274        .is_some_and(|ct| ct.starts_with("text/event-stream") || ct.starts_with("application/grpc"))
275}
276
277fn is_chunked_encoding(headers: &header::HeaderMap) -> bool {
278    !headers.contains_key(header::CONTENT_LENGTH)
279        || headers
280            .typed_get::<TransferEncoding>()
281            .map(|te| te.is_chunked())
282            .unwrap_or_default()
283}
rama_http/layer/compression/stream/service.rs

rama_http/layer/compression/stream/
service.rs
