rama_http/layer/compression/
predicate.rs

1//! Predicates for influencing compression of responses.
2
3use rama_core::extensions::{Extension, Extensions, ExtensionsRef};
4use rama_http_types::{HeaderMap, StatusCode, StreamingBody, Version, header};
5use rama_utils::str::arcstr::{ArcStr, arcstr};
6
7use crate::headers::encoding::Encoding;
8use crate::layer::decompression::DecompressedFrom;
9
10/// Predicate used to determine if a response should be compressed or not.
11pub trait Predicate: Clone {
12    /// Should this response be compressed or not?
13    ///
14    /// The response is mutable so predicates can attach extensions used later
15    /// during response-time compression negotiation.
16    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
17    where
18        B: StreamingBody;
19
20    /// Combine two predicates into one.
21    ///
22    /// The resulting predicate enables compression if both inner predicates do.
23    fn and<Other>(self, other: Other) -> And<Self, Other>
24    where
25        Self: Sized,
26        Other: Predicate,
27    {
28        And {
29            lhs: self,
30            rhs: other,
31        }
32    }
33}
34
35impl<F> Predicate for F
36where
37    F: Fn(StatusCode, Version, &HeaderMap, &Extensions) -> bool + Clone,
38{
39    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
40    where
41        B: StreamingBody,
42    {
43        let status = response.status();
44        let version = response.version();
45        let headers = response.headers().clone();
46        let extensions = response.extensions();
47        self(status, version, &headers, extensions)
48    }
49}
50
51/// Predicate to _always_ compress.
52#[derive(Debug, Clone, Default, Copy)]
53#[non_exhaustive]
54pub struct Always;
55
56impl Always {
57    #[must_use]
58    pub fn new() -> Self {
59        Self
60    }
61}
62
63impl Predicate for Always {
64    fn should_compress<B>(&self, _response: &mut rama_http_types::Response<B>) -> bool
65    where
66        B: StreamingBody,
67    {
68        true
69    }
70}
71
72impl<T> Predicate for Option<T>
73where
74    T: Predicate,
75{
76    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
77    where
78        B: StreamingBody,
79    {
80        self.as_ref()
81            .map(|inner| inner.should_compress(response))
82            .unwrap_or(true)
83    }
84}
85
86/// Two predicates combined into one.
87///
88/// Created with [`Predicate::and`]
89#[derive(Debug, Clone, Default, Copy)]
90pub struct And<Lhs, Rhs> {
91    lhs: Lhs,
92    rhs: Rhs,
93}
94
95impl<Lhs, Rhs> Predicate for And<Lhs, Rhs>
96where
97    Lhs: Predicate,
98    Rhs: Predicate,
99{
100    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
101    where
102        B: StreamingBody,
103    {
104        self.lhs.should_compress(response) && self.rhs.should_compress(response)
105    }
106}
107
108/// The default predicate used by [`Compression`] and [`CompressionLayer`].
109///
110/// This will compress responses unless:
111///
112/// - They're gRPC, which has its own protocol specific compression scheme.
113/// - It's an image as determined by the `content-type` starting with `image/`.
114/// - They're Server-Sent Events (SSE) as determined by the `content-type` being `text/event-stream`.
115/// - The response is less than 32 bytes.
116///
117/// # Configuring the defaults
118///
119/// `DefaultPredicate` doesn't support any configuration. Instead you can build your own predicate
120/// by combining types in this module:
121///
122/// ```rust
123/// use rama_utils::str::arcstr::arcstr;
124/// use rama_http::layer::compression::predicate::{SizeAbove, NotForContentType, Predicate};
125///
126/// // slightly large min size than the default 32
127/// let predicate = SizeAbove::new(256)
128///     // still don't compress gRPC
129///     .and(NotForContentType::GRPC)
130///     // still don't compress images
131///     .and(NotForContentType::IMAGES)
132///     // also don't compress JSON
133///     .and(NotForContentType::new(arcstr!("application/json")));
134/// ```
135///
136/// [`Compression`]: super::Compression
137/// [`CompressionLayer`]: super::CompressionLayer
138#[derive(Debug, Clone)]
139pub struct DefaultPredicate(
140    And<And<And<SizeAbove, NotForContentType>, NotForContentType>, NotForContentType>,
141);
142
143impl DefaultPredicate {
144    /// Create a new `DefaultPredicate`.
145    #[must_use]
146    pub fn new() -> Self {
147        let inner = SizeAbove::new(SizeAbove::DEFAULT_MIN_SIZE)
148            .and(NotForContentType::GRPC)
149            .and(NotForContentType::IMAGES)
150            .and(NotForContentType::SSE);
151        Self(inner)
152    }
153}
154
155impl Default for DefaultPredicate {
156    fn default() -> Self {
157        Self::new()
158    }
159}
160
161impl Predicate for DefaultPredicate {
162    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
163    where
164        B: StreamingBody,
165    {
166        self.0.should_compress(response)
167    }
168}
169
170#[derive(Debug, Clone)]
171pub struct DefaultStreamPredicate(And<SizeAbove, NotForContentType>);
172
173impl DefaultStreamPredicate {
174    /// Create a new `DefaultStreamPredicate`.
175    #[must_use]
176    pub fn new() -> Self {
177        let inner = SizeAbove::new(SizeAbove::DEFAULT_MIN_SIZE).and(NotForContentType::IMAGES);
178        Self(inner)
179    }
180}
181
182impl Default for DefaultStreamPredicate {
183    fn default() -> Self {
184        Self::new()
185    }
186}
187
188impl Predicate for DefaultStreamPredicate {
189    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
190    where
191        B: StreamingBody,
192    {
193        self.0.should_compress(response)
194    }
195}
196
197/// Preferred response encoding requested by a compression predicate.
198#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Extension)]
199#[extension(tags(http))]
200pub enum PreferredEncoding {
201    #[default]
202    Gzip,
203    Deflate,
204    Brotli,
205    Zstd,
206}
207
208impl PreferredEncoding {
209    #[must_use]
210    pub const fn as_encoding(self) -> Encoding {
211        match self {
212            Self::Gzip => Encoding::Gzip,
213            Self::Deflate => Encoding::Deflate,
214            Self::Brotli => Encoding::Brotli,
215            Self::Zstd => Encoding::Zstd,
216        }
217    }
218}
219
220/// [`Predicate`] that enables compression only for responses previously
221/// decompressed by Rama's [`crate::layer::decompression::DecompressionLayer`],
222/// while preferring the original upstream encoding for recompression.
223#[derive(Debug, Clone, Copy, Default)]
224#[non_exhaustive]
225pub struct MirrorDecompressed;
226
227impl MirrorDecompressed {
228    #[must_use]
229    pub fn new() -> Self {
230        Self
231    }
232}
233
234impl Predicate for MirrorDecompressed {
235    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
236    where
237        B: StreamingBody,
238    {
239        let preferred = match response.extensions().get_ref::<DecompressedFrom>() {
240            Some(DecompressedFrom::Gzip) => PreferredEncoding::Gzip,
241            Some(DecompressedFrom::Deflate) => PreferredEncoding::Deflate,
242            Some(DecompressedFrom::Brotli) => PreferredEncoding::Brotli,
243            Some(DecompressedFrom::Zstd) => PreferredEncoding::Zstd,
244            None => return false,
245        };
246
247        response.extensions().insert(preferred);
248        true
249    }
250}
251
252/// [`Predicate`] that will only allow compression of responses above a certain size.
253#[derive(Clone, Copy, Debug)]
254pub struct SizeAbove(u64);
255
256impl SizeAbove {
257    pub(crate) const DEFAULT_MIN_SIZE: u64 = 32;
258
259    /// Create a new `SizeAbove` predicate that will only compress responses larger than
260    /// `min_size_bytes`.
261    ///
262    /// The response will be compressed if the exact size cannot be determined through either the
263    /// `content-length` header or [`StreamingBody::size_hint`].
264    #[must_use]
265    pub const fn new(min_size_bytes: u64) -> Self {
266        Self(min_size_bytes)
267    }
268}
269
270impl Default for SizeAbove {
271    fn default() -> Self {
272        Self(Self::DEFAULT_MIN_SIZE)
273    }
274}
275
276impl Predicate for SizeAbove {
277    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
278    where
279        B: StreamingBody,
280    {
281        let content_size = response.body().size_hint().exact().or_else(|| {
282            response
283                .headers()
284                .get(header::CONTENT_LENGTH)
285                .and_then(|h| h.to_str().ok())
286                .and_then(|val| val.parse().ok())
287        });
288
289        match content_size {
290            Some(size) => size >= self.0,
291            _ => true,
292        }
293    }
294}
295
296/// Predicate that won't allow responses with a specific `content-type` to be compressed.
297#[derive(Clone, Debug)]
298pub struct NotForContentType {
299    content_type: ArcStr,
300    exception: Option<ArcStr>,
301}
302
303impl NotForContentType {
304    /// Predicate that won't compress gRPC responses.
305    pub const GRPC: Self = Self {
306        content_type: arcstr!("application/grpc"),
307        exception: Some(arcstr!("application/grpc-web")),
308    };
309
310    /// Predicate that won't compress images.
311    pub const IMAGES: Self = Self {
312        content_type: arcstr!("image/"),
313        exception: Some(arcstr!("image/svg+xml")),
314    };
315
316    /// Predicate that won't compress Server-Sent Events (SSE) responses.
317    pub const SSE: Self = Self::new(arcstr!("text/event-stream"));
318
319    /// Create a new `NotForContentType`.
320    #[must_use]
321    pub const fn new(content_type: ArcStr) -> Self {
322        Self {
323            content_type,
324            exception: None,
325        }
326    }
327}
328
329impl Predicate for NotForContentType {
330    fn should_compress<B>(&self, response: &mut rama_http_types::Response<B>) -> bool
331    where
332        B: StreamingBody,
333    {
334        let cty = content_type(response);
335        if let Some(except) = &self.exception
336            && cty.starts_with(except.as_str())
337        {
338            return true;
339        }
340
341        !cty.starts_with(self.content_type.as_str())
342    }
343}
344
345fn content_type<B>(response: &rama_http_types::Response<B>) -> &str {
346    response
347        .headers()
348        .get(header::CONTENT_TYPE)
349        .and_then(|h| h.to_str().ok())
350        .unwrap_or_default()
351}
rama_http/layer/compression/predicate.rs

rama_http/layer/compression/
predicate.rs
