Skip to main content

pg_proto/
server_component.rs

1//! Reusable construction and establishment for the client-facing server role.
2
3use std::{fmt, future::Future, io, pin::Pin, sync::Arc};
4
5use bytes::Bytes;
6use rand::RngExt as _;
7use rustls::{ServerConfig, pki_types::CertificateDer};
8use tokio::io::{AsyncRead, AsyncWrite};
9
10use crate::ServerMiddleware as _;
11use crate::{
12    Conn,
13    auth::{Ready, TlsServerEndPoint},
14    codec::{
15        Backend, BackendMessage, DEFAULT_MAX_FRAME_LEN, Direction as _, Frontend, FrontendMessage,
16    },
17    pre_startup::{DEFAULT_MAX_PRE_STARTUP_PACKET_LEN, PreStartupOffer},
18    server_auth::ServerProtocolOffer,
19    startup::{ProtocolVersion, StartupMessage},
20    tls::ServerTls,
21    transport::Buffered,
22};
23
24/// Async startup routing hook used by the intermediary facade.
25#[allow(clippy::type_complexity)]
26pub(crate) trait StartupResolver<State, Peer, Identity> {
27    type Route;
28    type Error;
29
30    /// Whether the intermediary must insert startup messages before readiness.
31    fn defer_ready(&self) -> bool {
32        false
33    }
34
35    fn resolve<'a>(
36        &'a mut self,
37        startup: &'a StartupMessage,
38        context: &'a ServerConnectionContext<Peer, Identity>,
39        state: &'a mut State,
40    ) -> Pin<Box<dyn Future<Output = Result<Self::Route, Self::Error>> + 'a>>;
41}
42
43/// Failure from either server establishment or application startup routing.
44#[derive(Debug)]
45pub(crate) enum RoutedAcceptError<TlsError, AuthenticationError, RouteError> {
46    Accept(AcceptError<TlsError, AuthenticationError>),
47    Route(RouteError),
48}
49
50impl<TlsError, AuthenticationError, RouteError> From<AcceptError<TlsError, AuthenticationError>>
51    for RoutedAcceptError<TlsError, AuthenticationError, RouteError>
52{
53    fn from(error: AcceptError<TlsError, AuthenticationError>) -> Self {
54        Self::Accept(error)
55    }
56}
57
58struct NoStartupRoute;
59
60impl<State, Peer, Identity> StartupResolver<State, Peer, Identity> for NoStartupRoute {
61    type Route = ();
62    type Error = std::convert::Infallible;
63
64    fn resolve<'a>(
65        &'a mut self,
66        _startup: &'a StartupMessage,
67        _context: &'a ServerConnectionContext<Peer, Identity>,
68        _state: &'a mut State,
69    ) -> Pin<Box<dyn Future<Output = Result<(), Self::Error>> + 'a>> {
70        Box::pin(async { Ok(()) })
71    }
72}
73
74/// A reloadable server identity resolved for each TLS connection.
75#[derive(Clone)]
76pub struct ServerIdentity {
77    config: Arc<ServerConfig>,
78    leaf_certificate: CertificateDer<'static>,
79}
80
81impl ServerIdentity {
82    /// Creates an identity from a rustls configuration and its leaf certificate.
83    ///
84    /// `leaf_certificate` must be the certificate selected by `config`; it is
85    /// retained separately because rustls does not expose configured resolver
86    /// certificates for PostgreSQL channel-binding derivation.
87    #[must_use]
88    pub const fn new(config: Arc<ServerConfig>, leaf_certificate: CertificateDer<'static>) -> Self {
89        Self {
90            config,
91            leaf_certificate,
92        }
93    }
94}
95
96impl fmt::Debug for ServerIdentity {
97    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
98        formatter.write_str("ServerIdentity([REDACTED])")
99    }
100}
101
102/// Application-owned source of the current TLS identity.
103pub trait ServerIdentityProvider {
104    /// Failure returned while resolving the current identity.
105    type Error;
106
107    /// Resolves the identity to use for one new connection.
108    ///
109    /// # Errors
110    ///
111    /// Returns the provider's error when no current identity is available.
112    fn resolve(&self) -> Result<ServerIdentity, Self::Error>;
113}
114
115/// TLS facts recorded after pre-startup negotiation.
116#[derive(Clone, Debug, Eq, PartialEq)]
117pub enum NegotiatedServerTls {
118    /// The connection is intentionally plaintext.
119    Plaintext,
120    /// TLS was negotiated and terminated by this server component.
121    Tls {
122        /// RFC 5929 `tls-server-end-point` channel-binding bytes.
123        server_end_point: Bytes,
124    },
125}
126
127/// Immutable inputs available to one authentication session.
128pub struct ServerAuthenticationRequest<'a, Peer> {
129    startup: &'a StartupMessage,
130    tls: &'a NegotiatedServerTls,
131    peer: &'a Peer,
132}
133
134impl<Peer> Copy for ServerAuthenticationRequest<'_, Peer> {}
135
136impl<Peer> Clone for ServerAuthenticationRequest<'_, Peer> {
137    fn clone(&self) -> Self {
138        *self
139    }
140}
141
142impl<Peer> ServerAuthenticationRequest<'_, Peer> {
143    /// Returns the accepted startup message.
144    #[must_use]
145    pub const fn startup(&self) -> &StartupMessage {
146        self.startup
147    }
148
149    /// Returns the negotiated transport security fact.
150    #[must_use]
151    pub const fn tls(&self) -> &NegotiatedServerTls {
152        self.tls
153    }
154
155    /// Returns immutable caller-supplied peer facts.
156    #[must_use]
157    pub const fn peer(&self) -> &Peer {
158        self.peer
159    }
160}
161
162/// The next protocol action selected by application authentication policy.
163#[derive(Clone, Debug, Eq, PartialEq)]
164pub enum ServerAuthenticationAction<Identity> {
165    /// Authentication is complete with typed identity evidence.
166    Accept(Identity),
167    /// Request a PostgreSQL cleartext password response.
168    CleartextPassword,
169    /// Request a PostgreSQL MD5 password response using the supplied salt.
170    Md5Password {
171        /// Four-byte server challenge salt.
172        salt: [u8; 4],
173    },
174    /// Offer SASL mechanisms and receive the client's initial response.
175    Sasl {
176        /// Mechanism names offered in preference order.
177        mechanisms: Vec<Bytes>,
178    },
179    /// Send a recursive SASL challenge.
180    SaslContinue(Bytes),
181    /// Send SASL server-final data and complete with typed identity evidence.
182    SaslFinal {
183        /// Verified mechanism-specific server-final data.
184        server_final: Bytes,
185        /// Typed identity evidence produced by the policy.
186        identity: Identity,
187    },
188    /// Request a Kerberos V5 response token.
189    KerberosV5,
190    /// Request a GSSAPI response token.
191    Gss,
192    /// Request an SSPI response token.
193    Sspi,
194    /// Send a recursive GSS continuation token.
195    GssContinue(Bytes),
196}
197
198/// Owned client response supplied to application authentication policy.
199#[derive(Clone, Debug, Eq, PartialEq)]
200pub enum ServerAuthenticationResponse {
201    /// Cleartext or MD5 password response body.
202    Password(Bytes),
203    /// SASL mechanism selection and optional initial response.
204    SaslInitial {
205        /// Mechanism selected by the client.
206        mechanism: Bytes,
207        /// Optional mechanism-specific initial data.
208        response: Option<Bytes>,
209    },
210    /// Recursive SASL response body.
211    Sasl(Bytes),
212    /// Kerberos, GSSAPI, or SSPI response token.
213    Token(Bytes),
214}
215
216/// Per-connection asynchronous authentication policy.
217///
218/// Authentication futures are not required to be [`Send`].
219#[allow(async_fn_in_trait)]
220pub trait ServerAuthentication<Peer> {
221    /// Typed evidence produced by successful authentication.
222    type Identity;
223    /// Application-defined authentication failure.
224    type Error;
225
226    /// Starts the application-driven authentication conversation.
227    async fn start(
228        &mut self,
229        request: ServerAuthenticationRequest<'_, Peer>,
230    ) -> Result<ServerAuthenticationAction<Self::Identity>, Self::Error>;
231
232    /// Advances the conversation after one protocol response.
233    async fn respond(
234        &mut self,
235        request: ServerAuthenticationRequest<'_, Peer>,
236        response: ServerAuthenticationResponse,
237    ) -> Result<ServerAuthenticationAction<Self::Identity>, Self::Error>;
238}
239
240/// Factory creating one isolated authentication policy per connection.
241pub trait ServerAuthenticationProvider {
242    /// Per-connection policy type.
243    type Authentication;
244
245    /// Creates a fresh policy instance.
246    fn create(&self) -> Self::Authentication;
247}
248
249/// Static username/password verification using PostgreSQL MD5 authentication.
250#[derive(Clone, Debug, Eq, PartialEq)]
251pub struct StaticMd5ServerCredentials {
252    username: Bytes,
253    password: Bytes,
254}
255
256impl StaticMd5ServerCredentials {
257    /// Creates a reusable credential provider.
258    #[must_use]
259    pub fn new(username: impl Into<Bytes>, password: impl Into<Bytes>) -> Self {
260        Self {
261            username: username.into(),
262            password: password.into(),
263        }
264    }
265}
266
267/// Isolated MD5 exchange for one connection.
268#[derive(Clone, Debug)]
269pub struct StaticMd5ServerCredentialSession {
270    username: Bytes,
271    password: Bytes,
272    salt: [u8; 4],
273}
274
275impl ServerAuthenticationProvider for StaticMd5ServerCredentials {
276    type Authentication = StaticMd5ServerCredentialSession;
277
278    fn create(&self) -> Self::Authentication {
279        let mut salt = [0; 4];
280        rand::rng().fill(&mut salt);
281        StaticMd5ServerCredentialSession {
282            username: self.username.clone(),
283            password: self.password.clone(),
284            salt,
285        }
286    }
287}
288
289impl<Peer> ServerAuthentication<Peer> for StaticMd5ServerCredentialSession {
290    type Identity = ();
291    type Error = crate::StaticCredentialError;
292
293    async fn start(
294        &mut self,
295        _: ServerAuthenticationRequest<'_, Peer>,
296    ) -> Result<ServerAuthenticationAction<()>, Self::Error> {
297        Ok(ServerAuthenticationAction::Md5Password { salt: self.salt })
298    }
299
300    async fn respond(
301        &mut self,
302        _: ServerAuthenticationRequest<'_, Peer>,
303        response: ServerAuthenticationResponse,
304    ) -> Result<ServerAuthenticationAction<()>, Self::Error> {
305        let ServerAuthenticationResponse::Password(received) = response else {
306            return Err(crate::StaticCredentialError::AuthenticationFailed);
307        };
308        if !crate::credentials::verify_md5_response(
309            &received,
310            &self.username,
311            &self.password,
312            self.salt,
313        ) {
314            return Err(crate::StaticCredentialError::AuthenticationFailed);
315        }
316        Ok(ServerAuthenticationAction::Accept(()))
317    }
318}
319
320/// Typed evidence for an explicitly trusted connection.
321#[derive(Clone, Copy, Debug, Eq, PartialEq)]
322pub struct TrustIdentity;
323
324/// Deterministic failures while constructing a reusable server component.
325#[derive(Clone, Copy, Debug, Eq, PartialEq)]
326pub enum BuildServerError {
327    /// No TLS posture was selected.
328    MissingTlsPolicy,
329    /// No authentication posture was selected.
330    MissingAuthenticationPolicy,
331    /// A protocol limit cannot support server-role establishment.
332    InvalidProtocolLimits,
333}
334
335impl fmt::Display for BuildServerError {
336    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
337        formatter.write_str(match self {
338            Self::MissingTlsPolicy => "server TLS policy is required",
339            Self::MissingAuthenticationPolicy => "server authentication policy is required",
340            Self::InvalidProtocolLimits => {
341                "server protocol limits cannot support connection establishment"
342            }
343        })
344    }
345}
346
347impl std::error::Error for BuildServerError {}
348
349/// Failures while establishing a live server-role connection.
350#[derive(Debug)]
351pub enum AcceptError<TlsError = NoServerIdentity, AuthenticationError = std::convert::Infallible> {
352    /// The transport failed or the peer sent invalid wire data.
353    Io(io::Error),
354    /// The startup packet requested an unsupported protocol major version.
355    UnsupportedProtocolVersion,
356    /// The configured policy requires TLS before startup.
357    TlsRequired,
358    /// The current TLS identity could not be resolved.
359    TlsIdentity(TlsError),
360    /// Application authentication rejected the connection.
361    Authentication(AuthenticationError),
362    /// The client sent a message invalid for the selected authentication mechanism.
363    AuthenticationProtocol,
364}
365
366impl<TlsError, AuthenticationError> fmt::Display for AcceptError<TlsError, AuthenticationError> {
367    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
368        match self {
369            Self::Io(error) => error.fmt(formatter),
370            Self::UnsupportedProtocolVersion => {
371                formatter.write_str("unsupported PostgreSQL protocol version")
372            }
373            Self::TlsRequired => formatter.write_str("TLS is required before startup"),
374            Self::TlsIdentity(_) => formatter.write_str("server TLS identity is unavailable"),
375            Self::Authentication(_) => formatter.write_str("authentication rejected"),
376            Self::AuthenticationProtocol => formatter.write_str("invalid authentication response"),
377        }
378    }
379}
380
381impl<TlsError, AuthenticationError> std::error::Error for AcceptError<TlsError, AuthenticationError>
382where
383    TlsError: std::error::Error + 'static,
384    AuthenticationError: std::error::Error + 'static,
385{
386    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
387        match self {
388            Self::Io(error) => Some(error),
389            Self::TlsIdentity(error) => Some(error),
390            Self::Authentication(error) => Some(error),
391            Self::UnsupportedProtocolVersion | Self::TlsRequired | Self::AuthenticationProtocol => {
392                None
393            }
394        }
395    }
396}
397
398/// Namespace for explicit server-side TLS policy values.
399#[derive(Clone, Copy, Debug, Eq, PartialEq)]
400pub struct ServerTlsPolicy;
401
402impl ServerTlsPolicy {
403    /// Deliberately serve plaintext and decline encryption negotiation.
404    #[allow(non_upper_case_globals)]
405    pub const Disabled: DisabledServerTls = DisabledServerTls;
406
407    /// Accepts plaintext or terminates TLS using identities from `provider`.
408    #[allow(non_snake_case)]
409    pub const fn Optional<Provider>(provider: Provider) -> OptionalServerTls<Provider> {
410        OptionalServerTls(provider)
411    }
412
413    /// Requires TLS using identities from `provider`.
414    #[allow(non_snake_case)]
415    pub const fn Required<Provider>(provider: Provider) -> RequiredServerTls<Provider> {
416        RequiredServerTls(provider)
417    }
418}
419
420/// Explicit plaintext-only server TLS policy.
421#[derive(Clone, Copy, Debug, Eq, PartialEq)]
422pub struct DisabledServerTls;
423
424/// Optional server TLS termination backed by a reloadable identity provider.
425#[derive(Clone)]
426pub struct OptionalServerTls<Provider>(Provider);
427
428/// Required server TLS termination backed by a reloadable identity provider.
429#[derive(Clone)]
430pub struct RequiredServerTls<Provider>(Provider);
431
432impl<Provider> fmt::Debug for OptionalServerTls<Provider> {
433    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
434        formatter.write_str("OptionalServerTls([REDACTED])")
435    }
436}
437
438impl<Provider> fmt::Debug for RequiredServerTls<Provider> {
439    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
440        formatter.write_str("RequiredServerTls([REDACTED])")
441    }
442}
443
444/// Marker provider used by the disabled TLS policy.
445#[derive(Clone, Copy, Debug, Eq, PartialEq)]
446pub struct NoServerIdentityProvider;
447
448/// Error returned by the marker provider, which has no identity.
449#[derive(Clone, Copy, Debug, Eq, PartialEq)]
450pub struct NoServerIdentity;
451
452impl fmt::Display for NoServerIdentity {
453    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
454        formatter.write_str("disabled TLS has no identity")
455    }
456}
457
458impl std::error::Error for NoServerIdentity {}
459
460impl ServerIdentityProvider for NoServerIdentityProvider {
461    type Error = NoServerIdentity;
462
463    fn resolve(&self) -> Result<ServerIdentity, Self::Error> {
464        Err(NoServerIdentity)
465    }
466}
467
468mod sealed {
469    pub trait Sealed {}
470}
471
472/// TLS configuration implemented by the facade policy values.
473#[doc(hidden)]
474pub trait ServerTlsConfiguration: sealed::Sealed {
475    /// Identity provider associated with this policy.
476    type Provider: ServerIdentityProvider;
477    /// Returns the provider when this policy can terminate TLS.
478    fn provider(&self) -> Option<&Self::Provider>;
479    /// Reports whether plaintext startup must be rejected.
480    fn required(&self) -> bool;
481    /// Returns a non-sensitive structural category for diagnostics.
482    fn category(&self) -> &'static str;
483}
484
485impl sealed::Sealed for DisabledServerTls {}
486impl<Provider> sealed::Sealed for OptionalServerTls<Provider> {}
487impl<Provider> sealed::Sealed for RequiredServerTls<Provider> {}
488
489impl ServerTlsConfiguration for DisabledServerTls {
490    type Provider = NoServerIdentityProvider;
491    fn provider(&self) -> Option<&Self::Provider> {
492        None
493    }
494    fn required(&self) -> bool {
495        false
496    }
497    fn category(&self) -> &'static str {
498        "disabled"
499    }
500}
501
502impl<Provider: ServerIdentityProvider> ServerTlsConfiguration for OptionalServerTls<Provider> {
503    type Provider = Provider;
504    fn provider(&self) -> Option<&Self::Provider> {
505        Some(&self.0)
506    }
507    fn required(&self) -> bool {
508        false
509    }
510    fn category(&self) -> &'static str {
511        "optional"
512    }
513}
514
515impl<Provider: ServerIdentityProvider> ServerTlsConfiguration for RequiredServerTls<Provider> {
516    type Provider = Provider;
517    fn provider(&self) -> Option<&Self::Provider> {
518        Some(&self.0)
519    }
520    fn required(&self) -> bool {
521        true
522    }
523    fn category(&self) -> &'static str {
524        "required"
525    }
526}
527
528/// Explicit trust authentication, which accepts every protocol-compatible client.
529#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
530pub struct TrustServerAuthentication;
531
532impl ServerAuthenticationProvider for TrustServerAuthentication {
533    type Authentication = Self;
534
535    fn create(&self) -> Self::Authentication {
536        *self
537    }
538}
539
540impl<Peer> ServerAuthentication<Peer> for TrustServerAuthentication {
541    type Identity = TrustIdentity;
542    type Error = std::convert::Infallible;
543
544    async fn start(
545        &mut self,
546        _request: ServerAuthenticationRequest<'_, Peer>,
547    ) -> Result<ServerAuthenticationAction<Self::Identity>, Self::Error> {
548        Ok(ServerAuthenticationAction::Accept(TrustIdentity))
549    }
550
551    async fn respond(
552        &mut self,
553        _request: ServerAuthenticationRequest<'_, Peer>,
554        _response: ServerAuthenticationResponse,
555    ) -> Result<ServerAuthenticationAction<Self::Identity>, Self::Error> {
556        unreachable!("trust authentication accepts before a response")
557    }
558}
559
560/// Conservative allocation limits applied to newly accepted transports.
561#[derive(Clone, Copy, Debug, Eq, PartialEq)]
562pub struct ServerProtocolLimits {
563    max_frame_len: usize,
564    max_pre_startup_packet_len: usize,
565}
566
567impl ServerProtocolLimits {
568    /// Returns limits with a different maximum tagged-frame length.
569    #[must_use]
570    pub const fn with_max_frame_len(mut self, bytes: usize) -> Self {
571        self.max_frame_len = bytes;
572        self
573    }
574
575    /// Returns limits with a different maximum untagged startup-packet length.
576    #[must_use]
577    pub const fn with_max_pre_startup_packet_len(mut self, bytes: usize) -> Self {
578        self.max_pre_startup_packet_len = bytes;
579        self
580    }
581
582    const fn is_valid(self) -> bool {
583        self.max_frame_len >= 9
584            && self.max_frame_len <= i32::MAX as usize
585            && self.max_pre_startup_packet_len >= 8
586            && self.max_pre_startup_packet_len <= i32::MAX as usize
587    }
588}
589
590impl Default for ServerProtocolLimits {
591    fn default() -> Self {
592        Self {
593            max_frame_len: DEFAULT_MAX_FRAME_LEN,
594            max_pre_startup_packet_len: DEFAULT_MAX_PRE_STARTUP_PACKET_LEN,
595        }
596    }
597}
598
599/// Reusable client-facing PostgreSQL server component.
600#[derive(Clone)]
601pub struct Server<
602    Tls = DisabledServerTls,
603    Authentication = TrustServerAuthentication,
604    Middleware = IdentityServerHandler,
605> {
606    tls: Tls,
607    authentication: Authentication,
608    limits: ServerProtocolLimits,
609    middleware: Middleware,
610}
611
612impl Server {
613    /// Starts configuration of a reusable server component.
614    #[must_use]
615    pub fn builder() -> ServerBuilder {
616        ServerBuilder::default()
617    }
618}
619
620impl<Tls: ServerTlsConfiguration, Authentication, Middleware> fmt::Debug
621    for Server<Tls, Authentication, Middleware>
622{
623    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
624        formatter
625            .debug_struct("Server")
626            .field("tls", &self.tls.category())
627            .field("authentication", &"<redacted>")
628            .field("limits", &self.limits)
629            .finish_non_exhaustive()
630    }
631}
632
633/// Builder for a reusable [`Server`].
634#[derive(Clone)]
635pub struct ServerBuilder<Tls = (), Authentication = (), Middleware = IdentityServerHandler> {
636    tls: Option<Tls>,
637    authentication: Option<Authentication>,
638    limits: ServerProtocolLimits,
639    middleware: Middleware,
640}
641
642impl<Tls, Authentication> fmt::Debug for ServerBuilder<Tls, Authentication> {
643    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
644        formatter
645            .debug_struct("ServerBuilder")
646            .field("tls_configured", &self.tls.is_some())
647            .field("authentication_configured", &self.authentication.is_some())
648            .field("limits", &self.limits)
649            .finish()
650    }
651}
652
653impl Default for ServerBuilder {
654    fn default() -> Self {
655        Self {
656            tls: None,
657            authentication: None,
658            limits: ServerProtocolLimits::default(),
659            middleware: IdentityServerHandler,
660        }
661    }
662}
663
664impl<Tls, Authentication, Middleware> ServerBuilder<Tls, Authentication, Middleware> {
665    /// Selects the client-facing TLS posture explicitly.
666    #[must_use]
667    pub fn tls<Next>(self, policy: Next) -> ServerBuilder<Next, Authentication, Middleware> {
668        ServerBuilder {
669            tls: Some(policy),
670            authentication: self.authentication,
671            limits: self.limits,
672            middleware: self.middleware,
673        }
674    }
675
676    /// Replaces the authentication policy used for each accepted connection.
677    #[must_use]
678    pub fn authentication<Next>(self, policy: Next) -> ServerBuilder<Tls, Next, Middleware> {
679        ServerBuilder {
680            tls: self.tls,
681            authentication: Some(policy),
682            limits: self.limits,
683            middleware: self.middleware,
684        }
685    }
686
687    /// Replaces the conservative protocol limits.
688    #[must_use]
689    pub fn limits(mut self, limits: ServerProtocolLimits) -> Self {
690        self.limits = limits;
691        self
692    }
693
694    /// Appends a synchronous, infallible per-connection middleware factory.
695    #[must_use]
696    pub fn middleware<Next>(
697        self,
698        factory: Next,
699    ) -> ServerBuilder<Tls, Authentication, crate::MiddlewareChain<Middleware, Next>> {
700        ServerBuilder {
701            tls: self.tls,
702            authentication: self.authentication,
703            limits: self.limits,
704            middleware: crate::MiddlewareChain(self.middleware, factory),
705        }
706    }
707
708    /// Validates configuration and creates an immutable reusable component.
709    ///
710    /// # Errors
711    ///
712    /// Returns an error when either security policy is omitted or a protocol
713    /// limit cannot be represented by the PostgreSQL wire format.
714    pub fn build(self) -> Result<Server<Tls, Authentication, Middleware>, BuildServerError> {
715        let tls = self.tls.ok_or(BuildServerError::MissingTlsPolicy)?;
716        let authentication = self
717            .authentication
718            .ok_or(BuildServerError::MissingAuthenticationPolicy)?;
719        if !self.limits.is_valid() {
720            return Err(BuildServerError::InvalidProtocolLimits);
721        }
722        Ok(Server {
723            tls,
724            authentication,
725            limits: self.limits,
726            middleware: self.middleware,
727        })
728    }
729}
730
731/// Immutable facts known about a client-facing connection.
732#[derive(Clone, Debug, Eq, PartialEq)]
733pub struct ServerConnectionContext<Peer, Identity> {
734    peer: Peer,
735    tls: Option<NegotiatedServerTls>,
736    identity: Option<Identity>,
737}
738
739impl<Peer, Identity> ServerConnectionContext<Peer, Identity> {
740    /// Returns caller-provided peer metadata.
741    #[must_use]
742    pub const fn peer(&self) -> &Peer {
743        &self.peer
744    }
745
746    /// Returns the negotiated TLS fact.
747    ///
748    /// # Panics
749    ///
750    /// Panics when called before pre-startup negotiation completes.
751    #[must_use]
752    pub const fn tls(&self) -> &NegotiatedServerTls {
753        match &self.tls {
754            Some(tls) => tls,
755            None => panic!("TLS is not known before pre-startup negotiation"),
756        }
757    }
758
759    /// Returns TLS evidence only after pre-startup negotiation has completed.
760    #[must_use]
761    pub const fn tls_if_known(&self) -> Option<&NegotiatedServerTls> {
762        self.tls.as_ref()
763    }
764
765    /// Returns typed evidence from application authentication.
766    ///
767    /// # Panics
768    ///
769    /// Panics when called before authentication completes.
770    #[must_use]
771    pub const fn identity(&self) -> &Identity {
772        match &self.identity {
773            Some(identity) => identity,
774            None => panic!("identity is not known before authentication"),
775        }
776    }
777
778    /// Returns identity evidence only after authentication has completed.
779    #[must_use]
780    pub const fn identity_if_known(&self) -> Option<&Identity> {
781        self.identity.as_ref()
782    }
783}
784
785/// Identity handler used until contextual middleware is configured.
786#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
787pub struct IdentityServerHandler;
788impl<C> crate::MiddlewareFactory<C> for IdentityServerHandler {
789    type Handler = Self;
790    fn create(&self, _: &C) -> Self {
791        *self
792    }
793}
794impl<S, C> crate::ServerMiddleware<S, C> for IdentityServerHandler {}
795
796/// A decoded out-of-band PostgreSQL cancellation request.
797#[derive(Clone, Eq, PartialEq)]
798pub struct CancellationRequest {
799    process_id: u32,
800    secret_key: Bytes,
801}
802
803impl fmt::Debug for CancellationRequest {
804    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
805        formatter
806            .debug_struct("CancellationRequest")
807            .field("process_id", &self.process_id)
808            .field("secret_key", &"[REDACTED]")
809            .finish()
810    }
811}
812
813impl CancellationRequest {
814    /// Returns the backend process identifier supplied by the client.
815    #[must_use]
816    pub const fn process_id(&self) -> u32 {
817        self.process_id
818    }
819
820    /// Returns the opaque cancellation key supplied by the client.
821    #[must_use]
822    pub fn secret_key(&self) -> &[u8] {
823        &self.secret_key
824    }
825}
826
827/// Result of accepting one caller-established transport.
828#[derive(Debug)]
829#[allow(clippy::large_enum_variant)]
830pub enum ServerAccept<
831    Transport,
832    State,
833    Peer,
834    Identity = TrustIdentity,
835    Handler = IdentityServerHandler,
836> {
837    /// Authentication completed and the connection is operational.
838    Session(ServerConnection<Transport, State, Peer, Identity, Handler>),
839    /// The first packet was an out-of-band cancellation request.
840    Cancellation(ServerCancellation<Transport, State, Peer, Handler>),
841}
842
843/// Non-`Send` future returned while accepting one server-role connection.
844pub type ServerAcceptFuture<
845    'a,
846    Transport,
847    State,
848    Peer,
849    Identity,
850    Handler,
851    TlsError,
852    AuthenticationError,
853> = Pin<
854    Box<
855        dyn Future<
856                Output = Result<
857                    ServerAccept<Transport, State, Peer, Identity, Handler>,
858                    AcceptError<TlsError, AuthenticationError>,
859                >,
860            > + 'a,
861    >,
862>;
863
864/// An operational server-role connection with all per-connection ownership.
865#[derive(Debug)]
866pub struct ServerConnection<
867    Transport,
868    State,
869    Peer,
870    Identity = TrustIdentity,
871    Handler = IdentityServerHandler,
872> {
873    core: ServerConnectionCore<Transport, Peer, Identity, Handler>,
874    state: State,
875}
876
877#[derive(Debug)]
878pub(crate) struct ServerConnectionCore<Transport, Peer, Identity, Handler> {
879    conn: ServerConnectionInner<Transport>,
880    startup: StartupMessage,
881    handler: Handler,
882    context: ServerConnectionContext<Peer, Identity>,
883}
884
885#[derive(Debug)]
886enum ServerConnectionInner<Transport> {
887    Plaintext(Box<Conn<Buffered<Transport, Frontend>, Ready>>),
888    Tls(Box<Conn<Buffered<ServerTls<Transport>, Frontend>, Ready>>),
889}
890
891/// Transport recovered when a server connection is explicitly torn down.
892#[derive(Debug)]
893pub enum AcceptedServerTransport<Transport> {
894    /// The original plaintext transport.
895    Plaintext(Transport),
896    /// A TLS stream over the original transport.
897    Tls(Box<ServerTls<Transport>>),
898}
899
900impl<Transport, State, Peer, Identity, Handler>
901    ServerConnection<Transport, State, Peer, Identity, Handler>
902{
903    /// Returns immutable connection facts.
904    #[must_use]
905    pub const fn context(&self) -> &ServerConnectionContext<Peer, Identity> {
906        &self.core.context
907    }
908
909    /// Returns the caller-owned connection state.
910    ///
911    #[must_use]
912    pub const fn state(&self) -> &State {
913        &self.state
914    }
915
916    pub(crate) fn into_core_and_state(
917        self,
918    ) -> (
919        ServerConnectionCore<Transport, Peer, Identity, Handler>,
920        State,
921    ) {
922        (self.core, self.state)
923    }
924
925    /// Returns the accepted startup parameters.
926    #[must_use]
927    pub const fn startup(&self) -> &StartupMessage {
928        &self.core.startup
929    }
930
931    /// Receives one operational frontend wire message without advancing the
932    /// typed session projection.
933    ///
934    /// This is the inspection boundary: application policy may inspect or
935    /// rewrite the owned message before a later facade operation projects it.
936    ///
937    /// # Errors
938    ///
939    /// Returns a transport, decoding, or configured frame-limit error.
940    ///
941    pub async fn receive_wire(&mut self) -> io::Result<FrontendMessage>
942    where
943        Transport: AsyncRead + AsyncWrite + Unpin,
944        Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
945    {
946        let message = self.core.receive_wire_raw().await?;
947        Ok(self.core.intercept_frontend(&mut self.state, message))
948    }
949
950    /// Sends one operational backend message after middleware interception.
951    ///
952    /// The replacement returned by middleware is the value encoded on the wire.
953    ///
954    /// # Errors
955    ///
956    /// Returns an encoding, configured frame-limit, or transport error.
957    ///
958    pub async fn send_wire(&mut self, message: BackendMessage) -> io::Result<()>
959    where
960        Transport: AsyncRead + AsyncWrite + Unpin,
961        Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
962    {
963        let message = self.core.intercept_backend(&mut self.state, message);
964        self.core.send_wire_raw(message).await
965    }
966
967    /// Deliberately ends typed ownership and recovers every connection part.
968    ///
969    #[must_use]
970    pub fn teardown(
971        self,
972    ) -> (
973        AcceptedServerTransport<Transport>,
974        State,
975        Handler,
976        ServerConnectionContext<Peer, Identity>,
977    ) {
978        let (transport, handler, context) = self.core.into_parts();
979        (transport, self.state, handler, context)
980    }
981}
982
983impl<Transport, State, Peer, Identity, Handler>
984    ServerConnection<Transport, State, Peer, Identity, Handler>
985where
986    Transport: AsyncRead + AsyncWrite + Unpin,
987    Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
988{
989    pub(crate) async fn send_generated_error(&mut self, message: BackendMessage) -> io::Result<()> {
990        let message = self.core.intercept_backend(&mut self.state, message);
991        if !matches!(message, BackendMessage::ErrorResponse(_)) {
992            return Err(io::Error::new(
993                io::ErrorKind::InvalidData,
994                "middleware rejected generated diagnostic",
995            ));
996        }
997        self.core.send_wire_raw(message).await
998    }
999}
1000
1001impl<Transport, Peer, Identity, Handler> ServerConnectionCore<Transport, Peer, Identity, Handler> {
1002    pub(crate) const fn context(&self) -> &ServerConnectionContext<Peer, Identity> {
1003        &self.context
1004    }
1005
1006    pub(crate) async fn receive_wire_raw(&mut self) -> io::Result<FrontendMessage>
1007    where
1008        Transport: AsyncRead + AsyncWrite + Unpin,
1009    {
1010        match &mut self.conn {
1011            ServerConnectionInner::Plaintext(conn) => conn.receive_frontend_wire().await,
1012            ServerConnectionInner::Tls(conn) => conn.receive_frontend_wire().await,
1013        }
1014    }
1015
1016    pub(crate) fn intercept_frontend<State>(
1017        &mut self,
1018        state: &mut State,
1019        message: FrontendMessage,
1020    ) -> FrontendMessage
1021    where
1022        Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
1023    {
1024        self.handler.frontend(&self.context, state, message)
1025    }
1026
1027    pub(crate) fn intercept_backend<State>(
1028        &mut self,
1029        state: &mut State,
1030        message: BackendMessage,
1031    ) -> BackendMessage
1032    where
1033        Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
1034    {
1035        self.handler.backend(&self.context, state, message)
1036    }
1037
1038    pub(crate) async fn send_wire_raw(&mut self, message: BackendMessage) -> io::Result<()>
1039    where
1040        Transport: AsyncRead + AsyncWrite + Unpin,
1041    {
1042        let frame = message.to_frame()?;
1043        match &mut self.conn {
1044            ServerConnectionInner::Plaintext(conn) => {
1045                conn.push_frame(frame)?;
1046                conn.flush().await
1047            }
1048            ServerConnectionInner::Tls(conn) => {
1049                conn.push_frame(frame)?;
1050                conn.flush().await
1051            }
1052        }
1053    }
1054
1055    pub(crate) fn into_parts(
1056        self,
1057    ) -> (
1058        AcceptedServerTransport<Transport>,
1059        Handler,
1060        ServerConnectionContext<Peer, Identity>,
1061    ) {
1062        let transport = match self.conn {
1063            ServerConnectionInner::Plaintext(conn) => {
1064                AcceptedServerTransport::Plaintext(conn.into_transport().into_inner())
1065            }
1066            ServerConnectionInner::Tls(conn) => {
1067                AcceptedServerTransport::Tls(Box::new(conn.into_transport().into_inner()))
1068            }
1069        };
1070        (transport, self.handler, self.context)
1071    }
1072}
1073
1074/// A cancellation branch retaining all caller and handler ownership.
1075#[derive(Debug)]
1076pub struct ServerCancellation<Transport, State, Peer, Handler = IdentityServerHandler> {
1077    transport: AcceptedServerTransport<Transport>,
1078    request: CancellationRequest,
1079    state: State,
1080    handler: Handler,
1081    context: ServerConnectionContext<Peer, ()>,
1082}
1083
1084impl<Transport, State, Peer, Handler> ServerCancellation<Transport, State, Peer, Handler> {
1085    /// Returns the decoded request.
1086    #[must_use]
1087    pub const fn request(&self) -> &CancellationRequest {
1088        &self.request
1089    }
1090
1091    /// Recovers every owned cancellation-connection part.
1092    #[must_use]
1093    pub fn teardown(
1094        self,
1095    ) -> (
1096        AcceptedServerTransport<Transport>,
1097        CancellationRequest,
1098        State,
1099        Handler,
1100        ServerConnectionContext<Peer, ()>,
1101    ) {
1102        (
1103            self.transport,
1104            self.request,
1105            self.state,
1106            self.handler,
1107            self.context,
1108        )
1109    }
1110}
1111
1112impl<Tls, Authentication, Middleware> Server<Tls, Authentication, Middleware>
1113where
1114    Tls: ServerTlsConfiguration,
1115    Authentication: ServerAuthenticationProvider,
1116{
1117    /// Accepts one transport through TLS negotiation and application authentication.
1118    ///
1119    /// The caller retains listener and task ownership. `state` and `peer` become
1120    /// owned parts of either returned branch.
1121    ///
1122    /// # Errors
1123    ///
1124    /// Returns [`AcceptError::Io`] for transport or wire failures,
1125    /// [`AcceptError::UnsupportedProtocolVersion`] for unsupported startup,
1126    /// [`AcceptError::TlsRequired`] when required TLS was bypassed,
1127    /// [`AcceptError::TlsIdentity`] with the provider's typed error when the
1128    /// current identity cannot be resolved, [`AcceptError::Authentication`]
1129    /// with the policy's typed error when authentication rejects the client,
1130    /// or [`AcceptError::AuthenticationProtocol`] for an invalid response to
1131    /// the selected wire mechanism.
1132    #[allow(clippy::type_complexity)]
1133    pub fn accept<'a, Transport, State, Peer>(
1134        &'a self,
1135        transport: Transport,
1136        peer: Peer,
1137        state: State,
1138    ) -> ServerAcceptFuture<
1139        'a,
1140        Transport,
1141        State,
1142        Peer,
1143        <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1144        <Middleware as crate::MiddlewareFactory<
1145            ServerConnectionContext<
1146                Peer,
1147                <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1148            >,
1149        >>::Handler,
1150        <Tls::Provider as ServerIdentityProvider>::Error,
1151        <Authentication::Authentication as ServerAuthentication<Peer>>::Error,
1152    >
1153    where
1154        Transport: AsyncRead + AsyncWrite + Unpin + 'a,
1155        State: 'a,
1156        Peer: 'a,
1157        Authentication::Authentication: ServerAuthentication<Peer>,
1158        Middleware: crate::MiddlewareFactory<
1159                ServerConnectionContext<
1160                    Peer,
1161                    <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1162                >,
1163            >,
1164        <Middleware as crate::MiddlewareFactory<
1165            ServerConnectionContext<
1166                Peer,
1167                <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1168            >,
1169        >>::Handler: crate::ServerMiddleware<
1170                State,
1171                ServerConnectionContext<
1172                    Peer,
1173                    <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1174                >,
1175            >,
1176    {
1177        Box::pin(async move {
1178            let mut resolver = NoStartupRoute;
1179            self.accept_routed(transport, peer, state, &mut resolver)
1180                .await
1181                .map(|(accepted, _)| accepted)
1182                .map_err(|error| match error {
1183                    RoutedAcceptError::Accept(error) => error,
1184                    RoutedAcceptError::Route(never) => match never {},
1185                })
1186        })
1187    }
1188
1189    #[allow(clippy::too_many_lines)]
1190    pub(crate) async fn accept_routed<Transport, State, Peer, Resolver>(
1191        &self,
1192        transport: Transport,
1193        peer: Peer,
1194        mut state: State,
1195        resolver: &mut Resolver,
1196    ) -> Result<
1197        (
1198            ServerAccept<
1199                Transport,
1200                State,
1201                Peer,
1202                <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1203                <Middleware as crate::MiddlewareFactory<
1204                    ServerConnectionContext<
1205                        Peer,
1206                        <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1207                    >,
1208                >>::Handler,
1209            >,
1210            Option<Resolver::Route>,
1211        ),
1212        RoutedAcceptError<
1213            <Tls::Provider as ServerIdentityProvider>::Error,
1214            <Authentication::Authentication as ServerAuthentication<Peer>>::Error,
1215            Resolver::Error,
1216        >,
1217    >
1218    where
1219        Transport: AsyncRead + AsyncWrite + Unpin,
1220        Authentication::Authentication: ServerAuthentication<Peer>,
1221        Middleware: crate::MiddlewareFactory<
1222                ServerConnectionContext<
1223                    Peer,
1224                    <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1225                >,
1226            >,
1227        <Middleware as crate::MiddlewareFactory<
1228            ServerConnectionContext<
1229                Peer,
1230                <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1231            >,
1232        >>::Handler: crate::ServerMiddleware<
1233                State,
1234                ServerConnectionContext<
1235                    Peer,
1236                    <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1237                >,
1238            >,
1239        Resolver: StartupResolver<
1240                State,
1241                Peer,
1242                <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1243            >,
1244    {
1245        let mut context = ServerConnectionContext {
1246            peer,
1247            tls: None,
1248            identity: None,
1249        };
1250        let mut handler = self.middleware.create(&context);
1251        let buffered = self
1252            .buffer_transport(transport)
1253            .map_err(AcceptError::Io)
1254            .map_err(RoutedAcceptError::Accept)?;
1255        let mut conn = Conn::new(buffered);
1256
1257        loop {
1258            let message = match conn.receive_pre_startup_wire().await {
1259                Ok(message) => message,
1260                Err(error) => {
1261                    let _ = conn.into_transport();
1262                    return Err(RoutedAcceptError::Accept(AcceptError::Io(error)));
1263                }
1264            };
1265            let message = handler.pre_startup(&context, &mut state, message);
1266            match conn.offer_pre_startup(message) {
1267                PreStartupOffer::Ssl(decision) => match self.tls.provider() {
1268                    None => {
1269                        conn = decision.decline_ssl();
1270                        conn = flush_or_abort(conn).await?;
1271                    }
1272                    Some(provider) => {
1273                        let identity = match provider.resolve() {
1274                            Ok(identity) => identity,
1275                            Err(error) => {
1276                                let _ = decision.into_transport();
1277                                return Err(RoutedAcceptError::Accept(AcceptError::TlsIdentity(
1278                                    error,
1279                                )));
1280                            }
1281                        };
1282                        let handshake = decision.approve_ssl();
1283                        let handshake = flush_or_abort(handshake).await?;
1284                        let encrypted = handshake
1285                            .accept_tls(identity.config, identity.leaf_certificate)
1286                            .await
1287                            .map_err(AcceptError::Io)?;
1288                        return accept_encrypted(
1289                            encrypted,
1290                            context,
1291                            state,
1292                            handler,
1293                            &self.authentication,
1294                            resolver,
1295                        )
1296                        .await;
1297                    }
1298                },
1299                PreStartupOffer::Gss(decision) => {
1300                    conn = decision.decline_gss();
1301                    conn = flush_or_abort(conn).await?;
1302                }
1303                PreStartupOffer::Cancel {
1304                    conn: terminal,
1305                    process_id,
1306                    secret_key,
1307                } => {
1308                    context.tls = Some(NegotiatedServerTls::Plaintext);
1309                    let request = handler.cancellation(
1310                        &context,
1311                        &mut state,
1312                        CancellationRequest {
1313                            process_id,
1314                            secret_key,
1315                        },
1316                    );
1317                    return Ok((
1318                        ServerAccept::Cancellation(ServerCancellation {
1319                            transport: AcceptedServerTransport::Plaintext(
1320                                terminal.into_transport().into_inner(),
1321                            ),
1322                            request,
1323                            state,
1324                            handler,
1325                            context: ServerConnectionContext {
1326                                peer: context.peer,
1327                                tls: Some(NegotiatedServerTls::Plaintext),
1328                                identity: None,
1329                            },
1330                        }),
1331                        None,
1332                    ));
1333                }
1334                PreStartupOffer::Startup {
1335                    conn: startup_conn,
1336                    message,
1337                } => {
1338                    if self.tls.required() {
1339                        let _ = startup_conn.into_transport();
1340                        return Err(RoutedAcceptError::Accept(AcceptError::TlsRequired));
1341                    }
1342                    context.tls = Some(NegotiatedServerTls::Plaintext);
1343                    let message = handler.startup(&context, &mut state, message);
1344                    let route = resolver
1345                        .resolve(&message, &context, &mut state)
1346                        .await
1347                        .map_err(RoutedAcceptError::Route)?;
1348                    let ready = complete_auth(
1349                        startup_conn,
1350                        &message,
1351                        &self.authentication,
1352                        &mut context,
1353                        &mut state,
1354                        &mut handler,
1355                        resolver.defer_ready(),
1356                    )
1357                    .await?;
1358                    return Ok((
1359                        ServerAccept::Session(ServerConnection {
1360                            core: ServerConnectionCore {
1361                                conn: ServerConnectionInner::Plaintext(Box::new(ready)),
1362                                startup: message,
1363                                handler,
1364                                context,
1365                            },
1366                            state,
1367                        }),
1368                        Some(route),
1369                    ));
1370                }
1371            }
1372        }
1373    }
1374
1375    fn buffer_transport<Transport>(
1376        &self,
1377        transport: Transport,
1378    ) -> io::Result<Buffered<Transport, Frontend>> {
1379        Buffered::with_limits_frontend(
1380            transport,
1381            self.limits.max_frame_len,
1382            self.limits.max_pre_startup_packet_len,
1383        )
1384    }
1385}
1386
1387async fn accept_encrypted<Transport, State, Peer, Authentication, TlsError, Handler, Resolver>(
1388    mut conn: Conn<Buffered<ServerTls<Transport>, Frontend>, crate::pre_startup::PreStartup>,
1389    mut context: ServerConnectionContext<
1390        Peer,
1391        <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1392    >,
1393    mut state: State,
1394    mut handler: Handler,
1395    authentication: &Authentication,
1396    resolver: &mut Resolver,
1397) -> Result<
1398    (
1399        ServerAccept<
1400            Transport,
1401            State,
1402            Peer,
1403            <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1404            Handler,
1405        >,
1406        Option<Resolver::Route>,
1407    ),
1408    RoutedAcceptError<
1409        TlsError,
1410        <Authentication::Authentication as ServerAuthentication<Peer>>::Error,
1411        Resolver::Error,
1412    >,
1413>
1414where
1415    Transport: AsyncRead + AsyncWrite + Unpin,
1416    Authentication: ServerAuthenticationProvider,
1417    Authentication::Authentication: ServerAuthentication<Peer>,
1418    Handler: crate::ServerMiddleware<
1419            State,
1420            ServerConnectionContext<
1421                Peer,
1422                <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1423            >,
1424        >,
1425    Resolver: StartupResolver<
1426            State,
1427            Peer,
1428            <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1429        >,
1430{
1431    let negotiated_tls = NegotiatedServerTls::Tls {
1432        server_end_point: Bytes::copy_from_slice(conn.transport().get_ref().tls_server_end_point()),
1433    };
1434    context.tls = Some(negotiated_tls.clone());
1435    loop {
1436        let message = match conn.receive_pre_startup_wire().await {
1437            Ok(message) => message,
1438            Err(error) => {
1439                let _ = conn.into_transport();
1440                return Err(RoutedAcceptError::Accept(AcceptError::Io(error)));
1441            }
1442        };
1443        let message = handler.pre_startup(&context, &mut state, message);
1444        match conn.offer_pre_startup(message) {
1445            PreStartupOffer::Ssl(decision) => {
1446                conn = decision.decline_ssl();
1447                conn = flush_or_abort(conn).await?;
1448            }
1449            PreStartupOffer::Gss(decision) => {
1450                conn = decision.decline_gss();
1451                conn = flush_or_abort(conn).await?;
1452            }
1453            PreStartupOffer::Cancel {
1454                conn: terminal,
1455                process_id,
1456                secret_key,
1457            } => {
1458                let request = handler.cancellation(
1459                    &context,
1460                    &mut state,
1461                    CancellationRequest {
1462                        process_id,
1463                        secret_key,
1464                    },
1465                );
1466                return Ok((
1467                    ServerAccept::Cancellation(ServerCancellation {
1468                        transport: AcceptedServerTransport::Tls(Box::new(
1469                            terminal.into_transport().into_inner(),
1470                        )),
1471                        request,
1472                        state,
1473                        handler,
1474                        context: ServerConnectionContext {
1475                            peer: context.peer,
1476                            tls: context.tls,
1477                            identity: None,
1478                        },
1479                    }),
1480                    None,
1481                ));
1482            }
1483            PreStartupOffer::Startup {
1484                conn: startup_conn,
1485                message,
1486            } => {
1487                let message = handler.startup(&context, &mut state, message);
1488                let route = resolver
1489                    .resolve(&message, &context, &mut state)
1490                    .await
1491                    .map_err(RoutedAcceptError::Route)?;
1492                let ready = complete_auth(
1493                    startup_conn,
1494                    &message,
1495                    authentication,
1496                    &mut context,
1497                    &mut state,
1498                    &mut handler,
1499                    resolver.defer_ready(),
1500                )
1501                .await?;
1502                return Ok((
1503                    ServerAccept::Session(ServerConnection {
1504                        core: ServerConnectionCore {
1505                            conn: ServerConnectionInner::Tls(Box::new(ready)),
1506                            startup: message,
1507                            handler,
1508                            context,
1509                        },
1510                        state,
1511                    }),
1512                    Some(route),
1513                ));
1514            }
1515        }
1516    }
1517}
1518
1519#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
1520async fn complete_auth<I, Authentication, Peer, TlsError, State, Handler>(
1521    startup_conn: Conn<Buffered<I, Frontend>, crate::pre_startup::Startup>,
1522    message: &StartupMessage,
1523    provider: &Authentication,
1524    context: &mut ServerConnectionContext<
1525        Peer,
1526        <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1527    >,
1528    state: &mut State,
1529    handler: &mut Handler,
1530    defer_ready: bool,
1531) -> Result<
1532    Conn<Buffered<I, Frontend>, Ready>,
1533    AcceptError<TlsError, <Authentication::Authentication as ServerAuthentication<Peer>>::Error>,
1534>
1535where
1536    I: AsyncRead + AsyncWrite + Unpin,
1537    Authentication: ServerAuthenticationProvider,
1538    Authentication::Authentication: ServerAuthentication<Peer>,
1539    Handler: crate::ServerMiddleware<
1540            State,
1541            ServerConnectionContext<
1542                Peer,
1543                <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1544            >,
1545        >,
1546{
1547    let validated = match startup_conn.validate_protocol(message.clone(), ProtocolVersion::V3_2) {
1548        ServerProtocolOffer::Supported { conn, .. } => conn,
1549        ServerProtocolOffer::Rejected { conn, .. } => {
1550            let _ = conn.into_transport();
1551            return Err(AcceptError::UnsupportedProtocolVersion);
1552        }
1553    };
1554    let mut policy = provider.create();
1555    let auth = validated.begin_server_auth();
1556    let request = ServerAuthenticationRequest {
1557        startup: message,
1558        tls: context.tls(),
1559        peer: context.peer(),
1560    };
1561    let action = match policy.start(request).await {
1562        Ok(action) => action,
1563        Err(error) => {
1564            let _ = auth.into_transport();
1565            return Err(AcceptError::Authentication(error));
1566        }
1567    };
1568    let (auth, identity, final_frame) = match action {
1569        ServerAuthenticationAction::Accept(identity) => (auth, identity, None),
1570        action @ (ServerAuthenticationAction::CleartextPassword
1571        | ServerAuthenticationAction::Md5Password { .. }) => {
1572            let (waiting, frame) = match action {
1573                ServerAuthenticationAction::CleartextPassword => auth.request_cleartext(),
1574                ServerAuthenticationAction::Md5Password { salt } => auth.request_md5(salt),
1575                _ => unreachable!("matched password action"),
1576            }
1577            .map_err(AcceptError::Io)?;
1578            let frame = intercept_server_backend(handler, context, state, frame)
1579                .map_err(AcceptError::Io)?;
1580            let waiting = push_or_abort(waiting, frame)?;
1581            let waiting = flush_or_abort(waiting).await?;
1582            let (waiting, wire) = receive_frontend_or_abort(waiting).await?;
1583            let wire = handler.frontend(context, state, wire);
1584            let (auth, credential) = match waiting.receive_password(wire) {
1585                Ok(response) => response,
1586                Err(rejected) => {
1587                    let (waiting, _) = *rejected;
1588                    let _ = waiting.into_transport();
1589                    return Err(AcceptError::AuthenticationProtocol);
1590                }
1591            };
1592            match policy
1593                .respond(request, ServerAuthenticationResponse::Password(credential))
1594                .await
1595            {
1596                Ok(ServerAuthenticationAction::Accept(identity)) => (auth, identity, None),
1597                Ok(_) => {
1598                    let _ = auth.into_transport();
1599                    return Err(AcceptError::AuthenticationProtocol);
1600                }
1601                Err(error) => {
1602                    let _ = auth.into_transport();
1603                    return Err(AcceptError::Authentication(error));
1604                }
1605            }
1606        }
1607        ServerAuthenticationAction::Sasl { mechanisms } => {
1608            authenticate_sasl(
1609                auth,
1610                mechanisms,
1611                &mut policy,
1612                request,
1613                context,
1614                state,
1615                handler,
1616            )
1617            .await?
1618        }
1619        action @ (ServerAuthenticationAction::KerberosV5
1620        | ServerAuthenticationAction::Gss
1621        | ServerAuthenticationAction::Sspi) => {
1622            authenticate_token(auth, action, &mut policy, request, context, state, handler).await?
1623        }
1624        ServerAuthenticationAction::SaslContinue(_)
1625        | ServerAuthenticationAction::SaslFinal { .. }
1626        | ServerAuthenticationAction::GssContinue(_) => {
1627            let _ = auth.into_transport();
1628            return Err(AcceptError::AuthenticationProtocol);
1629        }
1630    };
1631    context.identity = Some(identity);
1632    let (mut startup_ready, _authentication_ok) =
1633        auth.authentication_ok().map_err(AcceptError::Io)?;
1634    if let Some(final_frame) = final_frame {
1635        let final_frame = intercept_server_backend(handler, context, state, final_frame)
1636            .map_err(AcceptError::Io)?;
1637        startup_ready = push_or_abort(startup_ready, final_frame)?;
1638    }
1639    let authentication_ok = handler
1640        .backend(
1641            context,
1642            state,
1643            BackendMessage::Authentication(crate::codec::Authentication::Ok),
1644        )
1645        .to_frame()
1646        .map_err(AcceptError::Io)?;
1647    let startup_ready = push_or_abort(startup_ready, authentication_ok)?;
1648    let (ready, _ready_frame) = startup_ready.ready().map_err(AcceptError::Io)?;
1649    let ready = if defer_ready {
1650        ready
1651    } else {
1652        let ready_frame = handler
1653            .backend(
1654                context,
1655                state,
1656                BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
1657            )
1658            .to_frame()
1659            .map_err(AcceptError::Io)?;
1660        push_or_abort(ready, ready_frame)?
1661    };
1662    let ready = flush_or_abort(ready).await?;
1663    Ok(ready)
1664}
1665
1666async fn authenticate_sasl<I, Policy, Peer, TlsError, State, Handler>(
1667    auth: Conn<Buffered<I, Frontend>, crate::server_auth::ServerAuth>,
1668    mechanisms: Vec<Bytes>,
1669    policy: &mut Policy,
1670    request: ServerAuthenticationRequest<'_, Peer>,
1671    context: &ServerConnectionContext<Peer, Policy::Identity>,
1672    state: &mut State,
1673    handler: &mut Handler,
1674) -> Result<
1675    (
1676        Conn<Buffered<I, Frontend>, crate::server_auth::ServerAuth>,
1677        Policy::Identity,
1678        Option<crate::codec::Frame>,
1679    ),
1680    AcceptError<TlsError, Policy::Error>,
1681>
1682where
1683    I: AsyncRead + AsyncWrite + Unpin,
1684    Policy: ServerAuthentication<Peer>,
1685    Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Policy::Identity>>,
1686{
1687    if mechanisms.iter().any(|mechanism| mechanism.contains(&0)) {
1688        let _ = auth.into_transport();
1689        return Err(AcceptError::Io(io::Error::new(
1690            io::ErrorKind::InvalidInput,
1691            "SASL mechanism contains NUL",
1692        )));
1693    }
1694    let (initial, frame) = auth.request_sasl(mechanisms).map_err(AcceptError::Io)?;
1695    let frame =
1696        intercept_server_backend(handler, context, state, frame).map_err(AcceptError::Io)?;
1697    let initial = push_or_abort(initial, frame)?;
1698    let initial = flush_or_abort(initial).await?;
1699    let (initial, wire) = receive_frontend_or_abort(initial).await?;
1700    let wire = handler.frontend(context, state, wire);
1701    let (mut sasl, initial_response) = match initial.receive_initial(wire) {
1702        Ok(response) => response,
1703        Err(rejected) => {
1704            let (initial, _) = *rejected;
1705            let _ = initial.into_transport();
1706            return Err(AcceptError::AuthenticationProtocol);
1707        }
1708    };
1709    let mut action = match policy
1710        .respond(
1711            request,
1712            ServerAuthenticationResponse::SaslInitial {
1713                mechanism: initial_response.mechanism,
1714                response: initial_response.response,
1715            },
1716        )
1717        .await
1718    {
1719        Ok(action) => action,
1720        Err(error) => {
1721            let _ = sasl.into_transport();
1722            return Err(AcceptError::Authentication(error));
1723        }
1724    };
1725    loop {
1726        match action {
1727            ServerAuthenticationAction::SaslContinue(challenge) => {
1728                let (waiting, frame) = sasl.continue_with(challenge).map_err(AcceptError::Io)?;
1729                let frame = intercept_server_backend(handler, context, state, frame)
1730                    .map_err(AcceptError::Io)?;
1731                let waiting = push_or_abort(waiting, frame)?;
1732                let waiting = flush_or_abort(waiting).await?;
1733                let (waiting, wire) = receive_frontend_or_abort(waiting).await?;
1734                let wire = handler.frontend(context, state, wire);
1735                let (next, response) = match waiting.receive_response(wire) {
1736                    Ok(response) => response,
1737                    Err(rejected) => {
1738                        let (waiting, _) = *rejected;
1739                        let _ = waiting.into_transport();
1740                        return Err(AcceptError::AuthenticationProtocol);
1741                    }
1742                };
1743                sasl = next;
1744                action = match policy
1745                    .respond(request, ServerAuthenticationResponse::Sasl(response))
1746                    .await
1747                {
1748                    Ok(action) => action,
1749                    Err(error) => {
1750                        let _ = sasl.into_transport();
1751                        return Err(AcceptError::Authentication(error));
1752                    }
1753                };
1754            }
1755            ServerAuthenticationAction::SaslFinal {
1756                server_final,
1757                identity,
1758            } => {
1759                let (auth, frame) = sasl.finish(server_final).map_err(AcceptError::Io)?;
1760                return Ok((auth, identity, Some(frame)));
1761            }
1762            _ => {
1763                let _ = sasl.into_transport();
1764                return Err(AcceptError::AuthenticationProtocol);
1765            }
1766        }
1767    }
1768}
1769
1770async fn authenticate_token<I, Policy, Peer, TlsError, State, Handler>(
1771    auth: Conn<Buffered<I, Frontend>, crate::server_auth::ServerAuth>,
1772    initial_action: ServerAuthenticationAction<Policy::Identity>,
1773    policy: &mut Policy,
1774    request: ServerAuthenticationRequest<'_, Peer>,
1775    context: &ServerConnectionContext<Peer, Policy::Identity>,
1776    state: &mut State,
1777    handler: &mut Handler,
1778) -> Result<
1779    (
1780        Conn<Buffered<I, Frontend>, crate::server_auth::ServerAuth>,
1781        Policy::Identity,
1782        Option<crate::codec::Frame>,
1783    ),
1784    AcceptError<TlsError, Policy::Error>,
1785>
1786where
1787    I: AsyncRead + AsyncWrite + Unpin,
1788    Policy: ServerAuthentication<Peer>,
1789    Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Policy::Identity>>,
1790{
1791    let (waiting, frame) = match initial_action {
1792        ServerAuthenticationAction::KerberosV5 => auth.request_kerberos_v5(),
1793        ServerAuthenticationAction::Gss => auth.request_gss(),
1794        ServerAuthenticationAction::Sspi => auth.request_sspi(),
1795        _ => unreachable!("matched initial token action"),
1796    }
1797    .map_err(AcceptError::Io)?;
1798    let frame =
1799        intercept_server_backend(handler, context, state, frame).map_err(AcceptError::Io)?;
1800    let waiting = push_or_abort(waiting, frame)?;
1801    let mut waiting = flush_or_abort(waiting).await?;
1802    loop {
1803        let received = receive_frontend_or_abort(waiting).await?;
1804        waiting = received.0;
1805        let wire = handler.frontend(context, state, received.1);
1806        let (decision, token) = match waiting.receive_response(wire) {
1807            Ok(response) => response,
1808            Err(rejected) => {
1809                let (waiting, _) = *rejected;
1810                let _ = waiting.into_transport();
1811                return Err(AcceptError::AuthenticationProtocol);
1812            }
1813        };
1814        let action = match policy
1815            .respond(request, ServerAuthenticationResponse::Token(token))
1816            .await
1817        {
1818            Ok(action) => action,
1819            Err(error) => {
1820                let _ = decision.into_transport();
1821                return Err(AcceptError::Authentication(error));
1822            }
1823        };
1824        match action {
1825            ServerAuthenticationAction::Accept(identity) => {
1826                return Ok((decision.verified(), identity, None));
1827            }
1828            ServerAuthenticationAction::GssContinue(token) => {
1829                let (next, frame) = decision.continue_gss(token).map_err(AcceptError::Io)?;
1830                let frame = intercept_server_backend(handler, context, state, frame)
1831                    .map_err(AcceptError::Io)?;
1832                let next = push_or_abort(next, frame)?;
1833                waiting = flush_or_abort(next).await?;
1834            }
1835            _ => {
1836                let _ = decision.into_transport();
1837                return Err(AcceptError::AuthenticationProtocol);
1838            }
1839        }
1840    }
1841}
1842
1843fn intercept_server_backend<State, Context, Handler>(
1844    handler: &mut Handler,
1845    context: &Context,
1846    state: &mut State,
1847    frame: crate::codec::Frame,
1848) -> io::Result<crate::codec::Frame>
1849where
1850    Handler: crate::ServerMiddleware<State, Context>,
1851{
1852    handler
1853        .backend(context, state, Backend::decode(frame)?)
1854        .to_frame()
1855}
1856
1857async fn flush_or_abort<I, D, Phase, TlsError, AuthenticationError>(
1858    mut conn: Conn<Buffered<I, D>, Phase>,
1859) -> Result<Conn<Buffered<I, D>, Phase>, AcceptError<TlsError, AuthenticationError>>
1860where
1861    I: AsyncWrite + Unpin,
1862{
1863    if let Err(error) = conn.flush().await {
1864        let _ = conn.into_transport();
1865        return Err(AcceptError::Io(error));
1866    }
1867    Ok(conn)
1868}
1869
1870fn push_or_abort<I, D, Phase, TlsError, AuthenticationError>(
1871    mut conn: Conn<Buffered<I, D>, Phase>,
1872    frame: crate::codec::Frame,
1873) -> Result<Conn<Buffered<I, D>, Phase>, AcceptError<TlsError, AuthenticationError>> {
1874    if let Err(error) = conn.push_frame(frame) {
1875        let _ = conn.into_transport();
1876        return Err(AcceptError::Io(error));
1877    }
1878    Ok(conn)
1879}
1880
1881async fn receive_frontend_or_abort<I, Phase, TlsError, AuthenticationError>(
1882    mut conn: Conn<Buffered<I, Frontend>, Phase>,
1883) -> Result<
1884    (Conn<Buffered<I, Frontend>, Phase>, FrontendMessage),
1885    AcceptError<TlsError, AuthenticationError>,
1886>
1887where
1888    I: AsyncRead + Unpin,
1889{
1890    match conn.receive_frontend_wire().await {
1891        Ok(message) => Ok((conn, message)),
1892        Err(error) => {
1893            let _ = conn.into_transport();
1894            Err(AcceptError::Io(error))
1895        }
1896    }
1897}