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
205fn contains_ambiguous_path_escape(path: &str) -> bool {
206    let bytes = path.as_bytes();
207    let mut index = 0;
208    while index + 2 < bytes.len() {
209        if bytes[index] == b'%' {
210            if let (Some(high), Some(low)) =
211                (hex_value(bytes[index + 1]), hex_value(bytes[index + 2]))
212            {
213                let decoded = high << 4 | low;
214                if matches!(decoded, b'.' | b'/' | b'\\' | b'%') {
215                    return true;
216                }
217                index += 3;
218                continue;
219            }
220        }
221        index += 1;
222    }
223    false
224}
225
226fn hex_value(byte: u8) -> Option<u8> {
227    match byte {
228        b'0'..=b'9' => Some(byte - b'0'),
229        b'a'..=b'f' => Some(byte - b'a' + 10),
230        b'A'..=b'F' => Some(byte - b'A' + 10),
231        _ => None,
232    }
233}
234/// Default url query getter. This getter just return the query string from request uri.
235pub fn default_url_query_getter(req: &Request, _depot: &Depot) -> Option<String> {
236    req.uri().query().map(Into::into)
237}
238
239/// Default host header getter. This getter will get the host header from request uri
240pub fn default_host_header_getter(
241    forward_uri: &Uri,
242    _req: &Request,
243    _depot: &Depot,
244) -> Option<String> {
245    if let Some(host) = forward_uri.host() {
246        return Some(String::from(host));
247    }
248
249    None
250}
251
252/// RFC2616 complieant host header getter. This getter will get the host header from request uri,
253/// and add port if it's not default port. Falls back to default upon any forward URI parse error.
254pub fn rfc2616_host_header_getter(
255    forward_uri: &Uri,
256    req: &Request,
257    _depot: &Depot,
258) -> Option<String> {
259    let mut parts: Vec<String> = Vec::with_capacity(2);
260
261    if let Some(host) = forward_uri.host() {
262        parts.push(host.to_owned());
263
264        if let Some(scheme) = forward_uri.scheme_str()
265            && let Some(port) = forward_uri.port_u16()
266            && (scheme == "http" && port != 80 || scheme == "https" && port != 443)
267        {
268            parts.push(port.to_string());
269        }
270    }
271
272    if parts.is_empty() {
273        default_host_header_getter(forward_uri, req, _depot)
274    } else {
275        Some(parts.join(":"))
276    }
277}
278
279/// Preserve original host header getter. Propagates the original request host header to the proxied
280/// request.
281pub fn preserve_original_host_header_getter(
282    forward_uri: &Uri,
283    req: &Request,
284    _depot: &Depot,
285) -> Option<String> {
286    if let Some(host_header) = req.headers().get(HOST)
287        && let Ok(host) = host_header.to_str()
288    {
289        return Some(host.to_owned());
290    }
291
292    default_host_header_getter(forward_uri, req, _depot)
293}
294
295/// Handler that can proxy request to other server.
296#[non_exhaustive]
297pub struct Proxy<U, C>
298where
299    U: Upstreams,
300    C: Client,
301{
302    /// Upstreams list.
303    pub upstreams: U,
304    /// [`Client`] for proxy.
305    pub client: C,
306    /// Url path getter.
307    pub url_path_getter: UrlPartGetter,
308    /// Url query getter.
309    pub url_query_getter: UrlPartGetter,
310    /// Host header getter
311    pub host_header_getter: HostHeaderGetter,
312    /// Flag to enable x-forwarded-for header.
313    pub client_ip_forwarding_enabled: bool,
314    /// Flag to reject ambiguous percent-encoded path characters before proxying.
315    pub strict_path_normalization_enabled: bool,
316}
317
318impl<U, C> Debug for Proxy<U, C>
319where
320    U: Upstreams,
321    C: Client,
322{
323    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
324        f.debug_struct("Proxy").finish()
325    }
326}
327
328impl<U, C> Proxy<U, C>
329where
330    U: Upstreams,
331    U::Error: Into<BoxedError>,
332    C: Client,
333{
334    /// Create new `Proxy` with upstreams list.
335    #[must_use]
336    pub fn new(upstreams: U, client: C) -> Self {
337        Self {
338            upstreams,
339            client,
340            url_path_getter: Box::new(default_url_path_getter),
341            url_query_getter: Box::new(default_url_query_getter),
342            host_header_getter: Box::new(default_host_header_getter),
343            client_ip_forwarding_enabled: false,
344            strict_path_normalization_enabled: false,
345        }
346    }
347
348    /// Create new `Proxy` with upstreams list and enable x-forwarded-for header.
349    pub fn with_client_ip_forwarding(upstreams: U, client: C) -> Self {
350        Self {
351            upstreams,
352            client,
353            url_path_getter: Box::new(default_url_path_getter),
354            url_query_getter: Box::new(default_url_query_getter),
355            host_header_getter: Box::new(default_host_header_getter),
356            client_ip_forwarding_enabled: true,
357            strict_path_normalization_enabled: false,
358        }
359    }
360
361    /// Set url path getter.
362    #[inline]
363    #[must_use]
364    pub fn url_path_getter<G>(mut self, url_path_getter: G) -> Self
365    where
366        G: Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static,
367    {
368        self.url_path_getter = Box::new(url_path_getter);
369        self
370    }
371
372    /// Set url query getter.
373    #[inline]
374    #[must_use]
375    pub fn url_query_getter<G>(mut self, url_query_getter: G) -> Self
376    where
377        G: Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static,
378    {
379        self.url_query_getter = Box::new(url_query_getter);
380        self
381    }
382
383    /// Set host header query getter.
384    #[inline]
385    #[must_use]
386    pub fn host_header_getter<G>(mut self, host_header_getter: G) -> Self
387    where
388        G: Fn(&Uri, &Request, &Depot) -> Option<String> + Send + Sync + 'static,
389    {
390        self.host_header_getter = Box::new(host_header_getter);
391        self
392    }
393
394    /// Enable or disable strict path normalization.
395    ///
396    /// When enabled, the proxy rejects paths that still contain percent-encoded `.`, `/`, `\`,
397    /// or `%` characters after Salvo routing has extracted the path tail. This is useful when the
398    /// proxy is used as a security boundary and the upstream server may perform another decode
399    /// pass.
400    #[inline]
401    #[must_use]
402    pub fn strict_path_normalization(mut self, enable: bool) -> Self {
403        self.strict_path_normalization_enabled = enable;
404        self
405    }
406
407    /// Get upstreams list.
408    #[inline]
409    pub fn upstreams(&self) -> &U {
410        &self.upstreams
411    }
412    /// Get upstreams mutable list.
413    #[inline]
414    pub fn upstreams_mut(&mut self) -> &mut U {
415        &mut self.upstreams
416    }
417
418    /// Get client reference.
419    #[inline]
420    pub fn client(&self) -> &C {
421        &self.client
422    }
423    /// Get client mutable reference.
424    #[inline]
425    pub fn client_mut(&mut self) -> &mut C {
426        &mut self.client
427    }
428
429    /// Enable x-forwarded-for header prepending.
430    #[inline]
431    #[must_use]
432    pub fn client_ip_forwarding(mut self, enable: bool) -> Self {
433        self.client_ip_forwarding_enabled = enable;
434        self
435    }
436
437    async fn build_proxied_request(
438        &self,
439        req: &mut Request,
440        depot: &Depot,
441    ) -> Result<HyperRequest, Error> {
442        let upstream = self
443            .upstreams
444            .elect(req, depot)
445            .await
446            .map_err(Error::other)?;
447
448        if upstream.is_empty() {
449            tracing::error!("upstreams is empty");
450            return Err(Error::other("upstreams is empty"));
451        }
452
453        let path = (self.url_path_getter)(req, depot).unwrap_or_default();
454        if self.strict_path_normalization_enabled && contains_ambiguous_path_escape(&path) {
455            return Err(Error::other("ambiguous percent-encoded path"));
456        }
457        let path = encode_url_path(&normalize_url_path(&path));
458        let query = (self.url_query_getter)(req, depot);
459        let rest = if let Some(query) = query {
460            if let Some(stripped) = query.strip_prefix('?') {
461                format!("{path}?{}", utf8_percent_encode(stripped, QUERY_ENCODE_SET))
462            } else {
463                format!("{path}?{}", utf8_percent_encode(&query, QUERY_ENCODE_SET))
464            }
465        } else {
466            path
467        };
468        let forward_url = if upstream.ends_with('/') && rest.starts_with('/') {
469            format!("{}{}", upstream.trim_end_matches('/'), rest)
470        } else if upstream.ends_with('/') || rest.starts_with('/') {
471            format!("{upstream}{rest}")
472        } else if rest.is_empty() {
473            upstream.to_owned()
474        } else {
475            format!("{upstream}/{rest}")
476        };
477        let forward_url: Uri = TryFrom::try_from(forward_url).map_err(Error::other)?;
478        let mut build = hyper::Request::builder()
479            .method(req.method())
480            .uri(&forward_url);
481        let connection_headers = connection_header_names(req.headers());
482        let upgrade_type = get_upgrade_type(req.headers()).map(str::to_owned);
483        for (key, value) in req.headers() {
484            if key != HOST && !is_hop_by_hop_header(key, &connection_headers) {
485                build = build.header(key, value);
486            }
487        }
488        if let Some(upgrade_type) = upgrade_type {
489            build = build.header(CONNECTION, HeaderValue::from_static("upgrade"));
490            match HeaderValue::from_str(&upgrade_type) {
491                Ok(upgrade_type) => {
492                    build = build.header(UPGRADE, upgrade_type);
493                }
494                Err(e) => {
495                    tracing::error!(error = ?e, "invalid upgrade header value");
496                }
497            }
498        }
499        if let Some(host_value) = (self.host_header_getter)(&forward_url, req, depot) {
500            match HeaderValue::from_str(&host_value) {
501                Ok(host_value) => {
502                    build = build.header(HOST, host_value);
503                }
504                Err(e) => {
505                    tracing::error!(error = ?e, "invalid host header value");
506                }
507            }
508        }
509
510        if self.client_ip_forwarding_enabled {
511            let xff_header_name = HeaderName::from_static(X_FORWARDER_FOR_HEADER_NAME);
512            let current_xff = req.headers().get(&xff_header_name);
513
514            #[cfg(test)]
515            let system_ip_addr = match req.remote_addr() {
516                SocketAddr::IPv6(_) => Some(IpAddr::from(Ipv6Addr::new(
517                    0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8,
518                ))),
519                _ => Some(IpAddr::from(Ipv4Addr::new(101, 102, 103, 104))),
520            };
521
522            #[cfg(not(test))]
523            let system_ip_addr = match req.remote_addr() {
524                SocketAddr::IPv6(_) => local_ipv6().ok(),
525                _ => local_ip().ok(),
526            };
527
528            if let Some(system_ip_addr) = system_ip_addr {
529                let forwarded_addr = system_ip_addr.to_string();
530
531                let xff_value = match current_xff {
532                    Some(current_xff) => match current_xff.to_str() {
533                        Ok(current_xff) => format!("{forwarded_addr}, {current_xff}"),
534                        _ => forwarded_addr.clone(),
535                    },
536                    None => forwarded_addr.clone(),
537                };
538
539                let xff_header_halue = match HeaderValue::from_str(xff_value.as_str()) {
540                    Ok(xff_header_halue) => Some(xff_header_halue),
541                    Err(_) => match HeaderValue::from_str(forwarded_addr.as_str()) {
542                        Ok(xff_header_halue) => Some(xff_header_halue),
543                        Err(e) => {
544                            tracing::error!(error = ?e, "invalid x-forwarded-for header value");
545                            None
546                        }
547                    },
548                };
549
550                if let Some(xff) = xff_header_halue
551                    && let Some(headers) = build.headers_mut()
552                {
553                    headers.insert(&xff_header_name, xff);
554                }
555            }
556        }
557
558        build.body(req.take_body()).map_err(Error::other)
559    }
560}
561
562#[async_trait]
563impl<U, C> Handler for Proxy<U, C>
564where
565    U: Upstreams,
566    U::Error: Into<BoxedError>,
567    C: Client,
568{
569    async fn handle(
570        &self,
571        req: &mut Request,
572        depot: &mut Depot,
573        res: &mut Response,
574        _ctrl: &mut FlowCtrl,
575    ) {
576        match self.build_proxied_request(req, depot).await {
577            Ok(proxied_request) => {
578                match self
579                    .client
580                    .execute(proxied_request, req.extensions_mut().remove())
581                    .await
582                {
583                    Ok(response) => {
584                        let (
585                            salvo_core::http::response::Parts {
586                                status,
587                                // version,
588                                headers,
589                                // extensions,
590                                ..
591                            },
592                            body,
593                        ) = response.into_parts();
594                        res.status_code(status);
595                        for name in headers.keys() {
596                            for value in headers.get_all(name) {
597                                res.headers.append(name, value.to_owned());
598                            }
599                        }
600                        res.body(body);
601                    }
602                    Err(e) => {
603                        tracing::error!( error = ?e, uri = ?req.uri(), "get response data failed: {}", e);
604                        res.status_code(StatusCode::INTERNAL_SERVER_ERROR);
605                    }
606                }
607            }
608            Err(e) => {
609                tracing::error!(error = ?e, "build proxied request failed");
610                res.status_code(StatusCode::BAD_REQUEST);
611            }
612        }
613    }
614}
615
616fn connection_header_names(headers: &HeaderMap) -> Vec<HeaderName> {
617    headers
618        .get_all(CONNECTION)
619        .iter()
620        .filter_map(|value| value.to_str().ok())
621        .flat_map(|value| value.split(','))
622        .filter_map(|name| HeaderName::from_bytes(name.trim().as_bytes()).ok())
623        .collect()
624}
625
626fn is_hop_by_hop_header(name: &HeaderName, connection_headers: &[HeaderName]) -> bool {
627    HOP_BY_HOP_HEADERS
628        .iter()
629        .any(|hop_header| name.as_str().eq_ignore_ascii_case(hop_header))
630        || connection_headers.iter().any(|header| header == name)
631}
632
633#[inline]
634#[allow(dead_code)]
635fn get_upgrade_type(headers: &HeaderMap) -> Option<&str> {
636    if connection_header_names(headers)
637        .iter()
638        .any(|name| name == UPGRADE)
639        && let Some(upgrade_value) = headers.get(&UPGRADE)
640    {
641        tracing::debug!(
642            "found upgrade header with value: {:?}",
643            upgrade_value.to_str()
644        );
645        return upgrade_value.to_str().ok();
646    }
647
648    None
649}
650
651// Unit tests for Proxy
652#[cfg(test)]
653mod tests {
654    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
655    use std::str::FromStr;
656
657    use futures_util::{SinkExt, StreamExt};
658    use salvo_core::conn::{Acceptor, Listener};
659    use salvo_core::prelude::{Router, Server, StatusError, TcpListener, handler};
660    use salvo_extra::websocket::WebSocketUpgrade;
661    use tokio::io::{AsyncReadExt, AsyncWriteExt};
662    use tokio_tungstenite::tungstenite::Message;
663    use tokio_tungstenite::tungstenite::protocol::Role;
664
665    use super::*;
666
667    #[handler]
668    async fn websocket_echo(req: &mut Request, res: &mut Response) -> Result<(), StatusError> {
669        WebSocketUpgrade::new()
670            .upgrade(req, res, |mut ws| async move {
671                while let Some(message) = ws.recv().await {
672                    let Ok(message) = message else {
673                        return;
674                    };
675                    if ws.send(message).await.is_err() {
676                        return;
677                    }
678                }
679            })
680            .await
681    }
682
683    async fn spawn_server(router: Router) -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) {
684        let acceptor = TcpListener::new("127.0.0.1:0").bind().await;
685        let addr = acceptor.holdings()[0]
686            .local_addr
687            .clone()
688            .into_std()
689            .unwrap();
690        let handle = tokio::spawn(async move {
691            Server::new(acceptor).serve(router).await;
692        });
693        (addr, handle)
694    }
695
696    #[test]
697    fn test_encode_url_path() {
698        let path = "/test/path";
699        let encoded_path = encode_url_path(path);
700        assert_eq!(encoded_path, "/test/path");
701    }
702
703    #[test]
704    fn test_default_url_path_getter_uses_raw_tail() {
705        let mut request = Request::new();
706        request
707            .params_mut()
708            .insert("**rest", "guide/../index.html".to_owned());
709        let depot = Depot::new();
710
711        assert_eq!(
712            default_url_path_getter(&request, &depot).as_deref(),
713            Some("guide/../index.html")
714        );
715    }
716
717    #[test]
718    fn test_contains_ambiguous_path_escape() {
719        assert!(contains_ambiguous_path_escape("%2e%2e/admin"));
720        assert!(contains_ambiguous_path_escape("api%2Fadmin"));
721        assert!(contains_ambiguous_path_escape("api%5cadmin"));
722        assert!(contains_ambiguous_path_escape("%252e%252e/admin"));
723        assert!(!contains_ambiguous_path_escape("guide.v1/index.html"));
724        assert!(!contains_ambiguous_path_escape("files/%20space"));
725    }
726
727    #[test]
728    fn test_get_upgrade_type() {
729        let mut headers = HeaderMap::new();
730        headers.insert(CONNECTION, HeaderValue::from_static("upgrade"));
731        headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
732        let upgrade_type = get_upgrade_type(&headers);
733        assert_eq!(upgrade_type, Some("websocket"));
734    }
735
736    #[test]
737    fn test_get_upgrade_type_checks_all_connection_headers() {
738        let mut headers = HeaderMap::new();
739        headers.append(CONNECTION, HeaderValue::from_static("keep-alive"));
740        headers.append(CONNECTION, HeaderValue::from_static("Upgrade"));
741        headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
742
743        let upgrade_type = get_upgrade_type(&headers);
744
745        assert_eq!(upgrade_type, Some("websocket"));
746    }
747
748    #[test]
749    fn test_connection_header_names() {
750        let mut headers = HeaderMap::new();
751        headers.append(CONNECTION, HeaderValue::from_static("keep-alive, x-remove"));
752        headers.append(CONNECTION, HeaderValue::from_static("x-second"));
753
754        let names = connection_header_names(&headers);
755        assert!(names.contains(&HeaderName::from_static("keep-alive")));
756        assert!(names.contains(&HeaderName::from_static("x-remove")));
757        assert!(names.contains(&HeaderName::from_static("x-second")));
758    }
759
760    #[test]
761    fn test_host_header_handling() {
762        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
763        let uri = Uri::from_str("http://host.tld/test").unwrap();
764        let mut req = Request::new();
765        let depot = Depot::new();
766
767        assert_eq!(
768            default_host_header_getter(&uri, &req, &depot),
769            Some("host.tld".to_owned())
770        );
771
772        let uri_with_port = Uri::from_str("http://host.tld:8080/test").unwrap();
773        assert_eq!(
774            rfc2616_host_header_getter(&uri_with_port, &req, &depot),
775            Some("host.tld:8080".to_owned())
776        );
777
778        let uri_with_http_port = Uri::from_str("http://host.tld:80/test").unwrap();
779        assert_eq!(
780            rfc2616_host_header_getter(&uri_with_http_port, &req, &depot),
781            Some("host.tld".to_owned())
782        );
783
784        let uri_with_https_port = Uri::from_str("https://host.tld:443/test").unwrap();
785        assert_eq!(
786            rfc2616_host_header_getter(&uri_with_https_port, &req, &depot),
787            Some("host.tld".to_owned())
788        );
789
790        let uri_with_non_https_scheme_and_https_port =
791            Uri::from_str("http://host.tld:443/test").unwrap();
792        assert_eq!(
793            rfc2616_host_header_getter(&uri_with_non_https_scheme_and_https_port, &req, &depot),
794            Some("host.tld:443".to_owned())
795        );
796
797        req.headers_mut()
798            .insert(HOST, HeaderValue::from_static("test.host.tld"));
799        assert_eq!(
800            preserve_original_host_header_getter(&uri, &req, &depot),
801            Some("test.host.tld".to_owned())
802        );
803    }
804
805    #[tokio::test]
806    async fn test_build_proxied_request_strips_hop_by_hop_headers() {
807        let proxy = Proxy::new(vec!["http://example.com"], HyperClient::default());
808        let mut request = Request::new();
809        let depot = Depot::new();
810
811        request
812            .headers_mut()
813            .insert(HOST, HeaderValue::from_static("client.example"));
814        request
815            .headers_mut()
816            .insert(CONNECTION, HeaderValue::from_static("keep-alive, x-remove"));
817        request.headers_mut().insert(
818            HeaderName::from_static("keep-alive"),
819            HeaderValue::from_static("timeout=5"),
820        );
821        request.headers_mut().insert(
822            HeaderName::from_static("x-remove"),
823            HeaderValue::from_static("secret"),
824        );
825        request.headers_mut().insert(
826            HeaderName::from_static("te"),
827            HeaderValue::from_static("trailers"),
828        );
829        request.headers_mut().insert(
830            HeaderName::from_static("transfer-encoding"),
831            HeaderValue::from_static("chunked"),
832        );
833        request.headers_mut().insert(
834            HeaderName::from_static("x-keep"),
835            HeaderValue::from_static("ok"),
836        );
837
838        let proxied = proxy
839            .build_proxied_request(&mut request, &depot)
840            .await
841            .unwrap();
842
843        assert!(proxied.headers().get(CONNECTION).is_none());
844        assert!(
845            proxied
846                .headers()
847                .get(HeaderName::from_static("keep-alive"))
848                .is_none()
849        );
850        assert!(
851            proxied
852                .headers()
853                .get(HeaderName::from_static("x-remove"))
854                .is_none()
855        );
856        assert!(
857            proxied
858                .headers()
859                .get(HeaderName::from_static("te"))
860                .is_none()
861        );
862        assert!(
863            proxied
864                .headers()
865                .get(HeaderName::from_static("transfer-encoding"))
866                .is_none()
867        );
868        assert_eq!(
869            proxied.headers().get(HeaderName::from_static("x-keep")),
870            Some(&HeaderValue::from_static("ok"))
871        );
872    }
873
874    #[tokio::test]
875    async fn test_build_proxied_request_regenerates_upgrade_headers() {
876        let proxy = Proxy::new(vec!["http://example.com"], HyperClient::default());
877        let mut request = Request::new();
878        let depot = Depot::new();
879
880        request
881            .headers_mut()
882            .insert(CONNECTION, HeaderValue::from_static("x-remove, Upgrade"));
883        request
884            .headers_mut()
885            .insert(UPGRADE, HeaderValue::from_static("websocket"));
886        request.headers_mut().insert(
887            HeaderName::from_static("x-remove"),
888            HeaderValue::from_static("secret"),
889        );
890
891        let proxied = proxy
892            .build_proxied_request(&mut request, &depot)
893            .await
894            .unwrap();
895
896        assert_eq!(
897            proxied.headers().get(CONNECTION),
898            Some(&HeaderValue::from_static("upgrade"))
899        );
900        assert_eq!(
901            proxied.headers().get(UPGRADE),
902            Some(&HeaderValue::from_static("websocket"))
903        );
904        assert!(
905            proxied
906                .headers()
907                .get(HeaderName::from_static("x-remove"))
908                .is_none()
909        );
910    }
911
912    #[tokio::test]
913    async fn test_proxy_websocket_connection_with_split_connection_headers() {
914        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
915
916        let upstream_router = Router::with_path("ws").goal(websocket_echo);
917        let (upstream_addr, upstream_server) = spawn_server(upstream_router).await;
918
919        let proxy_router = Router::with_path("{**rest}").goal(Proxy::new(
920            vec![format!("http://{upstream_addr}")],
921            HyperClient::default(),
922        ));
923        let (proxy_addr, proxy_server) = spawn_server(proxy_router).await;
924
925        let mut stream = tokio::net::TcpStream::connect(proxy_addr).await.unwrap();
926        let request = format!(
927            "\
928GET /ws HTTP/1.1\r\n\
929Host: {proxy_addr}\r\n\
930Connection: keep-alive\r\n\
931Connection: Upgrade\r\n\
932Upgrade: websocket\r\n\
933Sec-WebSocket-Version: 13\r\n\
934Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
935\r\n"
936        );
937        stream.write_all(request.as_bytes()).await.unwrap();
938
939        let mut response = Vec::new();
940        let mut buffer = [0; 1024];
941        let header_end = loop {
942            let read = stream.read(&mut buffer).await.unwrap();
943            assert_ne!(
944                read, 0,
945                "server closed before websocket handshake completed"
946            );
947            response.extend_from_slice(&buffer[..read]);
948            if let Some(position) = response.windows(4).position(|window| window == b"\r\n\r\n") {
949                break position + 4;
950            }
951        };
952        let extra = response.split_off(header_end);
953        let response_head = String::from_utf8_lossy(&response);
954        assert!(
955            response_head.starts_with("HTTP/1.1 101"),
956            "unexpected websocket handshake response: {response_head}"
957        );
958
959        let mut websocket = tokio_tungstenite::WebSocketStream::from_partially_read(
960            stream,
961            extra,
962            Role::Client,
963            None,
964        )
965        .await;
966
967        websocket
968            .send(Message::text("proxied websocket"))
969            .await
970            .unwrap();
971        let echoed = websocket.next().await.unwrap().unwrap();
972        assert_eq!(echoed.into_text().unwrap(), "proxied websocket");
973
974        websocket.close(None).await.unwrap();
975        proxy_server.abort();
976        upstream_server.abort();
977    }
978
979    #[tokio::test]
980    async fn test_client_ip_forwarding() {
981        let xff_header_name = HeaderName::from_static(X_FORWARDER_FOR_HEADER_NAME);
982
983        let mut request = Request::new();
984        let depot = Depot::new();
985
986        // Test functionality not broken
987        let proxy_without_forwarding =
988            Proxy::new(vec!["http://example.com"], HyperClient::default());
989
990        assert!(!proxy_without_forwarding.client_ip_forwarding_enabled);
991
992        let proxy_with_forwarding = proxy_without_forwarding.client_ip_forwarding(true);
993
994        assert!(proxy_with_forwarding.client_ip_forwarding_enabled);
995
996        let proxy =
997            Proxy::with_client_ip_forwarding(vec!["http://example.com"], HyperClient::default());
998        assert!(proxy.client_ip_forwarding_enabled);
999
1000        match proxy.build_proxied_request(&mut request, &depot).await {
1001            Ok(req) => assert_eq!(
1002                req.headers().get(&xff_header_name),
1003                Some(&HeaderValue::from_static("101.102.103.104"))
1004            ),
1005            _ => panic!("expected Ok"),
1006        }
1007
1008        // Test choosing correct IP version depending on remote address
1009        *request.remote_addr_mut() =
1010            SocketAddr::from(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 12345, 0, 0));
1011
1012        match proxy.build_proxied_request(&mut request, &depot).await {
1013            Ok(req) => assert_eq!(
1014                req.headers().get(&xff_header_name),
1015                Some(&HeaderValue::from_static("1:2:3:4:5:6:7:8"))
1016            ),
1017            _ => panic!("expected Ok"),
1018        }
1019
1020        *request.remote_addr_mut() = SocketAddr::Unknown;
1021
1022        match proxy.build_proxied_request(&mut request, &depot).await {
1023            Ok(req) => assert_eq!(
1024                req.headers().get(&xff_header_name),
1025                Some(&HeaderValue::from_static("101.102.103.104"))
1026            ),
1027            _ => panic!("expected Ok"),
1028        }
1029
1030        // Test IP prepending when XFF header already exists in initial request.
1031        request.headers_mut().insert(
1032            &xff_header_name,
1033            HeaderValue::from_static("10.72.0.1, 127.0.0.1"),
1034        );
1035        *request.remote_addr_mut() =
1036            SocketAddr::from(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12345));
1037
1038        match proxy.build_proxied_request(&mut request, &depot).await {
1039            Ok(req) => assert_eq!(
1040                req.headers().get(&xff_header_name),
1041                Some(&HeaderValue::from_static(
1042                    "101.102.103.104, 10.72.0.1, 127.0.0.1"
1043                ))
1044            ),
1045            _ => panic!("expected Ok"),
1046        }
1047    }
1048
1049    #[tokio::test]
1050    async fn test_build_proxied_request_unsafe_tail() {
1051        let mut request = Request::new();
1052        request.params_mut().insert("**rest", "../admin".to_owned());
1053        let depot = Depot::new();
1054        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
1055
1056        let req = proxy
1057            .build_proxied_request(&mut request, &depot)
1058            .await
1059            .unwrap();
1060        assert_eq!(req.uri().to_string(), "http://example.com/api/admin");
1061    }
1062
1063    #[tokio::test]
1064    async fn test_build_proxied_request_normalizes_safe_tail() {
1065        let mut request = Request::new();
1066        request
1067            .params_mut()
1068            .insert("**rest", "guide\\index.html".to_owned());
1069        let depot = Depot::new();
1070        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
1071
1072        let proxied_request = proxy
1073            .build_proxied_request(&mut request, &depot)
1074            .await
1075            .unwrap();
1076        assert_eq!(
1077            proxied_request.uri().to_string(),
1078            "http://example.com/api/guide/index.html"
1079        );
1080    }
1081
1082    #[tokio::test]
1083    async fn test_build_proxied_request_preserves_encoded_tail_by_default() {
1084        let mut request = Request::new();
1085        request
1086            .params_mut()
1087            .insert("**rest", "%2e%2e/secrets/.env".to_owned());
1088        let depot = Depot::new();
1089        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
1090
1091        let proxied_request = proxy
1092            .build_proxied_request(&mut request, &depot)
1093            .await
1094            .unwrap();
1095        assert_eq!(
1096            proxied_request.uri().to_string(),
1097            "http://example.com/api/%2e%2e/secrets/.env"
1098        );
1099    }
1100
1101    #[tokio::test]
1102    async fn test_build_proxied_request_strict_path_normalization_rejects_ambiguous_escapes() {
1103        for path in [
1104            "%2e%2e/secrets/.env",
1105            "api%2fadmin",
1106            "api%5cadmin",
1107            "%252e%252e/secrets/.env",
1108        ] {
1109            let mut request = Request::new();
1110            request.params_mut().insert("**rest", path.to_owned());
1111            let depot = Depot::new();
1112            let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default())
1113                .strict_path_normalization(true);
1114
1115            let err = proxy.build_proxied_request(&mut request, &depot).await;
1116            assert!(err.is_err(), "path should be rejected: {path}");
1117        }
1118    }
1119}
salvo_proxy/lib.rs

salvo_proxy/
lib.rs
