salvo_proxy/
lib.rs

1#![cfg_attr(test, allow(clippy::unwrap_used))]
2//! Provide HTTP proxy capabilities for the Salvo web framework.
3//!
4//! This crate allows you to easily forward requests to upstream servers,
5//! supporting both HTTP and HTTPS protocols. It's useful for creating API gateways,
6//! load balancers, and reverse proxies.
7//!
8//! # Example
9//!
10//! In this example, requests to different hosts are proxied to different upstream servers:
11//! - Requests to <http://127.0.0.1:8698/> are proxied to <https://www.rust-lang.org>
12//! - Requests to <http://localhost:8698/> are proxied to <https://crates.io>
13//!
14//! ```no_run
15//! use salvo_core::prelude::*;
16//! use salvo_proxy::Proxy;
17//!
18//! #[tokio::main]
19//! async fn main() {
20//!     let router = Router::new()
21//!         .push(
22//!             Router::new()
23//!                 .host("127.0.0.1")
24//!                 .path("{**rest}")
25//!                 .goal(Proxy::use_hyper_client("https://www.rust-lang.org")),
26//!         )
27//!         .push(
28//!             Router::new()
29//!                 .host("localhost")
30//!                 .path("{**rest}")
31//!                 .goal(Proxy::use_hyper_client("https://crates.io")),
32//!         );
33//!
34//!     let acceptor = TcpListener::new("0.0.0.0:8698").bind().await;
35//!     Server::new(acceptor).serve(router).await;
36//! }
37//! ```
38#![doc(html_favicon_url = "https://salvo.rs/favicon-32x32.png")]
39#![doc(html_logo_url = "https://salvo.rs/images/logo.svg")]
40#![cfg_attr(docsrs, feature(doc_cfg))]
41
42use std::convert::Infallible;
43use std::error::Error as StdError;
44use std::fmt::{self, Debug, Formatter};
45#[cfg(test)]
46use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
47
48use hyper::upgrade::OnUpgrade;
49#[cfg(not(test))]
50use local_ip_address::{local_ip, local_ipv6};
51use percent_encoding::{AsciiSet, CONTROLS, utf8_percent_encode};
52use salvo_core::conn::SocketAddr;
53use salvo_core::http::header::{CONNECTION, HOST, HeaderMap, HeaderName, HeaderValue, UPGRADE};
54use salvo_core::http::uri::Uri;
55use salvo_core::http::{ReqBody, ResBody, StatusCode};
56use salvo_core::routing::normalize_url_path;
57use salvo_core::{BoxedError, Depot, Error, FlowCtrl, Handler, Request, Response, async_trait};
58
59#[macro_use]
60mod cfg;
61
62cfg_feature! {
63    #![feature = "hyper-client"]
64    mod hyper_client;
65    pub use hyper_client::*;
66}
67cfg_feature! {
68    #![feature = "reqwest-client"]
69    mod reqwest_client;
70    pub use reqwest_client::*;
71}
72
73cfg_feature! {
74    #![feature = "unix-sock-client"]
75    #[cfg(unix)]
76    mod unix_sock_client;
77    #[cfg(unix)]
78    pub use unix_sock_client::*;
79}
80
81type HyperRequest = hyper::Request<ReqBody>;
82type HyperResponse = hyper::Response<ResBody>;
83
84const X_FORWARDER_FOR_HEADER_NAME: &str = "x-forwarded-for";
85const HOP_BY_HOP_HEADERS: &[&str] = &[
86    "connection",
87    "keep-alive",
88    "proxy-authenticate",
89    "proxy-authorization",
90    "te",
91    "trailer",
92    "transfer-encoding",
93    "upgrade",
94];
95
96const QUERY_ENCODE_SET: &AsciiSet = &CONTROLS
97    .add(b' ')
98    .add(b'"')
99    .add(b'#')
100    .add(b'<')
101    .add(b'>')
102    .add(b'`');
103const PATH_ENCODE_SET: &AsciiSet = &QUERY_ENCODE_SET
104    .add(b'?')
105    .add(b'^')
106    .add(b'`')
107    .add(b'{')
108    .add(b'}');
109
110/// Encode url path. This can be used when build your custom url path getter.
111#[inline]
112pub(crate) fn encode_url_path(path: &str) -> String {
113    path.split('/')
114        .map(|s| utf8_percent_encode(s, PATH_ENCODE_SET).to_string())
115        .collect::<Vec<_>>()
116        .join("/")
117}
118
119/// Client trait for implementing different HTTP clients for proxying.
120///
121/// Implement this trait to create custom proxy clients with different
122/// backends or configurations.
123pub trait Client: Send + Sync + 'static {
124    /// Error type returned by the client.
125    type Error: StdError + Send + Sync + 'static;
126
127    /// Execute a request through the proxy client.
128    fn execute(
129        &self,
130        req: HyperRequest,
131        upgraded: Option<OnUpgrade>,
132    ) -> impl Future<Output = Result<HyperResponse, Self::Error>> + Send;
133}
134
135/// Upstreams trait for selecting target servers.
136///
137/// Implement this trait to customize how target servers are selected
138/// for proxying requests. This can be used to implement load balancing,
139/// failover, or other server selection strategies.
140pub trait Upstreams: Send + Sync + 'static {
141    /// Error type returned when selecting a server fails.
142    type Error: StdError + Send + Sync + 'static;
143
144    /// Elect a server to handle the current request.
145    fn elect(
146        &self,
147        req: &Request,
148        depot: &Depot,
149    ) -> impl Future<Output = Result<&str, Self::Error>> + Send;
150}
151impl Upstreams for &'static str {
152    type Error = Infallible;
153
154    async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
155        Ok(*self)
156    }
157}
158impl Upstreams for String {
159    type Error = Infallible;
160    async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
161        Ok(self.as_str())
162    }
163}
164
165impl<const N: usize> Upstreams for [&'static str; N] {
166    type Error = Error;
167    async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
168        if self.is_empty() {
169            return Err(Error::other("upstreams is empty"));
170        }
171        let index = fastrand::usize(..self.len());
172        Ok(self[index])
173    }
174}
175
176impl<T> Upstreams for Vec<T>
177where
178    T: AsRef<str> + Send + Sync + 'static,
179{
180    type Error = Error;
181    async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
182        if self.is_empty() {
183            return Err(Error::other("upstreams is empty"));
184        }
185        let index = fastrand::usize(..self.len());
186        Ok(self[index].as_ref())
187    }
188}
189
190/// Url part getter. You can use this to get the proxied url path or query.
191pub type UrlPartGetter = Box<dyn Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static>;
192
193/// Host header getter. You can use this to get the host header for the proxied request.
194pub type HostHeaderGetter =
195    Box<dyn Fn(&Uri, &Request, &Depot) -> Option<String> + Send + Sync + 'static>;
196
197/// Default url path getter.
198///
199/// This getter will get the last param as the rest url path from request.
200/// In most case you should use wildcard param, like `{**rest}`, `{*+rest}`.
201pub fn default_url_path_getter(req: &Request, _depot: &Depot) -> Option<String> {
202    req.params().tail().map(str::to_owned)
203}
204/// Default url query getter. This getter just return the query string from request uri.
205pub fn default_url_query_getter(req: &Request, _depot: &Depot) -> Option<String> {
206    req.uri().query().map(Into::into)
207}
208
209/// Default host header getter. This getter will get the host header from request uri
210pub fn default_host_header_getter(
211    forward_uri: &Uri,
212    _req: &Request,
213    _depot: &Depot,
214) -> Option<String> {
215    if let Some(host) = forward_uri.host() {
216        return Some(String::from(host));
217    }
218
219    None
220}
221
222/// RFC2616 complieant host header getter. This getter will get the host header from request uri,
223/// and add port if it's not default port. Falls back to default upon any forward URI parse error.
224pub fn rfc2616_host_header_getter(
225    forward_uri: &Uri,
226    req: &Request,
227    _depot: &Depot,
228) -> Option<String> {
229    let mut parts: Vec<String> = Vec::with_capacity(2);
230
231    if let Some(host) = forward_uri.host() {
232        parts.push(host.to_owned());
233
234        if let Some(scheme) = forward_uri.scheme_str()
235            && let Some(port) = forward_uri.port_u16()
236            && (scheme == "http" && port != 80 || scheme == "https" && port != 443)
237        {
238            parts.push(port.to_string());
239        }
240    }
241
242    if parts.is_empty() {
243        default_host_header_getter(forward_uri, req, _depot)
244    } else {
245        Some(parts.join(":"))
246    }
247}
248
249/// Preserve original host header getter. Propagates the original request host header to the proxied
250/// request.
251pub fn preserve_original_host_header_getter(
252    forward_uri: &Uri,
253    req: &Request,
254    _depot: &Depot,
255) -> Option<String> {
256    if let Some(host_header) = req.headers().get(HOST)
257        && let Ok(host) = host_header.to_str()
258    {
259        return Some(host.to_owned());
260    }
261
262    default_host_header_getter(forward_uri, req, _depot)
263}
264
265/// Handler that can proxy request to other server.
266#[non_exhaustive]
267pub struct Proxy<U, C>
268where
269    U: Upstreams,
270    C: Client,
271{
272    /// Upstreams list.
273    pub upstreams: U,
274    /// [`Client`] for proxy.
275    pub client: C,
276    /// Url path getter.
277    pub url_path_getter: UrlPartGetter,
278    /// Url query getter.
279    pub url_query_getter: UrlPartGetter,
280    /// Host header getter
281    pub host_header_getter: HostHeaderGetter,
282    /// Flag to enable x-forwarded-for header.
283    pub client_ip_forwarding_enabled: bool,
284}
285
286impl<U, C> Debug for Proxy<U, C>
287where
288    U: Upstreams,
289    C: Client,
290{
291    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
292        f.debug_struct("Proxy").finish()
293    }
294}
295
296impl<U, C> Proxy<U, C>
297where
298    U: Upstreams,
299    U::Error: Into<BoxedError>,
300    C: Client,
301{
302    /// Create new `Proxy` with upstreams list.
303    #[must_use]
304    pub fn new(upstreams: U, client: C) -> Self {
305        Self {
306            upstreams,
307            client,
308            url_path_getter: Box::new(default_url_path_getter),
309            url_query_getter: Box::new(default_url_query_getter),
310            host_header_getter: Box::new(default_host_header_getter),
311            client_ip_forwarding_enabled: false,
312        }
313    }
314
315    /// Create new `Proxy` with upstreams list and enable x-forwarded-for header.
316    pub fn with_client_ip_forwarding(upstreams: U, client: C) -> Self {
317        Self {
318            upstreams,
319            client,
320            url_path_getter: Box::new(default_url_path_getter),
321            url_query_getter: Box::new(default_url_query_getter),
322            host_header_getter: Box::new(default_host_header_getter),
323            client_ip_forwarding_enabled: true,
324        }
325    }
326
327    /// Set url path getter.
328    #[inline]
329    #[must_use]
330    pub fn url_path_getter<G>(mut self, url_path_getter: G) -> Self
331    where
332        G: Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static,
333    {
334        self.url_path_getter = Box::new(url_path_getter);
335        self
336    }
337
338    /// Set url query getter.
339    #[inline]
340    #[must_use]
341    pub fn url_query_getter<G>(mut self, url_query_getter: G) -> Self
342    where
343        G: Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static,
344    {
345        self.url_query_getter = Box::new(url_query_getter);
346        self
347    }
348
349    /// Set host header query getter.
350    #[inline]
351    #[must_use]
352    pub fn host_header_getter<G>(mut self, host_header_getter: G) -> Self
353    where
354        G: Fn(&Uri, &Request, &Depot) -> Option<String> + Send + Sync + 'static,
355    {
356        self.host_header_getter = Box::new(host_header_getter);
357        self
358    }
359
360    /// Get upstreams list.
361    #[inline]
362    pub fn upstreams(&self) -> &U {
363        &self.upstreams
364    }
365    /// Get upstreams mutable list.
366    #[inline]
367    pub fn upstreams_mut(&mut self) -> &mut U {
368        &mut self.upstreams
369    }
370
371    /// Get client reference.
372    #[inline]
373    pub fn client(&self) -> &C {
374        &self.client
375    }
376    /// Get client mutable reference.
377    #[inline]
378    pub fn client_mut(&mut self) -> &mut C {
379        &mut self.client
380    }
381
382    /// Enable x-forwarded-for header prepending.
383    #[inline]
384    #[must_use]
385    pub fn client_ip_forwarding(mut self, enable: bool) -> Self {
386        self.client_ip_forwarding_enabled = enable;
387        self
388    }
389
390    async fn build_proxied_request(
391        &self,
392        req: &mut Request,
393        depot: &Depot,
394    ) -> Result<HyperRequest, Error> {
395        let upstream = self
396            .upstreams
397            .elect(req, depot)
398            .await
399            .map_err(Error::other)?;
400
401        if upstream.is_empty() {
402            tracing::error!("upstreams is empty");
403            return Err(Error::other("upstreams is empty"));
404        }
405
406        let path = (self.url_path_getter)(req, depot).unwrap_or_default();
407        let path = encode_url_path(&normalize_url_path(&path));
408        let query = (self.url_query_getter)(req, depot);
409        let rest = if let Some(query) = query {
410            if let Some(stripped) = query.strip_prefix('?') {
411                format!("{path}?{}", utf8_percent_encode(stripped, QUERY_ENCODE_SET))
412            } else {
413                format!("{path}?{}", utf8_percent_encode(&query, QUERY_ENCODE_SET))
414            }
415        } else {
416            path
417        };
418        let forward_url = if upstream.ends_with('/') && rest.starts_with('/') {
419            format!("{}{}", upstream.trim_end_matches('/'), rest)
420        } else if upstream.ends_with('/') || rest.starts_with('/') {
421            format!("{upstream}{rest}")
422        } else if rest.is_empty() {
423            upstream.to_owned()
424        } else {
425            format!("{upstream}/{rest}")
426        };
427        let forward_url: Uri = TryFrom::try_from(forward_url).map_err(Error::other)?;
428        let mut build = hyper::Request::builder()
429            .method(req.method())
430            .uri(&forward_url);
431        let connection_headers = connection_header_names(req.headers());
432        let upgrade_type = get_upgrade_type(req.headers()).map(str::to_owned);
433        for (key, value) in req.headers() {
434            if key != HOST && !is_hop_by_hop_header(key, &connection_headers) {
435                build = build.header(key, value);
436            }
437        }
438        if let Some(upgrade_type) = upgrade_type {
439            build = build.header(CONNECTION, HeaderValue::from_static("upgrade"));
440            match HeaderValue::from_str(&upgrade_type) {
441                Ok(upgrade_type) => {
442                    build = build.header(UPGRADE, upgrade_type);
443                }
444                Err(e) => {
445                    tracing::error!(error = ?e, "invalid upgrade header value");
446                }
447            }
448        }
449        if let Some(host_value) = (self.host_header_getter)(&forward_url, req, depot) {
450            match HeaderValue::from_str(&host_value) {
451                Ok(host_value) => {
452                    build = build.header(HOST, host_value);
453                }
454                Err(e) => {
455                    tracing::error!(error = ?e, "invalid host header value");
456                }
457            }
458        }
459
460        if self.client_ip_forwarding_enabled {
461            let xff_header_name = HeaderName::from_static(X_FORWARDER_FOR_HEADER_NAME);
462            let current_xff = req.headers().get(&xff_header_name);
463
464            #[cfg(test)]
465            let system_ip_addr = match req.remote_addr() {
466                SocketAddr::IPv6(_) => Some(IpAddr::from(Ipv6Addr::new(
467                    0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8,
468                ))),
469                _ => Some(IpAddr::from(Ipv4Addr::new(101, 102, 103, 104))),
470            };
471
472            #[cfg(not(test))]
473            let system_ip_addr = match req.remote_addr() {
474                SocketAddr::IPv6(_) => local_ipv6().ok(),
475                _ => local_ip().ok(),
476            };
477
478            if let Some(system_ip_addr) = system_ip_addr {
479                let forwarded_addr = system_ip_addr.to_string();
480
481                let xff_value = match current_xff {
482                    Some(current_xff) => match current_xff.to_str() {
483                        Ok(current_xff) => format!("{forwarded_addr}, {current_xff}"),
484                        _ => forwarded_addr.clone(),
485                    },
486                    None => forwarded_addr.clone(),
487                };
488
489                let xff_header_halue = match HeaderValue::from_str(xff_value.as_str()) {
490                    Ok(xff_header_halue) => Some(xff_header_halue),
491                    Err(_) => match HeaderValue::from_str(forwarded_addr.as_str()) {
492                        Ok(xff_header_halue) => Some(xff_header_halue),
493                        Err(e) => {
494                            tracing::error!(error = ?e, "invalid x-forwarded-for header value");
495                            None
496                        }
497                    },
498                };
499
500                if let Some(xff) = xff_header_halue
501                    && let Some(headers) = build.headers_mut()
502                {
503                    headers.insert(&xff_header_name, xff);
504                }
505            }
506        }
507
508        build.body(req.take_body()).map_err(Error::other)
509    }
510}
511
512#[async_trait]
513impl<U, C> Handler for Proxy<U, C>
514where
515    U: Upstreams,
516    U::Error: Into<BoxedError>,
517    C: Client,
518{
519    async fn handle(
520        &self,
521        req: &mut Request,
522        depot: &mut Depot,
523        res: &mut Response,
524        _ctrl: &mut FlowCtrl,
525    ) {
526        match self.build_proxied_request(req, depot).await {
527            Ok(proxied_request) => {
528                match self
529                    .client
530                    .execute(proxied_request, req.extensions_mut().remove())
531                    .await
532                {
533                    Ok(response) => {
534                        let (
535                            salvo_core::http::response::Parts {
536                                status,
537                                // version,
538                                headers,
539                                // extensions,
540                                ..
541                            },
542                            body,
543                        ) = response.into_parts();
544                        res.status_code(status);
545                        for name in headers.keys() {
546                            for value in headers.get_all(name) {
547                                res.headers.append(name, value.to_owned());
548                            }
549                        }
550                        res.body(body);
551                    }
552                    Err(e) => {
553                        tracing::error!( error = ?e, uri = ?req.uri(), "get response data failed: {}", e);
554                        res.status_code(StatusCode::INTERNAL_SERVER_ERROR);
555                    }
556                }
557            }
558            Err(e) => {
559                tracing::error!(error = ?e, "build proxied request failed");
560                res.status_code(StatusCode::BAD_REQUEST);
561            }
562        }
563    }
564}
565
566fn connection_header_names(headers: &HeaderMap) -> Vec<HeaderName> {
567    headers
568        .get_all(CONNECTION)
569        .iter()
570        .filter_map(|value| value.to_str().ok())
571        .flat_map(|value| value.split(','))
572        .filter_map(|name| HeaderName::from_bytes(name.trim().as_bytes()).ok())
573        .collect()
574}
575
576fn is_hop_by_hop_header(name: &HeaderName, connection_headers: &[HeaderName]) -> bool {
577    HOP_BY_HOP_HEADERS
578        .iter()
579        .any(|hop_header| name.as_str().eq_ignore_ascii_case(hop_header))
580        || connection_headers.iter().any(|header| header == name)
581}
582
583#[inline]
584#[allow(dead_code)]
585fn get_upgrade_type(headers: &HeaderMap) -> Option<&str> {
586    if connection_header_names(headers)
587        .iter()
588        .any(|name| name == UPGRADE)
589        && let Some(upgrade_value) = headers.get(&UPGRADE)
590    {
591        tracing::debug!(
592            "found upgrade header with value: {:?}",
593            upgrade_value.to_str()
594        );
595        return upgrade_value.to_str().ok();
596    }
597
598    None
599}
600
601// Unit tests for Proxy
602#[cfg(test)]
603mod tests {
604    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
605    use std::str::FromStr;
606
607    use futures_util::{SinkExt, StreamExt};
608    use salvo_core::conn::{Acceptor, Listener};
609    use salvo_core::prelude::{Router, Server, StatusError, TcpListener, handler};
610    use salvo_extra::websocket::WebSocketUpgrade;
611    use tokio::io::{AsyncReadExt, AsyncWriteExt};
612    use tokio_tungstenite::tungstenite::Message;
613    use tokio_tungstenite::tungstenite::protocol::Role;
614
615    use super::*;
616
617    #[handler]
618    async fn websocket_echo(req: &mut Request, res: &mut Response) -> Result<(), StatusError> {
619        WebSocketUpgrade::new()
620            .upgrade(req, res, |mut ws| async move {
621                while let Some(message) = ws.recv().await {
622                    let Ok(message) = message else {
623                        return;
624                    };
625                    if ws.send(message).await.is_err() {
626                        return;
627                    }
628                }
629            })
630            .await
631    }
632
633    async fn spawn_server(router: Router) -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) {
634        let acceptor = TcpListener::new("127.0.0.1:0").bind().await;
635        let addr = acceptor.holdings()[0]
636            .local_addr
637            .clone()
638            .into_std()
639            .unwrap();
640        let handle = tokio::spawn(async move {
641            Server::new(acceptor).serve(router).await;
642        });
643        (addr, handle)
644    }
645
646    #[test]
647    fn test_encode_url_path() {
648        let path = "/test/path";
649        let encoded_path = encode_url_path(path);
650        assert_eq!(encoded_path, "/test/path");
651    }
652
653    #[test]
654    fn test_default_url_path_getter_uses_raw_tail() {
655        let mut request = Request::new();
656        request
657            .params_mut()
658            .insert("**rest", "guide/../index.html".to_owned());
659        let depot = Depot::new();
660
661        assert_eq!(
662            default_url_path_getter(&request, &depot).as_deref(),
663            Some("guide/../index.html")
664        );
665    }
666
667    #[test]
668    fn test_get_upgrade_type() {
669        let mut headers = HeaderMap::new();
670        headers.insert(CONNECTION, HeaderValue::from_static("upgrade"));
671        headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
672        let upgrade_type = get_upgrade_type(&headers);
673        assert_eq!(upgrade_type, Some("websocket"));
674    }
675
676    #[test]
677    fn test_get_upgrade_type_checks_all_connection_headers() {
678        let mut headers = HeaderMap::new();
679        headers.append(CONNECTION, HeaderValue::from_static("keep-alive"));
680        headers.append(CONNECTION, HeaderValue::from_static("Upgrade"));
681        headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
682
683        let upgrade_type = get_upgrade_type(&headers);
684
685        assert_eq!(upgrade_type, Some("websocket"));
686    }
687
688    #[test]
689    fn test_connection_header_names() {
690        let mut headers = HeaderMap::new();
691        headers.append(CONNECTION, HeaderValue::from_static("keep-alive, x-remove"));
692        headers.append(CONNECTION, HeaderValue::from_static("x-second"));
693
694        let names = connection_header_names(&headers);
695        assert!(names.contains(&HeaderName::from_static("keep-alive")));
696        assert!(names.contains(&HeaderName::from_static("x-remove")));
697        assert!(names.contains(&HeaderName::from_static("x-second")));
698    }
699
700    #[test]
701    fn test_host_header_handling() {
702        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
703        let uri = Uri::from_str("http://host.tld/test").unwrap();
704        let mut req = Request::new();
705        let depot = Depot::new();
706
707        assert_eq!(
708            default_host_header_getter(&uri, &req, &depot),
709            Some("host.tld".to_owned())
710        );
711
712        let uri_with_port = Uri::from_str("http://host.tld:8080/test").unwrap();
713        assert_eq!(
714            rfc2616_host_header_getter(&uri_with_port, &req, &depot),
715            Some("host.tld:8080".to_owned())
716        );
717
718        let uri_with_http_port = Uri::from_str("http://host.tld:80/test").unwrap();
719        assert_eq!(
720            rfc2616_host_header_getter(&uri_with_http_port, &req, &depot),
721            Some("host.tld".to_owned())
722        );
723
724        let uri_with_https_port = Uri::from_str("https://host.tld:443/test").unwrap();
725        assert_eq!(
726            rfc2616_host_header_getter(&uri_with_https_port, &req, &depot),
727            Some("host.tld".to_owned())
728        );
729
730        let uri_with_non_https_scheme_and_https_port =
731            Uri::from_str("http://host.tld:443/test").unwrap();
732        assert_eq!(
733            rfc2616_host_header_getter(&uri_with_non_https_scheme_and_https_port, &req, &depot),
734            Some("host.tld:443".to_owned())
735        );
736
737        req.headers_mut()
738            .insert(HOST, HeaderValue::from_static("test.host.tld"));
739        assert_eq!(
740            preserve_original_host_header_getter(&uri, &req, &depot),
741            Some("test.host.tld".to_owned())
742        );
743    }
744
745    #[tokio::test]
746    async fn test_build_proxied_request_strips_hop_by_hop_headers() {
747        let proxy = Proxy::new(vec!["http://example.com"], HyperClient::default());
748        let mut request = Request::new();
749        let depot = Depot::new();
750
751        request
752            .headers_mut()
753            .insert(HOST, HeaderValue::from_static("client.example"));
754        request
755            .headers_mut()
756            .insert(CONNECTION, HeaderValue::from_static("keep-alive, x-remove"));
757        request.headers_mut().insert(
758            HeaderName::from_static("keep-alive"),
759            HeaderValue::from_static("timeout=5"),
760        );
761        request.headers_mut().insert(
762            HeaderName::from_static("x-remove"),
763            HeaderValue::from_static("secret"),
764        );
765        request.headers_mut().insert(
766            HeaderName::from_static("te"),
767            HeaderValue::from_static("trailers"),
768        );
769        request.headers_mut().insert(
770            HeaderName::from_static("transfer-encoding"),
771            HeaderValue::from_static("chunked"),
772        );
773        request.headers_mut().insert(
774            HeaderName::from_static("x-keep"),
775            HeaderValue::from_static("ok"),
776        );
777
778        let proxied = proxy
779            .build_proxied_request(&mut request, &depot)
780            .await
781            .unwrap();
782
783        assert!(proxied.headers().get(CONNECTION).is_none());
784        assert!(
785            proxied
786                .headers()
787                .get(HeaderName::from_static("keep-alive"))
788                .is_none()
789        );
790        assert!(
791            proxied
792                .headers()
793                .get(HeaderName::from_static("x-remove"))
794                .is_none()
795        );
796        assert!(
797            proxied
798                .headers()
799                .get(HeaderName::from_static("te"))
800                .is_none()
801        );
802        assert!(
803            proxied
804                .headers()
805                .get(HeaderName::from_static("transfer-encoding"))
806                .is_none()
807        );
808        assert_eq!(
809            proxied.headers().get(HeaderName::from_static("x-keep")),
810            Some(&HeaderValue::from_static("ok"))
811        );
812    }
813
814    #[tokio::test]
815    async fn test_build_proxied_request_regenerates_upgrade_headers() {
816        let proxy = Proxy::new(vec!["http://example.com"], HyperClient::default());
817        let mut request = Request::new();
818        let depot = Depot::new();
819
820        request
821            .headers_mut()
822            .insert(CONNECTION, HeaderValue::from_static("x-remove, Upgrade"));
823        request
824            .headers_mut()
825            .insert(UPGRADE, HeaderValue::from_static("websocket"));
826        request.headers_mut().insert(
827            HeaderName::from_static("x-remove"),
828            HeaderValue::from_static("secret"),
829        );
830
831        let proxied = proxy
832            .build_proxied_request(&mut request, &depot)
833            .await
834            .unwrap();
835
836        assert_eq!(
837            proxied.headers().get(CONNECTION),
838            Some(&HeaderValue::from_static("upgrade"))
839        );
840        assert_eq!(
841            proxied.headers().get(UPGRADE),
842            Some(&HeaderValue::from_static("websocket"))
843        );
844        assert!(
845            proxied
846                .headers()
847                .get(HeaderName::from_static("x-remove"))
848                .is_none()
849        );
850    }
851
852    #[tokio::test]
853    async fn test_proxy_websocket_connection_with_split_connection_headers() {
854        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
855
856        let upstream_router = Router::with_path("ws").goal(websocket_echo);
857        let (upstream_addr, upstream_server) = spawn_server(upstream_router).await;
858
859        let proxy_router = Router::with_path("{**rest}").goal(Proxy::new(
860            vec![format!("http://{upstream_addr}")],
861            HyperClient::default(),
862        ));
863        let (proxy_addr, proxy_server) = spawn_server(proxy_router).await;
864
865        let mut stream = tokio::net::TcpStream::connect(proxy_addr).await.unwrap();
866        let request = format!(
867            "\
868GET /ws HTTP/1.1\r\n\
869Host: {proxy_addr}\r\n\
870Connection: keep-alive\r\n\
871Connection: Upgrade\r\n\
872Upgrade: websocket\r\n\
873Sec-WebSocket-Version: 13\r\n\
874Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
875\r\n"
876        );
877        stream.write_all(request.as_bytes()).await.unwrap();
878
879        let mut response = Vec::new();
880        let mut buffer = [0; 1024];
881        let header_end = loop {
882            let read = stream.read(&mut buffer).await.unwrap();
883            assert_ne!(
884                read, 0,
885                "server closed before websocket handshake completed"
886            );
887            response.extend_from_slice(&buffer[..read]);
888            if let Some(position) = response.windows(4).position(|window| window == b"\r\n\r\n") {
889                break position + 4;
890            }
891        };
892        let extra = response.split_off(header_end);
893        let response_head = String::from_utf8_lossy(&response);
894        assert!(
895            response_head.starts_with("HTTP/1.1 101"),
896            "unexpected websocket handshake response: {response_head}"
897        );
898
899        let mut websocket = tokio_tungstenite::WebSocketStream::from_partially_read(
900            stream,
901            extra,
902            Role::Client,
903            None,
904        )
905        .await;
906
907        websocket
908            .send(Message::text("proxied websocket"))
909            .await
910            .unwrap();
911        let echoed = websocket.next().await.unwrap().unwrap();
912        assert_eq!(echoed.into_text().unwrap(), "proxied websocket");
913
914        websocket.close(None).await.unwrap();
915        proxy_server.abort();
916        upstream_server.abort();
917    }
918
919    #[tokio::test]
920    async fn test_client_ip_forwarding() {
921        let xff_header_name = HeaderName::from_static(X_FORWARDER_FOR_HEADER_NAME);
922
923        let mut request = Request::new();
924        let depot = Depot::new();
925
926        // Test functionality not broken
927        let proxy_without_forwarding =
928            Proxy::new(vec!["http://example.com"], HyperClient::default());
929
930        assert!(!proxy_without_forwarding.client_ip_forwarding_enabled);
931
932        let proxy_with_forwarding = proxy_without_forwarding.client_ip_forwarding(true);
933
934        assert!(proxy_with_forwarding.client_ip_forwarding_enabled);
935
936        let proxy =
937            Proxy::with_client_ip_forwarding(vec!["http://example.com"], HyperClient::default());
938        assert!(proxy.client_ip_forwarding_enabled);
939
940        match proxy.build_proxied_request(&mut request, &depot).await {
941            Ok(req) => assert_eq!(
942                req.headers().get(&xff_header_name),
943                Some(&HeaderValue::from_static("101.102.103.104"))
944            ),
945            _ => panic!("expected Ok"),
946        }
947
948        // Test choosing correct IP version depending on remote address
949        *request.remote_addr_mut() =
950            SocketAddr::from(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 12345, 0, 0));
951
952        match proxy.build_proxied_request(&mut request, &depot).await {
953            Ok(req) => assert_eq!(
954                req.headers().get(&xff_header_name),
955                Some(&HeaderValue::from_static("1:2:3:4:5:6:7:8"))
956            ),
957            _ => panic!("expected Ok"),
958        }
959
960        *request.remote_addr_mut() = SocketAddr::Unknown;
961
962        match proxy.build_proxied_request(&mut request, &depot).await {
963            Ok(req) => assert_eq!(
964                req.headers().get(&xff_header_name),
965                Some(&HeaderValue::from_static("101.102.103.104"))
966            ),
967            _ => panic!("expected Ok"),
968        }
969
970        // Test IP prepending when XFF header already exists in initial request.
971        request.headers_mut().insert(
972            &xff_header_name,
973            HeaderValue::from_static("10.72.0.1, 127.0.0.1"),
974        );
975        *request.remote_addr_mut() =
976            SocketAddr::from(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12345));
977
978        match proxy.build_proxied_request(&mut request, &depot).await {
979            Ok(req) => assert_eq!(
980                req.headers().get(&xff_header_name),
981                Some(&HeaderValue::from_static(
982                    "101.102.103.104, 10.72.0.1, 127.0.0.1"
983                ))
984            ),
985            _ => panic!("expected Ok"),
986        }
987    }
988
989    #[tokio::test]
990    async fn test_build_proxied_request_unsafe_tail() {
991        let mut request = Request::new();
992        request.params_mut().insert("**rest", "../admin".to_owned());
993        let depot = Depot::new();
994        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
995
996        let req = proxy
997            .build_proxied_request(&mut request, &depot)
998            .await
999            .unwrap();
1000        assert_eq!(req.uri().to_string(), "http://example.com/api/admin");
1001    }
1002
1003    #[tokio::test]
1004    async fn test_build_proxied_request_normalizes_safe_tail() {
1005        let mut request = Request::new();
1006        request
1007            .params_mut()
1008            .insert("**rest", "guide\\index.html".to_owned());
1009        let depot = Depot::new();
1010        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
1011
1012        let proxied_request = proxy
1013            .build_proxied_request(&mut request, &depot)
1014            .await
1015            .unwrap();
1016        assert_eq!(
1017            proxied_request.uri().to_string(),
1018            "http://example.com/api/guide/index.html"
1019        );
1020    }
1021}
salvo_proxy/lib.rs

salvo_proxy/
lib.rs
