Skip to main content

pg_proto/
intermediary_component.rs

1//! Builder-centred composition of the client-facing and PostgreSQL-facing roles.
2
3use std::{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/// Middleware at the forwarding boundary between the two role components.
223pub trait IntermediaryMiddleware<State, ServerContext, ClientContext> {
224    /// Intercepts a client-originated message after server-role middleware and
225    /// before client-role middleware.
226    fn frontend(
227        &mut self,
228        _server: &ServerContext,
229        _client: &ClientContext,
230        _state: &mut State,
231        message: crate::codec::FrontendMessage,
232    ) -> crate::codec::FrontendMessage {
233        message
234    }
235
236    /// Intercepts a PostgreSQL-originated message after client-role middleware
237    /// and before server-role middleware.
238    fn backend(
239        &mut self,
240        _server: &ServerContext,
241        _client: &ClientContext,
242        _state: &mut State,
243        message: crate::codec::BackendMessage,
244    ) -> crate::codec::BackendMessage {
245        message
246    }
247}
248
249/// Identity forwarding-boundary middleware.
250#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
251pub struct IdentityIntermediaryMiddleware;
252
253impl<State, ServerContext, ClientContext>
254    IntermediaryMiddleware<State, ServerContext, ClientContext> for IdentityIntermediaryMiddleware
255{
256}
257
258/// Creates fresh forwarding-boundary middleware for one established pair.
259pub trait IntermediaryMiddlewareFactory<ServerContext, ClientContext> {
260    /// Per-connection boundary handler.
261    type Handler;
262    /// Creates an isolated handler from both distinct role contexts.
263    fn create(&self, server: &ServerContext, client: &ClientContext) -> Self::Handler;
264}
265
266impl<ServerContext, ClientContext, Handler, Factory>
267    IntermediaryMiddlewareFactory<ServerContext, ClientContext> for Factory
268where
269    Factory: Fn(&ServerContext, &ClientContext) -> Handler,
270{
271    type Handler = Handler;
272    fn create(&self, server: &ServerContext, client: &ClientContext) -> Handler {
273        self(server, client)
274    }
275}
276
277impl<ServerContext, ClientContext> IntermediaryMiddlewareFactory<ServerContext, ClientContext>
278    for IdentityIntermediaryMiddleware
279{
280    type Handler = Self;
281    fn create(&self, _server: &ServerContext, _client: &ClientContext) -> Self {
282        *self
283    }
284}
285
286/// A reusable operational intermediary configuration.
287pub struct Intermediary<
288    Server = (),
289    Client = (),
290    Resolver = (),
291    Route = AllowAuthenticatedRoute,
292    Policy = NoPipeline,
293    Boundary = IdentityIntermediaryMiddleware,
294    Cancellation = RejectCancellation,
295> {
296    pub(crate) server: Server,
297    pub(crate) client: Client,
298    pub(crate) resolver: Resolver,
299    pub(crate) route: Route,
300    pub(crate) pipeline: Policy,
301    pub(crate) boundary: Boundary,
302    pub(crate) cancellation: CancellationPolicy,
303    pub(crate) cancellation_registry: Cancellation,
304    pub(crate) failure_policy: EstablishmentFailurePolicy,
305}
306
307impl Intermediary<()> {
308    /// Starts composition of the two complete role configurations.
309    #[must_use]
310    pub fn builder() -> IntermediaryBuilder {
311        IntermediaryBuilder::default()
312    }
313}
314
315impl<S, C, R, A, P, B, K> fmt::Debug for Intermediary<S, C, R, A, P, B, K> {
316    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
317        formatter
318            .debug_struct("Intermediary")
319            .field("server", &"<configured>")
320            .field("client", &"<configured>")
321            .field("resolver", &"<redacted>")
322            .field("authenticated_route", &"<redacted>")
323            .field("cancellation", &self.cancellation)
324            .finish_non_exhaustive()
325    }
326}
327
328/// Progressive builder for [`Intermediary`].
329pub struct IntermediaryBuilder<
330    Server = (),
331    Client = (),
332    Resolver = (),
333    Route = AllowAuthenticatedRoute,
334    Policy = NoPipeline,
335    Boundary = IdentityIntermediaryMiddleware,
336    Cancellation = RejectCancellation,
337> {
338    server: Option<Server>,
339    client: Option<Client>,
340    resolver: Option<Resolver>,
341    route: Route,
342    pipeline: Policy,
343    boundary: Boundary,
344    cancellation: Option<CancellationPolicy>,
345    cancellation_registry: Cancellation,
346    failure_policy: EstablishmentFailurePolicy,
347}
348
349impl Default for IntermediaryBuilder {
350    fn default() -> Self {
351        Self {
352            server: None,
353            client: None,
354            resolver: None,
355            route: AllowAuthenticatedRoute,
356            pipeline: NoPipeline,
357            boundary: IdentityIntermediaryMiddleware,
358            cancellation: None,
359            cancellation_registry: RejectCancellation,
360            failure_policy: EstablishmentFailurePolicy::Close,
361        }
362    }
363}
364
365impl<S, C, R, A, P, B, K> IntermediaryBuilder<S, C, R, A, P, B, K> {
366    /// Supplies the complete client-facing role configuration.
367    #[must_use]
368    pub fn server<Next>(self, server: Next) -> IntermediaryBuilder<Next, C, R, A, P, B, K> {
369        IntermediaryBuilder {
370            server: Some(server),
371            client: self.client,
372            resolver: self.resolver,
373            route: self.route,
374            pipeline: self.pipeline,
375            boundary: self.boundary,
376            cancellation: self.cancellation,
377            cancellation_registry: self.cancellation_registry,
378            failure_policy: self.failure_policy,
379        }
380    }
381
382    /// Supplies the complete PostgreSQL-facing role configuration.
383    #[must_use]
384    pub fn client<Next>(self, client: Next) -> IntermediaryBuilder<S, Next, R, A, P, B, K> {
385        IntermediaryBuilder {
386            server: self.server,
387            client: Some(client),
388            resolver: self.resolver,
389            route: self.route,
390            pipeline: self.pipeline,
391            boundary: self.boundary,
392            cancellation: self.cancellation,
393            cancellation_registry: self.cancellation_registry,
394            failure_policy: self.failure_policy,
395        }
396    }
397
398    /// Supplies the required asynchronous startup resolver.
399    #[must_use]
400    pub fn startup_resolver<Next>(
401        self,
402        resolver: Next,
403    ) -> IntermediaryBuilder<S, C, Next, A, P, B, K> {
404        IntermediaryBuilder {
405            server: self.server,
406            client: self.client,
407            resolver: Some(resolver),
408            route: self.route,
409            pipeline: self.pipeline,
410            boundary: self.boundary,
411            cancellation: self.cancellation,
412            cancellation_registry: self.cancellation_registry,
413            failure_policy: self.failure_policy,
414        }
415    }
416
417    /// Supplies optional post-authentication routing policy.
418    #[must_use]
419    pub fn authenticated_route<Next>(
420        self,
421        route: Next,
422    ) -> IntermediaryBuilder<S, C, R, Next, P, B, K> {
423        IntermediaryBuilder {
424            server: self.server,
425            client: self.client,
426            resolver: self.resolver,
427            route,
428            pipeline: self.pipeline,
429            boundary: self.boundary,
430            cancellation: self.cancellation,
431            cancellation_registry: self.cancellation_registry,
432            failure_policy: self.failure_policy,
433        }
434    }
435
436    /// Selects lock-step or bounded request pipelining.
437    #[must_use]
438    pub fn pipeline<Next: PipelinePolicy>(
439        self,
440        pipeline: Next,
441    ) -> IntermediaryBuilder<S, C, R, A, Next, B, K> {
442        IntermediaryBuilder {
443            server: self.server,
444            client: self.client,
445            resolver: self.resolver,
446            route: self.route,
447            pipeline,
448            boundary: self.boundary,
449            cancellation: self.cancellation,
450            cancellation_registry: self.cancellation_registry,
451            failure_policy: self.failure_policy,
452        }
453    }
454
455    /// Supplies middleware for the forwarding boundary.
456    #[must_use]
457    pub fn middleware<Next>(self, boundary: Next) -> IntermediaryBuilder<S, C, R, A, P, Next, K> {
458        IntermediaryBuilder {
459            server: self.server,
460            client: self.client,
461            resolver: self.resolver,
462            route: self.route,
463            pipeline: self.pipeline,
464            boundary,
465            cancellation: self.cancellation,
466            cancellation_registry: self.cancellation_registry,
467            failure_policy: self.failure_policy,
468        }
469    }
470
471    /// Selects an explicit cancellation posture.
472    #[must_use]
473    pub fn cancellation(mut self, cancellation: CancellationPolicy) -> Self {
474        self.cancellation = match cancellation {
475            CancellationPolicy::Reject => Some(CancellationPolicy::Reject),
476            CancellationPolicy::Forward => None,
477        };
478        self
479    }
480
481    /// Selects conservative close or one fixed safe diagnostic on establishment failure.
482    #[must_use]
483    pub fn establishment_failure(mut self, policy: EstablishmentFailurePolicy) -> Self {
484        self.failure_policy = policy;
485        self
486    }
487
488    /// Enables forwarding through an application-owned concurrent registry.
489    #[must_use]
490    pub fn cancellation_registry<Next>(
491        self,
492        registry: Next,
493    ) -> IntermediaryBuilder<S, C, R, A, P, B, Next> {
494        IntermediaryBuilder {
495            server: self.server,
496            client: self.client,
497            resolver: self.resolver,
498            route: self.route,
499            pipeline: self.pipeline,
500            boundary: self.boundary,
501            cancellation: Some(CancellationPolicy::Forward),
502            cancellation_registry: registry,
503            failure_policy: self.failure_policy,
504        }
505    }
506
507    /// Validates composition and creates a reusable component.
508    ///
509    /// # Errors
510    ///
511    /// Returns the first missing mandatory role, resolver, or cancellation configuration.
512    #[allow(clippy::type_complexity)]
513    pub fn build(self) -> Result<Intermediary<S, C, R, A, P, B, K>, IntermediaryBuildError> {
514        Ok(Intermediary {
515            server: self.server.ok_or(IntermediaryBuildError::MissingServer)?,
516            client: self.client.ok_or(IntermediaryBuildError::MissingClient)?,
517            resolver: self
518                .resolver
519                .ok_or(IntermediaryBuildError::MissingStartupResolver)?,
520            route: self.route,
521            pipeline: self.pipeline,
522            boundary: self.boundary,
523            cancellation: self
524                .cancellation
525                .ok_or(IntermediaryBuildError::MissingCancellationPolicy)?,
526            cancellation_registry: self.cancellation_registry,
527            failure_policy: self.failure_policy,
528        })
529    }
530}
531
532struct StartupResolverAdapter<'a, Resolver> {
533    resolver: &'a Resolver,
534}
535
536/// Failure while decoding or resolving a startup route.
537#[derive(Debug)]
538pub enum StartupResolutionError<Error> {
539    /// A startup parameter was not representable by the structured facade.
540    Parameters(io::Error),
541    /// The application resolver rejected the route.
542    Resolver(Error),
543}
544
545impl<Error: fmt::Display> fmt::Display for StartupResolutionError<Error> {
546    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
547        match self {
548            Self::Parameters(error) => error.fmt(formatter),
549            Self::Resolver(error) => error.fmt(formatter),
550        }
551    }
552}
553
554impl<Error: std::error::Error + 'static> std::error::Error for StartupResolutionError<Error> {
555    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
556        match self {
557            Self::Parameters(error) => Some(error),
558            Self::Resolver(error) => Some(error),
559        }
560    }
561}
562
563impl<Resolver, State, Peer, Identity>
564    crate::server_component::StartupResolver<State, Peer, Identity>
565    for StartupResolverAdapter<'_, Resolver>
566where
567    Resolver: StartupRouteResolver<Peer>,
568{
569    type Route = ConnectTarget;
570    type Error = StartupResolutionError<Resolver::Error>;
571
572    fn defer_ready(&self) -> bool {
573        true
574    }
575
576    fn resolve<'a>(
577        &'a mut self,
578        startup: &'a crate::startup::StartupMessage,
579        context: &'a crate::ServerConnectionContext<Peer, Identity>,
580        _state: &'a mut State,
581    ) -> Pin<Box<dyn Future<Output = Result<Self::Route, Self::Error>> + 'a>> {
582        let parameters = StartupParameters::from_wire(startup);
583        let initial = context
584            .tls_if_known()
585            .map(|tls| InitialServerContext::new(context.peer(), tls));
586        let resolver = self.resolver;
587        Box::pin(async move {
588            let parameters = parameters.map_err(StartupResolutionError::Parameters)?;
589            let initial = initial.expect("startup routing runs after TLS negotiation");
590            resolver
591                .resolve(parameters, initial)
592                .await
593                .map_err(StartupResolutionError::Resolver)
594        })
595    }
596}
597
598/// Failure while establishing both independently authenticated roles.
599pub enum IntermediaryAcceptError<
600    ServerError,
601    ResolverError,
602    RouteError,
603    ClientError,
604    RegistryError = std::convert::Infallible,
605    CancellationError = std::convert::Infallible,
606> {
607    /// Client-facing TLS, startup, or authentication failed.
608    Server(ServerError),
609    /// Startup routing failed before client-facing authentication.
610    StartupRoute(StartupResolutionError<ResolverError>),
611    /// The explicit cancellation posture rejected an out-of-band request.
612    CancellationRejected,
613    /// Authenticated routing rejected or failed to refine the destination.
614    AuthenticatedRoute(RouteError),
615    /// PostgreSQL-facing connection, TLS, startup, or authentication failed.
616    Client(ClientError),
617    /// Cancellation-key allocation, collision detection, or storage failed.
618    CancellationRegistry(RegistryError),
619    /// A generated establishment message could not be written downstream.
620    ServerOutput(io::Error),
621    /// Opening or writing the one-shot upstream cancellation connection failed.
622    Cancellation(CancellationError),
623}
624
625impl<S, R, A, C, K, X> fmt::Debug for IntermediaryAcceptError<S, R, A, C, K, X> {
626    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
627        formatter.write_str(match self {
628            Self::Server(_) => "IntermediaryAcceptError::Server([REDACTED])",
629            Self::StartupRoute(_) => "IntermediaryAcceptError::StartupRoute([REDACTED])",
630            Self::CancellationRejected => "IntermediaryAcceptError::CancellationRejected",
631            Self::AuthenticatedRoute(_) => {
632                "IntermediaryAcceptError::AuthenticatedRoute([REDACTED])"
633            }
634            Self::Client(_) => "IntermediaryAcceptError::Client([REDACTED])",
635            Self::CancellationRegistry(_) => {
636                "IntermediaryAcceptError::CancellationRegistry([REDACTED])"
637            }
638            Self::ServerOutput(_) => "IntermediaryAcceptError::ServerOutput([REDACTED])",
639            Self::Cancellation(_) => "IntermediaryAcceptError::Cancellation([REDACTED])",
640        })
641    }
642}
643
644impl<S, R, A, C, K, X> fmt::Display for IntermediaryAcceptError<S, R, A, C, K, X> {
645    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
646        match self {
647            Self::Server(_) => formatter.write_str("client-facing establishment failed"),
648            Self::StartupRoute(_) => formatter.write_str("startup routing failed"),
649            Self::CancellationRejected => {
650                formatter.write_str("cancellation is explicitly rejected")
651            }
652            Self::AuthenticatedRoute(_) => formatter.write_str("authenticated routing failed"),
653            Self::Client(_) => formatter.write_str("PostgreSQL-facing establishment failed"),
654            Self::CancellationRegistry(_) => {
655                formatter.write_str("cancellation registration failed")
656            }
657            Self::ServerOutput(_) => {
658                formatter.write_str("client-facing establishment output failed")
659            }
660            Self::Cancellation(_) => formatter.write_str("cancellation forwarding failed"),
661        }
662    }
663}
664
665impl<S, R, A, C, K, X> std::error::Error for IntermediaryAcceptError<S, R, A, C, K, X>
666where
667    S: std::error::Error + 'static,
668    R: std::error::Error + 'static,
669    A: std::error::Error + 'static,
670    C: std::error::Error + 'static,
671    K: std::error::Error + 'static,
672    X: std::error::Error + 'static,
673{
674    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
675        match self {
676            Self::Server(error) => Some(error),
677            Self::StartupRoute(error) => Some(error),
678            Self::CancellationRejected => None,
679            Self::AuthenticatedRoute(error) => Some(error),
680            Self::Client(error) => Some(error),
681            Self::CancellationRegistry(error) => Some(error),
682            Self::ServerOutput(error) => Some(error),
683            Self::Cancellation(error) => Some(error),
684        }
685    }
686}
687
688/// Both role contexts recovered during deliberate intermediary teardown.
689#[derive(Debug)]
690pub struct IntermediaryContexts<ServerContext, ClientContext> {
691    server: ServerContext,
692    client: ClientContext,
693}
694
695impl<ServerContext, ClientContext> IntermediaryContexts<ServerContext, ClientContext> {
696    /// Returns the client-facing role context.
697    #[must_use]
698    pub const fn server(&self) -> &ServerContext {
699        &self.server
700    }
701    /// Returns the PostgreSQL-facing role context.
702    #[must_use]
703    pub const fn client(&self) -> &ClientContext {
704        &self.client
705    }
706}
707
708/// One operational, independently authenticated intermediary session.
709pub struct IntermediaryConnection<
710    DT,
711    UT,
712    State,
713    Peer,
714    ServerIdentity,
715    ClientEvidence,
716    ServerHandler,
717    ClientHandler,
718    Boundary,
719    Policy,
720    Cancellation = RejectCancellation,
721> {
722    downstream:
723        crate::server_component::ServerConnectionCore<DT, Peer, ServerIdentity, ServerHandler>,
724    upstream: crate::client_component::ClientConnectionCore<
725        crate::ClientTransport<UT>,
726        crate::Pristine,
727        ClientEvidence,
728        ClientHandler,
729    >,
730    state: State,
731    boundary: Boundary,
732    pipeline: Pipeline<Policy>,
733    target: ConnectTarget,
734    pending_frontend: Option<crate::codec::FrontendMessage>,
735    cancellation_registry: Cancellation,
736    client_cancel_key: Option<crate::demux::CancelKey>,
737}
738
739/// Result of accepting either an ordinary session or an out-of-band request.
740#[derive(Debug)]
741pub enum IntermediaryAccept<Connection> {
742    /// A fully established, independently authenticated session pair.
743    Session(Connection),
744    /// The resolved cancellation packet was rewritten and forwarded.
745    CancellationForwarded,
746}
747
748impl<Connection> IntermediaryAccept<Connection> {
749    /// Extracts the ordinary session branch.
750    ///
751    /// # Panics
752    /// Panics when the accepted connection was cancellation-only.
753    #[must_use]
754    pub fn into_session(self) -> Connection {
755        match self {
756            Self::Session(connection) => connection,
757            Self::CancellationForwarded => panic!("accepted cancellation has no session"),
758        }
759    }
760}
761
762/// Direction selected by one cancellation-safe duplex forwarding step.
763#[derive(Debug)]
764pub enum ForwardedMessage {
765    /// A client-originated message was forwarded to PostgreSQL.
766    Frontend(crate::codec::FrontendMessage),
767    /// A PostgreSQL-originated message was forwarded to the client.
768    Backend(crate::codec::BackendMessage),
769}
770
771impl<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
772    IntermediaryConnection<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
773where
774    Policy: PipelinePolicy,
775{
776    /// Returns the authoritative destination selected for the client component.
777    #[must_use]
778    pub const fn target(&self) -> &ConnectTarget {
779        &self.target
780    }
781    /// Returns the single caller-owned state shared by all three middleware layers.
782    #[must_use]
783    pub const fn state(&self) -> &State {
784        &self.state
785    }
786    /// Returns the proxy-issued cancellation key for this live session.
787    #[must_use]
788    pub const fn cancellation_key(&self) -> Option<&crate::demux::CancelKey> {
789        self.client_cancel_key.as_ref()
790    }
791
792    /// Detaches this session's cancellation mapping explicitly.
793    pub fn detach_cancellation(&mut self) -> Option<CancellationRoute>
794    where
795        K: IntermediaryCancellationRegistry,
796    {
797        self.client_cancel_key
798            .take()
799            .and_then(|key| self.cancellation_registry.detach(&key))
800    }
801}
802
803impl<
804    DT,
805    UT,
806    State,
807    Peer,
808    ServerIdentity,
809    ClientEvidence,
810    ServerHandler,
811    ClientHandler,
812    Boundary,
813    Policy,
814    K,
815>
816    IntermediaryConnection<
817        DT,
818        UT,
819        State,
820        Peer,
821        ServerIdentity,
822        ClientEvidence,
823        ServerHandler,
824        ClientHandler,
825        Boundary,
826        Policy,
827        K,
828    >
829where
830    DT: AsyncRead + AsyncWrite + Unpin,
831    UT: AsyncRead + AsyncWrite + Unpin,
832    ServerHandler:
833        crate::ServerMiddleware<State, crate::ServerConnectionContext<Peer, ServerIdentity>>,
834    ClientHandler: crate::ClientMiddleware<State, crate::ClientConnectionContext<ClientEvidence>>,
835    Boundary: IntermediaryMiddleware<
836            State,
837            crate::ServerConnectionContext<Peer, ServerIdentity>,
838            crate::ClientConnectionContext<ClientEvidence>,
839        >,
840    Policy: PipelinePolicy,
841    K: IntermediaryCancellationRegistry,
842{
843    /// Receives one legal client message and forwards it upstream in
844    /// source-role, boundary, destination-role middleware order.
845    ///
846    /// # Errors
847    ///
848    /// Returns transport, framing, protocol-legality, or capacity failures.
849    pub async fn forward_frontend(
850        &mut self,
851    ) -> Result<crate::codec::FrontendMessage, ForwardError> {
852        if let Some(message) = self.pending_frontend.take() {
853            self.process_frontend(message, false).await
854        } else {
855            let message = self.downstream.receive_wire_raw().await?;
856            self.process_frontend(message, true).await
857        }
858    }
859
860    async fn process_frontend(
861        &mut self,
862        message: crate::codec::FrontendMessage,
863        intercept_source_and_boundary: bool,
864    ) -> Result<crate::codec::FrontendMessage, ForwardError> {
865        let message = if intercept_source_and_boundary {
866            let message = self.downstream.intercept_frontend(&mut self.state, message);
867            let message = self.boundary.frontend(
868                self.downstream.context(),
869                self.upstream.context(),
870                &mut self.state,
871                message,
872            );
873            self.upstream.intercept_frontend(&mut self.state, message)
874        } else {
875            message
876        };
877        let admission = match self
878            .pipeline
879            .accept_frontend(message.clone(), FrontendHandling::Forward)
880        {
881            Ok(admission) => admission,
882            Err(error) => {
883                self.pending_frontend = Some(message);
884                return Err(ForwardError::Frontend(error));
885            }
886        };
887        let FrontendAction::Forward { message, .. } = admission.into_action() else {
888            unreachable!()
889        };
890        self.upstream.send_wire_raw(message.clone()).await?;
891        Ok(message)
892    }
893
894    /// Receives one legal PostgreSQL response and forwards it downstream in
895    /// source-role, boundary, destination-role middleware order.
896    ///
897    /// # Errors
898    ///
899    /// Returns transport, framing, ordering, or protocol-legality failures.
900    pub async fn forward_backend(&mut self) -> Result<crate::codec::BackendMessage, ForwardError> {
901        let message = self.upstream.receive_wire_raw().await?;
902        self.process_backend(message).await
903    }
904
905    async fn process_backend(
906        &mut self,
907        message: crate::codec::BackendMessage,
908    ) -> Result<crate::codec::BackendMessage, ForwardError> {
909        let message = self.upstream.intercept_backend(&mut self.state, message);
910        let message = self.boundary.backend(
911            self.downstream.context(),
912            self.upstream.context(),
913            &mut self.state,
914            message,
915        );
916        let message = self.downstream.intercept_backend(&mut self.state, message);
917        let message = match self
918            .pipeline
919            .accept_backend(message)
920            .map_err(ForwardError::Backend)?
921        {
922            BackendAction::Emit(message) => message,
923            BackendAction::Deferred(message) => return Err(ForwardError::Deferred(message)),
924        };
925        self.downstream.send_wire_raw(message.clone()).await?;
926        Ok(message)
927    }
928
929    /// Waits on both transports and forwards whichever legal message becomes
930    /// available first. This is the duplex driver for asynchronous traffic,
931    /// COPY BOTH, and physical replication.
932    ///
933    /// When frontend capacity is exhausted, the unchanged pending request is
934    /// retained and only backend progress is polled until capacity recovers.
935    ///
936    /// # Errors
937    ///
938    /// Returns transport, framing, ordering, protocol-legality, or capacity failures.
939    pub async fn forward_next(&mut self) -> Result<ForwardedMessage, ForwardError> {
940        if self.pending_frontend.is_some() {
941            let message = self.upstream.receive_wire_raw().await?;
942            return self
943                .process_backend(message)
944                .await
945                .map(ForwardedMessage::Backend);
946        }
947        tokio::select! {
948            result = self.downstream.receive_wire_raw() => {
949                let message = result?;
950                self.process_frontend(message, true).await.map(ForwardedMessage::Frontend)
951            }
952            result = self.upstream.receive_wire_raw() => {
953                let message = result?;
954                self.process_backend(message).await.map(ForwardedMessage::Backend)
955            }
956        }
957    }
958
959    /// Deliberately tears down both roles and recovers transports, handlers,
960    /// contexts, boundary middleware, and the sole connection state.
961    #[allow(clippy::type_complexity)]
962    pub fn teardown(
963        mut self,
964    ) -> (
965        crate::AcceptedServerTransport<DT>,
966        crate::ClientTransport<UT>,
967        State,
968        Boundary,
969        (ServerHandler, ClientHandler),
970        IntermediaryContexts<
971            crate::ServerConnectionContext<Peer, ServerIdentity>,
972            crate::ClientConnectionContext<ClientEvidence>,
973        >,
974    ) {
975        let _ = self.detach_cancellation();
976        let (downstream, server_handler, server_context) = self.downstream.into_parts();
977        let (upstream, client_handler, client_context) = self.upstream.into_parts();
978        (
979            downstream,
980            upstream,
981            self.state,
982            self.boundary,
983            (server_handler, client_handler),
984            IntermediaryContexts {
985                server: server_context,
986                client: client_context,
987            },
988        )
989    }
990}
991
992/// Operational forwarding or pipeline projection failure.
993#[derive(Debug)]
994pub enum ForwardError {
995    /// Transport, decoding, or encoding failure.
996    Io(io::Error),
997    /// Frontend backpressure or protocol-legality rejection.
998    Frontend(crate::pipeline::FrontendProjectionError),
999    /// Backend protocol-legality rejection.
1000    Backend(crate::pipeline::BackendProjectionError),
1001    /// A bounded response arrived before its operation became emittable.
1002    Deferred(crate::codec::BackendMessage),
1003}
1004
1005impl fmt::Display for ForwardError {
1006    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1007        match self {
1008            Self::Io(error) => error.fmt(formatter),
1009            Self::Frontend(_) => {
1010                formatter.write_str("frontend message violates pipeline legality or capacity")
1011            }
1012            Self::Backend(_) => formatter.write_str("backend message violates pipeline legality"),
1013            Self::Deferred(_) => formatter.write_str("backend response is not yet emittable"),
1014        }
1015    }
1016}
1017
1018impl std::error::Error for ForwardError {
1019    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1020        match self {
1021            Self::Io(error) => Some(error),
1022            _ => None,
1023        }
1024    }
1025}
1026
1027impl From<io::Error> for ForwardError {
1028    fn from(error: io::Error) -> Self {
1029        Self::Io(error)
1030    }
1031}
1032
1033impl<ST, SA, SM, Connector, CT, CA, CM, Resolver, Route, Policy, Boundary, K>
1034    Intermediary<
1035        crate::Server<ST, SA, SM>,
1036        crate::Client<Connector, CT, CA, CM>,
1037        Resolver,
1038        Route,
1039        Policy,
1040        Boundary,
1041        K,
1042    >
1043where
1044    ST: crate::ServerTlsConfiguration,
1045    SA: crate::ServerAuthenticationProvider,
1046    CT: crate::client_component::ClientTlsConfiguration,
1047    CA: crate::ClientAuthentication,
1048    CM: crate::MiddlewareFactory<crate::ClientInitialContext>,
1049    Policy: PipelinePolicy,
1050    K: IntermediaryCancellationRegistry + Clone,
1051{
1052    /// Establishes both independently authenticated roles around one shared state.
1053    ///
1054    /// # Errors
1055    ///
1056    /// Returns the typed failure from either role or routing policy, or explicit
1057    /// cancellation rejection.
1058    #[allow(clippy::type_complexity, clippy::too_many_lines)]
1059    pub async fn accept<DT, State, Peer, CW, UT, CE>(
1060        &self,
1061        transport: DT,
1062        peer: Peer,
1063        state: State,
1064    ) -> Result<
1065        IntermediaryAccept<
1066            IntermediaryConnection<
1067                DT,
1068                UT,
1069                State,
1070                Peer,
1071                <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1072                CA::Evidence,
1073                <SM as crate::MiddlewareFactory<
1074                    crate::ServerConnectionContext<
1075                        Peer,
1076                        <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1077                    >,
1078                >>::Handler,
1079                <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler,
1080                <Boundary as IntermediaryMiddlewareFactory<
1081                    crate::ServerConnectionContext<
1082                        Peer,
1083                        <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1084                    >,
1085                    crate::ClientConnectionContext<CA::Evidence>,
1086                >>::Handler,
1087                Policy,
1088                K,
1089            >,
1090        >,
1091        IntermediaryAcceptError<
1092            crate::AcceptError<
1093                <ST::Provider as crate::ServerIdentityProvider>::Error,
1094                <SA::Authentication as crate::ServerAuthentication<Peer>>::Error,
1095            >,
1096            Resolver::Error,
1097            Route::Error,
1098            crate::ConnectError<
1099                CE,
1100                crate::ClientTlsError<<CT::Provider as crate::ClientTlsProvider>::Error>,
1101                crate::ClientAuthenticationError<CA::Error>,
1102            >,
1103            K::Error,
1104            crate::CancelError<CE>,
1105        >,
1106    >
1107    where
1108        DT: AsyncRead + AsyncWrite + Unpin,
1109        UT: AsyncRead + AsyncWrite + Unpin,
1110        SA::Authentication: crate::ServerAuthentication<Peer>,
1111        SM: crate::MiddlewareFactory<
1112                crate::ServerConnectionContext<
1113                    Peer,
1114                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1115                >,
1116            >,
1117        <SM as crate::MiddlewareFactory<
1118            crate::ServerConnectionContext<
1119                Peer,
1120                <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1121            >,
1122        >>::Handler: crate::ServerMiddleware<
1123                State,
1124                crate::ServerConnectionContext<
1125                    Peer,
1126                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1127                >,
1128            >,
1129        Resolver: StartupRouteResolver<Peer>,
1130        Connector: Fn(&ConnectTarget) -> CW,
1131        CW: Future<Output = Result<UT, CE>>,
1132        <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler:
1133            crate::ClientMiddleware<State, crate::ClientConnectionContext<CA::Evidence>>,
1134        Route: AuthenticatedRoutePolicy<
1135                Peer,
1136                <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1137            >,
1138        Boundary: IntermediaryMiddlewareFactory<
1139                crate::ServerConnectionContext<
1140                    Peer,
1141                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1142                >,
1143                crate::ClientConnectionContext<CA::Evidence>,
1144            >,
1145        <Boundary as IntermediaryMiddlewareFactory<
1146            crate::ServerConnectionContext<
1147                Peer,
1148                <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1149            >,
1150            crate::ClientConnectionContext<CA::Evidence>,
1151        >>::Handler: IntermediaryMiddleware<
1152                State,
1153                crate::ServerConnectionContext<
1154                    Peer,
1155                    <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1156                >,
1157                crate::ClientConnectionContext<CA::Evidence>,
1158            >,
1159    {
1160        let mut resolver = StartupResolverAdapter {
1161            resolver: &self.resolver,
1162        };
1163        let (accepted, selected) = self
1164            .server
1165            .accept_routed(transport, peer, state, &mut resolver)
1166            .await
1167            .map_err(|error| match error {
1168                crate::server_component::RoutedAcceptError::Accept(error) => {
1169                    IntermediaryAcceptError::Server(error)
1170                }
1171                crate::server_component::RoutedAcceptError::Route(error) => {
1172                    IntermediaryAcceptError::StartupRoute(error)
1173                }
1174            })?;
1175        let mut downstream = match accepted {
1176            crate::ServerAccept::Session(downstream) => downstream,
1177            crate::ServerAccept::Cancellation(cancellation) => {
1178                if self.cancellation == CancellationPolicy::Reject {
1179                    let _ = cancellation.teardown();
1180                    return Err(IntermediaryAcceptError::CancellationRejected);
1181                }
1182                let request = cancellation.request();
1183                let client_key = crate::demux::CancelKey {
1184                    process_id: request.process_id(),
1185                    secret_key: bytes::Bytes::copy_from_slice(request.secret_key()),
1186                };
1187                let Some(route) = self.cancellation_registry.resolve(&client_key) else {
1188                    let _ = cancellation.teardown();
1189                    return Err(IntermediaryAcceptError::CancellationRejected);
1190                };
1191                if let Err(error) = self
1192                    .client
1193                    .cancel(route.target(), route.upstream_key())
1194                    .await
1195                {
1196                    let _ = cancellation.teardown();
1197                    return Err(IntermediaryAcceptError::Cancellation(error));
1198                }
1199                let _ = cancellation.teardown();
1200                return Ok(IntermediaryAccept::CancellationForwarded);
1201            }
1202        };
1203        let startup = match StartupParameters::from_wire(downstream.startup()) {
1204            Ok(startup) => startup,
1205            Err(error) => {
1206                if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1207                    let _ = downstream
1208                        .send_generated_error(safe_establishment_diagnostic())
1209                        .await;
1210                }
1211                let _ = downstream.teardown();
1212                return Err(IntermediaryAcceptError::StartupRoute(
1213                    StartupResolutionError::Parameters(error),
1214                ));
1215            }
1216        };
1217        let context = AuthenticatedRouteContext {
1218            peer: downstream.context().peer(),
1219            identity: downstream.context().identity(),
1220        };
1221        let Some(selected) = selected else {
1222            let _ = downstream.teardown();
1223            return Err(IntermediaryAcceptError::CancellationRejected);
1224        };
1225        let selected = match self.route.route(selected, context).await {
1226            Ok(target) => target,
1227            Err(error) => {
1228                if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1229                    let _ = downstream
1230                        .send_generated_error(safe_establishment_diagnostic())
1231                        .await;
1232                }
1233                let _ = downstream.teardown();
1234                return Err(IntermediaryAcceptError::AuthenticatedRoute(error));
1235            }
1236        };
1237        let (mut downstream, mut state) = downstream.into_core_and_state();
1238        let upstream = match self
1239            .client
1240            .connect_core(selected.clone(), startup, &mut state)
1241            .await
1242        {
1243            Ok(upstream) => upstream,
1244            Err(error) => {
1245                if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1246                    let diagnostic = safe_establishment_diagnostic();
1247                    let diagnostic = downstream.intercept_backend(&mut state, diagnostic);
1248                    if matches!(diagnostic, crate::codec::BackendMessage::ErrorResponse(_)) {
1249                        // A failed encode/write is a terminal close; do not recursively
1250                        // invoke failure handling or middleware.
1251                        let _ = downstream.send_wire_raw(diagnostic).await;
1252                    }
1253                }
1254                let _ = downstream.into_parts();
1255                return Err(IntermediaryAcceptError::Client(error));
1256            }
1257        };
1258        let boundary = self
1259            .boundary
1260            .create(downstream.context(), upstream.context());
1261        let (client_cancel_key, backend_key_message) =
1262            match (self.cancellation, upstream.context().backend_key().cloned()) {
1263                (CancellationPolicy::Forward, Some(upstream_key)) => {
1264                    let client_key = match self
1265                        .cancellation_registry
1266                        .register(CancellationRoute::new(selected.clone(), upstream_key))
1267                    {
1268                        Ok(key) => key,
1269                        Err(error) => {
1270                            if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1271                                let diagnostic = downstream
1272                                    .intercept_backend(&mut state, safe_establishment_diagnostic());
1273                                if matches!(
1274                                    diagnostic,
1275                                    crate::codec::BackendMessage::ErrorResponse(_)
1276                                ) {
1277                                    let _ = downstream.send_wire_raw(diagnostic).await;
1278                                }
1279                            }
1280                            let _ = downstream.into_parts();
1281                            let _ = upstream.into_parts();
1282                            return Err(IntermediaryAcceptError::CancellationRegistry(error));
1283                        }
1284                    };
1285                    let message = crate::codec::BackendMessage::BackendKeyData {
1286                        process_id: client_key.process_id,
1287                        secret_key: client_key.secret_key.clone(),
1288                    };
1289                    (Some(client_key), Some(message))
1290                }
1291                _ => (None, None),
1292            };
1293        let mut connection = IntermediaryConnection {
1294            downstream,
1295            upstream,
1296            state,
1297            boundary,
1298            pipeline: Pipeline::new(self.pipeline),
1299            target: selected,
1300            pending_frontend: None,
1301            cancellation_registry: self.cancellation_registry.clone(),
1302            client_cancel_key,
1303        };
1304        if let Some(message) = backend_key_message {
1305            let expected = message.clone();
1306            let message = connection.boundary.backend(
1307                connection.downstream.context(),
1308                connection.upstream.context(),
1309                &mut connection.state,
1310                message,
1311            );
1312            let message = connection
1313                .downstream
1314                .intercept_backend(&mut connection.state, message);
1315            if message != expected {
1316                let _ = connection.detach_cancellation();
1317                let _ = connection.teardown();
1318                return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
1319                    io::ErrorKind::InvalidData,
1320                    "middleware rejected generated cancellation key",
1321                )));
1322            }
1323            if let Err(error) = connection.downstream.send_wire_raw(message).await {
1324                let _ = connection.detach_cancellation();
1325                let _ = connection.teardown();
1326                return Err(IntermediaryAcceptError::ServerOutput(error));
1327            }
1328        }
1329        let ready = connection.downstream.intercept_backend(
1330            &mut connection.state,
1331            crate::codec::BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
1332        );
1333        if !matches!(ready, crate::codec::BackendMessage::ReadyForQuery(_)) {
1334            let _ = connection.detach_cancellation();
1335            let _ = connection.teardown();
1336            return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
1337                io::ErrorKind::InvalidData,
1338                "middleware rejected generated readiness",
1339            )));
1340        }
1341        if let Err(error) = connection.downstream.send_wire_raw(ready).await {
1342            let _ = connection.detach_cancellation();
1343            let _ = connection.teardown();
1344            return Err(IntermediaryAcceptError::ServerOutput(error));
1345        }
1346        Ok(IntermediaryAccept::Session(connection))
1347    }
1348}