Skip to main content

rama_http/layer/decompression/request/
service.rs

1use crate::StreamingBody;
2use crate::headers::encoding::{AcceptEncoding, SupportedEncodings};
3use crate::layer::{
4    decompression::DecompressionBody,
5    decompression::body::BodyInner,
6    util::compression::{CompressionLevel, WrapBody},
7};
8use crate::{HeaderValue, Request, Response, StatusCode, header};
9use rama_core::Service;
10use rama_core::bytes::Bytes;
11use rama_core::error::{BoxError, ErrorContext as _};
12use rama_http_types::Body;
13use rama_utils::macros::define_inner_service_accessors;
14
15/// Decompresses request bodies and calls its underlying service.
16///
17/// Transparently decompresses request bodies based on the `Content-Encoding` header.
18/// When the encoding in the `Content-Encoding` header is not accepted an `Unsupported Media Type`
19/// status code will be returned with the accepted encodings in the `Accept-Encoding` header.
20///
21/// Enabling pass-through of unaccepted encodings will not return an `Unsupported Media Type` but
22/// will call the underlying service with the unmodified request if the encoding is not supported.
23/// This is disabled by default.
24///
25/// See the [module docs](crate::layer::decompression) for more details.
26#[derive(Debug, Clone)]
27pub struct RequestDecompression<S> {
28    pub(super) inner: S,
29    pub(super) accept: AcceptEncoding,
30    pub(super) pass_through_unaccepted: bool,
31}
32
33impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RequestDecompression<S>
34where
35    S: Service<
36            Request<DecompressionBody<ReqBody>>,
37            Output = Response<ResBody>,
38            Error: Into<BoxError>,
39        >,
40    ReqBody: StreamingBody + Send + 'static,
41    ResBody: StreamingBody<Data = Bytes, Error: Into<BoxError>> + Send + Sync + 'static,
42{
43    type Output = Response;
44    type Error = BoxError;
45
46    async fn serve(&self, req: Request<ReqBody>) -> Result<Self::Output, Self::Error> {
47        let (mut parts, body) = req.into_parts();
48
49        let body =
50            if let header::Entry::Occupied(entry) = parts.headers.entry(header::CONTENT_ENCODING) {
51                match entry.get().as_bytes() {
52                    b"gzip" if self.accept.gzip() => {
53                        entry.remove();
54                        parts.headers.remove(header::CONTENT_LENGTH);
55                        BodyInner::gzip(WrapBody::new(body, CompressionLevel::default()))
56                    }
57                    b"deflate" if self.accept.deflate() => {
58                        entry.remove();
59                        parts.headers.remove(header::CONTENT_LENGTH);
60                        BodyInner::deflate(WrapBody::new(body, CompressionLevel::default()))
61                    }
62                    b"br" if self.accept.br() => {
63                        entry.remove();
64                        parts.headers.remove(header::CONTENT_LENGTH);
65                        BodyInner::brotli(WrapBody::new(body, CompressionLevel::default()))
66                    }
67                    b"zstd" if self.accept.zstd() => {
68                        entry.remove();
69                        parts.headers.remove(header::CONTENT_LENGTH);
70                        BodyInner::zstd(WrapBody::new(body, CompressionLevel::default()))
71                    }
72                    b"identity" => BodyInner::identity(body),
73                    _ if self.pass_through_unaccepted => BodyInner::identity(body),
74                    _ => return unsupported_encoding(self.accept),
75                }
76            } else {
77                BodyInner::identity(body)
78            };
79        let body = DecompressionBody::new(body);
80        let req = Request::from_parts(parts, body);
81        self.inner
82            .serve(req)
83            .await
84            .map(|res| res.map(Body::new))
85            .into_box_error()
86    }
87}
88
89fn unsupported_encoding(accept: AcceptEncoding) -> Result<Response, BoxError> {
90    let res = Response::builder()
91        .header(
92            header::ACCEPT_ENCODING,
93            accept
94                .maybe_to_header_value()
95                .unwrap_or(HeaderValue::from_static("identity")),
96        )
97        .status(StatusCode::UNSUPPORTED_MEDIA_TYPE)
98        .body(Body::empty())?;
99    Ok(res)
100}
101
102impl<S> RequestDecompression<S> {
103    /// Creates a new `RequestDecompression` wrapping the `service`.
104    pub fn new(service: S) -> Self {
105        Self {
106            inner: service,
107            accept: AcceptEncoding::default(),
108            pass_through_unaccepted: false,
109        }
110    }
111
112    define_inner_service_accessors!();
113
114    rama_utils::macros::generate_set_and_with! {
115        /// Passes through the request even when the encoding is not supported.
116        ///
117        /// By default pass-through is disabled.
118        pub fn pass_through_unaccepted(mut self, enabled: bool) -> Self {
119            self.pass_through_unaccepted = enabled;
120            self
121        }
122    }
123
124    rama_utils::macros::generate_set_and_with! {
125        /// Sets whether to support gzip encoding.
126        pub fn gzip(mut self, enable: bool) -> Self {
127            self.accept.set_gzip(enable);
128            self
129        }
130    }
131
132    rama_utils::macros::generate_set_and_with! {
133        /// Sets whether to support Deflate encoding.
134        pub fn deflate(mut self, enable: bool) -> Self {
135            self.accept.set_deflate(enable);
136            self
137        }
138    }
139
140    rama_utils::macros::generate_set_and_with! {
141        /// Sets whether to support Brotli encoding.
142        pub fn br(mut self, enable: bool) -> Self {
143            self.accept.set_br(enable);
144            self
145        }
146    }
147
148    rama_utils::macros::generate_set_and_with! {
149        /// Sets whether to support Zstd encoding.
150        pub fn zstd(mut self, enable: bool) -> Self {
151            self.accept.set_zstd(enable);
152            self
153        }
154    }
155}