1use crate::client::{
2 ConnectionError, ConnectionErrorDomain, ConnectionErrorKind, ConnectorService,
3 EstablishedClientConnection,
4};
5use crate::std::sync::Arc;
6use core::{fmt, time::Duration};
7use rama_core::{
8 Fork, Layer, Service,
9 error::{BoxError, BoxErrorExt as _},
10 extensions::{ExtensionsRef, FromExtensions},
11 telemetry::tracing,
12};
13use rama_utils::macros::{define_inner_service_accessors, generate_set_and_with};
14use tokio::time::Instant;
15
16use super::{ProxyRoute, ProxyRouteIndex, ProxyRoutes};
17
18const DIRECT_PROXY_ROUTES: [ProxyRoute; 1] = [ProxyRoute::Direct];
19
20fn routes_or_direct(routes: &[ProxyRoute]) -> &[ProxyRoute] {
21 if routes.is_empty() {
22 &DIRECT_PROXY_ROUTES
23 } else {
24 routes
25 }
26}
27
28fn route_error_context(
29 error: ConnectionError,
30 route: &ProxyRoute,
31 index: usize,
32) -> ConnectionError {
33 let error = error.context_field("proxy_route_index", index);
34 match route {
35 ProxyRoute::Direct => error.context_field("proxy_route", "DIRECT"),
36 ProxyRoute::Proxy(proxy) => error
37 .context_field("proxy_route", "PROXY")
38 .context_field("proxy_host", proxy.address.host.clone())
39 .context_field("proxy_port", proxy.address.port),
40 }
41}
42
43fn should_try_next_route(error: &ConnectionError) -> bool {
44 error.domain() == ConnectionErrorDomain::Transport
45 && matches!(
46 error.kind(),
47 ConnectionErrorKind::Unavailable
48 | ConnectionErrorKind::Timeout
49 | ConnectionErrorKind::Rejected
50 | ConnectionErrorKind::Protocol
51 | ConnectionErrorKind::Other
52 )
53}
54
55pub struct ProxyRouteConnectError {
64 failures: Box<[ConnectionError]>,
65}
66
67impl ProxyRouteConnectError {
68 fn new(failures: Vec<ConnectionError>) -> Self {
69 debug_assert!(failures.len() > 1);
70 Self {
71 failures: failures.into_boxed_slice(),
72 }
73 }
74
75 pub fn failures(&self) -> &[ConnectionError] {
77 &self.failures
78 }
79
80 #[must_use]
82 pub fn into_failures(self) -> Box<[ConnectionError]> {
83 self.failures
84 }
85}
86
87impl fmt::Debug for ProxyRouteConnectError {
88 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
89 f.debug_struct("ProxyRouteConnectError")
90 .field("failure_count", &self.failures.len())
91 .field(
92 "final_classification",
93 &self
94 .failures
95 .last()
96 .map(|error| (error.domain(), error.kind())),
97 )
98 .finish()
99 }
100}
101
102impl fmt::Display for ProxyRouteConnectError {
103 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
104 write!(
105 f,
106 "all {} attempted proxy routes failed",
107 self.failures.len()
108 )
109 }
110}
111
112impl core::error::Error for ProxyRouteConnectError {
113 fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
114 self.failures.last().map(|error| error as _)
115 }
116}
117
118#[derive(FromExtensions)]
119enum ProxyRouteSelection {
120 Route(Arc<ProxyRoute>),
121 Routes(Arc<ProxyRoutes>),
122}
123
124#[derive(Debug, Clone, Default)]
141#[non_exhaustive]
142pub struct ProxyRoutesLayer {
143 routes: Option<Arc<ProxyRoutes>>,
144 overwrite: bool,
145}
146
147impl ProxyRoutesLayer {
148 #[must_use]
150 pub const fn new() -> Self {
151 Self {
152 routes: None,
153 overwrite: false,
154 }
155 }
156
157 #[must_use]
162 pub fn with_routes(routes: impl Into<ProxyRoutes>) -> Self {
163 Self {
164 routes: Some(Arc::new(routes.into())),
165 overwrite: false,
166 }
167 }
168
169 generate_set_and_with! {
170 pub const fn overwrite(mut self, overwrite: bool) -> Self {
173 self.overwrite = overwrite;
174 self
175 }
176 }
177}
178
179impl<S> Layer<S> for ProxyRoutesLayer {
180 type Service = ProxyRoutesService<S>;
181
182 fn layer(&self, inner: S) -> Self::Service {
183 ProxyRoutesService {
184 inner,
185 routes: self.routes.clone(),
186 overwrite: self.overwrite,
187 }
188 }
189
190 fn into_layer(self, inner: S) -> Self::Service {
191 ProxyRoutesService {
192 inner,
193 routes: self.routes,
194 overwrite: self.overwrite,
195 }
196 }
197}
198
199#[derive(Debug, Clone)]
201pub struct ProxyRoutesService<S> {
202 inner: S,
203 routes: Option<Arc<ProxyRoutes>>,
204 overwrite: bool,
205}
206
207impl<S> ProxyRoutesService<S> {
208 pub const fn new(inner: S) -> Self {
210 Self {
211 inner,
212 routes: None,
213 overwrite: false,
214 }
215 }
216
217 define_inner_service_accessors!();
218}
219
220impl<S, Input> Service<Input> for ProxyRoutesService<S>
221where
222 S: Service<Input>,
223 Input: ExtensionsRef + Send + 'static,
224{
225 type Output = S::Output;
226 type Error = S::Error;
227
228 fn serve(
229 &self,
230 input: Input,
231 ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send + '_ {
232 let extensions = input.extensions();
233 let selected = if self.overwrite {
234 self.routes
235 .clone()
236 .map(ProxyRouteSelection::Routes)
237 .or_else(|| ProxyRouteSelection::from_extensions(extensions))
238 } else {
239 ProxyRouteSelection::from_extensions(extensions)
240 .or_else(|| self.routes.clone().map(ProxyRouteSelection::Routes))
241 };
242 match selected {
243 Some(ProxyRouteSelection::Routes(routes)) => {
244 let singular = routes.as_slice().len() == 1;
245 let mut route = routes
246 .as_slice()
247 .first()
248 .cloned()
249 .unwrap_or(ProxyRoute::Direct);
250 if singular {
251 if let Some(route_extensions) = routes.route_extensions(0) {
252 extensions.extend(route_extensions);
253 }
254 } else if let ProxyRoute::Proxy(address) = &mut route {
255 address.credential = None;
256 }
257 extensions.insert(route);
258
259 extensions.insert_arc(routes);
262 }
263 Some(ProxyRouteSelection::Route(_)) | None => {}
264 }
265 self.inner.serve(input)
266 }
267}
268
269#[derive(Debug, Clone)]
286#[non_exhaustive]
287pub struct ProxyRoutesConnector<S> {
288 inner: S,
289 timeout: Option<Duration>,
290}
291
292impl<S> ProxyRoutesConnector<S> {
293 #[must_use]
295 pub const fn new(inner: S) -> Self {
296 Self {
297 inner,
298 timeout: None,
299 }
300 }
301
302 generate_set_and_with! {
303 pub fn timeout(mut self, timeout: Duration) -> Self {
311 self.timeout = Some(timeout);
312 self
313 }
314 }
315
316 define_inner_service_accessors!();
317
318 async fn connect_routes<Input>(
319 &self,
320 input: Input,
321 routes: &[ProxyRoute],
322 route_contexts: Option<&ProxyRoutes>,
323 deadline: Option<Instant>,
324 ) -> Result<EstablishedClientConnection<S::Connection, Input>, ConnectionError>
325 where
326 S: ConnectorService<Input>,
327 Input: Fork + ExtensionsRef + Send + 'static,
328 {
329 let mut failures = Vec::new();
330 for (index, route) in routes.iter().enumerate() {
331 let attempt = input.fork();
332 if let Some(extensions) =
333 route_contexts.and_then(|contexts| contexts.route_extensions(index))
334 {
335 attempt.extensions().extend(extensions);
336 }
337 attempt.extensions().insert(route.clone());
338 attempt.extensions().insert(ProxyRouteIndex::new(index));
339
340 let result = match deadline {
341 Some(deadline) => {
342 match tokio::time::timeout_at(deadline, self.inner.connect(attempt)).await {
343 Ok(result) => result,
344 Err(error) => {
345 let error = route_error_context(
346 ConnectionError::local(error, ConnectionErrorKind::Timeout)
347 .context("proxy route connector: overall timeout"),
348 route,
349 index,
350 );
351 if failures.is_empty() {
352 return Err(error);
353 }
354
355 failures.push(error);
356 return Err(ConnectionError::new(
357 ProxyRouteConnectError::new(failures),
358 ConnectionErrorDomain::Local,
359 ConnectionErrorKind::Timeout,
360 ));
361 }
362 }
363 }
364 None => self.inner.connect(attempt).await,
365 };
366
367 match result {
368 Ok(established) => return Ok(established),
369 Err(error) => {
370 let error = route_error_context(error, route, index);
371 let try_next = should_try_next_route(&error) && index + 1 < routes.len();
372
373 if try_next {
374 match route {
375 ProxyRoute::Direct => tracing::debug!(
376 route.index = index,
377 route.kind = "direct",
378 error = ?error,
379 "proxy route failed; trying next route",
380 ),
381 ProxyRoute::Proxy(proxy) => tracing::debug!(
382 route.index = index,
383 route.kind = "proxy",
384 server.address = %proxy.address.host,
385 server.port = proxy.address.port,
386 error = ?error,
387 "proxy route failed; trying next route",
388 ),
389 }
390 failures.push(error);
391 continue;
392 }
393
394 if failures.is_empty() {
395 return Err(error);
396 }
397
398 let domain = error.domain();
399 let kind = error.kind();
400 failures.push(error);
401 return Err(ConnectionError::new(
402 ProxyRouteConnectError::new(failures),
403 domain,
404 kind,
405 ));
406 }
407 }
408 }
409
410 Err(ConnectionError::local(
411 BoxError::from_static_str("proxy route resolution produced no attempts"),
412 ConnectionErrorKind::Internal,
413 ))
414 }
415
416 async fn connect_routes_with_timeout<Input>(
417 &self,
418 input: Input,
419 routes: &[ProxyRoute],
420 route_contexts: Option<&ProxyRoutes>,
421 ) -> Result<EstablishedClientConnection<S::Connection, Input>, ConnectionError>
422 where
423 S: ConnectorService<Input>,
424 Input: Fork + ExtensionsRef + Send + 'static,
425 {
426 let deadline = self.timeout.map(|timeout| Instant::now() + timeout);
427 self.connect_routes(input, routes, route_contexts, deadline)
428 .await
429 }
430}
431
432impl<S, Input> Service<Input> for ProxyRoutesConnector<S>
433where
434 S: ConnectorService<Input>,
435 Input: Fork + ExtensionsRef + Send + 'static,
436{
437 type Output = EstablishedClientConnection<S::Connection, Input>;
438 type Error = ConnectionError;
439
440 async fn serve(&self, input: Input) -> Result<Self::Output, Self::Error> {
441 match ProxyRouteSelection::from_extensions(input.extensions()) {
442 Some(ProxyRouteSelection::Route(route)) => {
443 return self
444 .connect_routes_with_timeout(input, core::slice::from_ref(route.as_ref()), None)
445 .await;
446 }
447 Some(ProxyRouteSelection::Routes(routes)) => {
448 return self
449 .connect_routes_with_timeout(
450 input,
451 routes_or_direct(routes.as_slice()),
452 Some(routes.as_ref()),
453 )
454 .await;
455 }
456 None => {}
457 }
458
459 self.connect_routes_with_timeout(input, &DIRECT_PROXY_ROUTES, None)
460 .await
461 }
462}
463
464#[derive(Debug, Clone, Default)]
466#[non_exhaustive]
467pub struct ProxyRoutesConnectorLayer {
468 timeout: Option<Duration>,
469}
470
471impl ProxyRoutesConnectorLayer {
472 #[must_use]
474 pub const fn new() -> Self {
475 Self { timeout: None }
476 }
477
478 generate_set_and_with! {
479 pub const fn timeout(mut self, timeout: Duration) -> Self {
482 self.timeout = Some(timeout);
483 self
484 }
485 }
486}
487
488impl<S> Layer<S> for ProxyRoutesConnectorLayer {
489 type Service = ProxyRoutesConnector<S>;
490
491 fn layer(&self, inner: S) -> Self::Service {
492 ProxyRoutesConnector {
493 inner,
494 timeout: self.timeout,
495 }
496 }
497
498 fn into_layer(self, inner: S) -> Self::Service {
499 ProxyRoutesConnector {
500 inner,
501 timeout: self.timeout,
502 }
503 }
504}
505
506#[cfg(test)]
507mod tests {
508 use core::{
509 convert::Infallible,
510 sync::atomic::{AtomicUsize, Ordering},
511 };
512
513 use parking_lot::Mutex;
514 use rama_core::{
515 ServiceInput,
516 error::BoxError,
517 extensions::{Extension, Extensions},
518 layer::TimeoutLayer,
519 service::service_fn,
520 };
521
522 use crate::{
523 address::{HostWithPort, ProxyAddress},
524 client::{ConnectRequest, ConnectionErrorKind},
525 };
526
527 use super::*;
528
529 fn proxy(name: &str) -> ProxyRoute {
530 ProxyRoute::Proxy(
531 format!("http://{name}.example:8080")
532 .parse::<ProxyAddress>()
533 .unwrap(),
534 )
535 }
536
537 fn route_name(route: &ProxyRoute) -> String {
538 match route {
539 ProxyRoute::Direct => "DIRECT".to_owned(),
540 ProxyRoute::Proxy(address) => address.address.host.to_string(),
541 }
542 }
543
544 #[tokio::test]
545 async fn route_layer_materializes_plans_without_replacing_singular_routes() {
546 let service =
547 ProxyRoutesLayer::new().into_layer(service_fn(|input: ServiceInput<()>| async move {
548 Ok::<_, Infallible>(input)
549 }));
550
551 let singular = ServiceInput::new(());
552 singular.extensions.insert(proxy("selected"));
553 let singular = service.serve(singular).await.unwrap();
554 assert_eq!(
555 route_name(singular.extensions.get_ref::<ProxyRoute>().unwrap()),
556 "selected.example"
557 );
558 assert!(!singular.extensions.contains::<ProxyRoutes>());
559
560 let authoritative = ServiceInput::new(());
561 authoritative.extensions.insert(proxy("stale"));
562 authoritative
563 .extensions
564 .insert(ProxyRoutes::new([proxy("primary"), ProxyRoute::Direct]));
565 let authoritative = service.serve(authoritative).await.unwrap();
566 assert_eq!(
567 route_name(authoritative.extensions.get_ref::<ProxyRoute>().unwrap()),
568 "primary.example"
569 );
570 assert_eq!(
571 authoritative
572 .extensions
573 .get_ref::<ProxyRoutes>()
574 .unwrap()
575 .as_slice()
576 .len(),
577 2
578 );
579
580 let route_extensions = Extensions::new();
581 route_extensions.insert(RoutePreference("singleton"));
582 let singleton = ServiceInput::new(());
583 singleton.extensions.insert(
584 [(proxy("only"), route_extensions)]
585 .into_iter()
586 .collect::<ProxyRoutes>(),
587 );
588 let singleton = service.serve(singleton).await.unwrap();
589 assert_eq!(
590 singleton.extensions.get_ref::<RoutePreference>(),
591 Some(&RoutePreference("singleton"))
592 );
593 }
594
595 #[derive(Debug, Clone, PartialEq, Eq, Extension)]
596 struct RoutePreference(&'static str);
597
598 #[tokio::test]
599 async fn retries_transport_failures_in_order() {
600 let attempts = Arc::new(Mutex::new(Vec::new()));
601 let inner = service_fn({
602 let attempts = attempts.clone();
603 move |input: ConnectRequest| {
604 let attempts = attempts.clone();
605 async move {
606 let route = input.extensions.get_ref::<ProxyRoute>().unwrap();
607 attempts.lock().push(route_name(route));
608 if attempts.lock().len() < 3 {
609 Err(ConnectionError::transport(
610 BoxError::from_static_str("route unavailable"),
611 ConnectionErrorKind::Unavailable,
612 ))
613 } else {
614 Ok(EstablishedClientConnection {
615 input,
616 conn: ServiceInput::new(()),
617 })
618 }
619 }
620 }
621 });
622 let connector = ProxyRoutesConnector::new(inner);
623 let input = ConnectRequest::new(HostWithPort::example_domain_https());
624 input
625 .extensions
626 .insert(ProxyRoutes::new([proxy("a"), proxy("b"), proxy("c")]));
627
628 let established = connector.serve(input).await.unwrap();
629 assert_eq!(
630 attempts.lock().as_slice(),
631 ["a.example", "b.example", "c.example"]
632 );
633 assert_eq!(
634 route_name(
635 established
636 .input
637 .extensions
638 .get_ref::<ProxyRoute>()
639 .unwrap()
640 ),
641 "c.example"
642 );
643 assert_eq!(
644 established
645 .input
646 .extensions
647 .get_ref::<ProxyRouteIndex>()
648 .copied()
649 .map(ProxyRouteIndex::get),
650 Some(2)
651 );
652 }
653
654 #[tokio::test]
655 async fn route_layer_keeps_extensions_isolated_per_attempt() {
656 let attempts = Arc::new(AtomicUsize::new(0));
657 let inner = service_fn({
658 let attempts = attempts.clone();
659 move |input: ConnectRequest| {
660 let attempts = attempts.clone();
661 async move {
662 let attempt = attempts.fetch_add(1, Ordering::SeqCst);
663 let preference = input
664 .extensions
665 .get_ref::<RoutePreference>()
666 .expect("route preference");
667 assert_eq!(preference.0, if attempt == 0 { "first" } else { "second" });
668 assert_eq!(
669 input
670 .extensions
671 .get_ref::<ProxyRoute>()
672 .and_then(ProxyRoute::proxy_address)
673 .map(|address| address.address.host.to_string()),
674 Some(if attempt == 0 {
675 "first.example".to_owned()
676 } else {
677 "second.example".to_owned()
678 })
679 );
680 assert_eq!(
681 input
682 .extensions
683 .get_ref::<ProxyRouteIndex>()
684 .copied()
685 .map(ProxyRouteIndex::get),
686 Some(attempt)
687 );
688
689 if attempt == 0 {
690 Err(ConnectionError::transport(
691 BoxError::from_static_str("first route unavailable"),
692 ConnectionErrorKind::Unavailable,
693 ))
694 } else {
695 Ok(EstablishedClientConnection {
696 input,
697 conn: ServiceInput::new(()),
698 })
699 }
700 }
701 }
702 });
703 let first_extensions = Extensions::new();
704 first_extensions.insert(RoutePreference("first"));
705 first_extensions.insert(proxy("hidden-first"));
706 first_extensions.insert(ProxyRouteIndex::new(99));
707 let second_extensions = Extensions::new();
708 second_extensions.insert(RoutePreference("second"));
709 second_extensions.insert(proxy("hidden-second"));
710 second_extensions.insert(ProxyRouteIndex::new(99));
711 let routes = [
712 (proxy("first"), first_extensions),
713 (proxy("second"), second_extensions),
714 ]
715 .into_iter()
716 .collect::<ProxyRoutes>();
717 let connector = ProxyRoutesLayer::new().into_layer(ProxyRoutesConnector::new(inner));
718 let input = ConnectRequest::new(HostWithPort::example_domain_https());
719 let original_extensions = input.extensions.clone();
720 input.extensions.insert(routes);
721
722 let established = connector.serve(input).await.unwrap();
723
724 assert_eq!(attempts.load(Ordering::SeqCst), 2);
725 assert_eq!(
726 established.input.extensions.get_ref::<RoutePreference>(),
727 Some(&RoutePreference("second"))
728 );
729 assert_eq!(
730 established
731 .input
732 .extensions
733 .iter_ref::<RoutePreference>()
734 .count(),
735 1
736 );
737 assert!(original_extensions.get_ref::<RoutePreference>().is_none());
738 }
739
740 #[tokio::test]
741 async fn retryable_transport_kinds_advance_to_next_route() {
742 for kind in [
743 ConnectionErrorKind::Unavailable,
744 ConnectionErrorKind::Timeout,
745 ConnectionErrorKind::Rejected,
746 ConnectionErrorKind::Protocol,
747 ConnectionErrorKind::Other,
748 ] {
749 let attempts = Arc::new(AtomicUsize::new(0));
750 let inner = service_fn({
751 let attempts = attempts.clone();
752 move |input: ConnectRequest| {
753 let attempts = attempts.clone();
754 async move {
755 if attempts.fetch_add(1, Ordering::SeqCst) == 0 {
756 Err(ConnectionError::transport(
757 BoxError::from_static_str("try another route"),
758 kind,
759 ))
760 } else {
761 Ok(EstablishedClientConnection {
762 input,
763 conn: ServiceInput::new(()),
764 })
765 }
766 }
767 }
768 });
769 let connector = ProxyRoutesConnector::new(inner);
770 let input = ConnectRequest::new(HostWithPort::example_domain_https());
771 input
772 .extensions
773 .insert(ProxyRoutes::new([proxy("a"), proxy("b")]));
774
775 connector.serve(input).await.unwrap();
776 assert_eq!(attempts.load(Ordering::SeqCst), 2, "kind: {kind}");
777 }
778 }
779
780 #[tokio::test(start_paused = true)]
781 async fn timeout_layer_failure_advances_to_next_route() {
782 let attempts = Arc::new(AtomicUsize::new(0));
783 let inner = TimeoutLayer::new(Duration::from_secs(1)).into_layer(service_fn({
784 let attempts = attempts.clone();
785 move |input: ConnectRequest| {
786 let attempts = attempts.clone();
787 async move {
788 if attempts.fetch_add(1, Ordering::SeqCst) == 0 {
789 core::future::pending::<()>().await;
790 }
791 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
792 input,
793 conn: ServiceInput::new(()),
794 })
795 }
796 }
797 }));
798 let connector = ProxyRoutesConnector::new(inner);
799 let input = ConnectRequest::new(HostWithPort::example_domain_https());
800 input
801 .extensions
802 .insert(ProxyRoutes::new([proxy("a"), proxy("b")]));
803
804 connector.serve(input).await.unwrap();
805 assert_eq!(attempts.load(Ordering::SeqCst), 2);
806 }
807
808 #[tokio::test(start_paused = true)]
809 async fn overall_timeout_limits_all_route_attempts() {
810 let attempts = Arc::new(AtomicUsize::new(0));
811 let inner = service_fn({
812 let attempts = attempts.clone();
813 move |_input: ConnectRequest| {
814 let attempts = attempts.clone();
815 async move {
816 attempts.fetch_add(1, Ordering::SeqCst);
817 tokio::time::sleep(Duration::from_secs(10)).await;
818 Err::<EstablishedClientConnection<ServiceInput<()>, _>, _>(
819 ConnectionError::transport(
820 BoxError::from_static_str("route unavailable"),
821 ConnectionErrorKind::Unavailable,
822 ),
823 )
824 }
825 }
826 });
827 let connector = ProxyRoutesConnector::new(inner).with_timeout(Duration::from_secs(15));
828 let input = ConnectRequest::new(HostWithPort::example_domain_https());
829 input
830 .extensions
831 .insert(ProxyRoutes::new([proxy("a"), proxy("b"), proxy("c")]));
832
833 let error = connector.serve(input).await.unwrap_err();
834 assert_eq!(error.domain(), ConnectionErrorDomain::Local);
835 assert_eq!(error.kind(), ConnectionErrorKind::Timeout);
836 assert_eq!(attempts.load(Ordering::SeqCst), 2);
837 let aggregate = error
838 .get_ref()
839 .downcast_ref::<ProxyRouteConnectError>()
840 .unwrap();
841 assert_eq!(aggregate.failures().len(), 2);
842 assert_eq!(
843 aggregate.failures()[0].kind(),
844 ConnectionErrorKind::Unavailable
845 );
846 assert_eq!(aggregate.failures()[1].kind(), ConnectionErrorKind::Timeout);
847 }
848
849 #[tokio::test]
850 async fn application_failure_stops_route_fallback() {
851 let attempts = Arc::new(AtomicUsize::new(0));
852 let inner = service_fn({
853 let attempts = attempts.clone();
854 move |_input: ConnectRequest| {
855 let attempts = attempts.clone();
856 async move {
857 attempts.fetch_add(1, Ordering::SeqCst);
858 Err::<EstablishedClientConnection<ServiceInput<()>, _>, _>(
859 ConnectionError::application(
860 BoxError::from_static_str("origin handshake failed"),
861 ConnectionErrorKind::Protocol,
862 ),
863 )
864 }
865 }
866 });
867 let connector = ProxyRoutesConnector::new(inner);
868 let input = ConnectRequest::new(HostWithPort::example_domain_https());
869 input
870 .extensions
871 .insert(ProxyRoutes::new([proxy("a"), proxy("b")]));
872
873 let error = connector.serve(input).await.unwrap_err();
874 assert_eq!(error.domain(), ConnectionErrorDomain::Application);
875 assert_eq!(attempts.load(Ordering::SeqCst), 1);
876 }
877
878 #[tokio::test]
879 async fn unsafe_failure_classifications_stop_route_fallback() {
880 for (domain, kind) in [
881 (
882 ConnectionErrorDomain::Transport,
883 ConnectionErrorKind::Authentication,
884 ),
885 (
886 ConnectionErrorDomain::Transport,
887 ConnectionErrorKind::InvalidInput,
888 ),
889 (
890 ConnectionErrorDomain::Transport,
891 ConnectionErrorKind::Internal,
892 ),
893 (ConnectionErrorDomain::Local, ConnectionErrorKind::Internal),
894 (ConnectionErrorDomain::Unknown, ConnectionErrorKind::Other),
895 ] {
896 let attempts = Arc::new(AtomicUsize::new(0));
897 let inner = service_fn({
898 let attempts = attempts.clone();
899 move |_input: ConnectRequest| {
900 let attempts = attempts.clone();
901 async move {
902 attempts.fetch_add(1, Ordering::SeqCst);
903 Err::<EstablishedClientConnection<ServiceInput<()>, _>, _>(
904 ConnectionError::new(
905 BoxError::from_static_str("do not retry"),
906 domain,
907 kind,
908 ),
909 )
910 }
911 }
912 });
913 let connector = ProxyRoutesConnector::new(inner);
914 let input = ConnectRequest::new(HostWithPort::example_domain_https());
915 input
916 .extensions
917 .insert(ProxyRoutes::new([proxy("a"), proxy("b")]));
918
919 let error = connector.serve(input).await.unwrap_err();
920 assert_eq!(error.domain(), domain);
921 assert_eq!(error.kind(), kind);
922 assert_eq!(attempts.load(Ordering::SeqCst), 1);
923 }
924 }
925
926 #[tokio::test]
927 async fn exhaustion_retains_ordered_route_failures_without_credentials() {
928 let inner = service_fn(async |_input: ConnectRequest| {
929 Err::<EstablishedClientConnection<ServiceInput<()>, _>, _>(ConnectionError::transport(
930 BoxError::from_static_str("route unavailable"),
931 ConnectionErrorKind::Unavailable,
932 ))
933 });
934 let connector = ProxyRoutesConnector::new(inner);
935 let input = ConnectRequest::new(HostWithPort::example_domain_https());
936 input.extensions.insert(ProxyRoutes::new([
937 ProxyRoute::Proxy("http://alice:first-secret@a.example:8080".parse().unwrap()),
938 ProxyRoute::Proxy("http://bob:second-secret@b.example:8080".parse().unwrap()),
939 ]));
940
941 let error = connector.serve(input).await.unwrap_err();
942 assert_eq!(error.domain(), ConnectionErrorDomain::Transport);
943 assert_eq!(error.kind(), ConnectionErrorKind::Unavailable);
944 assert_eq!(error.to_string(), "all 2 attempted proxy routes failed");
945
946 let aggregate = error
947 .get_ref()
948 .downcast_ref::<ProxyRouteConnectError>()
949 .unwrap();
950 let failures = aggregate.failures();
951 assert_eq!(failures.len(), 2);
952 for (index, host) in ["a.example", "b.example"].into_iter().enumerate() {
953 let message = failures[index].to_string();
954 assert!(
955 message.contains(&format!("proxy_route_index=\"{index}\"")),
956 "{message}"
957 );
958 assert!(
959 message.contains(&format!("proxy_host=\"{host}\"")),
960 "{message}"
961 );
962 assert!(message.contains("proxy_port=\"8080\""), "{message}");
963 }
964
965 let formatted = format!("{aggregate:?} {aggregate}");
966 assert!(formatted.contains("ProxyRouteConnectError"), "{formatted}");
967 assert!(formatted.contains("failure_count: 2"), "{formatted}");
968 assert!(!formatted.contains("first-secret"), "{formatted}");
969 assert!(!formatted.contains("second-secret"), "{formatted}");
970 assert!(!formatted.contains("alice"), "{formatted}");
971 assert!(!formatted.contains("bob"), "{formatted}");
972
973 let final_source = core::error::Error::source(aggregate).unwrap();
974 assert!(
975 final_source.to_string().contains("proxy_route_index=\"1\""),
976 "{final_source}"
977 );
978 }
979
980 #[derive(Debug, Extension)]
981 struct FailedAttemptMarker;
982
983 #[tokio::test]
984 async fn failed_attempt_extensions_do_not_leak() {
985 let attempts = Arc::new(AtomicUsize::new(0));
986 let inner = service_fn({
987 let attempts = attempts.clone();
988 move |input: ConnectRequest| {
989 let attempts = attempts.clone();
990 async move {
991 let index = attempts.fetch_add(1, Ordering::SeqCst);
992 if index == 0 {
993 input.extensions.insert(FailedAttemptMarker);
994 Err(ConnectionError::transport(
995 BoxError::from_static_str("first route failed"),
996 ConnectionErrorKind::Unavailable,
997 ))
998 } else {
999 assert!(!input.extensions.contains::<FailedAttemptMarker>());
1000 Ok(EstablishedClientConnection {
1001 input,
1002 conn: ServiceInput::new(()),
1003 })
1004 }
1005 }
1006 }
1007 });
1008 let connector = ProxyRoutesConnector::new(inner);
1009 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1010 input
1011 .extensions
1012 .insert(ProxyRoutes::new([proxy("a"), proxy("b")]));
1013
1014 connector.serve(input).await.unwrap();
1015 }
1016
1017 #[tokio::test]
1018 async fn concurrent_route_plans_remain_isolated() {
1019 let inner = service_fn(async |input: ConnectRequest| {
1020 let route = route_name(input.extensions.get_ref::<ProxyRoute>().unwrap());
1021 tokio::task::yield_now().await;
1022 if route.contains("first") {
1023 Err(ConnectionError::transport(
1024 BoxError::from_static_str("first route unavailable"),
1025 ConnectionErrorKind::Unavailable,
1026 ))
1027 } else {
1028 Ok(EstablishedClientConnection {
1029 input,
1030 conn: ServiceInput::new(()),
1031 })
1032 }
1033 });
1034 let connector = ProxyRoutesConnector::new(inner);
1035 let first = ConnectRequest::new(HostWithPort::example_domain_https());
1036 first
1037 .extensions
1038 .insert(ProxyRoutes::new([proxy("a-first"), proxy("a-second")]));
1039 let second = ConnectRequest::new(HostWithPort::example_domain_https());
1040 second
1041 .extensions
1042 .insert(ProxyRoutes::new([proxy("b-first"), proxy("b-second")]));
1043
1044 let (first, second) = tokio::join!(connector.serve(first), connector.serve(second));
1045 let first = first.unwrap();
1046 let second = second.unwrap();
1047
1048 assert_eq!(
1049 route_name(first.input.extensions.get_ref::<ProxyRoute>().unwrap()),
1050 "a-second.example"
1051 );
1052 assert_eq!(
1053 route_name(second.input.extensions.get_ref::<ProxyRoute>().unwrap()),
1054 "b-second.example"
1055 );
1056 }
1057
1058 #[tokio::test]
1059 async fn empty_routes_mean_direct() {
1060 let inner = service_fn(async |input: ConnectRequest| {
1061 assert_eq!(
1062 input.extensions.get_ref::<ProxyRoute>(),
1063 Some(&ProxyRoute::Direct)
1064 );
1065 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
1066 input,
1067 conn: ServiceInput::new(()),
1068 })
1069 });
1070 let connector = ProxyRoutesConnector::new(inner);
1071 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1072 input.extensions.insert(ProxyRoutes::default());
1073
1074 connector.serve(input).await.unwrap();
1075 }
1076
1077 #[tokio::test]
1078 async fn singular_route_overrides_context_routes() {
1079 let inner = service_fn(async |input: ConnectRequest| {
1080 assert_eq!(
1081 route_name(input.extensions.get_ref::<ProxyRoute>().unwrap()),
1082 "selected.example"
1083 );
1084 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
1085 input,
1086 conn: ServiceInput::new(()),
1087 })
1088 });
1089 let connector = ProxyRoutesConnector::new(inner);
1090 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1091 input.extensions.insert(ProxyRoutes::from(proxy("planned")));
1092 input.extensions.insert(proxy("selected"));
1093
1094 connector.serve(input).await.unwrap();
1095 }
1096
1097 #[tokio::test]
1098 async fn newer_context_routes_override_a_singular_route() {
1099 let inner = service_fn(async |input: ConnectRequest| {
1100 assert_eq!(
1101 route_name(input.extensions.get_ref::<ProxyRoute>().unwrap()),
1102 "planned.example"
1103 );
1104 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
1105 input,
1106 conn: ServiceInput::new(()),
1107 })
1108 });
1109 let connector = ProxyRoutesConnector::new(inner);
1110 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1111 input.extensions.insert(proxy("selected"));
1112 input.extensions.insert(ProxyRoutes::from(proxy("planned")));
1113
1114 connector.serve(input).await.unwrap();
1115 }
1116
1117 #[tokio::test]
1118 async fn input_plan_overrides_configured_default_routes() {
1119 let inner = service_fn(async |input: ConnectRequest| {
1120 assert_eq!(
1121 route_name(input.extensions.get_ref::<ProxyRoute>().unwrap()),
1122 "context.example"
1123 );
1124 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
1125 input,
1126 conn: ServiceInput::new(()),
1127 })
1128 });
1129 let connector = ProxyRoutesLayer::with_routes(proxy("fixed"))
1130 .into_layer(ProxyRoutesConnector::new(inner));
1131 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1132 input.extensions.insert(ProxyRoutes::from(proxy("context")));
1133
1134 connector.serve(input).await.unwrap();
1135 }
1136
1137 #[tokio::test]
1138 async fn singular_route_overrides_configured_default_routes() {
1139 let inner = service_fn(async |input: ConnectRequest| {
1140 assert_eq!(
1141 route_name(input.extensions.get_ref::<ProxyRoute>().unwrap()),
1142 "selected.example"
1143 );
1144 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
1145 input,
1146 conn: ServiceInput::new(()),
1147 })
1148 });
1149 let connector = ProxyRoutesLayer::with_routes(proxy("fixed"))
1150 .into_layer(ProxyRoutesConnector::new(inner));
1151 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1152 input.extensions.insert(proxy("selected"));
1153
1154 connector.serve(input).await.unwrap();
1155 }
1156
1157 #[tokio::test]
1158 async fn route_layer_can_overwrite_a_singular_route() {
1159 let inner = service_fn(async |input: ConnectRequest| {
1160 assert_eq!(
1161 route_name(input.extensions.get_ref::<ProxyRoute>().unwrap()),
1162 "fixed.example"
1163 );
1164 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
1165 input,
1166 conn: ServiceInput::new(()),
1167 })
1168 });
1169 let connector = ProxyRoutesLayer::with_routes(proxy("fixed"))
1170 .with_overwrite(true)
1171 .into_layer(ProxyRoutesConnector::new(inner));
1172 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1173 input.extensions.insert(proxy("selected"));
1174
1175 connector.serve(input).await.unwrap();
1176 }
1177
1178 #[tokio::test]
1179 async fn configured_routes_are_used_without_an_input_decision() {
1180 let inner = service_fn(async |input: ConnectRequest| {
1181 assert_eq!(
1182 route_name(input.extensions.get_ref::<ProxyRoute>().unwrap()),
1183 "fixed.example"
1184 );
1185 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
1186 input,
1187 conn: ServiceInput::new(()),
1188 })
1189 });
1190 let connector = ProxyRoutesLayer::with_routes(proxy("fixed"))
1191 .into_layer(ProxyRoutesConnector::new(inner));
1192 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1193
1194 connector.serve(input).await.unwrap();
1195 }
1196
1197 #[tokio::test]
1198 async fn authoritative_configured_routes_override_an_input_plan() {
1199 let inner = service_fn(async |input: ConnectRequest| {
1200 assert_eq!(
1201 route_name(input.extensions.get_ref::<ProxyRoute>().unwrap()),
1202 "fixed.example"
1203 );
1204 assert!(input.extensions.get_ref::<RoutePreference>().is_none());
1205 Ok::<_, core::convert::Infallible>(EstablishedClientConnection {
1206 input,
1207 conn: ServiceInput::new(()),
1208 })
1209 });
1210 let connector = ProxyRoutesLayer::with_routes(proxy("fixed"))
1211 .with_overwrite(true)
1212 .into_layer(ProxyRoutesConnector::new(inner));
1213 let input = ConnectRequest::new(HostWithPort::example_domain_https());
1214 let route_extensions = Extensions::new();
1215 route_extensions.insert(RoutePreference("context"));
1216 input.extensions.insert(
1217 [(proxy("context"), route_extensions)]
1218 .into_iter()
1219 .collect::<ProxyRoutes>(),
1220 );
1221
1222 connector.serve(input).await.unwrap();
1223 }
1224}