rama_http/layer/forwarded/
set_forwarded_multi.rs

1#![expect(
2    clippy::allow_attributes,
3    reason = "macro-generated `#[allow]` attributes whose underlying lints fire only for some expansions"
4)]
5
6use crate::Request;
7use crate::headers::HeaderMapExt;
8use crate::headers::forwarded::ForwardHeader;
9use rama_core::error::{BoxError, BoxErrorExt as _, ErrorContext as _};
10use rama_core::{Layer, Service, extensions::ExtensionsRef};
11use rama_net::address::Domain;
12use rama_net::forwarded::{Forwarded, ForwardedElement, NodeId};
13use rama_net::stream::SocketInfo;
14use rama_net::{AuthorityInputExt, Protocol, ProtocolInputExt};
15use rama_utils::macros::all_the_tuples_no_last_special_case;
16use std::fmt;
17use std::marker::PhantomData;
18
19/// Layer to write [`Forwarded`] information for this proxy,
20/// added to the end of the chain of forwarded information already known.
21///
22/// Existing wire headers are not parsed by this layer. Use the matching
23/// [`GetForwardedHeadersLayer`](super::GetForwardedHeadersLayer) earlier in the
24/// request path when inbound chains should be preserved and extended.
25///
26/// Use [`super::SetForwardedHeaderLayer`] if you only need a single a header.
27///
28/// This layer can set any headers as long as you have a [`ForwardHeader`] implementation
29/// for the headers you want to set. You can pass it as the type to the layer when creating
30/// the layer using [`SetForwardedHeadersLayer::new`], with the headers in a single tuple.
31pub struct SetForwardedHeadersLayer<T = Forwarded> {
32    by_node: NodeId,
33    _headers: PhantomData<fn() -> T>,
34}
35
36impl<T: fmt::Debug> fmt::Debug for SetForwardedHeadersLayer<T> {
37    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
38        f.debug_struct("SetForwardedHeadersLayer")
39            .field("by_node", &self.by_node)
40            .field(
41                "_headers",
42                &format_args!("{}", std::any::type_name::<fn() -> T>()),
43            )
44            .finish()
45    }
46}
47
48impl<T: Clone> Clone for SetForwardedHeadersLayer<T> {
49    fn clone(&self) -> Self {
50        Self {
51            by_node: self.by_node.clone(),
52            _headers: PhantomData,
53        }
54    }
55}
56
57impl<T> Default for SetForwardedHeadersLayer<T> {
58    #[inline]
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64impl<T> SetForwardedHeadersLayer<T> {
65    /// Create a new `SetForwardedHeadersLayer` for the specified headers `T`.
66    #[must_use]
67    pub fn new() -> Self {
68        Self {
69            by_node: Domain::from_static("rama").into(),
70            _headers: PhantomData,
71        }
72    }
73}
74
75impl<H, S> Layer<S> for SetForwardedHeadersLayer<H> {
76    type Service = SetForwardedHeadersService<S, H>;
77
78    fn layer(&self, inner: S) -> Self::Service {
79        Self::Service {
80            inner,
81            by_node: self.by_node.clone(),
82            _headers: PhantomData,
83        }
84    }
85
86    fn into_layer(self, inner: S) -> Self::Service {
87        Self::Service {
88            inner,
89            by_node: self.by_node,
90            _headers: PhantomData,
91        }
92    }
93}
94
95/// Middleware [`Service`] to write [`Forwarded`] information for this proxy,
96/// added to the end of the chain of forwarded information already known.
97///
98/// See [`SetForwardedHeadersLayer`] for more information.
99pub struct SetForwardedHeadersService<S, T = Forwarded> {
100    inner: S,
101    by_node: NodeId,
102    _headers: PhantomData<fn() -> T>,
103}
104
105impl<S: fmt::Debug, T> fmt::Debug for SetForwardedHeadersService<S, T> {
106    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
107        f.debug_struct("SetForwardedHeadersService")
108            .field("inner", &self.inner)
109            .field("by_node", &self.by_node)
110            .field(
111                "_headers",
112                &format_args!("{}", std::any::type_name::<fn() -> T>()),
113            )
114            .finish()
115    }
116}
117
118impl<S: Clone, T> Clone for SetForwardedHeadersService<S, T> {
119    fn clone(&self) -> Self {
120        Self {
121            inner: self.inner.clone(),
122            by_node: self.by_node.clone(),
123            _headers: PhantomData,
124        }
125    }
126}
127
128impl<S, T> SetForwardedHeadersService<S, T> {
129    /// Create a new `SetForwardedHeadersService` for the specified headers `T`.
130    pub fn new(inner: S) -> Self {
131        Self {
132            inner,
133            by_node: Domain::from_static("rama").into(),
134            _headers: PhantomData,
135        }
136    }
137}
138
139macro_rules! set_forwarded_service_for_tuple {
140    ( $($ty:ident),* $(,)? ) => {
141        #[allow(non_snake_case)]
142        impl<S, $($ty),* , Body> Service<Request<Body>> for SetForwardedHeadersService<S, ($($ty,)*)>
143        where
144            $( $ty: ForwardHeader + Send + Sync + 'static, )*
145            S: Service<Request<Body>, Error: Into<BoxError>>,
146            Body: Send + 'static,
147        {
148            type Output = S::Output;
149            type Error = BoxError;
150
151            async fn serve(
152                &self,
153                mut req: Request<Body>,
154            ) -> Result<Self::Output, Self::Error> {
155                let forwarded: Option<Forwarded> = req.extensions().get_ref().cloned();
156
157                let mut forwarded_element = ForwardedElement::new_forwarded_by(self.by_node.clone());
158
159                if let Some(peer_addr) = req
160                    .extensions()
161                    .ingress()
162                    .and_then(|ext|ext.get_ref::<SocketInfo>())
163                    .map(|socket| socket.peer_addr())
164                {
165
166                    forwarded_element.set_forwarded_for(peer_addr);
167                }
168
169                let authority = req
170                    .authority()
171                    .ok_or_else(|| BoxError::from_static_str("set forwarded: no authority"))?;
172
173                forwarded_element.set_forwarded_host(authority);
174
175                let protocol = req.protocol().unwrap_or(&Protocol::HTTP);
176                if let Ok(forwarded_proto) = protocol.try_into() {
177                    forwarded_element.set_forwarded_proto(forwarded_proto);
178                }
179
180                let forwarded = match forwarded {
181                    None => Some(Forwarded::new(forwarded_element)),
182                    Some(mut forwarded) => {
183                        forwarded.append(forwarded_element);
184                        Some(forwarded)
185                    }
186                };
187
188                if let Some(forwarded) = forwarded {
189                    $(
190                        if let Some(header) = $ty::try_from_forwarded(forwarded.iter()) {
191                            req.headers_mut().typed_insert(header);
192                        }
193                    )*
194                }
195
196                self.inner.serve(req).await.into_box_error()
197            }
198        }
199    };
200}
201all_the_tuples_no_last_special_case!(set_forwarded_service_for_tuple);
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use crate::{
207        Response, StatusCode,
208        headers::forwarded::{TrueClientIp, XClientIp, XRealIp},
209        service::web::response::IntoResponse,
210    };
211    use rama_core::{Layer, error::BoxError, service::service_fn};
212    use rama_http_headers::forwarded::XForwardedProto;
213    use std::convert::Infallible;
214
215    fn assert_is_service<T: Service<Request<()>>>(_: T) {}
216
217    async fn dummy_service_fn() -> Result<Response, BoxError> {
218        Ok(StatusCode::OK.into_response())
219    }
220
221    #[test]
222    fn test_set_forwarded_service_is_service() {
223        assert_is_service(SetForwardedHeadersService::<_, (TrueClientIp,)>::new(
224            service_fn(dummy_service_fn),
225        ));
226        assert_is_service(
227            SetForwardedHeadersService::<_, (TrueClientIp, XClientIp)>::new(service_fn(
228                dummy_service_fn,
229            )),
230        );
231        assert_is_service(
232            SetForwardedHeadersLayer::<(XRealIp, XForwardedProto)>::new()
233                .into_layer(service_fn(dummy_service_fn)),
234        );
235    }
236
237    #[tokio::test]
238    async fn test_set_forwarded_service_forwarded() {
239        async fn svc(request: Request<()>) -> Result<(), Infallible> {
240            assert_eq!(
241                request.headers().get("Forwarded").unwrap(),
242                "by=rama;host=\"example.com:80\";proto=http"
243            );
244            Ok(())
245        }
246
247        let service =
248            SetForwardedHeadersService::<_, (rama_http_headers::forwarded::Forwarded,)>::new(
249                service_fn(svc),
250            );
251        let req = Request::builder()
252            .uri("http://example.com")
253            .body(())
254            .unwrap();
255        service.serve(req).await.unwrap();
256    }
257}
rama_http/layer/forwarded/set_forwarded_multi.rs

rama_http/layer/forwarded/
set_forwarded_multi.rs
