rama_http/layer/remove_header/
request.rs1use 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)]
37pub 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 pub fn prefix(prefix: impl Into<SmolStr>) -> Self {
58 Self {
59 mode: RemoveRequestHeaderMode::Prefix(prefix.into()),
60 }
61 }
62
63 pub fn exact(header: HeaderName) -> Self {
67 Self {
68 mode: RemoveRequestHeaderMode::Exact(header),
69 }
70 }
71
72 #[must_use]
76 pub fn hop_by_hop() -> Self {
77 Self {
78 mode: RemoveRequestHeaderMode::Hop,
79 }
80 }
81
82 #[must_use]
87 pub fn forwarding_metadata() -> Self {
88 Self {
89 mode: RemoveRequestHeaderMode::ForwardingMetadata,
90 }
91 }
92
93 #[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#[derive(Debug, Clone)]
124pub struct RemoveRequestHeader<S> {
125 inner: S,
126 mode: RemoveRequestHeaderMode,
127}
128
129impl<S> RemoveRequestHeader<S> {
130 pub fn prefix(prefix: impl Into<SmolStr>, inner: S) -> Self {
134 RemoveRequestHeaderLayer::prefix(prefix.into()).into_layer(inner)
135 }
136
137 pub fn exact(header: HeaderName, inner: S) -> Self {
141 RemoveRequestHeaderLayer::exact(header).into_layer(inner)
142 }
143
144 pub fn hop_by_hop(inner: S) -> Self {
148 RemoveRequestHeaderLayer::hop_by_hop().into_layer(inner)
149 }
150
151 pub fn forwarding_metadata(inner: S) -> Self {
156 RemoveRequestHeaderLayer::forwarding_metadata().into_layer(inner)
157 }
158
159 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}