pg_proto/
intermediary_component.rs

1//! Builder-centred composition of the client-facing and PostgreSQL-facing roles.
2
3use std::{collections::VecDeque, fmt, future::Future, io, pin::Pin};
4
5use tokio::io::{AsyncRead, AsyncWrite};
6
7use crate::{
8    ConnectTarget, NoPipeline, StartupParameters,
9    pipeline::{BackendAction, FrontendAction, FrontendHandling, Pipeline, PipelinePolicy},
10};
11
12/// Required posture for out-of-band cancellation connections.
13///
14/// Forwarding cancellation is implemented by issue #36. Until then the only
15/// safe operational posture is an explicit rejection.
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub enum CancellationPolicy {
18    /// Reject cancellation packets instead of silently routing them.
19    Reject,
20    /// Resolve and forward cancellation using the configured registry.
21    Forward,
22}
23
24/// Disclosure-safe handling for failures after a downstream connection exists.
25#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
26pub enum EstablishmentFailurePolicy {
27    /// Silently close without exposing internal failure details.
28    #[default]
29    Close,
30    /// Send one fixed, non-disclosing PostgreSQL diagnostic and then close.
31    SafeDiagnostic,
32}
33
34fn safe_establishment_diagnostic() -> crate::codec::BackendMessage {
35    crate::codec::BackendMessage::ErrorResponse(crate::codec::DiagnosticResponse {
36        fields: vec![
37            crate::codec::DiagnosticField {
38                code: b'S',
39                value: bytes::Bytes::from_static(b"ERROR"),
40            },
41            crate::codec::DiagnosticField {
42                code: b'M',
43                value: bytes::Bytes::from_static(b"connection establishment failed"),
44            },
45        ],
46    })
47}
48
49/// A destination and upstream key retained independently of startup routing.
50#[derive(Clone, Debug, Eq, PartialEq)]
51pub struct CancellationRoute {
52    target: ConnectTarget,
53    upstream: crate::demux::CancelKey,
54}
55
56impl CancellationRoute {
57    /// Creates a cancellation route.
58    #[must_use]
59    pub const fn new(target: ConnectTarget, upstream: crate::demux::CancelKey) -> Self {
60        Self { target, upstream }
61    }
62    /// Returns the original destination, including application metadata.
63    #[must_use]
64    pub const fn target(&self) -> &ConnectTarget {
65        &self.target
66    }
67    /// Returns the upstream cancellation key.
68    #[must_use]
69    pub const fn upstream_key(&self) -> &crate::demux::CancelKey {
70        &self.upstream
71    }
72}
73
74/// Application-owned concurrent cancellation mapping and key allocator.
75///
76/// Methods take `&self` so implementations can use an application-selected
77/// lock, actor, shared store, or other concurrency mechanism. No global
78/// `Send`, `Sync`, or `'static` requirement is imposed.
79pub trait IntermediaryCancellationRegistry {
80    /// Collision, allocation, or storage failure.
81    type Error;
82    /// Records a live route and returns the proxy key exposed downstream.
83    ///
84    /// # Errors
85    ///
86    /// Returns an application-defined allocation, collision, or storage error.
87    fn register(&self, route: CancellationRoute) -> Result<crate::demux::CancelKey, Self::Error>;
88    /// Resolves a later out-of-band request without consulting startup routing.
89    fn resolve(&self, client: &crate::demux::CancelKey) -> Option<CancellationRoute>;
90    /// Explicitly detaches a live client key.
91    fn detach(&self, client: &crate::demux::CancelKey) -> Option<CancellationRoute>;
92}
93
94/// Marker registry used by explicit cancellation rejection.
95#[derive(Clone, Copy, Debug, Default)]
96pub struct RejectCancellation;
97impl IntermediaryCancellationRegistry for RejectCancellation {
98    type Error = std::convert::Infallible;
99    fn register(&self, _: CancellationRoute) -> Result<crate::demux::CancelKey, Self::Error> {
100        unreachable!()
101    }
102    fn resolve(&self, _: &crate::demux::CancelKey) -> Option<CancellationRoute> {
103        None
104    }
105    fn detach(&self, _: &crate::demux::CancelKey) -> Option<CancellationRoute> {
106        None
107    }
108}
109
110/// Deterministic failure while assembling an intermediary component.
111#[derive(Clone, Copy, Debug, Eq, PartialEq)]
112pub enum IntermediaryBuildError {
113    /// No complete client-facing server role was supplied.
114    MissingServer,
115    /// No complete PostgreSQL-facing client role was supplied.
116    MissingClient,
117    /// No asynchronous startup resolver was supplied.
118    MissingStartupResolver,
119    /// Cancellation behavior was not selected explicitly.
120    MissingCancellationPolicy,
121}
122
123impl fmt::Display for IntermediaryBuildError {
124    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
125        formatter.write_str(match self {
126            Self::MissingServer => "an intermediary server component is required",
127            Self::MissingClient => "an intermediary client component is required",
128            Self::MissingStartupResolver => "an asynchronous startup resolver is required",
129            Self::MissingCancellationPolicy => "an explicit cancellation policy is required",
130        })
131    }
132}
133
134impl std::error::Error for IntermediaryBuildError {}
135
136/// Immutable server-side facts available before authentication begins.
137#[derive(Clone, Copy, Debug)]
138pub struct InitialServerContext<'a, Peer> {
139    peer: &'a Peer,
140    tls: &'a crate::NegotiatedServerTls,
141}
142
143impl<'a, Peer> InitialServerContext<'a, Peer> {
144    pub(crate) const fn new(peer: &'a Peer, tls: &'a crate::NegotiatedServerTls) -> Self {
145        Self { peer, tls }
146    }
147
148    /// Returns application-supplied peer metadata.
149    #[must_use]
150    pub const fn peer(&self) -> &Peer {
151        self.peer
152    }
153
154    /// Returns transport security negotiated on the client-facing side.
155    #[must_use]
156    pub const fn tls(&self) -> &crate::NegotiatedServerTls {
157        self.tls
158    }
159}
160
161/// Required asynchronous startup routing policy.
162pub trait StartupRouteResolver<Peer> {
163    /// Application resolver failure.
164    type Error;
165
166    /// Selects a destination before client-facing authentication begins.
167    fn resolve<'a>(
168        &'a self,
169        startup: StartupParameters,
170        context: InitialServerContext<'a, Peer>,
171    ) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>>;
172}
173
174/// Optional policy that validates or refines a destination after authentication.
175pub trait AuthenticatedRoutePolicy<Peer, Identity> {
176    /// Application policy failure.
177    type Error;
178    /// Validates or refines the startup-selected target using typed identity evidence.
179    fn route<'a>(
180        &'a self,
181        target: ConnectTarget,
182        context: AuthenticatedRouteContext<'a, Peer, Identity>,
183    ) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>>;
184}
185
186/// Borrowed facts passed to authenticated route policy.
187#[derive(Clone, Copy, Debug)]
188pub struct AuthenticatedRouteContext<'a, Peer, Identity> {
189    peer: &'a Peer,
190    identity: &'a Identity,
191}
192
193impl<Peer, Identity> AuthenticatedRouteContext<'_, Peer, Identity> {
194    /// Returns application-supplied peer metadata.
195    #[must_use]
196    pub const fn peer(&self) -> &Peer {
197        self.peer
198    }
199
200    /// Returns independently verified client-facing identity evidence.
201    #[must_use]
202    pub const fn identity(&self) -> &Identity {
203        self.identity
204    }
205}
206
207/// Identity authenticated-route policy.
208#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
209pub struct AllowAuthenticatedRoute;
210
211impl<Peer, Identity> AuthenticatedRoutePolicy<Peer, Identity> for AllowAuthenticatedRoute {
212    type Error = std::convert::Infallible;
213    fn route<'a>(
214        &'a self,
215        target: ConnectTarget,
216        _context: AuthenticatedRouteContext<'a, Peer, Identity>,
217    ) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>> {
218        Box::pin(async move { Ok(target) })
219    }
220}
221
222/// Result of asynchronously intercepting a client-originated message.
223#[derive(Debug, Eq, PartialEq)]
224pub enum FrontendMiddlewareOutput {
225    /// Forward this owned message to PostgreSQL.
226    Forward(crate::codec::FrontendMessage),
227    /// Consume the message without forwarding it or registering a response.
228    Suppress(crate::codec::FrontendMessage),
229    /// Handle the request locally and emit these responses in protocol order.
230    Respond {
231        /// Consumed client request.
232        request: crate::codec::FrontendMessage,
233        /// Responses generated for this request, in emission order.
234        responses: Vec<crate::codec::BackendMessage>,
235    },
236}
237
238/// Result of asynchronously intercepting a PostgreSQL-originated message.
239#[derive(Debug, Eq, PartialEq)]
240pub enum BackendMiddlewareOutput {
241    /// Forward this owned message to the client.
242    Forward(crate::codec::BackendMessage),
243    /// Consume the response after advancing protocol and pipeline state.
244    Suppress(crate::codec::BackendMessage),
245}
246
247/// Asynchronous, fallible middleware at the forwarding boundary.
248pub trait IntermediaryMiddleware<State, ServerContext, ClientContext> {
249    /// Application-defined interception failure.
250    type Error;
251
252    /// Intercepts a client-originated message after server-role middleware and
253    /// before client-role middleware.
254    fn frontend<'a>(
255        &'a mut self,
256        _server: &'a ServerContext,
257        _client: &'a ClientContext,
258        _state: &'a mut State,
259        message: crate::codec::FrontendMessage,
260    ) -> Pin<Box<dyn Future<Output = Result<FrontendMiddlewareOutput, Self::Error>> + 'a>> {
261        Box::pin(async move { Ok(FrontendMiddlewareOutput::Forward(message)) })
262    }
263
264    /// Intercepts a PostgreSQL-originated message after client-role middleware
265    /// and before server-role middleware.
266    fn backend<'a>(
267        &'a mut self,
268        _server: &'a ServerContext,
269        _client: &'a ClientContext,
270        _state: &'a mut State,
271        message: crate::codec::BackendMessage,
272    ) -> Pin<Box<dyn Future<Output = Result<BackendMiddlewareOutput, Self::Error>> + 'a>> {
273        Box::pin(async move { Ok(BackendMiddlewareOutput::Forward(message)) })
274    }
275}
276
277/// Identity forwarding-boundary middleware.
278#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
279pub struct IdentityIntermediaryMiddleware;
280
281impl<State, ServerContext, ClientContext>
282    IntermediaryMiddleware<State, ServerContext, ClientContext> for IdentityIntermediaryMiddleware
283{
284    type Error = std::convert::Infallible;
285}
286
287/// Creates fresh forwarding-boundary middleware for one established pair.
288pub trait IntermediaryMiddlewareFactory<ServerContext, ClientContext> {
289    /// Per-connection boundary handler.
290    type Handler;
291    /// Creates an isolated handler from both distinct role contexts.
292    fn create(&self, server: &ServerContext, client: &ClientContext) -> Self::Handler;
293}
294
295impl<ServerContext, ClientContext, Handler, Factory>
296    IntermediaryMiddlewareFactory<ServerContext, ClientContext> for Factory
297where
298    Factory: Fn(&ServerContext, &ClientContext) -> Handler,
299{
300    type Handler = Handler;
301    fn create(&self, server: &ServerContext, client: &ClientContext) -> Handler {
302        self(server, client)
303    }
304}
305
306impl<ServerContext, ClientContext> IntermediaryMiddlewareFactory<ServerContext, ClientContext>
307    for IdentityIntermediaryMiddleware
308{
309    type Handler = Self;
310    fn create(&self, _server: &ServerContext, _client: &ClientContext) -> Self {
311        *self
312    }
313}
314
315/// A reusable operational intermediary configuration.
316pub struct Intermediary<
317    Server = (),
318    Client = (),
319    Resolver = (),
320    Route = AllowAuthenticatedRoute,
321    Policy = NoPipeline,
322    Boundary = IdentityIntermediaryMiddleware,
323    Cancellation = RejectCancellation,
324> {
325    pub(crate) server: Server,
326    pub(crate) client: Client,
327    pub(crate) resolver: Resolver,
328    pub(crate) route: Route,
329    pub(crate) pipeline: Policy,
330    pub(crate) boundary: Boundary,
331    pub(crate) cancellation: CancellationPolicy,
332    pub(crate) cancellation_registry: Cancellation,
333    pub(crate) failure_policy: EstablishmentFailurePolicy,
334}
335
336impl Intermediary<()> {
337    /// Starts composition of the two complete role configurations.
338    #[must_use]
339    pub fn builder() -> IntermediaryBuilder {
340        IntermediaryBuilder::default()
341    }
342}
343
344impl<S, C, R, A, P, B, K> fmt::Debug for Intermediary<S, C, R, A, P, B, K> {
345    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
346        formatter
347            .debug_struct("Intermediary")
348            .field("server", &"<configured>")
349            .field("client", &"<configured>")
350            .field("resolver", &"<redacted>")
351            .field("authenticated_route", &"<redacted>")
352            .field("cancellation", &self.cancellation)
353            .finish_non_exhaustive()
354    }
355}
356
357/// Progressive builder for [`Intermediary`].
358pub struct IntermediaryBuilder<
359    Server = (),
360    Client = (),
361    Resolver = (),
362    Route = AllowAuthenticatedRoute,
363    Policy = NoPipeline,
364    Boundary = IdentityIntermediaryMiddleware,
365    Cancellation = RejectCancellation,
366> {
367    server: Option<Server>,
368    client: Option<Client>,
369    resolver: Option<Resolver>,
370    route: Route,
371    pipeline: Policy,
372    boundary: Boundary,
373    cancellation: Option<CancellationPolicy>,
374    cancellation_registry: Cancellation,
375    failure_policy: EstablishmentFailurePolicy,
376}
377
378impl Default for IntermediaryBuilder {
379    fn default() -> Self {
380        Self {
381            server: None,
382            client: None,
383            resolver: None,
384            route: AllowAuthenticatedRoute,
385            pipeline: NoPipeline,
386            boundary: IdentityIntermediaryMiddleware,
387            cancellation: None,
388            cancellation_registry: RejectCancellation,
389            failure_policy: EstablishmentFailurePolicy::Close,
390        }
391    }
392}
393
394impl<S, C, R, A, P, B, K> IntermediaryBuilder<S, C, R, A, P, B, K> {
395    /// Supplies the complete client-facing role configuration.
396    #[must_use]
397    pub fn server<Next>(self, server: Next) -> IntermediaryBuilder<Next, C, R, A, P, B, K> {
398        IntermediaryBuilder {
399            server: Some(server),
400            client: self.client,
401            resolver: self.resolver,
402            route: self.route,
403            pipeline: self.pipeline,
404            boundary: self.boundary,
405            cancellation: self.cancellation,
406            cancellation_registry: self.cancellation_registry,
407            failure_policy: self.failure_policy,
408        }
409    }
410
411    /// Supplies the complete PostgreSQL-facing role configuration.
412    #[must_use]
413    pub fn client<Next>(self, client: Next) -> IntermediaryBuilder<S, Next, R, A, P, B, K> {
414        IntermediaryBuilder {
415            server: self.server,
416            client: Some(client),
417            resolver: self.resolver,
418            route: self.route,
419            pipeline: self.pipeline,
420            boundary: self.boundary,
421            cancellation: self.cancellation,
422            cancellation_registry: self.cancellation_registry,
423            failure_policy: self.failure_policy,
424        }
425    }
426
427    /// Supplies the required asynchronous startup resolver.
428    #[must_use]
429    pub fn startup_resolver<Next>(
430        self,
431        resolver: Next,
432    ) -> IntermediaryBuilder<S, C, Next, A, P, B, K> {
433        IntermediaryBuilder {
434            server: self.server,
435            client: self.client,
436            resolver: Some(resolver),
437            route: self.route,
438            pipeline: self.pipeline,
439            boundary: self.boundary,
440            cancellation: self.cancellation,
441            cancellation_registry: self.cancellation_registry,
442            failure_policy: self.failure_policy,
443        }
444    }
445
446    /// Supplies optional post-authentication routing policy.
447    #[must_use]
448    pub fn authenticated_route<Next>(
449        self,
450        route: Next,
451    ) -> IntermediaryBuilder<S, C, R, Next, P, B, K> {
452        IntermediaryBuilder {
453            server: self.server,
454            client: self.client,
455            resolver: self.resolver,
456            route,
457            pipeline: self.pipeline,
458            boundary: self.boundary,
459            cancellation: self.cancellation,
460            cancellation_registry: self.cancellation_registry,
461            failure_policy: self.failure_policy,
462        }
463    }
464
465    /// Selects lock-step or bounded request pipelining.
466    #[must_use]
467    pub fn pipeline<Next: PipelinePolicy>(
468        self,
469        pipeline: Next,
470    ) -> IntermediaryBuilder<S, C, R, A, Next, B, K> {
471        IntermediaryBuilder {
472            server: self.server,
473            client: self.client,
474            resolver: self.resolver,
475            route: self.route,
476            pipeline,
477            boundary: self.boundary,
478            cancellation: self.cancellation,
479            cancellation_registry: self.cancellation_registry,
480            failure_policy: self.failure_policy,
481        }
482    }
483
484    /// Supplies middleware for the forwarding boundary.
485    #[must_use]
486    pub fn middleware<Next>(self, boundary: Next) -> IntermediaryBuilder<S, C, R, A, P, Next, K> {
487        IntermediaryBuilder {
488            server: self.server,
489            client: self.client,
490            resolver: self.resolver,
491            route: self.route,
492            pipeline: self.pipeline,
493            boundary,
494            cancellation: self.cancellation,
495            cancellation_registry: self.cancellation_registry,
496            failure_policy: self.failure_policy,
497        }
498    }
499
500    /// Selects an explicit cancellation posture.
501    #[must_use]
502    pub fn cancellation(mut self, cancellation: CancellationPolicy) -> Self {
503        self.cancellation = match cancellation {
504            CancellationPolicy::Reject => Some(CancellationPolicy::Reject),
505            CancellationPolicy::Forward => None,
506        };
507        self
508    }
509
510    /// Selects conservative close or one fixed safe diagnostic on establishment failure.
511    #[must_use]
512    pub fn establishment_failure(mut self, policy: EstablishmentFailurePolicy) -> Self {
513        self.failure_policy = policy;
514        self
515    }
516
517    /// Enables forwarding through an application-owned concurrent registry.
518    #[must_use]
519    pub fn cancellation_registry<Next>(
520        self,
521        registry: Next,
522    ) -> IntermediaryBuilder<S, C, R, A, P, B, Next> {
523        IntermediaryBuilder {
524            server: self.server,
525            client: self.client,
526            resolver: self.resolver,
527            route: self.route,
528            pipeline: self.pipeline,
529            boundary: self.boundary,
530            cancellation: Some(CancellationPolicy::Forward),
531            cancellation_registry: registry,
532            failure_policy: self.failure_policy,
533        }
534    }
535
536    /// Validates composition and creates a reusable component.
537    ///
538    /// # Errors
539    ///
540    /// Returns the first missing mandatory role, resolver, or cancellation configuration.
541    #[allow(clippy::type_complexity)]
542    pub fn build(self) -> Result<Intermediary<S, C, R, A, P, B, K>, IntermediaryBuildError> {
543        Ok(Intermediary {
544            server: self.server.ok_or(IntermediaryBuildError::MissingServer)?,
545            client: self.client.ok_or(IntermediaryBuildError::MissingClient)?,
546            resolver: self
547                .resolver
548                .ok_or(IntermediaryBuildError::MissingStartupResolver)?,
549            route: self.route,
550            pipeline: self.pipeline,
551            boundary: self.boundary,
552            cancellation: self
553                .cancellation
554                .ok_or(IntermediaryBuildError::MissingCancellationPolicy)?,
555            cancellation_registry: self.cancellation_registry,
556            failure_policy: self.failure_policy,
557        })
558    }
559}
560
561struct StartupResolverAdapter<'a, Resolver> {
562    resolver: &'a Resolver,
563}
564
565/// Failure while decoding or resolving a startup route.
566#[derive(Debug)]
567pub enum StartupResolutionError<Error> {
568    /// A startup parameter was not representable by the structured facade.
569    Parameters(io::Error),
570    /// The application resolver rejected the route.
571    Resolver(Error),
572}
573
574impl<Error: fmt::Display> fmt::Display for StartupResolutionError<Error> {
575    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
576        match self {
577            Self::Parameters(error) => error.fmt(formatter),
578            Self::Resolver(error) => error.fmt(formatter),
579        }
580    }
581}
582
583impl<Error: std::error::Error + 'static> std::error::Error for StartupResolutionError<Error> {
584    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
585        match self {
586            Self::Parameters(error) => Some(error),
587            Self::Resolver(error) => Some(error),
588        }
589    }
590}
591
592impl<Resolver, State, Peer, Identity>
593    crate::server_component::StartupResolver<State, Peer, Identity>
594    for StartupResolverAdapter<'_, Resolver>
595where
596    Resolver: StartupRouteResolver<Peer>,
597{
598    type Route = ConnectTarget;
599    type Error = StartupResolutionError<Resolver::Error>;
600
601    fn defer_ready(&self) -> bool {
602        true
603    }
604
605    fn resolve<'a>(
606        &'a mut self,
607        startup: &'a crate::startup::StartupMessage,
608        context: &'a crate::ServerConnectionContext<Peer, Identity>,
609        _state: &'a mut State,
610    ) -> Pin<Box<dyn Future<Output = Result<Self::Route, Self::Error>> + 'a>> {
611        let parameters = StartupParameters::from_wire(startup);
612        let initial = context
613            .tls_if_known()
614            .map(|tls| InitialServerContext::new(context.peer(), tls));
615        let resolver = self.resolver;
616        Box::pin(async move {
617            let parameters = parameters.map_err(StartupResolutionError::Parameters)?;
618            let initial = initial.expect("startup routing runs after TLS negotiation");
619            resolver
620                .resolve(parameters, initial)
621                .await
622                .map_err(StartupResolutionError::Resolver)
623        })
624    }
625}
626
627/// Failure while establishing both independently authenticated roles.
628pub enum IntermediaryAcceptError<
629    ServerError,
630    ResolverError,
631    RouteError,
632    ClientError,
633    RegistryError = std::convert::Infallible,
634    CancellationError = std::convert::Infallible,
635    MiddlewareError = std::convert::Infallible,
636> {
637    /// Client-facing TLS, startup, or authentication failed.
638    Server(ServerError),
639    /// Startup routing failed before client-facing authentication.
640    StartupRoute(StartupResolutionError<ResolverError>),
641    /// The explicit cancellation posture rejected an out-of-band request.
642    CancellationRejected,
643    /// Authenticated routing rejected or failed to refine the destination.
644    AuthenticatedRoute(RouteError),
645    /// PostgreSQL-facing connection, TLS, startup, or authentication failed.
646    Client(ClientError),
647    /// Cancellation-key allocation, collision detection, or storage failed.
648    CancellationRegistry(RegistryError),
649    /// A generated establishment message could not be written downstream.
650    ServerOutput(io::Error),
651    /// Opening or writing the one-shot upstream cancellation connection failed.
652    Cancellation(CancellationError),
653    /// Forwarding-boundary middleware rejected generated establishment output.
654    Middleware(MiddlewareError),
655}
656
657impl<S, R, A, C, K, X, M> fmt::Debug for IntermediaryAcceptError<S, R, A, C, K, X, M> {
658    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
659        formatter.write_str(match self {
660            Self::Server(_) => "IntermediaryAcceptError::Server([REDACTED])",
661            Self::StartupRoute(_) => "IntermediaryAcceptError::StartupRoute([REDACTED])",
662            Self::CancellationRejected => "IntermediaryAcceptError::CancellationRejected",
663            Self::AuthenticatedRoute(_) => {
664                "IntermediaryAcceptError::AuthenticatedRoute([REDACTED])"
665            }
666            Self::Client(_) => "IntermediaryAcceptError::Client([REDACTED])",
667            Self::CancellationRegistry(_) => {
668                "IntermediaryAcceptError::CancellationRegistry([REDACTED])"
669            }
670            Self::ServerOutput(_) => "IntermediaryAcceptError::ServerOutput([REDACTED])",
671            Self::Cancellation(_) => "IntermediaryAcceptError::Cancellation([REDACTED])",
672            Self::Middleware(_) => "IntermediaryAcceptError::Middleware([REDACTED])",
673        })
674    }
675}
676
677impl<S, R, A, C, K, X, M> fmt::Display for IntermediaryAcceptError<S, R, A, C, K, X, M> {
678    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
679        match self {
680            Self::Server(_) => formatter.write_str("client-facing establishment failed"),
681            Self::StartupRoute(_) => formatter.write_str("startup routing failed"),
682            Self::CancellationRejected => {
683                formatter.write_str("cancellation is explicitly rejected")
684            }
685            Self::AuthenticatedRoute(_) => formatter.write_str("authenticated routing failed"),
686            Self::Client(_) => formatter.write_str("PostgreSQL-facing establishment failed"),
687            Self::CancellationRegistry(_) => {
688                formatter.write_str("cancellation registration failed")
689            }
690            Self::ServerOutput(_) => {
691                formatter.write_str("client-facing establishment output failed")
692            }
693            Self::Cancellation(_) => formatter.write_str("cancellation forwarding failed"),
694            Self::Middleware(_) => {
695                formatter.write_str("forwarding middleware rejected establishment output")
696            }
697        }
698    }
699}
700
701impl<S, R, A, C, K, X, M> std::error::Error for IntermediaryAcceptError<S, R, A, C, K, X, M>
702where
703    S: std::error::Error + 'static,
704    R: std::error::Error + 'static,
705    A: std::error::Error + 'static,
706    C: std::error::Error + 'static,
707    K: std::error::Error + 'static,
708    X: std::error::Error + 'static,
709    M: std::error::Error + 'static,
710{
711    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
712        match self {
713            Self::Server(error) => Some(error),
714            Self::StartupRoute(error) => Some(error),
715            Self::CancellationRejected => None,
716            Self::AuthenticatedRoute(error) => Some(error),
717            Self::Client(error) => Some(error),
718            Self::CancellationRegistry(error) => Some(error),
719            Self::ServerOutput(error) => Some(error),
720            Self::Cancellation(error) => Some(error),
721            Self::Middleware(error) => Some(error),
722        }
723    }
724}
725
726/// Both role contexts recovered during deliberate intermediary teardown.
727#[derive(Debug)]
728pub struct IntermediaryContexts<ServerContext, ClientContext> {
729    server: ServerContext,
730    client: ClientContext,
731}
732
733impl<ServerContext, ClientContext> IntermediaryContexts<ServerContext, ClientContext> {
734    /// Returns the client-facing role context.
735    #[must_use]
736    pub const fn server(&self) -> &ServerContext {
737        &self.server
738    }
739    /// Returns the PostgreSQL-facing role context.
740    #[must_use]
741    pub const fn client(&self) -> &ClientContext {
742        &self.client
743    }
744}
745
746/// One operational, independently authenticated intermediary session.
747pub struct IntermediaryConnection<
748    DT,
749    UT,
750    State,
751    Peer,
752    ServerIdentity,
753    ClientEvidence,
754    ServerHandler,
755    ClientHandler,
756    Boundary,
757    Policy,
758    Cancellation = RejectCancellation,
759> {
760    downstream:
761        crate::server_component::ServerConnectionCore<DT, Peer, ServerIdentity, ServerHandler>,
762    upstream: crate::client_component::ClientConnectionCore<
763        crate::ClientTransport<UT>,
764        crate::Pristine,
765        ClientEvidence,
766        ClientHandler,
767    >,
768    state: State,
769    boundary: Boundary,
770    pipeline: Pipeline<Policy>,
771    target: ConnectTarget,
772    pending_frontend: Option<crate::codec::FrontendMessage>,
773    pending_local: VecDeque<PendingLocalResponses>,
774    cancellation_registry: Cancellation,
775    client_cancel_key: Option<crate::demux::CancelKey>,
776}
777
778struct PendingLocalResponses {
779    operation: crate::pipeline::OperationId,
780    messages: VecDeque<crate::codec::BackendMessage>,
781}
782
783/// Result of accepting either an ordinary session or an out-of-band request.
784#[derive(Debug)]
785pub enum IntermediaryAccept<Connection> {
786    /// A fully established, independently authenticated session pair.
787    Session(Connection),
788    /// The resolved cancellation packet was rewritten and forwarded.
789    CancellationForwarded,
790}
791
792impl<Connection> IntermediaryAccept<Connection> {
793    /// Extracts the ordinary session branch.
794    ///
795    /// # Panics
796    /// Panics when the accepted connection was cancellation-only.
797    #[must_use]
798    pub fn into_session(self) -> Connection {
799        match self {
800            Self::Session(connection) => connection,
801            Self::CancellationForwarded => panic!("accepted cancellation has no session"),
802        }
803    }
804}
805
806/// Direction selected by one cancellation-safe duplex forwarding step.
807#[derive(Debug)]
808pub enum ForwardedMessage {
809    /// A client-originated message was forwarded to PostgreSQL.
810    Frontend(crate::codec::FrontendMessage),
811    /// A PostgreSQL-originated message was forwarded to the client.
812    Backend(crate::codec::BackendMessage),
813    /// A client-originated message was deliberately suppressed.
814    FrontendSuppressed(crate::codec::FrontendMessage),
815    /// A request was handled locally; responses were emitted or queued in order.
816    FrontendLocallyHandled(crate::codec::FrontendMessage),
817    /// A PostgreSQL-originated response advanced protocol state but was suppressed.
818    BackendSuppressed(crate::codec::BackendMessage),
819}
820
821/// Observable result of processing one client-originated message.
822#[derive(Debug, Eq, PartialEq)]
823pub enum FrontendForwarding {
824    /// The message was sent to PostgreSQL.
825    Forwarded(crate::codec::FrontendMessage),
826    /// The message was consumed without being sent.
827    Suppressed(crate::codec::FrontendMessage),
828    /// The request was admitted as local and its responses were emitted or queued.
829    LocallyHandled(crate::codec::FrontendMessage),
830}
831
832impl FrontendForwarding {
833    /// Returns the owned client message represented by this outcome.
834    #[must_use]
835    pub fn into_message(self) -> crate::codec::FrontendMessage {
836        match self {
837            Self::Forwarded(message)
838            | Self::Suppressed(message)
839            | Self::LocallyHandled(message) => message,
840        }
841    }
842}
843
844/// Observable result of processing one PostgreSQL-originated response.
845#[derive(Debug, Eq, PartialEq)]
846pub enum BackendForwarding {
847    /// The response was sent to the client.
848    Forwarded(crate::codec::BackendMessage),
849    /// The response advanced protocol state but was not sent.
850    Suppressed(crate::codec::BackendMessage),
851}
852
853impl BackendForwarding {
854    /// Returns the owned PostgreSQL response represented by this outcome.
855    #[must_use]
856    pub fn into_message(self) -> crate::codec::BackendMessage {
857        match self {
858            Self::Forwarded(message) | Self::Suppressed(message) => message,
859        }
860    }
861}
862
863impl<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
864    IntermediaryConnection<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
865where
866    Policy: PipelinePolicy,
867{
868    /// Returns the authoritative destination selected for the client component.
869    #[must_use]
870    pub const fn target(&self) -> &ConnectTarget {
871        &self.target
872    }
873    /// Returns the single caller-owned state shared by all three middleware layers.
874    #[must_use]
875    pub const fn state(&self) -> &State {
876        &self.state
877    }
878    /// Returns the proxy-issued cancellation key for this live session.
879    #[must_use]
880    pub const fn cancellation_key(&self) -> Option<&crate::demux::CancelKey> {
881        self.client_cancel_key.as_ref()
882    }
883
884    /// Detaches this session's cancellation mapping explicitly.
885    pub fn detach_cancellation(&mut self) -> Option<CancellationRoute>
886    where
887        K: IntermediaryCancellationRegistry,
888    {
889        self.client_cancel_key
890            .take()
891            .and_then(|key| self.cancellation_registry.detach(&key))
892    }
893}
894
895impl<
896    DT,
897    UT,
898    State,
899    Peer,
900    ServerIdentity,
901    ClientEvidence,
902    ServerHandler,
903    ClientHandler,
904    Boundary,
905    Policy,
906    K,
907>
908    IntermediaryConnection<
909        DT,
910        UT,
911        State,
912        Peer,
913        ServerIdentity,
914        ClientEvidence,
915        ServerHandler,
916        ClientHandler,
917        Boundary,
918        Policy,
919        K,
920    >
921where
922    DT: AsyncRead + AsyncWrite + Unpin,
923    UT: AsyncRead + AsyncWrite + Unpin,
924    ServerHandler:
925        crate::ServerMiddleware<State, crate::ServerConnectionContext<Peer, ServerIdentity>>,
926    ClientHandler: crate::ClientMiddleware<State, crate::ClientConnectionContext<ClientEvidence>>,
927    Boundary: IntermediaryMiddleware<
928            State,
929            crate::ServerConnectionContext<Peer, ServerIdentity>,
930            crate::ClientConnectionContext<ClientEvidence>,
931        >,
932    Policy: PipelinePolicy,
933    K: IntermediaryCancellationRegistry,
934{
935    /// Receives one client message and reports whether it was forwarded,
936    /// suppressed, or handled with pipeline-ordered local responses.
937    ///
938    /// # Errors
939    ///
940    /// Returns middleware, transport, protocol-legality, or capacity failures.
941    pub async fn forward_frontend(
942        &mut self,
943    ) -> Result<FrontendForwarding, ForwardError<Boundary::Error>> {
944        if let Some(message) = self.pending_frontend.take() {
945            self.process_frontend(message, false).await
946        } else {
947            let message = self.downstream.receive_wire_raw().await?;
948            self.process_frontend(message, true).await
949        }
950    }
951
952    async fn process_frontend(
953        &mut self,
954        message: crate::codec::FrontendMessage,
955        intercept_source_and_boundary: bool,
956    ) -> Result<FrontendForwarding, ForwardError<Boundary::Error>> {
957        let decision = if intercept_source_and_boundary {
958            let message = self.downstream.intercept_frontend(&mut self.state, message);
959            self.boundary
960                .frontend(
961                    self.downstream.context(),
962                    self.upstream.context(),
963                    &mut self.state,
964                    message,
965                )
966                .await
967                .map_err(ForwardError::Middleware)?
968        } else {
969            FrontendMiddlewareOutput::Forward(message)
970        };
971        let (message, handling) = match decision {
972            FrontendMiddlewareOutput::Forward(message) => {
973                let message = if intercept_source_and_boundary {
974                    self.upstream.intercept_frontend(&mut self.state, message)
975                } else {
976                    message
977                };
978                (message, FrontendHandling::Forward)
979            }
980            FrontendMiddlewareOutput::Suppress(message) => {
981                return Ok(FrontendForwarding::Suppressed(message));
982            }
983            FrontendMiddlewareOutput::Respond { request, responses } => {
984                let admission = self
985                    .pipeline
986                    .accept_frontend(request.clone(), FrontendHandling::Local)
987                    .map_err(ForwardError::Frontend)?;
988                let FrontendAction::Discard { id } = admission.into_action() else {
989                    unreachable!()
990                };
991                let messages = responses
992                    .into_iter()
993                    .map(|message| self.downstream.intercept_backend(&mut self.state, message))
994                    .collect();
995                self.pending_local.push_back(PendingLocalResponses {
996                    operation: id,
997                    messages,
998                });
999                self.flush_local_responses().await?;
1000                return Ok(FrontendForwarding::LocallyHandled(request));
1001            }
1002        };
1003        let admission = match self.pipeline.accept_frontend(message.clone(), handling) {
1004            Ok(admission) => admission,
1005            Err(error) => {
1006                self.pending_frontend = Some(message);
1007                return Err(ForwardError::Frontend(error));
1008            }
1009        };
1010        let FrontendAction::Forward { message, .. } = admission.into_action() else {
1011            unreachable!()
1012        };
1013        self.upstream.send_wire_raw(message.clone()).await?;
1014        Ok(FrontendForwarding::Forwarded(message))
1015    }
1016
1017    /// Receives one legal PostgreSQL response and forwards it downstream in
1018    /// source-role, boundary, destination-role middleware order.
1019    ///
1020    /// # Errors
1021    ///
1022    /// Returns transport, framing, ordering, or protocol-legality failures.
1023    /// Receives one PostgreSQL response and reports whether it was forwarded or suppressed.
1024    ///
1025    /// # Errors
1026    ///
1027    /// Returns middleware, transport, ordering, or protocol-legality failures.
1028    pub async fn forward_backend(
1029        &mut self,
1030    ) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
1031        let message = self.upstream.receive_wire_raw().await?;
1032        self.process_backend(message).await
1033    }
1034
1035    async fn process_backend(
1036        &mut self,
1037        message: crate::codec::BackendMessage,
1038    ) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
1039        let message = self.upstream.intercept_backend(&mut self.state, message);
1040        let decision = self
1041            .boundary
1042            .backend(
1043                self.downstream.context(),
1044                self.upstream.context(),
1045                &mut self.state,
1046                message,
1047            )
1048            .await
1049            .map_err(ForwardError::Middleware)?;
1050        let (message, suppress) = match decision {
1051            BackendMiddlewareOutput::Forward(message) => (
1052                self.downstream.intercept_backend(&mut self.state, message),
1053                false,
1054            ),
1055            BackendMiddlewareOutput::Suppress(message) => (message, true),
1056        };
1057        let message = match self
1058            .pipeline
1059            .accept_backend(message)
1060            .map_err(ForwardError::Backend)?
1061        {
1062            BackendAction::Emit(message) => message,
1063            BackendAction::Deferred(message) => return Err(ForwardError::Deferred(message)),
1064        };
1065        if suppress {
1066            self.flush_local_responses().await?;
1067            Ok(BackendForwarding::Suppressed(message))
1068        } else {
1069            self.downstream.send_wire_raw(message.clone()).await?;
1070            self.flush_local_responses().await?;
1071            Ok(BackendForwarding::Forwarded(message))
1072        }
1073    }
1074
1075    async fn flush_local_responses(&mut self) -> Result<(), ForwardError<Boundary::Error>> {
1076        loop {
1077            let Some(pending) = self.pending_local.front_mut() else {
1078                return Ok(());
1079            };
1080            let Some(message) = pending.messages.pop_front() else {
1081                self.pending_local.pop_front();
1082                continue;
1083            };
1084            match self.pipeline.try_emit_local(pending.operation, message) {
1085                Ok(BackendAction::Emit(message)) => {
1086                    self.downstream.send_wire_raw(message).await?;
1087                }
1088                Ok(BackendAction::Deferred(message)) => {
1089                    pending.messages.push_front(message);
1090                    return Ok(());
1091                }
1092                Err(error) => return Err(ForwardError::Backend(error)),
1093            }
1094        }
1095    }
1096
1097    /// Waits on both transports and forwards whichever legal message becomes
1098    /// available first. This is the duplex driver for asynchronous traffic,
1099    /// COPY BOTH, and physical replication.
1100    ///
1101    /// When frontend capacity is exhausted, the unchanged pending request is
1102    /// retained and only backend progress is polled until capacity recovers.
1103    ///
1104    /// # Errors
1105    ///
1106    /// Returns transport, framing, ordering, protocol-legality, or capacity failures.
1107    pub async fn forward_next(
1108        &mut self,
1109    ) -> Result<ForwardedMessage, ForwardError<Boundary::Error>> {
1110        if self.pending_frontend.is_some() {
1111            let message = self.upstream.receive_wire_raw().await?;
1112            return self
1113                .process_backend(message)
1114                .await
1115                .map(|outcome| match outcome {
1116                    BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
1117                    BackendForwarding::Suppressed(message) => {
1118                        ForwardedMessage::BackendSuppressed(message)
1119                    }
1120                });
1121        }
1122        tokio::select! {
1123            result = self.downstream.receive_wire_raw() => {
1124                let message = result?;
1125                self.process_frontend(message, true).await.map(|outcome| match outcome {
1126                    FrontendForwarding::Forwarded(message) => ForwardedMessage::Frontend(message),
1127                    FrontendForwarding::Suppressed(message) => ForwardedMessage::FrontendSuppressed(message),
1128                    FrontendForwarding::LocallyHandled(message) => ForwardedMessage::FrontendLocallyHandled(message),
1129                })
1130            }
1131            result = self.upstream.receive_wire_raw() => {
1132                let message = result?;
1133                self.process_backend(message).await.map(|outcome| match outcome {
1134                    BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
1135                    BackendForwarding::Suppressed(message) => ForwardedMessage::BackendSuppressed(message),
1136                })
1137            }
1138        }
1139    }
1140
1141    /// Deliberately tears down both roles and recovers transports, handlers,
1142    /// contexts, boundary middleware, and the sole connection state.
1143    #[allow(clippy::type_complexity)]
1144    pub fn teardown(
1145        mut self,
1146    ) -> (
1147        crate::AcceptedServerTransport<DT>,
1148        crate::ClientTransport<UT>,
1149        State,
1150        Boundary,
1151        (ServerHandler, ClientHandler),
1152        IntermediaryContexts<
1153            crate::ServerConnectionContext<Peer, ServerIdentity>,
1154            crate::ClientConnectionContext<ClientEvidence>,
1155        >,
1156    ) {
1157        let _ = self.detach_cancellation();
1158        let (downstream, server_handler, server_context) = self.downstream.into_parts();
1159        let (upstream, client_handler, client_context) = self.upstream.into_parts();
1160        (
1161            downstream,
1162            upstream,
1163            self.state,
1164            self.boundary,
1165            (server_handler, client_handler),
1166            IntermediaryContexts {
1167                server: server_context,
1168                client: client_context,
1169            },
1170        )
1171    }
1172}
1173
1174/// Operational forwarding or pipeline projection failure.
1175#[derive(Debug)]
1176pub enum ForwardError<MiddlewareError = std::convert::Infallible> {
1177    /// Transport, decoding, or encoding failure.
1178    Io(io::Error),
1179    /// Frontend backpressure or protocol-legality rejection.
1180    Frontend(crate::pipeline::FrontendProjectionError),
1181    /// Backend protocol-legality rejection.
1182    Backend(crate::pipeline::BackendProjectionError),
1183    /// A bounded response arrived before its operation became emittable.
1184    Deferred(crate::codec::BackendMessage),
1185    /// Forwarding-boundary middleware rejected a message.
1186    Middleware(MiddlewareError),
1187}
1188
1189impl<E> fmt::Display for ForwardError<E> {
1190    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1191        match self {
1192            Self::Io(error) => error.fmt(formatter),
1193            Self::Frontend(_) => {
1194                formatter.write_str("frontend message violates pipeline legality or capacity")
1195            }
1196            Self::Backend(_) => formatter.write_str("backend message violates pipeline legality"),
1197            Self::Deferred(_) => formatter.write_str("backend response is not yet emittable"),
1198            Self::Middleware(_) => formatter.write_str("forwarding middleware rejected a message"),
1199        }
1200    }
1201}
1202
1203impl<E> std::error::Error for ForwardError<E>
1204where
1205    E: std::error::Error + 'static,
1206{
1207    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1208        match self {
1209            Self::Io(error) => Some(error),
1210            Self::Middleware(error) => Some(error),
1211            Self::Frontend(_) | Self::Backend(_) | Self::Deferred(_) => None,
1212        }
1213    }
1214}
1215
1216impl<E> From<io::Error> for ForwardError<E> {
1217    fn from(error: io::Error) -> Self {
1218        Self::Io(error)
1219    }
1220}
1221
1222impl<ST, SA, SM, Connector, CT, CA, CM, Resolver, Route, Policy, Boundary, K>
1223    Intermediary<
1224        crate::Server<ST, SA, SM>,
1225        crate::Client<Connector, CT, CA, CM>,
1226        Resolver,
1227        Route,
1228        Policy,
1229        Boundary,
1230        K,
1231    >
1232where
1233    ST: crate::ServerTlsConfiguration,
1234    SA: crate::ServerAuthenticationProvider,
1235    CT: crate::client_component::ClientTlsConfiguration,
1236    CA: crate::ClientAuthentication,
1237    CM: crate::MiddlewareFactory<crate::ClientInitialContext>,
1238    Policy: PipelinePolicy,
1239    K: IntermediaryCancellationRegistry + Clone,
1240{
1241    /// Establishes both independently authenticated roles around one shared state.
1242    ///
1243    /// # Errors
1244    ///
1245    /// Returns the typed failure from either role or routing policy, or explicit
1246    /// cancellation rejection.
1247    #[allow(clippy::type_complexity, clippy::too_many_lines)]
1248    pub async fn accept<DT, State, Peer, CW, UT, CE>(
1249        &self,
1250        transport: DT,
1251        peer: Peer,
1252        state: State,
1253    ) -> Result<
1254        IntermediaryAccept<
1255            IntermediaryConnection<
1256                DT,
1257                UT,
1258                State,
1259                Peer,
1260                <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1261                CA::Evidence,
1262                <SM as crate::MiddlewareFactory<
1263                    crate::ServerConnectionContext<
1264                        Peer,
1265                        <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1266                    >,
1267                >>::Handler,
1268                <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler,
1269                <Boundary as IntermediaryMiddlewareFactory<
1270                    crate::ServerConnectionContext<
1271                        Peer,
1272                        <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1273                    >,
1274                    crate::ClientConnectionContext<CA::Evidence>,
1275                >>::Handler,
1276                Policy,
1277                K,
1278            >,
1279        >,
1280        IntermediaryAcceptError<
1281            crate::AcceptError<
1282                <ST::Provider as crate::ServerIdentityProvider>::Error,
1283                <SA::Authentication as crate::ServerAuthentication<Peer>>::Error,
1284            >,
1285            Resolver::Error,
1286            Route::Error,
1287            crate::ConnectError<
1288                CE,
1289                crate::ClientTlsError<<CT::Provider as crate::ClientTlsProvider>::Error>,
1290                crate::ClientAuthenticationError<CA::Error>,
1291            >,
1292            K::Error,
1293            crate::CancelError<CE>,
1294            <<Boundary as IntermediaryMiddlewareFactory<
1295                crate::ServerConnectionContext<
1296                    Peer,
1297                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1298                >,
1299                crate::ClientConnectionContext<CA::Evidence>,
1300            >>::Handler as IntermediaryMiddleware<
1301                State,
1302                crate::ServerConnectionContext<
1303                    Peer,
1304                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1305                >,
1306                crate::ClientConnectionContext<CA::Evidence>,
1307            >>::Error,
1308        >,
1309    >
1310    where
1311        DT: AsyncRead + AsyncWrite + Unpin,
1312        UT: AsyncRead + AsyncWrite + Unpin,
1313        SA::Authentication: crate::ServerAuthentication<Peer>,
1314        SM: crate::MiddlewareFactory<
1315                crate::ServerConnectionContext<
1316                    Peer,
1317                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1318                >,
1319            >,
1320        <SM as crate::MiddlewareFactory<
1321            crate::ServerConnectionContext<
1322                Peer,
1323                <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1324            >,
1325        >>::Handler: crate::ServerMiddleware<
1326                State,
1327                crate::ServerConnectionContext<
1328                    Peer,
1329                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1330                >,
1331            >,
1332        Resolver: StartupRouteResolver<Peer>,
1333        Connector: Fn(&ConnectTarget) -> CW,
1334        CW: Future<Output = Result<UT, CE>>,
1335        <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler:
1336            crate::ClientMiddleware<State, crate::ClientConnectionContext<CA::Evidence>>,
1337        Route: AuthenticatedRoutePolicy<
1338                Peer,
1339                <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1340            >,
1341        Boundary: IntermediaryMiddlewareFactory<
1342                crate::ServerConnectionContext<
1343                    Peer,
1344                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1345                >,
1346                crate::ClientConnectionContext<CA::Evidence>,
1347            >,
1348        <Boundary as IntermediaryMiddlewareFactory<
1349            crate::ServerConnectionContext<
1350                Peer,
1351                <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1352            >,
1353            crate::ClientConnectionContext<CA::Evidence>,
1354        >>::Handler: IntermediaryMiddleware<
1355                State,
1356                crate::ServerConnectionContext<
1357                    Peer,
1358                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1359                >,
1360                crate::ClientConnectionContext<CA::Evidence>,
1361            >,
1362    {
1363        let mut resolver = StartupResolverAdapter {
1364            resolver: &self.resolver,
1365        };
1366        let (accepted, selected) = self
1367            .server
1368            .accept_routed(transport, peer, state, &mut resolver)
1369            .await
1370            .map_err(|error| match error {
1371                crate::server_component::RoutedAcceptError::Accept(error) => {
1372                    IntermediaryAcceptError::Server(error)
1373                }
1374                crate::server_component::RoutedAcceptError::Route(error) => {
1375                    IntermediaryAcceptError::StartupRoute(error)
1376                }
1377            })?;
1378        let mut downstream = match accepted {
1379            crate::ServerAccept::Session(downstream) => downstream,
1380            crate::ServerAccept::Cancellation(cancellation) => {
1381                if self.cancellation == CancellationPolicy::Reject {
1382                    let _ = cancellation.teardown();
1383                    return Err(IntermediaryAcceptError::CancellationRejected);
1384                }
1385                let request = cancellation.request();
1386                let client_key = crate::demux::CancelKey {
1387                    process_id: request.process_id(),
1388                    secret_key: bytes::Bytes::copy_from_slice(request.secret_key()),
1389                };
1390                let Some(route) = self.cancellation_registry.resolve(&client_key) else {
1391                    let _ = cancellation.teardown();
1392                    return Err(IntermediaryAcceptError::CancellationRejected);
1393                };
1394                if let Err(error) = self
1395                    .client
1396                    .cancel(route.target(), route.upstream_key())
1397                    .await
1398                {
1399                    let _ = cancellation.teardown();
1400                    return Err(IntermediaryAcceptError::Cancellation(error));
1401                }
1402                let _ = cancellation.teardown();
1403                return Ok(IntermediaryAccept::CancellationForwarded);
1404            }
1405        };
1406        let startup = match StartupParameters::from_wire(downstream.startup()) {
1407            Ok(startup) => startup,
1408            Err(error) => {
1409                if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1410                    let _ = downstream
1411                        .send_generated_error(safe_establishment_diagnostic())
1412                        .await;
1413                }
1414                let _ = downstream.teardown();
1415                return Err(IntermediaryAcceptError::StartupRoute(
1416                    StartupResolutionError::Parameters(error),
1417                ));
1418            }
1419        };
1420        let context = AuthenticatedRouteContext {
1421            peer: downstream.context().peer(),
1422            identity: downstream.context().identity(),
1423        };
1424        let Some(selected) = selected else {
1425            let _ = downstream.teardown();
1426            return Err(IntermediaryAcceptError::CancellationRejected);
1427        };
1428        let selected = match self.route.route(selected, context).await {
1429            Ok(target) => target,
1430            Err(error) => {
1431                if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1432                    let _ = downstream
1433                        .send_generated_error(safe_establishment_diagnostic())
1434                        .await;
1435                }
1436                let _ = downstream.teardown();
1437                return Err(IntermediaryAcceptError::AuthenticatedRoute(error));
1438            }
1439        };
1440        let (mut downstream, mut state) = downstream.into_core_and_state();
1441        let upstream = match self
1442            .client
1443            .connect_core(selected.clone(), startup, &mut state)
1444            .await
1445        {
1446            Ok(upstream) => upstream,
1447            Err(error) => {
1448                if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1449                    let diagnostic = safe_establishment_diagnostic();
1450                    let diagnostic = downstream.intercept_backend(&mut state, diagnostic);
1451                    if matches!(diagnostic, crate::codec::BackendMessage::ErrorResponse(_)) {
1452                        // A failed encode/write is a terminal close; do not recursively
1453                        // invoke failure handling or middleware.
1454                        let _ = downstream.send_wire_raw(diagnostic).await;
1455                    }
1456                }
1457                let _ = downstream.into_parts();
1458                return Err(IntermediaryAcceptError::Client(error));
1459            }
1460        };
1461        let boundary = self
1462            .boundary
1463            .create(downstream.context(), upstream.context());
1464        let (client_cancel_key, backend_key_message) =
1465            match (self.cancellation, upstream.context().backend_key().cloned()) {
1466                (CancellationPolicy::Forward, Some(upstream_key)) => {
1467                    let client_key = match self
1468                        .cancellation_registry
1469                        .register(CancellationRoute::new(selected.clone(), upstream_key))
1470                    {
1471                        Ok(key) => key,
1472                        Err(error) => {
1473                            if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1474                                let diagnostic = downstream
1475                                    .intercept_backend(&mut state, safe_establishment_diagnostic());
1476                                if matches!(
1477                                    diagnostic,
1478                                    crate::codec::BackendMessage::ErrorResponse(_)
1479                                ) {
1480                                    let _ = downstream.send_wire_raw(diagnostic).await;
1481                                }
1482                            }
1483                            let _ = downstream.into_parts();
1484                            let _ = upstream.into_parts();
1485                            return Err(IntermediaryAcceptError::CancellationRegistry(error));
1486                        }
1487                    };
1488                    let message = crate::codec::BackendMessage::BackendKeyData {
1489                        process_id: client_key.process_id,
1490                        secret_key: client_key.secret_key.clone(),
1491                    };
1492                    (Some(client_key), Some(message))
1493                }
1494                _ => (None, None),
1495            };
1496        let mut connection = IntermediaryConnection {
1497            downstream,
1498            upstream,
1499            state,
1500            boundary,
1501            pipeline: Pipeline::new(self.pipeline),
1502            target: selected,
1503            pending_frontend: None,
1504            pending_local: VecDeque::new(),
1505            cancellation_registry: self.cancellation_registry.clone(),
1506            client_cancel_key,
1507        };
1508        if let Some(message) = backend_key_message {
1509            let expected = message.clone();
1510            let message = connection
1511                .boundary
1512                .backend(
1513                    connection.downstream.context(),
1514                    connection.upstream.context(),
1515                    &mut connection.state,
1516                    message,
1517                )
1518                .await;
1519            let message = match message {
1520                Ok(BackendMiddlewareOutput::Forward(message)) => message,
1521                Ok(BackendMiddlewareOutput::Suppress(_)) => {
1522                    let _ = connection.detach_cancellation();
1523                    let _ = connection.teardown();
1524                    return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
1525                        io::ErrorKind::InvalidData,
1526                        "middleware suppressed generated cancellation key",
1527                    )));
1528                }
1529                Err(error) => {
1530                    let _ = connection.detach_cancellation();
1531                    let _ = connection.teardown();
1532                    return Err(IntermediaryAcceptError::Middleware(error));
1533                }
1534            };
1535            let message = connection
1536                .downstream
1537                .intercept_backend(&mut connection.state, message);
1538            if message != expected {
1539                let _ = connection.detach_cancellation();
1540                let _ = connection.teardown();
1541                return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
1542                    io::ErrorKind::InvalidData,
1543                    "middleware rejected generated cancellation key",
1544                )));
1545            }
1546            if let Err(error) = connection.downstream.send_wire_raw(message).await {
1547                let _ = connection.detach_cancellation();
1548                let _ = connection.teardown();
1549                return Err(IntermediaryAcceptError::ServerOutput(error));
1550            }
1551        }
1552        let ready = connection.downstream.intercept_backend(
1553            &mut connection.state,
1554            crate::codec::BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
1555        );
1556        if !matches!(ready, crate::codec::BackendMessage::ReadyForQuery(_)) {
1557            let _ = connection.detach_cancellation();
1558            let _ = connection.teardown();
1559            return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
1560                io::ErrorKind::InvalidData,
1561                "middleware rejected generated readiness",
1562            )));
1563        }
1564        if let Err(error) = connection.downstream.send_wire_raw(ready).await {
1565            let _ = connection.detach_cancellation();
1566            let _ = connection.teardown();
1567            return Err(IntermediaryAcceptError::ServerOutput(error));
1568        }
1569        Ok(IntermediaryAccept::Session(connection))
1570    }
1571}
pg_proto/intermediary_component.rs

pg_proto/
intermediary_component.rs
