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
46use hyper::upgrade::OnUpgrade;
47use percent_encoding::{AsciiSet, CONTROLS, utf8_percent_encode};
48use salvo_core::http::header::{
49    AUTHORIZATION, CONNECTION, HOST, HeaderMap, HeaderName, HeaderValue, UPGRADE,
50};
51use salvo_core::http::uri::Uri;
52use salvo_core::http::{ReqBody, ResBody, StatusCode};
53use salvo_core::routing::normalize_url_path;
54use salvo_core::{
55    BoxedError, Depot, Error, FlowCtrl, Handler, Request, Response, async_trait, cfg_feature,
56};
57
58cfg_feature! {
59    #![feature = "hyper-client"]
60    mod hyper_client;
61    pub use hyper_client::*;
62}
63cfg_feature! {
64    #![feature = "reqwest-client"]
65    mod reqwest_client;
66    pub use reqwest_client::*;
67}
68
69cfg_feature! {
70    #![feature = "unix-sock-client"]
71    #[cfg(unix)]
72    mod unix_sock_client;
73    #[cfg(unix)]
74    pub use unix_sock_client::*;
75}
76
77type HyperRequest = hyper::Request<ReqBody>;
78type HyperResponse = hyper::Response<ResBody>;
79
80const X_FORWARDED_FOR_HEADER_NAME: &str = "x-forwarded-for";
81const HOP_BY_HOP_HEADERS: &[&str] = &[
82    "connection",
83    "keep-alive",
84    "proxy-authenticate",
85    "proxy-authorization",
86    "te",
87    "trailer",
88    "transfer-encoding",
89    "upgrade",
90];
91
92const QUERY_ENCODE_SET: &AsciiSet = &CONTROLS
93    .add(b' ')
94    .add(b'"')
95    .add(b'#')
96    .add(b'<')
97    .add(b'>')
98    .add(b'`');
99const PATH_ENCODE_SET: &AsciiSet = &QUERY_ENCODE_SET
100    .add(b'?')
101    .add(b'^')
102    .add(b'`')
103    .add(b'{')
104    .add(b'}');
105
106/// Encode URL path. This can be used when building a custom URL path getter.
107///
108/// Walks the segments separated by `/` and percent-encodes each one in place using
109/// [`PATH_ENCODE_SET`], without allocating an intermediate `Vec<String>` or a
110/// `String` per segment.
111#[inline]
112pub(crate) fn encode_url_path(path: &str) -> String {
113    use std::fmt::Write as _;
114
115    let mut out = String::with_capacity(path.len());
116    let mut first = true;
117    for segment in path.split('/') {
118        if first {
119            first = false;
120        } else {
121            out.push('/');
122        }
123        // `Display` for `PercentEncode` writes directly into the formatter, so this
124        // does not allocate a temporary `String` per segment.
125        let _ = write!(
126            &mut out,
127            "{}",
128            utf8_percent_encode(segment, PATH_ENCODE_SET)
129        );
130    }
131    out
132}
133
134/// Client trait for implementing different HTTP clients for proxying.
135///
136/// Implement this trait to create custom proxy clients with different
137/// backends or configurations.
138pub trait Client: Send + Sync + 'static {
139    /// Error type returned by the client.
140    type Error: StdError + Send + Sync + 'static;
141
142    /// Execute a request through the proxy client.
143    fn execute(
144        &self,
145        req: HyperRequest,
146        upgraded: Option<OnUpgrade>,
147    ) -> impl Future<Output = Result<HyperResponse, Self::Error>> + Send;
148}
149
150/// Upstreams trait for selecting target servers.
151///
152/// Implement this trait to customize how target servers are selected
153/// for proxying requests. This can be used to implement load balancing,
154/// failover, or other server selection strategies.
155pub trait Upstreams: Send + Sync + 'static {
156    /// Error type returned when selecting a server fails.
157    type Error: StdError + Send + Sync + 'static;
158
159    /// Elect a server to handle the current request.
160    fn elect(
161        &self,
162        req: &Request,
163        depot: &Depot,
164    ) -> impl Future<Output = Result<&str, Self::Error>> + Send;
165}
166impl Upstreams for &'static str {
167    type Error = Infallible;
168
169    async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
170        Ok(*self)
171    }
172}
173impl Upstreams for String {
174    type Error = Infallible;
175    async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
176        Ok(self.as_str())
177    }
178}
179
180impl<const N: usize> Upstreams for [&'static str; N] {
181    type Error = Error;
182    async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
183        if self.is_empty() {
184            return Err(Error::other("upstreams is empty"));
185        }
186        let index = fastrand::usize(..self.len());
187        Ok(self[index])
188    }
189}
190
191impl<T> Upstreams for Vec<T>
192where
193    T: AsRef<str> + Send + Sync + 'static,
194{
195    type Error = Error;
196    async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
197        if self.is_empty() {
198            return Err(Error::other("upstreams is empty"));
199        }
200        let index = fastrand::usize(..self.len());
201        Ok(self[index].as_ref())
202    }
203}
204
205/// URL part getter. You can use this to get the proxied URL path or query.
206pub type UrlPartGetter = Box<dyn Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static>;
207
208/// Host header getter. You can use this to get the host header for the proxied request.
209pub type HostHeaderGetter =
210    Box<dyn Fn(&Uri, &Request, &Depot) -> Option<String> + Send + Sync + 'static>;
211
212/// Default URL path getter.
213///
214/// This getter gets the last param as the rest URL path from the request.
215/// In most cases you should use a wildcard param, like `{**rest}`, `{*+rest}`.
216pub fn default_url_path_getter(req: &Request, _depot: &Depot) -> Option<String> {
217    req.params().tail().map(str::to_owned)
218}
219
220fn contains_ambiguous_path_escape(path: &str) -> bool {
221    let bytes = path.as_bytes();
222    let mut index = 0;
223    while index + 2 < bytes.len() {
224        if bytes[index] == b'%'
225            && let (Some(high), Some(low)) =
226                (hex_value(bytes[index + 1]), hex_value(bytes[index + 2]))
227        {
228            let decoded = high << 4 | low;
229            if matches!(decoded, b'.' | b'/' | b'\\' | b'%') {
230                return true;
231            }
232            index += 3;
233            continue;
234        }
235        index += 1;
236    }
237    false
238}
239
240fn contains_parent_dir_component(path: &str) -> bool {
241    path.split(['/', '\\']).any(|part| part == "..")
242}
243
244fn hex_value(byte: u8) -> Option<u8> {
245    match byte {
246        b'0'..=b'9' => Some(byte - b'0'),
247        b'a'..=b'f' => Some(byte - b'a' + 10),
248        b'A'..=b'F' => Some(byte - b'A' + 10),
249        _ => None,
250    }
251}
252/// Default URL query getter. This getter returns the query string from the request URI.
253pub fn default_url_query_getter(req: &Request, _depot: &Depot) -> Option<String> {
254    req.uri().query().map(Into::into)
255}
256
257/// Default host header getter.
258///
259/// This getter returns only the host name from the request URI. It does not include a non-default
260/// port. Use [`standard_host_header_getter`] when the forwarded `Host` header should include a
261/// non-default port from the upstream URI.
262pub fn default_host_header_getter(
263    forward_uri: &Uri,
264    _req: &Request,
265    _depot: &Depot,
266) -> Option<String> {
267    if let Some(host) = forward_uri.host() {
268        return Some(String::from(host));
269    }
270
271    None
272}
273
274/// Standards-compliant host header getter.
275///
276/// This getter gets the host header from the request URI and adds the port when it is not the
277/// default port for the scheme. This follows the Host header behavior specified by modern HTTP
278/// semantics (RFC 7230 and RFC 9110).
279pub fn standard_host_header_getter(
280    forward_uri: &Uri,
281    req: &Request,
282    _depot: &Depot,
283) -> Option<String> {
284    let mut parts: Vec<String> = Vec::with_capacity(2);
285
286    if let Some(host) = forward_uri.host() {
287        parts.push(host.to_owned());
288
289        if let Some(scheme) = forward_uri.scheme_str()
290            && let Some(port) = forward_uri.port_u16()
291            && (scheme == "http" && port != 80 || scheme == "https" && port != 443)
292        {
293            parts.push(port.to_string());
294        }
295    }
296
297    if parts.is_empty() {
298        default_host_header_getter(forward_uri, req, _depot)
299    } else {
300        Some(parts.join(":"))
301    }
302}
303
304/// Preserve original host header getter. Propagates the original request host header to the proxied
305/// request.
306pub fn preserve_original_host_header_getter(
307    forward_uri: &Uri,
308    req: &Request,
309    _depot: &Depot,
310) -> Option<String> {
311    if let Some(host_header) = req.headers().get(HOST)
312        && let Ok(host) = host_header.to_str()
313    {
314        return Some(host.to_owned());
315    }
316
317    default_host_header_getter(forward_uri, req, _depot)
318}
319
320/// Handler that can proxy request to other server.
321#[non_exhaustive]
322pub struct Proxy<U, C>
323where
324    U: Upstreams,
325    C: Client,
326{
327    /// Upstreams list.
328    pub upstreams: U,
329    /// [`Client`] for proxy.
330    pub client: C,
331    /// URL path getter.
332    pub url_path_getter: UrlPartGetter,
333    /// URL query getter.
334    pub url_query_getter: UrlPartGetter,
335    /// Host header getter
336    pub host_header_getter: HostHeaderGetter,
337    /// Flag to enable x-forwarded-for header.
338    pub client_ip_forwarding_enabled: bool,
339    /// Flag to reject ambiguous percent-encoded path characters before proxying.
340    pub strict_path_normalization_enabled: bool,
341    /// Flag to drop the inbound `Authorization` header before forwarding.
342    pub strip_authorization_header_enabled: bool,
343}
344
345impl<U, C> Debug for Proxy<U, C>
346where
347    U: Upstreams,
348    C: Client,
349{
350    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
351        f.debug_struct("Proxy")
352            .field(
353                "client_ip_forwarding_enabled",
354                &self.client_ip_forwarding_enabled,
355            )
356            .field(
357                "strict_path_normalization_enabled",
358                &self.strict_path_normalization_enabled,
359            )
360            .field(
361                "strip_authorization_header_enabled",
362                &self.strip_authorization_header_enabled,
363            )
364            .finish_non_exhaustive()
365    }
366}
367
368impl<U, C> Proxy<U, C>
369where
370    U: Upstreams,
371    U::Error: Into<BoxedError>,
372    C: Client,
373{
374    /// Creates a new `Proxy` with an upstream list.
375    ///
376    /// The default host header getter is [`standard_host_header_getter`], which forwards the
377    /// upstream host and includes a non-default port (RFC 7230 / RFC 9110). To forward only
378    /// the bare host name, configure [`default_host_header_getter`] via
379    /// [`Self::host_header_getter`].
380    #[must_use]
381    pub fn new(upstreams: U, client: C) -> Self {
382        Self {
383            upstreams,
384            client,
385            url_path_getter: Box::new(default_url_path_getter),
386            url_query_getter: Box::new(default_url_query_getter),
387            host_header_getter: Box::new(standard_host_header_getter),
388            client_ip_forwarding_enabled: false,
389            strict_path_normalization_enabled: true,
390            strip_authorization_header_enabled: false,
391        }
392    }
393
394    /// Creates a new `Proxy` with an upstream list and enables the x-forwarded-for header.
395    ///
396    /// Client IP forwarding overwrites any inbound `X-Forwarded-For` value with the direct
397    /// client IP from the connection.
398    pub fn with_client_ip_forwarding(upstreams: U, client: C) -> Self {
399        Self {
400            upstreams,
401            client,
402            url_path_getter: Box::new(default_url_path_getter),
403            url_query_getter: Box::new(default_url_query_getter),
404            host_header_getter: Box::new(standard_host_header_getter),
405            client_ip_forwarding_enabled: true,
406            strict_path_normalization_enabled: true,
407            strip_authorization_header_enabled: false,
408        }
409    }
410
411    /// Set URL path getter.
412    #[inline]
413    #[must_use]
414    pub fn url_path_getter<G>(mut self, url_path_getter: G) -> Self
415    where
416        G: Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static,
417    {
418        self.url_path_getter = Box::new(url_path_getter);
419        self
420    }
421
422    /// Set URL query getter.
423    #[inline]
424    #[must_use]
425    pub fn url_query_getter<G>(mut self, url_query_getter: G) -> Self
426    where
427        G: Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static,
428    {
429        self.url_query_getter = Box::new(url_query_getter);
430        self
431    }
432
433    /// Set host header getter.
434    #[inline]
435    #[must_use]
436    pub fn host_header_getter<G>(mut self, host_header_getter: G) -> Self
437    where
438        G: Fn(&Uri, &Request, &Depot) -> Option<String> + Send + Sync + 'static,
439    {
440        self.host_header_getter = Box::new(host_header_getter);
441        self
442    }
443
444    /// Enable or disable strict path normalization.
445    ///
446    /// When enabled, the proxy rejects literal parent-directory (`..`) path components and paths
447    /// that still contain percent-encoded `.`, `/`, `\`, or `%` characters after Salvo routing has
448    /// extracted the path tail. This is useful when the proxy is used as a security boundary and
449    /// the upstream server may perform another decode or normalization pass.
450    #[inline]
451    #[must_use]
452    pub fn strict_path_normalization(mut self, enable: bool) -> Self {
453        self.strict_path_normalization_enabled = enable;
454        self
455    }
456
457    /// Get upstreams list.
458    #[inline]
459    pub fn upstreams(&self) -> &U {
460        &self.upstreams
461    }
462    /// Get upstreams mutable list.
463    #[inline]
464    pub fn upstreams_mut(&mut self) -> &mut U {
465        &mut self.upstreams
466    }
467
468    /// Get client reference.
469    #[inline]
470    pub fn client(&self) -> &C {
471        &self.client
472    }
473    /// Get client mutable reference.
474    #[inline]
475    pub fn client_mut(&mut self) -> &mut C {
476        &mut self.client
477    }
478
479    /// Enable x-forwarded-for header forwarding.
480    ///
481    /// When enabled, the proxy overwrites any inbound `X-Forwarded-For` value with the direct
482    /// client IP from the connection instead of trusting a client-supplied forwarding chain.
483    #[inline]
484    #[must_use]
485    pub fn client_ip_forwarding(mut self, enable: bool) -> Self {
486        self.client_ip_forwarding_enabled = enable;
487        self
488    }
489
490    /// Drop the inbound `Authorization` header before forwarding the
491    /// request upstream.
492    ///
493    /// Disabled by default to preserve the typical reverse-proxy
494    /// pass-through behaviour. Enable this when the proxy sits in
495    /// front of services that should not see the caller's credentials
496    /// — for example a public CORS gateway, a third-party API
497    /// bridge, or any deployment where the upstream is in a different
498    /// trust domain than the inbound client.
499    #[inline]
500    #[must_use]
501    pub fn strip_authorization_header(mut self, enable: bool) -> Self {
502        self.strip_authorization_header_enabled = enable;
503        self
504    }
505
506    async fn build_proxied_request(
507        &self,
508        req: &mut Request,
509        depot: &Depot,
510    ) -> Result<HyperRequest, Error> {
511        let upstream = self
512            .upstreams
513            .elect(req, depot)
514            .await
515            .map_err(Error::other)?;
516
517        if upstream.is_empty() {
518            tracing::error!("upstreams is empty");
519            return Err(Error::other("upstreams is empty"));
520        }
521
522        let path = (self.url_path_getter)(req, depot).unwrap_or_else(|| {
523            // A `None` from the path getter means "no extra path"; the request is
524            // forwarded to the upstream root. Log it so a misconfigured custom getter
525            // (one that fails to extract the intended segment) is observable instead
526            // of silently proxying to the upstream root.
527            tracing::debug!("url_path_getter returned None; forwarding to upstream root path");
528            String::new()
529        });
530        if self.strict_path_normalization_enabled {
531            if contains_ambiguous_path_escape(&path) {
532                return Err(Error::other("ambiguous percent-encoded path"));
533            }
534            if contains_parent_dir_component(&path) {
535                return Err(Error::other("parent directory path segment"));
536            }
537        }
538        let path = encode_url_path(&normalize_url_path(&path));
539        let query = (self.url_query_getter)(req, depot);
540        let path_and_query = if let Some(query) = query {
541            if let Some(stripped) = query.strip_prefix('?') {
542                format!("{path}?{}", utf8_percent_encode(stripped, QUERY_ENCODE_SET))
543            } else {
544                format!("{path}?{}", utf8_percent_encode(&query, QUERY_ENCODE_SET))
545            }
546        } else {
547            path
548        };
549        let forward_url = if upstream.ends_with('/') && path_and_query.starts_with('/') {
550            format!("{}{}", upstream.trim_end_matches('/'), path_and_query)
551        } else if upstream.ends_with('/') || path_and_query.starts_with('/') {
552            format!("{upstream}{path_and_query}")
553        } else if path_and_query.is_empty() {
554            upstream.to_owned()
555        } else {
556            format!("{upstream}/{path_and_query}")
557        };
558        let forward_url: Uri = TryFrom::try_from(forward_url).map_err(Error::other)?;
559        let mut request_builder = hyper::Request::builder()
560            .method(req.method())
561            .uri(&forward_url);
562        let connection_headers = connection_header_names(req.headers());
563        let upgrade_type = get_upgrade_type(req.headers()).map(str::to_owned);
564        for (key, value) in req.headers() {
565            if key == HOST || is_hop_by_hop_header(key, &connection_headers) {
566                continue;
567            }
568            if self.strip_authorization_header_enabled && key == AUTHORIZATION {
569                continue;
570            }
571            request_builder = request_builder.header(key, value);
572        }
573        if let Some(upgrade_type) = upgrade_type {
574            request_builder =
575                request_builder.header(CONNECTION, HeaderValue::from_static("upgrade"));
576            match HeaderValue::from_str(&upgrade_type) {
577                Ok(upgrade_type) => {
578                    request_builder = request_builder.header(UPGRADE, upgrade_type);
579                }
580                Err(e) => {
581                    tracing::error!(error = ?e, "invalid upgrade header value");
582                }
583            }
584        }
585        if let Some(host_value) = (self.host_header_getter)(&forward_url, req, depot) {
586            match HeaderValue::from_str(&host_value) {
587                Ok(host_value) => {
588                    request_builder = request_builder.header(HOST, host_value);
589                }
590                Err(e) => {
591                    tracing::error!(error = ?e, "invalid host header value");
592                }
593            }
594        }
595
596        if self.client_ip_forwarding_enabled {
597            let xff_header_name = HeaderName::from_static(X_FORWARDED_FOR_HEADER_NAME);
598            if let Some(client_ip) = req.remote_addr().ip() {
599                match HeaderValue::from_str(&client_ip.to_string()) {
600                    Ok(xff) => {
601                        if let Some(headers) = request_builder.headers_mut() {
602                            headers.insert(&xff_header_name, xff);
603                        }
604                    }
605                    Err(e) => {
606                        tracing::error!(error = ?e, "invalid x-forwarded-for header value");
607                    }
608                }
609            }
610        }
611
612        request_builder.body(req.take_body()).map_err(Error::other)
613    }
614}
615
616#[async_trait]
617impl<U, C> Handler for Proxy<U, C>
618where
619    U: Upstreams,
620    U::Error: Into<BoxedError>,
621    C: Client,
622{
623    async fn handle(
624        &self,
625        req: &mut Request,
626        depot: &mut Depot,
627        res: &mut Response,
628        _ctrl: &mut FlowCtrl,
629    ) {
630        match self.build_proxied_request(req, depot).await {
631            Ok(proxied_request) => {
632                match self
633                    .client
634                    .execute(proxied_request, req.extensions_mut().remove())
635                    .await
636                {
637                    Ok(response) => {
638                        let (
639                            salvo_core::http::response::Parts {
640                                status,
641                                // version,
642                                headers,
643                                // extensions,
644                                ..
645                            },
646                            body,
647                        ) = response.into_parts();
648                        res.status_code(status);
649                        append_end_to_end_headers(res.headers_mut(), &headers, status);
650                        res.body(body);
651                    }
652                    Err(e) => {
653                        tracing::error!( error = ?e, uri = ?req.uri(), "get response data failed: {}", e);
654                        res.status_code(StatusCode::INTERNAL_SERVER_ERROR);
655                    }
656                }
657            }
658            Err(e) => {
659                tracing::error!(error = ?e, "build proxied request failed");
660                res.status_code(StatusCode::BAD_REQUEST);
661            }
662        }
663    }
664}
665
666fn connection_header_names(headers: &HeaderMap) -> Vec<HeaderName> {
667    headers
668        .get_all(CONNECTION)
669        .iter()
670        .filter_map(|value| value.to_str().ok())
671        .flat_map(|value| value.split(','))
672        .filter_map(|name| HeaderName::from_bytes(name.trim().as_bytes()).ok())
673        .collect()
674}
675
676fn is_hop_by_hop_header(name: &HeaderName, connection_headers: &[HeaderName]) -> bool {
677    HOP_BY_HOP_HEADERS
678        .iter()
679        .any(|hop_header| name.as_str().eq_ignore_ascii_case(hop_header))
680        || connection_headers.iter().any(|header| header == name)
681}
682
683fn append_end_to_end_headers(destination: &mut HeaderMap, source: &HeaderMap, status: StatusCode) {
684    let connection_headers = connection_header_names(source);
685    let upgrade_type = if status == StatusCode::SWITCHING_PROTOCOLS {
686        get_upgrade_type(source).map(str::to_owned)
687    } else {
688        None
689    };
690    for name in source.keys() {
691        if is_hop_by_hop_header(name, &connection_headers) {
692            continue;
693        }
694        for value in source.get_all(name) {
695            destination.append(name, value.to_owned());
696        }
697    }
698    if let Some(upgrade_type) = upgrade_type {
699        destination.append(CONNECTION, HeaderValue::from_static("upgrade"));
700        match HeaderValue::from_str(&upgrade_type) {
701            Ok(upgrade_type) => {
702                destination.append(UPGRADE, upgrade_type);
703            }
704            Err(e) => {
705                tracing::error!(error = ?e, "invalid upgrade header value");
706            }
707        }
708    }
709}
710
711#[inline]
712fn get_upgrade_type(headers: &HeaderMap) -> Option<&str> {
713    if connection_header_names(headers).contains(&UPGRADE)
714        && let Some(upgrade_value) = headers.get(&UPGRADE)
715    {
716        tracing::debug!(
717            "found upgrade header with value: {:?}",
718            upgrade_value.to_str()
719        );
720        return upgrade_value.to_str().ok();
721    }
722
723    None
724}
725
726/// Compare two `Upgrade` token candidates from request and response.
727///
728/// Per RFC 7230 §6.7 the upgrade token is case-insensitive (`websocket`,
729/// `WebSocket`, `WEBSOCKET` are equivalent), so this helper compares with
730/// `eq_ignore_ascii_case`. Two missing tokens compare equal (the previous
731/// case-sensitive check matched too) so callers behave the same when neither
732/// side advertises an upgrade.
733#[inline]
734pub(crate) fn upgrade_types_match(request: Option<&str>, response: Option<&str>) -> bool {
735    match (request, response) {
736        (Some(req), Some(resp)) => req.eq_ignore_ascii_case(resp),
737        (None, None) => true,
738        _ => false,
739    }
740}
741
742// Unit tests for Proxy
743#[cfg(test)]
744mod tests {
745    use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
746    use std::str::FromStr;
747
748    use futures_util::{SinkExt, StreamExt};
749    use salvo_core::conn::{Acceptor, Listener, SocketAddr};
750    use salvo_core::prelude::{Router, Server, StatusError, TcpListener, handler};
751    use salvo_extra::websocket::WebSocketUpgrade;
752    use tokio::io::{AsyncReadExt, AsyncWriteExt};
753    use tokio_tungstenite::tungstenite::Message;
754    use tokio_tungstenite::tungstenite::protocol::Role;
755
756    use super::*;
757
758    #[handler]
759    async fn websocket_echo(req: &mut Request, res: &mut Response) -> Result<(), StatusError> {
760        WebSocketUpgrade::new()
761            .upgrade(req, res, |mut ws| async move {
762                while let Some(message) = ws.recv().await {
763                    let Ok(message) = message else {
764                        return;
765                    };
766                    if ws.send(message).await.is_err() {
767                        return;
768                    }
769                }
770            })
771            .await
772    }
773
774    async fn spawn_server(router: Router) -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) {
775        let acceptor = TcpListener::new("127.0.0.1:0").bind().await;
776        let addr = acceptor.holdings()[0]
777            .local_addr
778            .clone()
779            .into_std()
780            .unwrap();
781        let handle = tokio::spawn(async move {
782            Server::new(acceptor).serve(router).await;
783        });
784        (addr, handle)
785    }
786
787    #[test]
788    fn test_encode_url_path() {
789        let path = "/test/path";
790        let encoded_path = encode_url_path(path);
791        assert_eq!(encoded_path, "/test/path");
792    }
793
794    #[test]
795    fn test_upgrade_types_match_is_case_insensitive() {
796        // RFC 7230 §6.7 — upgrade tokens are case-insensitive.
797        assert!(upgrade_types_match(Some("WebSocket"), Some("websocket")));
798        assert!(upgrade_types_match(Some("WEBSOCKET"), Some("WebSocket")));
799        assert!(upgrade_types_match(Some("h2c"), Some("h2c")));
800        assert!(upgrade_types_match(None, None));
801
802        assert!(!upgrade_types_match(Some("websocket"), Some("h2c")));
803        assert!(!upgrade_types_match(Some("websocket"), None));
804        assert!(!upgrade_types_match(None, Some("websocket")));
805    }
806
807    #[test]
808    fn test_encode_url_path_preserves_segments_and_escapes_unsafe_chars() {
809        // Empty input round-trips.
810        assert_eq!(encode_url_path(""), "");
811
812        // Leading and trailing slashes survive (they become empty segments).
813        assert_eq!(encode_url_path("/"), "/");
814        assert_eq!(encode_url_path("//a//b//"), "//a//b//");
815
816        // Unsafe path characters in `PATH_ENCODE_SET` get percent-encoded per
817        // segment; the slash separator itself is left intact.
818        assert_eq!(encode_url_path("a b/c d"), "a%20b/c%20d");
819        assert_eq!(encode_url_path("a/{b}/c"), "a/%7Bb%7D/c");
820    }
821
822    #[test]
823    fn test_default_url_path_getter_uses_raw_tail() {
824        let mut request = Request::new();
825        request
826            .params_mut()
827            .insert("**rest", "guide/../index.html".to_owned());
828        let depot = Depot::new();
829
830        assert_eq!(
831            default_url_path_getter(&request, &depot).as_deref(),
832            Some("guide/../index.html")
833        );
834    }
835
836    #[test]
837    fn test_contains_ambiguous_path_escape() {
838        assert!(contains_ambiguous_path_escape("%2e%2e/admin"));
839        assert!(contains_ambiguous_path_escape("api%2Fadmin"));
840        assert!(contains_ambiguous_path_escape("api%5cadmin"));
841        assert!(contains_ambiguous_path_escape("%252e%252e/admin"));
842        assert!(!contains_ambiguous_path_escape("guide.v1/index.html"));
843        assert!(!contains_ambiguous_path_escape("files/%20space"));
844    }
845
846    #[test]
847    fn test_contains_parent_dir_component() {
848        assert!(contains_parent_dir_component("../admin"));
849        assert!(contains_parent_dir_component("api/../admin"));
850        assert!(contains_parent_dir_component(r"api\..\admin"));
851        assert!(!contains_parent_dir_component("guide.v1/index.html"));
852        assert!(!contains_parent_dir_component("files/%2e%2e/admin"));
853        assert!(!contains_parent_dir_component("..hidden/admin"));
854    }
855
856    #[test]
857    fn test_get_upgrade_type() {
858        let mut headers = HeaderMap::new();
859        headers.insert(CONNECTION, HeaderValue::from_static("upgrade"));
860        headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
861        let upgrade_type = get_upgrade_type(&headers);
862        assert_eq!(upgrade_type, Some("websocket"));
863    }
864
865    #[test]
866    fn test_get_upgrade_type_checks_all_connection_headers() {
867        let mut headers = HeaderMap::new();
868        headers.append(CONNECTION, HeaderValue::from_static("keep-alive"));
869        headers.append(CONNECTION, HeaderValue::from_static("Upgrade"));
870        headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
871
872        let upgrade_type = get_upgrade_type(&headers);
873
874        assert_eq!(upgrade_type, Some("websocket"));
875    }
876
877    #[test]
878    fn test_connection_header_names() {
879        let mut headers = HeaderMap::new();
880        headers.append(CONNECTION, HeaderValue::from_static("keep-alive, x-remove"));
881        headers.append(CONNECTION, HeaderValue::from_static("x-second"));
882
883        let names = connection_header_names(&headers);
884        assert!(names.contains(&HeaderName::from_static("keep-alive")));
885        assert!(names.contains(&HeaderName::from_static("x-remove")));
886        assert!(names.contains(&HeaderName::from_static("x-second")));
887    }
888
889    #[test]
890    fn test_host_header_handling() {
891        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
892        let uri = Uri::from_str("http://host.tld/test").unwrap();
893        let mut req = Request::new();
894        let depot = Depot::new();
895
896        assert_eq!(
897            default_host_header_getter(&uri, &req, &depot),
898            Some("host.tld".to_owned())
899        );
900
901        let uri_with_port = Uri::from_str("http://host.tld:8080/test").unwrap();
902        assert_eq!(
903            default_host_header_getter(&uri_with_port, &req, &depot),
904            Some("host.tld".to_owned())
905        );
906        assert_eq!(
907            standard_host_header_getter(&uri_with_port, &req, &depot),
908            Some("host.tld:8080".to_owned())
909        );
910
911        let uri_with_http_port = Uri::from_str("http://host.tld:80/test").unwrap();
912        assert_eq!(
913            standard_host_header_getter(&uri_with_http_port, &req, &depot),
914            Some("host.tld".to_owned())
915        );
916
917        let uri_with_https_port = Uri::from_str("https://host.tld:443/test").unwrap();
918        assert_eq!(
919            standard_host_header_getter(&uri_with_https_port, &req, &depot),
920            Some("host.tld".to_owned())
921        );
922
923        let uri_with_non_https_scheme_and_https_port =
924            Uri::from_str("http://host.tld:443/test").unwrap();
925        assert_eq!(
926            standard_host_header_getter(&uri_with_non_https_scheme_and_https_port, &req, &depot),
927            Some("host.tld:443".to_owned())
928        );
929
930        req.headers_mut()
931            .insert(HOST, HeaderValue::from_static("test.host.tld"));
932        assert_eq!(
933            preserve_original_host_header_getter(&uri, &req, &depot),
934            Some("test.host.tld".to_owned())
935        );
936    }
937
938    #[test]
939    fn test_proxy_default_host_header_getter_includes_port() {
940        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
941        // `Proxy::new` now defaults to the standards-compliant getter, which keeps
942        // a non-default upstream port in the forwarded `Host`.
943        let proxy = Proxy::new(vec!["http://host.tld:8080"], HyperClient::default());
944        let uri = Uri::from_str("http://host.tld:8080/test").unwrap();
945        let req = Request::new();
946        let depot = Depot::new();
947        assert_eq!(
948            (proxy.host_header_getter)(&uri, &req, &depot),
949            Some("host.tld:8080".to_owned())
950        );
951    }
952
953    #[tokio::test]
954    async fn test_build_proxied_request_strips_hop_by_hop_headers() {
955        let proxy = Proxy::new(vec!["http://example.com"], HyperClient::default());
956        let mut request = Request::new();
957        let depot = Depot::new();
958
959        request
960            .headers_mut()
961            .insert(HOST, HeaderValue::from_static("client.example"));
962        request
963            .headers_mut()
964            .insert(CONNECTION, HeaderValue::from_static("keep-alive, x-remove"));
965        request.headers_mut().insert(
966            HeaderName::from_static("keep-alive"),
967            HeaderValue::from_static("timeout=5"),
968        );
969        request.headers_mut().insert(
970            HeaderName::from_static("x-remove"),
971            HeaderValue::from_static("secret"),
972        );
973        request.headers_mut().insert(
974            HeaderName::from_static("te"),
975            HeaderValue::from_static("trailers"),
976        );
977        request.headers_mut().insert(
978            HeaderName::from_static("transfer-encoding"),
979            HeaderValue::from_static("chunked"),
980        );
981        request.headers_mut().insert(
982            HeaderName::from_static("x-keep"),
983            HeaderValue::from_static("ok"),
984        );
985
986        let proxied = proxy
987            .build_proxied_request(&mut request, &depot)
988            .await
989            .unwrap();
990
991        assert!(proxied.headers().get(CONNECTION).is_none());
992        assert!(
993            proxied
994                .headers()
995                .get(HeaderName::from_static("keep-alive"))
996                .is_none()
997        );
998        assert!(
999            proxied
1000                .headers()
1001                .get(HeaderName::from_static("x-remove"))
1002                .is_none()
1003        );
1004        assert!(
1005            proxied
1006                .headers()
1007                .get(HeaderName::from_static("te"))
1008                .is_none()
1009        );
1010        assert!(
1011            proxied
1012                .headers()
1013                .get(HeaderName::from_static("transfer-encoding"))
1014                .is_none()
1015        );
1016        assert_eq!(
1017            proxied.headers().get(HeaderName::from_static("x-keep")),
1018            Some(&HeaderValue::from_static("ok"))
1019        );
1020    }
1021
1022    #[tokio::test]
1023    async fn test_build_proxied_request_regenerates_upgrade_headers() {
1024        let proxy = Proxy::new(vec!["http://example.com"], HyperClient::default());
1025        let mut request = Request::new();
1026        let depot = Depot::new();
1027
1028        request
1029            .headers_mut()
1030            .insert(CONNECTION, HeaderValue::from_static("x-remove, Upgrade"));
1031        request
1032            .headers_mut()
1033            .insert(UPGRADE, HeaderValue::from_static("websocket"));
1034        request.headers_mut().insert(
1035            HeaderName::from_static("x-remove"),
1036            HeaderValue::from_static("secret"),
1037        );
1038
1039        let proxied = proxy
1040            .build_proxied_request(&mut request, &depot)
1041            .await
1042            .unwrap();
1043
1044        assert_eq!(
1045            proxied.headers().get(CONNECTION),
1046            Some(&HeaderValue::from_static("upgrade"))
1047        );
1048        assert_eq!(
1049            proxied.headers().get(UPGRADE),
1050            Some(&HeaderValue::from_static("websocket"))
1051        );
1052        assert!(
1053            proxied
1054                .headers()
1055                .get(HeaderName::from_static("x-remove"))
1056                .is_none()
1057        );
1058    }
1059
1060    #[tokio::test]
1061    async fn test_build_proxied_request_forwards_authorization_by_default() {
1062        let proxy = Proxy::new(vec!["http://example.com"], HyperClient::default());
1063        let mut request = Request::new();
1064        let depot = Depot::new();
1065
1066        request
1067            .headers_mut()
1068            .insert(AUTHORIZATION, HeaderValue::from_static("Bearer secret"));
1069
1070        let proxied = proxy
1071            .build_proxied_request(&mut request, &depot)
1072            .await
1073            .unwrap();
1074
1075        assert_eq!(
1076            proxied.headers().get(AUTHORIZATION),
1077            Some(&HeaderValue::from_static("Bearer secret"))
1078        );
1079    }
1080
1081    #[tokio::test]
1082    async fn test_build_proxied_request_strips_authorization_when_enabled() {
1083        let proxy = Proxy::new(vec!["http://example.com"], HyperClient::default())
1084            .strip_authorization_header(true);
1085        let mut request = Request::new();
1086        let depot = Depot::new();
1087
1088        request
1089            .headers_mut()
1090            .insert(AUTHORIZATION, HeaderValue::from_static("Bearer secret"));
1091        request.headers_mut().insert(
1092            HeaderName::from_static("x-keep"),
1093            HeaderValue::from_static("ok"),
1094        );
1095
1096        let proxied = proxy
1097            .build_proxied_request(&mut request, &depot)
1098            .await
1099            .unwrap();
1100
1101        assert!(proxied.headers().get(AUTHORIZATION).is_none());
1102        assert_eq!(
1103            proxied.headers().get(HeaderName::from_static("x-keep")),
1104            Some(&HeaderValue::from_static("ok"))
1105        );
1106    }
1107
1108    #[tokio::test]
1109    async fn test_proxy_websocket_connection_with_split_connection_headers() {
1110        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
1111
1112        let upstream_router = Router::with_path("ws").goal(websocket_echo);
1113        let (upstream_addr, upstream_server) = spawn_server(upstream_router).await;
1114
1115        let proxy_router = Router::with_path("{**rest}").goal(Proxy::new(
1116            vec![format!("http://{upstream_addr}")],
1117            HyperClient::default(),
1118        ));
1119        let (proxy_addr, proxy_server) = spawn_server(proxy_router).await;
1120
1121        let mut stream = tokio::net::TcpStream::connect(proxy_addr).await.unwrap();
1122        let request = format!(
1123            "\
1124GET /ws HTTP/1.1\r\n\
1125Host: {proxy_addr}\r\n\
1126Connection: keep-alive\r\n\
1127Connection: Upgrade\r\n\
1128Upgrade: websocket\r\n\
1129Sec-WebSocket-Version: 13\r\n\
1130Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\
1131\r\n"
1132        );
1133        stream.write_all(request.as_bytes()).await.unwrap();
1134
1135        let mut response = Vec::new();
1136        let mut buffer = [0; 1024];
1137        let header_end = loop {
1138            let read = stream.read(&mut buffer).await.unwrap();
1139            assert_ne!(
1140                read, 0,
1141                "server closed before websocket handshake completed"
1142            );
1143            response.extend_from_slice(&buffer[..read]);
1144            if let Some(position) = response.windows(4).position(|window| window == b"\r\n\r\n") {
1145                break position + 4;
1146            }
1147        };
1148        let extra = response.split_off(header_end);
1149        let response_head = String::from_utf8_lossy(&response);
1150        assert!(
1151            response_head.starts_with("HTTP/1.1 101"),
1152            "unexpected websocket handshake response: {response_head}"
1153        );
1154        let response_head_lower = response_head.to_ascii_lowercase();
1155        assert!(
1156            response_head_lower.contains("\r\nconnection: upgrade\r\n"),
1157            "missing connection upgrade header: {response_head}"
1158        );
1159        assert!(
1160            response_head_lower.contains("\r\nupgrade: websocket\r\n"),
1161            "missing upgrade header: {response_head}"
1162        );
1163
1164        let mut websocket = tokio_tungstenite::WebSocketStream::from_partially_read(
1165            stream,
1166            extra,
1167            Role::Client,
1168            None,
1169        )
1170        .await;
1171
1172        websocket
1173            .send(Message::text("proxied websocket"))
1174            .await
1175            .unwrap();
1176        let echoed = websocket.next().await.unwrap().unwrap();
1177        assert_eq!(echoed.into_text().unwrap(), "proxied websocket");
1178
1179        websocket.close(None).await.unwrap();
1180        proxy_server.abort();
1181        upstream_server.abort();
1182    }
1183
1184    #[tokio::test]
1185    async fn test_client_ip_forwarding() {
1186        let xff_header_name = HeaderName::from_static(X_FORWARDED_FOR_HEADER_NAME);
1187
1188        let mut request = Request::new();
1189        let depot = Depot::new();
1190
1191        // Test functionality not broken
1192        let proxy_without_forwarding =
1193            Proxy::new(vec!["http://example.com"], HyperClient::default());
1194
1195        assert!(!proxy_without_forwarding.client_ip_forwarding_enabled);
1196
1197        let proxy_with_forwarding = proxy_without_forwarding.client_ip_forwarding(true);
1198
1199        assert!(proxy_with_forwarding.client_ip_forwarding_enabled);
1200
1201        let proxy =
1202            Proxy::with_client_ip_forwarding(vec!["http://example.com"], HyperClient::default());
1203        assert!(proxy.client_ip_forwarding_enabled);
1204
1205        match proxy.build_proxied_request(&mut request, &depot).await {
1206            Ok(req) => assert!(req.headers().get(&xff_header_name).is_none()),
1207            _ => panic!("expected Ok"),
1208        }
1209
1210        *request.remote_addr_mut() =
1211            SocketAddr::from(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12345));
1212
1213        match proxy.build_proxied_request(&mut request, &depot).await {
1214            Ok(req) => assert_eq!(
1215                req.headers().get(&xff_header_name),
1216                Some(&HeaderValue::from_static("127.0.0.1"))
1217            ),
1218            _ => panic!("expected Ok"),
1219        }
1220
1221        // Test choosing correct IP version depending on remote address.
1222        *request.remote_addr_mut() =
1223            SocketAddr::from(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 12345, 0, 0));
1224
1225        match proxy.build_proxied_request(&mut request, &depot).await {
1226            Ok(req) => assert_eq!(
1227                req.headers().get(&xff_header_name),
1228                Some(&HeaderValue::from_static("::1"))
1229            ),
1230            _ => panic!("expected Ok"),
1231        }
1232
1233        *request.remote_addr_mut() = SocketAddr::Unknown;
1234
1235        match proxy.build_proxied_request(&mut request, &depot).await {
1236            Ok(req) => assert!(req.headers().get(&xff_header_name).is_none()),
1237            _ => panic!("expected Ok"),
1238        }
1239
1240        // Test incoming XFF is overwritten instead of preserving untrusted client input.
1241        request.headers_mut().insert(
1242            &xff_header_name,
1243            HeaderValue::from_static("10.72.0.1, 127.0.0.1"),
1244        );
1245        *request.remote_addr_mut() =
1246            SocketAddr::from(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12345));
1247
1248        match proxy.build_proxied_request(&mut request, &depot).await {
1249            Ok(req) => assert_eq!(
1250                req.headers().get(&xff_header_name),
1251                Some(&HeaderValue::from_static("127.0.0.1"))
1252            ),
1253            _ => panic!("expected Ok"),
1254        }
1255    }
1256
1257    #[tokio::test]
1258    async fn test_build_proxied_request_rejects_parent_dir_tail_by_default() {
1259        let mut request = Request::new();
1260        request.params_mut().insert("**rest", "../admin".to_owned());
1261        let depot = Depot::new();
1262        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
1263
1264        assert!(
1265            proxy
1266                .build_proxied_request(&mut request, &depot)
1267                .await
1268                .is_err()
1269        );
1270    }
1271
1272    #[tokio::test]
1273    async fn test_build_proxied_request_can_opt_out_of_parent_dir_rejection() {
1274        let mut request = Request::new();
1275        request.params_mut().insert("**rest", "../admin".to_owned());
1276        let depot = Depot::new();
1277        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default())
1278            .strict_path_normalization(false);
1279
1280        let req = proxy
1281            .build_proxied_request(&mut request, &depot)
1282            .await
1283            .unwrap();
1284        assert_eq!(req.uri().to_string(), "http://example.com/api/admin");
1285    }
1286
1287    #[tokio::test]
1288    async fn test_build_proxied_request_normalizes_safe_tail() {
1289        let mut request = Request::new();
1290        request
1291            .params_mut()
1292            .insert("**rest", "guide\\index.html".to_owned());
1293        let depot = Depot::new();
1294        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
1295
1296        let proxied_request = proxy
1297            .build_proxied_request(&mut request, &depot)
1298            .await
1299            .unwrap();
1300        assert_eq!(
1301            proxied_request.uri().to_string(),
1302            "http://example.com/api/guide/index.html"
1303        );
1304    }
1305
1306    #[tokio::test]
1307    async fn test_build_proxied_request_rejects_ambiguous_encoded_tail_by_default() {
1308        let mut request = Request::new();
1309        request
1310            .params_mut()
1311            .insert("**rest", "%2e%2e/secrets/.env".to_owned());
1312        let depot = Depot::new();
1313        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
1314
1315        assert!(
1316            proxy
1317                .build_proxied_request(&mut request, &depot)
1318                .await
1319                .is_err()
1320        );
1321    }
1322
1323    #[tokio::test]
1324    async fn test_build_proxied_request_can_opt_out_of_strict_path_normalization() {
1325        let mut request = Request::new();
1326        request
1327            .params_mut()
1328            .insert("**rest", "%2e%2e/secrets/.env".to_owned());
1329        let depot = Depot::new();
1330        let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default())
1331            .strict_path_normalization(false);
1332
1333        let proxied_request = proxy
1334            .build_proxied_request(&mut request, &depot)
1335            .await
1336            .unwrap();
1337        assert_eq!(
1338            proxied_request.uri().to_string(),
1339            "http://example.com/api/%2e%2e/secrets/.env"
1340        );
1341    }
1342
1343    #[tokio::test]
1344    async fn test_build_proxied_request_strict_path_normalization_rejects_ambiguous_escapes() {
1345        for path in [
1346            "%2e%2e/secrets/.env",
1347            "api%2fadmin",
1348            "api%5cadmin",
1349            "%252e%252e/secrets/.env",
1350        ] {
1351            let mut request = Request::new();
1352            request.params_mut().insert("**rest", path.to_owned());
1353            let depot = Depot::new();
1354            let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
1355
1356            let err = proxy.build_proxied_request(&mut request, &depot).await;
1357            assert!(err.is_err(), "path should be rejected: {path}");
1358        }
1359    }
1360
1361    #[test]
1362    fn test_append_end_to_end_headers_strips_response_hop_by_hop_headers() {
1363        let mut source = HeaderMap::new();
1364        source.insert(CONNECTION, HeaderValue::from_static("x-remove"));
1365        source.insert(UPGRADE, HeaderValue::from_static("websocket"));
1366        source.insert(
1367            HeaderName::from_static("transfer-encoding"),
1368            HeaderValue::from_static("chunked"),
1369        );
1370        source.insert(
1371            HeaderName::from_static("x-remove"),
1372            HeaderValue::from_static("secret"),
1373        );
1374        source.insert(
1375            HeaderName::from_static("x-keep"),
1376            HeaderValue::from_static("ok"),
1377        );
1378
1379        let mut destination = HeaderMap::new();
1380        append_end_to_end_headers(&mut destination, &source, StatusCode::OK);
1381
1382        assert!(destination.get(CONNECTION).is_none());
1383        assert!(destination.get(UPGRADE).is_none());
1384        assert!(
1385            destination
1386                .get(HeaderName::from_static("transfer-encoding"))
1387                .is_none()
1388        );
1389        assert!(
1390            destination
1391                .get(HeaderName::from_static("x-remove"))
1392                .is_none()
1393        );
1394        assert_eq!(
1395            destination.get(HeaderName::from_static("x-keep")),
1396            Some(&HeaderValue::from_static("ok"))
1397        );
1398    }
1399
1400    #[test]
1401    fn test_append_end_to_end_headers_preserves_upgrade_handshake_on_101() {
1402        let mut source = HeaderMap::new();
1403        source.insert(CONNECTION, HeaderValue::from_static("Upgrade, x-remove"));
1404        source.insert(UPGRADE, HeaderValue::from_static("websocket"));
1405        source.insert(
1406            HeaderName::from_static("transfer-encoding"),
1407            HeaderValue::from_static("chunked"),
1408        );
1409        source.insert(
1410            HeaderName::from_static("x-remove"),
1411            HeaderValue::from_static("secret"),
1412        );
1413        source.insert(
1414            HeaderName::from_static("x-keep"),
1415            HeaderValue::from_static("ok"),
1416        );
1417
1418        let mut destination = HeaderMap::new();
1419        append_end_to_end_headers(&mut destination, &source, StatusCode::SWITCHING_PROTOCOLS);
1420
1421        assert_eq!(
1422            destination.get(CONNECTION),
1423            Some(&HeaderValue::from_static("upgrade"))
1424        );
1425        assert_eq!(
1426            destination.get(UPGRADE),
1427            Some(&HeaderValue::from_static("websocket"))
1428        );
1429        assert!(
1430            destination
1431                .get(HeaderName::from_static("transfer-encoding"))
1432                .is_none()
1433        );
1434        assert!(
1435            destination
1436                .get(HeaderName::from_static("x-remove"))
1437                .is_none()
1438        );
1439        assert_eq!(
1440            destination.get(HeaderName::from_static("x-keep")),
1441            Some(&HeaderValue::from_static("ok"))
1442        );
1443    }
1444}
salvo_proxy/lib.rs

salvo_proxy/
lib.rs
