Skip to main content

rama_http/layer/remove_header/
request.rs

1//! Remove headers from a request.
2//!
3//! # Example
4//!
5//! ```
6//! use rama_http::layer::remove_header::RemoveRequestHeaderLayer;
7//! use rama_http::{Body, Request, Response, header::{self, HeaderValue}};
8//! use rama_core::service::service_fn;
9//! use rama_core::{Service, Layer};
10//! use rama_core::error::BoxError;
11//!
12//! # #[tokio::main]
13//! # async fn main() -> Result<(), BoxError> {
14//! # let http_client = service_fn(async |_: Request| {
15//! #     Ok::<_, std::convert::Infallible>(Response::new(Body::empty()))
16//! # });
17//! #
18//! let mut svc = (
19//!     // Layer that removes all request headers with the prefix `x-foo`.`ac
20//!     RemoveRequestHeaderLayer::prefix("x-foo"),
21//! ).into_layer(http_client);
22//!
23//! let request = Request::new(Body::empty());
24//!
25//! let response = svc.serve(request).await?;
26//! #
27//! # Ok(())
28//! # }
29//! ```
30
31use crate::{HeaderName, Request, Response};
32use rama_core::{Layer, Service};
33use rama_utils::macros::define_inner_service_accessors;
34use rama_utils::str::smol_str::SmolStr;
35
36#[derive(Debug, Clone)]
37/// Layer that applies [`RemoveRequestHeader`] which removes request headers.
38///
39/// See [`RemoveRequestHeader`] for more details.
40pub struct RemoveRequestHeaderLayer {
41    mode: RemoveRequestHeaderMode,
42}
43
44#[derive(Debug, Clone)]
45enum RemoveRequestHeaderMode {
46    Prefix(SmolStr),
47    Exact(HeaderName),
48    Hop,
49    ForwardingMetadata,
50    Sensitive,
51}
52
53impl RemoveRequestHeaderLayer {
54    /// Create a new [`RemoveRequestHeaderLayer`].
55    ///
56    /// Removes request headers by prefix.
57    pub fn prefix(prefix: impl Into<SmolStr>) -> Self {
58        Self {
59            mode: RemoveRequestHeaderMode::Prefix(prefix.into()),
60        }
61    }
62
63    /// Create a new [`RemoveRequestHeaderLayer`].
64    ///
65    /// Removes the request header with the exact name.
66    pub fn exact(header: HeaderName) -> Self {
67        Self {
68            mode: RemoveRequestHeaderMode::Exact(header),
69        }
70    }
71
72    /// Create a new [`RemoveRequestHeaderLayer`].
73    ///
74    /// Removes all hop-by-hop request headers as specified in [RFC 9110](https://github.com/plabayo/rama/blob/main/rama-http-core/specifications/rfc9110.txt#section-7.6.1).
75    #[must_use]
76    pub fn hop_by_hop() -> Self {
77        Self {
78            mode: RemoveRequestHeaderMode::Hop,
79        }
80    }
81
82    /// Create a new [`RemoveRequestHeaderLayer`].
83    ///
84    /// Removes headers that describe the forwarding chain or client address using
85    /// [`remove_forwarding_metadata_request_headers`](super::remove_forwarding_metadata_request_headers).
86    #[must_use]
87    pub fn forwarding_metadata() -> Self {
88        Self {
89            mode: RemoveRequestHeaderMode::ForwardingMetadata,
90        }
91    }
92
93    /// Create a new [`RemoveRequestHeaderLayer`].
94    ///
95    /// Removes all sensitive headers.
96    #[must_use]
97    pub fn sensitive() -> Self {
98        Self {
99            mode: RemoveRequestHeaderMode::Sensitive,
100        }
101    }
102}
103
104impl<S> Layer<S> for RemoveRequestHeaderLayer {
105    type Service = RemoveRequestHeader<S>;
106
107    fn layer(&self, inner: S) -> Self::Service {
108        RemoveRequestHeader {
109            inner,
110            mode: self.mode.clone(),
111        }
112    }
113
114    fn into_layer(self, inner: S) -> Self::Service {
115        RemoveRequestHeader {
116            inner,
117            mode: self.mode,
118        }
119    }
120}
121
122/// Middleware that removes headers from a request.
123#[derive(Debug, Clone)]
124pub struct RemoveRequestHeader<S> {
125    inner: S,
126    mode: RemoveRequestHeaderMode,
127}
128
129impl<S> RemoveRequestHeader<S> {
130    /// Create a new [`RemoveRequestHeader`].
131    ///
132    /// Removes headers by prefix.
133    pub fn prefix(prefix: impl Into<SmolStr>, inner: S) -> Self {
134        RemoveRequestHeaderLayer::prefix(prefix.into()).into_layer(inner)
135    }
136
137    /// Create a new [`RemoveRequestHeader`].
138    ///
139    /// Removes the header with the exact name.
140    pub fn exact(header: HeaderName, inner: S) -> Self {
141        RemoveRequestHeaderLayer::exact(header).into_layer(inner)
142    }
143
144    /// Create a new [`RemoveRequestHeader`].
145    ///
146    /// Removes all hop-by-hop request headers as specified in [RFC 9110](https://github.com/plabayo/rama/blob/main/rama-http-core/specifications/rfc9110.txt#section-7.6.1).
147    pub fn hop_by_hop(inner: S) -> Self {
148        RemoveRequestHeaderLayer::hop_by_hop().into_layer(inner)
149    }
150
151    /// Create a new [`RemoveRequestHeader`].
152    ///
153    /// Removes headers that describe the forwarding chain or client address using
154    /// [`remove_forwarding_metadata_request_headers`](super::remove_forwarding_metadata_request_headers).
155    pub fn forwarding_metadata(inner: S) -> Self {
156        RemoveRequestHeaderLayer::forwarding_metadata().into_layer(inner)
157    }
158
159    /// Create a new [`RemoveRequestHeader`].
160    ///
161    /// Removes all sensitive headers.
162    pub fn sensitive(inner: S) -> Self {
163        RemoveRequestHeaderLayer::sensitive().into_layer(inner)
164    }
165
166    define_inner_service_accessors!();
167}
168
169impl<ReqBody, ResBody, S> Service<Request<ReqBody>> for RemoveRequestHeader<S>
170where
171    ReqBody: Send + 'static,
172    ResBody: Send + 'static,
173    S: Service<Request<ReqBody>, Output = Response<ResBody>>,
174{
175    type Output = S::Output;
176    type Error = S::Error;
177
178    fn serve(
179        &self,
180        mut req: Request<ReqBody>,
181    ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + '_ {
182        match &self.mode {
183            RemoveRequestHeaderMode::Hop => {
184                super::remove_hop_by_hop_request_headers(req.headers_mut())
185            }
186            RemoveRequestHeaderMode::ForwardingMetadata => {
187                super::remove_forwarding_metadata_request_headers(req.headers_mut())
188            }
189            RemoveRequestHeaderMode::Sensitive => {
190                super::remove_sensitive_request_headers(req.headers_mut())
191            }
192            RemoveRequestHeaderMode::Prefix(prefix) => {
193                super::remove_headers_by_prefix(req.headers_mut(), prefix)
194            }
195            RemoveRequestHeaderMode::Exact(header) => {
196                super::remove_headers_by_exact_name(req.headers_mut(), header)
197            }
198        }
199        self.inner.serve(req)
200    }
201}
202
203#[cfg(test)]
204mod test {
205    use super::*;
206    use crate::{Body, Response};
207    use rama_core::{Layer, Service, service::service_fn};
208    use std::convert::Infallible;
209
210    #[tokio::test]
211    async fn remove_request_header_prefix() {
212        let svc = RemoveRequestHeaderLayer::prefix("x-foo").into_layer(service_fn(
213            async |req: Request| {
214                assert!(req.headers().get("x-foo-bar").is_none());
215                assert_eq!(
216                    req.headers().get("foo").map(|v| v.to_str().unwrap()),
217                    Some("bar")
218                );
219                Ok::<_, Infallible>(Response::new(Body::empty()))
220            },
221        ));
222        let req = Request::builder()
223            .header("x-foo-bar", "baz")
224            .header("foo", "bar")
225            .body(Body::empty())
226            .unwrap();
227        _ = svc.serve(req).await.unwrap();
228    }
229
230    #[tokio::test]
231    async fn remove_request_header_exact() {
232        let svc = RemoveRequestHeaderLayer::exact(HeaderName::from_static("x-foo")).into_layer(
233            service_fn(async |req: Request| {
234                assert!(req.headers().get("x-foo").is_none());
235                assert_eq!(
236                    req.headers().get("x-foo-bar").map(|v| v.to_str().unwrap()),
237                    Some("baz")
238                );
239                Ok::<_, Infallible>(Response::new(Body::empty()))
240            }),
241        );
242        let req = Request::builder()
243            .header("x-foo", "baz")
244            .header("x-foo-bar", "baz")
245            .body(Body::empty())
246            .unwrap();
247        _ = svc.serve(req).await.unwrap();
248    }
249
250    #[tokio::test]
251    async fn remove_request_header_hop_by_hop() {
252        let svc =
253            RemoveRequestHeaderLayer::hop_by_hop().into_layer(service_fn(async |req: Request| {
254                assert!(req.headers().get("connection").is_none());
255                assert_eq!(
256                    req.headers().get("foo").map(|v| v.to_str().unwrap()),
257                    Some("bar")
258                );
259                Ok::<_, Infallible>(Response::new(Body::empty()))
260            }));
261        let req = Request::builder()
262            .header("connection", "close")
263            .header("foo", "bar")
264            .body(Body::empty())
265            .unwrap();
266        _ = svc.serve(req).await.unwrap();
267    }
268
269    #[tokio::test]
270    async fn remove_request_header_hop_by_hop_with_connection_list() {
271        let svc =
272            RemoveRequestHeaderLayer::hop_by_hop().into_layer(service_fn(async |req: Request| {
273                assert!(req.headers().get("connection").is_none());
274                assert!(req.headers().get("x-foo").is_none());
275                assert!(req.headers().get("x-bar").is_none());
276                assert!(req.headers().get("connection").is_none());
277                assert_eq!(req.headers().get("x-real-ip").unwrap(), "1.2.3.4");
278                assert_eq!(
279                    req.headers().get("foo").map(|v| v.to_str().unwrap()),
280                    Some("bar")
281                );
282                Ok::<_, Infallible>(Response::new(Body::empty()))
283            }));
284        let req = Request::builder()
285            .header("connection", "x-foo, x-bar")
286            .header("x-foo", "1")
287            .header("x-real-ip", "1.2.3.4")
288            .header("foo", "bar")
289            .body(Body::empty())
290            .unwrap();
291        _ = svc.serve(req).await.unwrap();
292    }
293
294    #[tokio::test]
295    async fn remove_request_header_forwarding_metadata() {
296        let svc = RemoveRequestHeaderLayer::forwarding_metadata().into_layer(service_fn(
297            async |req: Request| {
298                assert!(req.headers().get("x-forwarded-for").is_none());
299                assert!(req.headers().get("x-real-ip").is_none());
300                assert_eq!(req.headers().get("connection").unwrap(), "close");
301                Ok::<_, Infallible>(Response::new(Body::empty()))
302            },
303        ));
304        let req = Request::builder()
305            .header("x-forwarded-for", "1.2.3.4")
306            .header("x-real-ip", "1.2.3.4")
307            .header("connection", "close")
308            .body(Body::empty())
309            .unwrap();
310        _ = svc.serve(req).await.unwrap();
311    }
312}