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

pg_proto/
server_component.rs
