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

pg_proto/
server_component.rs
