1#![cfg_attr(test, allow(clippy::unwrap_used))]
2#![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";
85
86const QUERY_ENCODE_SET: &AsciiSet = &CONTROLS
87 .add(b' ')
88 .add(b'"')
89 .add(b'#')
90 .add(b'<')
91 .add(b'>')
92 .add(b'`');
93const PATH_ENCODE_SET: &AsciiSet = &QUERY_ENCODE_SET
94 .add(b'?')
95 .add(b'^')
96 .add(b'`')
97 .add(b'{')
98 .add(b'}');
99
100#[inline]
102pub(crate) fn encode_url_path(path: &str) -> String {
103 path.split('/')
104 .map(|s| utf8_percent_encode(s, PATH_ENCODE_SET).to_string())
105 .collect::<Vec<_>>()
106 .join("/")
107}
108
109pub trait Client: Send + Sync + 'static {
114 type Error: StdError + Send + Sync + 'static;
116
117 fn execute(
119 &self,
120 req: HyperRequest,
121 upgraded: Option<OnUpgrade>,
122 ) -> impl Future<Output = Result<HyperResponse, Self::Error>> + Send;
123}
124
125pub trait Upstreams: Send + Sync + 'static {
131 type Error: StdError + Send + Sync + 'static;
133
134 fn elect(
136 &self,
137 req: &Request,
138 depot: &Depot,
139 ) -> impl Future<Output = Result<&str, Self::Error>> + Send;
140}
141impl Upstreams for &'static str {
142 type Error = Infallible;
143
144 async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
145 Ok(*self)
146 }
147}
148impl Upstreams for String {
149 type Error = Infallible;
150 async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
151 Ok(self.as_str())
152 }
153}
154
155impl<const N: usize> Upstreams for [&'static str; N] {
156 type Error = Error;
157 async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
158 if self.is_empty() {
159 return Err(Error::other("upstreams is empty"));
160 }
161 let index = fastrand::usize(..self.len());
162 Ok(self[index])
163 }
164}
165
166impl<T> Upstreams for Vec<T>
167where
168 T: AsRef<str> + Send + Sync + 'static,
169{
170 type Error = Error;
171 async fn elect(&self, _: &Request, _: &Depot) -> Result<&str, Self::Error> {
172 if self.is_empty() {
173 return Err(Error::other("upstreams is empty"));
174 }
175 let index = fastrand::usize(..self.len());
176 Ok(self[index].as_ref())
177 }
178}
179
180pub type UrlPartGetter = Box<dyn Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static>;
182
183pub type HostHeaderGetter =
185 Box<dyn Fn(&Uri, &Request, &Depot) -> Option<String> + Send + Sync + 'static>;
186
187pub fn default_url_path_getter(req: &Request, _depot: &Depot) -> Option<String> {
192 req.params().tail().map(str::to_owned)
193}
194pub fn default_url_query_getter(req: &Request, _depot: &Depot) -> Option<String> {
196 req.uri().query().map(Into::into)
197}
198
199pub fn default_host_header_getter(
201 forward_uri: &Uri,
202 _req: &Request,
203 _depot: &Depot,
204) -> Option<String> {
205 if let Some(host) = forward_uri.host() {
206 return Some(String::from(host));
207 }
208
209 None
210}
211
212pub fn rfc2616_host_header_getter(
215 forward_uri: &Uri,
216 req: &Request,
217 _depot: &Depot,
218) -> Option<String> {
219 let mut parts: Vec<String> = Vec::with_capacity(2);
220
221 if let Some(host) = forward_uri.host() {
222 parts.push(host.to_owned());
223
224 if let Some(scheme) = forward_uri.scheme_str()
225 && let Some(port) = forward_uri.port_u16()
226 && (scheme == "http" && port != 80 || scheme == "https" && port != 443)
227 {
228 parts.push(port.to_string());
229 }
230 }
231
232 if parts.is_empty() {
233 default_host_header_getter(forward_uri, req, _depot)
234 } else {
235 Some(parts.join(":"))
236 }
237}
238
239pub fn preserve_original_host_header_getter(
242 forward_uri: &Uri,
243 req: &Request,
244 _depot: &Depot,
245) -> Option<String> {
246 if let Some(host_header) = req.headers().get(HOST)
247 && let Ok(host) = host_header.to_str()
248 {
249 return Some(host.to_owned());
250 }
251
252 default_host_header_getter(forward_uri, req, _depot)
253}
254
255#[non_exhaustive]
257pub struct Proxy<U, C>
258where
259 U: Upstreams,
260 C: Client,
261{
262 pub upstreams: U,
264 pub client: C,
266 pub url_path_getter: UrlPartGetter,
268 pub url_query_getter: UrlPartGetter,
270 pub host_header_getter: HostHeaderGetter,
272 pub client_ip_forwarding_enabled: bool,
274}
275
276impl<U, C> Debug for Proxy<U, C>
277where
278 U: Upstreams,
279 C: Client,
280{
281 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
282 f.debug_struct("Proxy").finish()
283 }
284}
285
286impl<U, C> Proxy<U, C>
287where
288 U: Upstreams,
289 U::Error: Into<BoxedError>,
290 C: Client,
291{
292 #[must_use]
294 pub fn new(upstreams: U, client: C) -> Self {
295 Self {
296 upstreams,
297 client,
298 url_path_getter: Box::new(default_url_path_getter),
299 url_query_getter: Box::new(default_url_query_getter),
300 host_header_getter: Box::new(default_host_header_getter),
301 client_ip_forwarding_enabled: false,
302 }
303 }
304
305 pub fn with_client_ip_forwarding(upstreams: U, client: C) -> Self {
307 Self {
308 upstreams,
309 client,
310 url_path_getter: Box::new(default_url_path_getter),
311 url_query_getter: Box::new(default_url_query_getter),
312 host_header_getter: Box::new(default_host_header_getter),
313 client_ip_forwarding_enabled: true,
314 }
315 }
316
317 #[inline]
319 #[must_use]
320 pub fn url_path_getter<G>(mut self, url_path_getter: G) -> Self
321 where
322 G: Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static,
323 {
324 self.url_path_getter = Box::new(url_path_getter);
325 self
326 }
327
328 #[inline]
330 #[must_use]
331 pub fn url_query_getter<G>(mut self, url_query_getter: G) -> Self
332 where
333 G: Fn(&Request, &Depot) -> Option<String> + Send + Sync + 'static,
334 {
335 self.url_query_getter = Box::new(url_query_getter);
336 self
337 }
338
339 #[inline]
341 #[must_use]
342 pub fn host_header_getter<G>(mut self, host_header_getter: G) -> Self
343 where
344 G: Fn(&Uri, &Request, &Depot) -> Option<String> + Send + Sync + 'static,
345 {
346 self.host_header_getter = Box::new(host_header_getter);
347 self
348 }
349
350 #[inline]
352 pub fn upstreams(&self) -> &U {
353 &self.upstreams
354 }
355 #[inline]
357 pub fn upstreams_mut(&mut self) -> &mut U {
358 &mut self.upstreams
359 }
360
361 #[inline]
363 pub fn client(&self) -> &C {
364 &self.client
365 }
366 #[inline]
368 pub fn client_mut(&mut self) -> &mut C {
369 &mut self.client
370 }
371
372 #[inline]
374 #[must_use]
375 pub fn client_ip_forwarding(mut self, enable: bool) -> Self {
376 self.client_ip_forwarding_enabled = enable;
377 self
378 }
379
380 async fn build_proxied_request(
381 &self,
382 req: &mut Request,
383 depot: &Depot,
384 ) -> Result<HyperRequest, Error> {
385 let upstream = self
386 .upstreams
387 .elect(req, depot)
388 .await
389 .map_err(Error::other)?;
390
391 if upstream.is_empty() {
392 tracing::error!("upstreams is empty");
393 return Err(Error::other("upstreams is empty"));
394 }
395
396 let path = (self.url_path_getter)(req, depot).unwrap_or_default();
397 let path = encode_url_path(&normalize_url_path(&path));
398 let query = (self.url_query_getter)(req, depot);
399 let rest = if let Some(query) = query {
400 if let Some(stripped) = query.strip_prefix('?') {
401 format!("{path}?{}", utf8_percent_encode(stripped, QUERY_ENCODE_SET))
402 } else {
403 format!("{path}?{}", utf8_percent_encode(&query, QUERY_ENCODE_SET))
404 }
405 } else {
406 path
407 };
408 let forward_url = if upstream.ends_with('/') && rest.starts_with('/') {
409 format!("{}{}", upstream.trim_end_matches('/'), rest)
410 } else if upstream.ends_with('/') || rest.starts_with('/') {
411 format!("{upstream}{rest}")
412 } else if rest.is_empty() {
413 upstream.to_owned()
414 } else {
415 format!("{upstream}/{rest}")
416 };
417 let forward_url: Uri = TryFrom::try_from(forward_url).map_err(Error::other)?;
418 let mut build = hyper::Request::builder()
419 .method(req.method())
420 .uri(&forward_url);
421 for (key, value) in req.headers() {
422 if key != HOST {
423 build = build.header(key, value);
424 }
425 }
426 if let Some(host_value) = (self.host_header_getter)(&forward_url, req, depot) {
427 match HeaderValue::from_str(&host_value) {
428 Ok(host_value) => {
429 build = build.header(HOST, host_value);
430 }
431 Err(e) => {
432 tracing::error!(error = ?e, "invalid host header value");
433 }
434 }
435 }
436
437 if self.client_ip_forwarding_enabled {
438 let xff_header_name = HeaderName::from_static(X_FORWARDER_FOR_HEADER_NAME);
439 let current_xff = req.headers().get(&xff_header_name);
440
441 #[cfg(test)]
442 let system_ip_addr = match req.remote_addr() {
443 SocketAddr::IPv6(_) => Some(IpAddr::from(Ipv6Addr::new(
444 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8,
445 ))),
446 _ => Some(IpAddr::from(Ipv4Addr::new(101, 102, 103, 104))),
447 };
448
449 #[cfg(not(test))]
450 let system_ip_addr = match req.remote_addr() {
451 SocketAddr::IPv6(_) => local_ipv6().ok(),
452 _ => local_ip().ok(),
453 };
454
455 if let Some(system_ip_addr) = system_ip_addr {
456 let forwarded_addr = system_ip_addr.to_string();
457
458 let xff_value = match current_xff {
459 Some(current_xff) => match current_xff.to_str() {
460 Ok(current_xff) => format!("{forwarded_addr}, {current_xff}"),
461 _ => forwarded_addr.clone(),
462 },
463 None => forwarded_addr.clone(),
464 };
465
466 let xff_header_halue = match HeaderValue::from_str(xff_value.as_str()) {
467 Ok(xff_header_halue) => Some(xff_header_halue),
468 Err(_) => match HeaderValue::from_str(forwarded_addr.as_str()) {
469 Ok(xff_header_halue) => Some(xff_header_halue),
470 Err(e) => {
471 tracing::error!(error = ?e, "invalid x-forwarded-for header value");
472 None
473 }
474 },
475 };
476
477 if let Some(xff) = xff_header_halue
478 && let Some(headers) = build.headers_mut()
479 {
480 headers.insert(&xff_header_name, xff);
481 }
482 }
483 }
484
485 build.body(req.take_body()).map_err(Error::other)
486 }
487}
488
489#[async_trait]
490impl<U, C> Handler for Proxy<U, C>
491where
492 U: Upstreams,
493 U::Error: Into<BoxedError>,
494 C: Client,
495{
496 async fn handle(
497 &self,
498 req: &mut Request,
499 depot: &mut Depot,
500 res: &mut Response,
501 _ctrl: &mut FlowCtrl,
502 ) {
503 match self.build_proxied_request(req, depot).await {
504 Ok(proxied_request) => {
505 match self
506 .client
507 .execute(proxied_request, req.extensions_mut().remove())
508 .await
509 {
510 Ok(response) => {
511 let (
512 salvo_core::http::response::Parts {
513 status,
514 headers,
516 ..
518 },
519 body,
520 ) = response.into_parts();
521 res.status_code(status);
522 for name in headers.keys() {
523 for value in headers.get_all(name) {
524 res.headers.append(name, value.to_owned());
525 }
526 }
527 res.body(body);
528 }
529 Err(e) => {
530 tracing::error!( error = ?e, uri = ?req.uri(), "get response data failed: {}", e);
531 res.status_code(StatusCode::INTERNAL_SERVER_ERROR);
532 }
533 }
534 }
535 Err(e) => {
536 tracing::error!(error = ?e, "build proxied request failed");
537 res.status_code(StatusCode::BAD_REQUEST);
538 }
539 }
540 }
541}
542#[inline]
543#[allow(dead_code)]
544fn get_upgrade_type(headers: &HeaderMap) -> Option<&str> {
545 if headers
546 .get(&CONNECTION)
547 .map(|value| {
548 value
549 .to_str()
550 .unwrap_or_default()
551 .split(',')
552 .any(|e| e.trim() == UPGRADE)
553 })
554 .unwrap_or(false)
555 && let Some(upgrade_value) = headers.get(&UPGRADE)
556 {
557 tracing::debug!(
558 "found upgrade header with value: {:?}",
559 upgrade_value.to_str()
560 );
561 return upgrade_value.to_str().ok();
562 }
563
564 None
565}
566
567#[cfg(test)]
569mod tests {
570 use std::net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
571 use std::str::FromStr;
572
573 use super::*;
574
575 #[test]
576 fn test_encode_url_path() {
577 let path = "/test/path";
578 let encoded_path = encode_url_path(path);
579 assert_eq!(encoded_path, "/test/path");
580 }
581
582 #[test]
583 fn test_default_url_path_getter_uses_raw_tail() {
584 let mut request = Request::new();
585 request
586 .params_mut()
587 .insert("**rest", "guide/../index.html".to_owned());
588 let depot = Depot::new();
589
590 assert_eq!(
591 default_url_path_getter(&request, &depot).as_deref(),
592 Some("guide/../index.html")
593 );
594 }
595
596 #[test]
597 fn test_get_upgrade_type() {
598 let mut headers = HeaderMap::new();
599 headers.insert(CONNECTION, HeaderValue::from_static("upgrade"));
600 headers.insert(UPGRADE, HeaderValue::from_static("websocket"));
601 let upgrade_type = get_upgrade_type(&headers);
602 assert_eq!(upgrade_type, Some("websocket"));
603 }
604
605 #[test]
606 fn test_host_header_handling() {
607 let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
608 let uri = Uri::from_str("http://host.tld/test").unwrap();
609 let mut req = Request::new();
610 let depot = Depot::new();
611
612 assert_eq!(
613 default_host_header_getter(&uri, &req, &depot),
614 Some("host.tld".to_owned())
615 );
616
617 let uri_with_port = Uri::from_str("http://host.tld:8080/test").unwrap();
618 assert_eq!(
619 rfc2616_host_header_getter(&uri_with_port, &req, &depot),
620 Some("host.tld:8080".to_owned())
621 );
622
623 let uri_with_http_port = Uri::from_str("http://host.tld:80/test").unwrap();
624 assert_eq!(
625 rfc2616_host_header_getter(&uri_with_http_port, &req, &depot),
626 Some("host.tld".to_owned())
627 );
628
629 let uri_with_https_port = Uri::from_str("https://host.tld:443/test").unwrap();
630 assert_eq!(
631 rfc2616_host_header_getter(&uri_with_https_port, &req, &depot),
632 Some("host.tld".to_owned())
633 );
634
635 let uri_with_non_https_scheme_and_https_port =
636 Uri::from_str("http://host.tld:443/test").unwrap();
637 assert_eq!(
638 rfc2616_host_header_getter(&uri_with_non_https_scheme_and_https_port, &req, &depot),
639 Some("host.tld:443".to_owned())
640 );
641
642 req.headers_mut()
643 .insert(HOST, HeaderValue::from_static("test.host.tld"));
644 assert_eq!(
645 preserve_original_host_header_getter(&uri, &req, &depot),
646 Some("test.host.tld".to_owned())
647 );
648 }
649
650 #[tokio::test]
651 async fn test_client_ip_forwarding() {
652 let xff_header_name = HeaderName::from_static(X_FORWARDER_FOR_HEADER_NAME);
653
654 let mut request = Request::new();
655 let depot = Depot::new();
656
657 let proxy_without_forwarding =
659 Proxy::new(vec!["http://example.com"], HyperClient::default());
660
661 assert!(!proxy_without_forwarding.client_ip_forwarding_enabled);
662
663 let proxy_with_forwarding = proxy_without_forwarding.client_ip_forwarding(true);
664
665 assert!(proxy_with_forwarding.client_ip_forwarding_enabled);
666
667 let proxy =
668 Proxy::with_client_ip_forwarding(vec!["http://example.com"], HyperClient::default());
669 assert!(proxy.client_ip_forwarding_enabled);
670
671 match proxy.build_proxied_request(&mut request, &depot).await {
672 Ok(req) => assert_eq!(
673 req.headers().get(&xff_header_name),
674 Some(&HeaderValue::from_static("101.102.103.104"))
675 ),
676 _ => panic!("expected Ok"),
677 }
678
679 *request.remote_addr_mut() =
681 SocketAddr::from(SocketAddrV6::new(Ipv6Addr::LOCALHOST, 12345, 0, 0));
682
683 match proxy.build_proxied_request(&mut request, &depot).await {
684 Ok(req) => assert_eq!(
685 req.headers().get(&xff_header_name),
686 Some(&HeaderValue::from_static("1:2:3:4:5:6:7:8"))
687 ),
688 _ => panic!("expected Ok"),
689 }
690
691 *request.remote_addr_mut() = SocketAddr::Unknown;
692
693 match proxy.build_proxied_request(&mut request, &depot).await {
694 Ok(req) => assert_eq!(
695 req.headers().get(&xff_header_name),
696 Some(&HeaderValue::from_static("101.102.103.104"))
697 ),
698 _ => panic!("expected Ok"),
699 }
700
701 request.headers_mut().insert(
703 &xff_header_name,
704 HeaderValue::from_static("10.72.0.1, 127.0.0.1"),
705 );
706 *request.remote_addr_mut() =
707 SocketAddr::from(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 12345));
708
709 match proxy.build_proxied_request(&mut request, &depot).await {
710 Ok(req) => assert_eq!(
711 req.headers().get(&xff_header_name),
712 Some(&HeaderValue::from_static(
713 "101.102.103.104, 10.72.0.1, 127.0.0.1"
714 ))
715 ),
716 _ => panic!("expected Ok"),
717 }
718 }
719
720 #[tokio::test]
721 async fn test_build_proxied_request_unsafe_tail() {
722 let mut request = Request::new();
723 request.params_mut().insert("**rest", "../admin".to_owned());
724 let depot = Depot::new();
725 let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
726
727 let req = proxy
728 .build_proxied_request(&mut request, &depot)
729 .await
730 .unwrap();
731 assert_eq!(req.uri().to_string(), "http://example.com/api/admin");
732 }
733
734 #[tokio::test]
735 async fn test_build_proxied_request_normalizes_safe_tail() {
736 let mut request = Request::new();
737 request
738 .params_mut()
739 .insert("**rest", "guide\\index.html".to_owned());
740 let depot = Depot::new();
741 let proxy = Proxy::new(vec!["http://example.com/api"], HyperClient::default());
742
743 let proxied_request = proxy
744 .build_proxied_request(&mut request, &depot)
745 .await
746 .unwrap();
747 assert_eq!(
748 proxied_request.uri().to_string(),
749 "http://example.com/api/guide/index.html"
750 );
751 }
752}