Skip to main content

rama_net/client/proxy/route/
routes.rs

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
55/// Errors produced by every attempted route of an unsuccessful connection.
56///
57/// The failures remain ordered by route preference and each one contains safe
58/// route metadata such as its index, kind, host and port. The aggregate's
59/// [`Display`](fmt::Display) and [`Debug`](fmt::Debug) implementations do not
60/// print route addresses or nested error messages, so they cannot newly expose
61/// proxy credentials. Callers that intentionally need the detailed causes can
62/// inspect [`Self::failures`].
63pub 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    /// Return the attempted route failures in route preference order.
76    pub fn failures(&self) -> &[ConnectionError] {
77        &self.failures
78    }
79
80    /// Consume the aggregate and return its ordered route failures.
81    #[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/// Resolve and materialize the route decision for downstream middleware.
125///
126/// Place every route-selection layer before this one, then place middleware
127/// that consumes [`ProxyRoute`] after it. This is the sole middleware boundary
128/// that needs to understand both route forms. The most recently inserted input
129/// decision wins. An optional configured plan supplies defaults unless
130/// [`Self::with_overwrite`] makes it authoritative. The selected plan publishes
131/// its first route (or direct for an empty plan) to middleware and remains
132/// authoritative for [`ProxyRoutesConnector`], which still owns fallback
133/// attempts.
134///
135/// Credentials and route-specific extensions are omitted from the middleware
136/// view of a multi-route plan. One HTTP header cannot safely authenticate every
137/// fallback and could leak to a later direct route, while publishing the first
138/// route's extensions would contaminate later attempts. The connector retains
139/// the original per-route state and applies it to isolated attempts.
140#[derive(Debug, Clone, Default)]
141#[non_exhaustive]
142pub struct ProxyRoutesLayer {
143    routes: Option<Arc<ProxyRoutes>>,
144    overwrite: bool,
145}
146
147impl ProxyRoutesLayer {
148    /// Create a route materialization layer without configured defaults.
149    #[must_use]
150    pub const fn new() -> Self {
151        Self {
152            routes: None,
153            overwrite: false,
154        }
155    }
156
157    /// Create a route materialization layer with an ordered default plan.
158    ///
159    /// An input [`ProxyRoute`] or `ProxyRoutes` decision takes precedence by
160    /// default. Use [`Self::with_overwrite`] when this plan is authoritative.
161    #[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        /// Let the configured route plan take precedence over an input route
171        /// decision.
172        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/// Service produced by [`ProxyRoutesLayer`].
200#[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    /// Create a service that materializes route plans for middleware.
209    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                // Keep the complete plan authoritative after publishing its
260                // middleware view as a singular route.
261                extensions.insert_arc(routes);
262            }
263            Some(ProxyRouteSelection::Route(_)) | None => {}
264        }
265        self.inner.serve(input)
266    }
267}
268
269/// Try ordered proxy routes until a connection is established.
270///
271/// Every route receives an isolated [`Fork`] of the original input with the
272/// selected [`ProxyRoute`] inserted into its extensions. A transport-domain
273/// failure advances to the next route only when its kind indicates that another
274/// route can plausibly help. Authentication, invalid-input and internal failures
275/// stop fallback even when reported by the transport domain. Application, local
276/// and unclassified failures also stop immediately because another transport
277/// route should not normally change their outcome. If multiple attempted routes
278/// fail, their contextualized errors are retained in a
279/// [`ProxyRouteConnectError`].
280///
281/// Input route decisions use extension insertion order: the most recently
282/// inserted [`ProxyRoute`] or [`ProxyRoutes`] wins. Configure default or
283/// authoritative plans on [`ProxyRoutesLayer`], before route-aware middleware,
284/// so every consumer observes the same selected route.
285#[derive(Debug, Clone)]
286#[non_exhaustive]
287pub struct ProxyRoutesConnector<S> {
288    inner: S,
289    timeout: Option<Duration>,
290}
291
292impl<S> ProxyRoutesConnector<S> {
293    /// Create a connector that reads routes from the input extensions.
294    #[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        /// Limit the complete ordered-route operation to `timeout`.
304        ///
305        /// This is distinct from a timeout applied to the inner connector: an
306        /// inner timeout applies to one route and can advance to the next route,
307        /// while this budget covers every route and stops the operation when it is
308        /// exhausted. Failures completed before the budget expired remain available
309        /// through [`ProxyRouteConnectError`].
310        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/// Layer that tries ordered proxy routes while establishing a connection.
465#[derive(Debug, Clone, Default)]
466#[non_exhaustive]
467pub struct ProxyRoutesConnectorLayer {
468    timeout: Option<Duration>,
469}
470
471impl ProxyRoutesConnectorLayer {
472    /// Create a layer that reads routes from input extensions.
473    #[must_use]
474    pub const fn new() -> Self {
475        Self { timeout: None }
476    }
477
478    generate_set_and_with! {
479        /// Limit the complete ordered-route operation to `timeout` while retaining
480        /// any route failures completed before the budget expires.
481        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}