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