pg_proto/
client_component.rs

1//! Builder-centred client-role component.
2
3use std::{
4    collections::BTreeMap,
5    convert::Infallible,
6    fmt,
7    future::Future,
8    io,
9    pin::Pin,
10    sync::Arc,
11    task::{Context, Poll},
12};
13
14use bytes::Bytes;
15use tokio::io::ReadBuf;
16use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _};
17
18use crate::ClientMiddleware as _;
19use crate::{
20    Conn, Pristine,
21    auth::{AuthEvent, AuthOffer, Ready},
22    codec::{Backend, FrontendMessage},
23    demux::SessionItem,
24    session::{ReadyState, SimpleTransition},
25    startup::{ProtocolVersion, StartupMessage},
26    transport::Buffered,
27};
28
29/// A deterministic client component configuration failure.
30#[derive(Clone, Copy, Debug, Eq, PartialEq)]
31pub enum BuildError {
32    /// No transport connector was configured.
33    MissingConnector,
34    /// No explicit TLS policy was configured.
35    MissingTls,
36    /// No explicit authentication policy was configured.
37    MissingAuthentication,
38}
39
40impl fmt::Display for BuildError {
41    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
42        formatter.write_str(match self {
43            Self::MissingConnector => "client connector is required",
44            Self::MissingTls => "an explicit client TLS policy is required",
45            Self::MissingAuthentication => "an explicit client authentication policy is required",
46        })
47    }
48}
49
50impl std::error::Error for BuildError {}
51
52/// Explicit libpq-compatible TLS policy and its reloadable configuration provider.
53#[derive(Clone, Copy, Debug, Eq, PartialEq)]
54pub enum ClientTlsPolicy {
55    /// Intentionally use plaintext transport.
56    Disabled,
57}
58
59impl ClientTlsPolicy {
60    /// Configures a libpq-compatible SSL mode and reloadable provider.
61    #[must_use]
62    pub fn libpq<Provider>(
63        mode: crate::pre_startup::SslMode,
64        provider: Provider,
65    ) -> ReloadableClientTls<Provider> {
66        ReloadableClientTls { mode, provider }
67    }
68}
69
70/// A libpq-compatible policy backed by application-owned reloadable TLS material.
71#[derive(Clone)]
72pub struct ReloadableClientTls<Provider> {
73    mode: crate::pre_startup::SslMode,
74    provider: Provider,
75}
76
77impl<Provider> fmt::Debug for ReloadableClientTls<Provider> {
78    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
79        formatter
80            .debug_struct("Libpq")
81            .field("mode", &self.mode)
82            .field("provider", &"<redacted>")
83            .finish()
84    }
85}
86
87/// Internal shape shared by disabled and reloadable TLS policies.
88pub trait ClientTlsConfiguration {
89    /// Reloadable provider type.
90    type Provider: ClientTlsProvider;
91
92    /// Returns the libpq mode and provider, or `None` for explicit plaintext.
93    fn configured(&self) -> Option<(crate::pre_startup::SslMode, &Self::Provider)>;
94}
95
96impl ClientTlsConfiguration for ClientTlsPolicy {
97    type Provider = ();
98
99    fn configured(&self) -> Option<(crate::pre_startup::SslMode, &Self::Provider)> {
100        None
101    }
102}
103
104impl<Provider: ClientTlsProvider> ClientTlsConfiguration for ReloadableClientTls<Provider> {
105    type Provider = Provider;
106
107    fn configured(&self) -> Option<(crate::pre_startup::SslMode, &Self::Provider)> {
108        Some((self.mode, &self.provider))
109    }
110}
111
112/// Application-owned TLS material resolved afresh for a connection attempt.
113///
114/// The application owns reload and rotation of destination names and trust
115/// anchors. `pg-proto` deliberately constructs the final rustls verifier from
116/// the selected [`SslMode`](crate::SslMode), so a provider cannot silently
117/// weaken `VerifyCa` or `VerifyFull`.
118#[derive(Clone)]
119pub struct ClientTlsConfig {
120    server_name: rustls::pki_types::ServerName<'static>,
121    roots: rustls::RootCertStore,
122}
123
124impl ClientTlsConfig {
125    /// Creates resolved TLS material.
126    #[must_use]
127    pub fn new(
128        server_name: rustls::pki_types::ServerName<'static>,
129        roots: rustls::RootCertStore,
130    ) -> Self {
131        Self { server_name, roots }
132    }
133}
134
135impl fmt::Debug for ClientTlsConfig {
136    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
137        formatter.write_str("ClientTlsConfig(<redacted>)")
138    }
139}
140
141/// Application-owned source of reloadable client TLS material.
142pub trait ClientTlsProvider {
143    /// Provider resolution failure.
144    type Error;
145
146    /// Resolves material for one connection attempt.
147    fn resolve<'a>(
148        &'a self,
149        target: &'a ConnectTarget,
150    ) -> Pin<Box<dyn Future<Output = Result<ClientTlsConfig, Self::Error>> + 'a>>;
151}
152
153/// Failure while resolving or establishing client TLS.
154#[derive(Debug)]
155pub enum ClientTlsError<ProviderError> {
156    /// The application-owned reloadable provider failed.
157    Provider(ProviderError),
158    /// PostgreSQL negotiation or the TLS handshake failed.
159    Handshake(io::Error),
160}
161
162impl<ProviderError: fmt::Display> fmt::Display for ClientTlsError<ProviderError> {
163    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
164        match self {
165            Self::Provider(error) => error.fmt(formatter),
166            Self::Handshake(error) => error.fmt(formatter),
167        }
168    }
169}
170
171impl<ProviderError: std::error::Error + 'static> std::error::Error
172    for ClientTlsError<ProviderError>
173{
174    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
175        match self {
176            Self::Provider(error) => Some(error),
177            Self::Handshake(error) => Some(error),
178        }
179    }
180}
181
182impl ClientTlsProvider for () {
183    type Error = Infallible;
184
185    fn resolve<'a>(
186        &'a self,
187        _target: &'a ConnectTarget,
188    ) -> Pin<Box<dyn Future<Output = Result<ClientTlsConfig, Self::Error>> + 'a>> {
189        Box::pin(async { unreachable!("disabled TLS never resolves a provider") })
190    }
191}
192
193/// Transport selected by libpq-compatible negotiation.
194#[derive(Debug)]
195pub enum ClientTransport<Transport> {
196    /// Unencrypted PostgreSQL transport.
197    Plain(Transport),
198    /// TLS-protected PostgreSQL transport.
199    Tls(Box<crate::tls::ClientTls<Transport>>),
200}
201
202impl<Transport: AsyncRead + AsyncWrite + Unpin> AsyncRead for ClientTransport<Transport> {
203    fn poll_read(
204        mut self: Pin<&mut Self>,
205        cx: &mut Context<'_>,
206        buffer: &mut ReadBuf<'_>,
207    ) -> Poll<io::Result<()>> {
208        match &mut *self {
209            Self::Plain(stream) => Pin::new(stream).poll_read(cx, buffer),
210            Self::Tls(stream) => Pin::new(stream).poll_read(cx, buffer),
211        }
212    }
213}
214
215impl<Transport: AsyncRead + AsyncWrite + Unpin> AsyncWrite for ClientTransport<Transport> {
216    fn poll_write(
217        mut self: Pin<&mut Self>,
218        cx: &mut Context<'_>,
219        buffer: &[u8],
220    ) -> Poll<io::Result<usize>> {
221        match &mut *self {
222            Self::Plain(stream) => Pin::new(stream).poll_write(cx, buffer),
223            Self::Tls(stream) => Pin::new(stream).poll_write(cx, buffer),
224        }
225    }
226
227    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
228        match &mut *self {
229            Self::Plain(stream) => Pin::new(stream).poll_flush(cx),
230            Self::Tls(stream) => Pin::new(stream).poll_flush(cx),
231        }
232    }
233
234    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
235        match &mut *self {
236            Self::Plain(stream) => Pin::new(stream).poll_shutdown(cx),
237            Self::Tls(stream) => Pin::new(stream).poll_shutdown(cx),
238        }
239    }
240}
241
242async fn negotiate_client_tls<Transport, Provider, State, Handler, Evidence>(
243    mut transport: Transport,
244    mode: crate::pre_startup::SslMode,
245    provider: &Provider,
246    target: &ConnectTarget,
247    context: &mut ClientConnectionContext<Evidence>,
248    state: &mut State,
249    handler: &mut Handler,
250) -> Result<ClientTransport<Transport>, ClientTlsError<Provider::Error>>
251where
252    Transport: AsyncRead + AsyncWrite + Unpin,
253    Provider: ClientTlsProvider,
254    Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
255{
256    if !mode.strategy().request_on_first_connection {
257        context.tls = Some(ClientTlsStatus::Plaintext);
258        return Ok(ClientTransport::Plain(transport));
259    }
260    let request = handler.pre_startup(
261        context,
262        state,
263        crate::pre_startup::PreStartupMessage::SslRequest,
264    );
265    let crate::pre_startup::PreStartupMessage::SslRequest = request else {
266        return Err(ClientTlsError::Handshake(io::Error::new(
267            io::ErrorKind::InvalidData,
268            "middleware replaced SSLRequest with an incompatible packet",
269        )));
270    };
271    transport
272        .write_all(&[0, 0, 0, 8, 4, 210, 22, 47])
273        .await
274        .map_err(ClientTlsError::Handshake)?;
275    transport.flush().await.map_err(ClientTlsError::Handshake)?;
276    match transport
277        .read_u8()
278        .await
279        .map_err(ClientTlsError::Handshake)?
280    {
281        b'S' => {
282            let resolved = provider
283                .resolve(target)
284                .await
285                .map_err(ClientTlsError::Provider)?;
286            let config = Arc::new(crate::tls::client_config(mode, resolved.roots));
287            let stream = crate::tls::connect(transport, resolved.server_name, config)
288                .await
289                .map_err(ClientTlsError::Handshake)?;
290            context.tls = Some(ClientTlsStatus::Encrypted);
291            Ok(ClientTransport::Tls(Box::new(stream)))
292        }
293        b'N' if mode.strategy().allow_server_rejection => {
294            context.tls = Some(ClientTlsStatus::Plaintext);
295            Ok(ClientTransport::Plain(transport))
296        }
297        b'N' => Err(ClientTlsError::Handshake(io::Error::new(
298            io::ErrorKind::PermissionDenied,
299            "server rejected required TLS",
300        ))),
301        b'E' => Err(ClientTlsError::Handshake(io::Error::new(
302            io::ErrorKind::ConnectionAborted,
303            "server terminated TLS negotiation",
304        ))),
305        _ => Err(ClientTlsError::Handshake(io::Error::new(
306            io::ErrorKind::InvalidData,
307            "invalid TLS negotiation response",
308        ))),
309    }
310}
311
312/// Explicit client authentication policy which accepts only `AuthenticationOk`.
313#[derive(Clone, Copy, Default, Eq, PartialEq)]
314pub struct TrustClientAuthentication;
315
316/// An authentication request offered by a PostgreSQL server.
317#[derive(Clone, Debug, Eq, PartialEq)]
318#[non_exhaustive]
319pub enum ClientAuthenticationChallenge {
320    /// The server requested a cleartext password response.
321    CleartextPassword,
322    /// The server requested a PostgreSQL MD5 password response.
323    Md5Password([u8; 4]),
324    /// The server offered SASL mechanisms.
325    Sasl(Vec<Bytes>),
326    /// The server supplied another SASL challenge.
327    SaslContinue(Bytes),
328    /// The server supplied the SASL verifier.
329    SaslFinal(Bytes),
330    /// The server requested an opaque GSS token.
331    Gss,
332    /// The server requested an opaque SSPI token.
333    Sspi,
334    /// The server requested Kerberos V5 authentication.
335    KerberosV5,
336    /// The server supplied another opaque token challenge.
337    TokenContinue(Bytes),
338}
339
340/// An application authentication policy's wire response.
341#[derive(Clone, Debug, Eq, PartialEq)]
342#[non_exhaustive]
343pub enum ClientAuthenticationResponse {
344    /// Send a cleartext or precomputed MD5 password.
345    Password(Bytes),
346    /// Begin SASL using the named mechanism and initial response.
347    SaslInitial {
348        /// Selected SASL mechanism name.
349        mechanism: Bytes,
350        /// Mechanism-specific initial response.
351        response: Bytes,
352    },
353    /// Continue a SASL exchange.
354    Sasl(Bytes),
355    /// Send an opaque GSS, SSPI, or Kerberos token.
356    Token(Bytes),
357    /// Accept a verified SASL final message without sending another frame.
358    Verified,
359}
360
361/// Factory for asynchronous, fallible per-connection authentication sessions.
362pub trait ClientAuthentication {
363    /// Typed identity evidence produced after server confirmation.
364    type Evidence;
365    /// Per-connection mutable authentication state.
366    type Session: ClientAuthenticationSession<Evidence = Self::Evidence, Error = Self::Error>;
367    /// Application authentication failure.
368    type Error;
369
370    /// Creates fresh authentication state using the selected route.
371    fn begin<'a>(
372        &'a self,
373        target: &'a ConnectTarget,
374    ) -> ClientAuthenticationFuture<'a, Self::Session, Self::Error>;
375}
376
377/// Boxed future used by application-defined authentication policies.
378pub type ClientAuthenticationFuture<'a, Output, Error> =
379    Pin<Box<dyn Future<Output = Result<Output, Error>> + 'a>>;
380
381/// Mutable authentication policy state owned by one connection attempt.
382pub trait ClientAuthenticationSession {
383    /// Typed identity evidence produced after server confirmation.
384    type Evidence;
385    /// Application authentication failure.
386    type Error;
387
388    /// Answers one server authentication challenge.
389    fn respond(
390        &mut self,
391        challenge: ClientAuthenticationChallenge,
392    ) -> ClientAuthenticationFuture<'_, ClientAuthenticationResponse, Self::Error>;
393
394    /// Produces identity evidence after `AuthenticationOk` was received.
395    fn authenticated(self) -> ClientAuthenticationFuture<'static, Self::Evidence, Self::Error>;
396}
397
398/// Failure while running an application authentication policy.
399#[derive(Debug)]
400pub enum ClientAuthenticationError<PolicyError> {
401    /// Application policy creation or evaluation failed.
402    Policy(PolicyError),
403    /// The PostgreSQL server rejected authentication.
404    Rejected,
405    /// The policy returned a response illegal for the active challenge.
406    InvalidResponse,
407}
408
409impl<PolicyError: fmt::Display> fmt::Display for ClientAuthenticationError<PolicyError> {
410    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
411        match self {
412            Self::Policy(error) => error.fmt(formatter),
413            Self::Rejected => formatter.write_str("server rejected authentication"),
414            Self::InvalidResponse => {
415                formatter.write_str("authentication policy rejected the credential challenge")
416            }
417        }
418    }
419}
420
421impl<PolicyError: std::error::Error + 'static> std::error::Error
422    for ClientAuthenticationError<PolicyError>
423{
424    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
425        match self {
426            Self::Policy(error) => Some(error),
427            Self::Rejected | Self::InvalidResponse => None,
428        }
429    }
430}
431
432enum AuthenticationDriveError<PolicyError> {
433    Policy(PolicyError),
434    Rejected,
435    InvalidResponse,
436    Protocol(io::Error),
437}
438
439impl ClientAuthentication for TrustClientAuthentication {
440    type Evidence = ();
441    type Session = Self;
442    type Error = Infallible;
443
444    fn begin<'a>(
445        &'a self,
446        _target: &'a ConnectTarget,
447    ) -> Pin<Box<dyn Future<Output = Result<Self::Session, Self::Error>> + 'a>> {
448        Box::pin(async { Ok(Self) })
449    }
450}
451
452impl ClientAuthenticationSession for TrustClientAuthentication {
453    type Evidence = ();
454    type Error = Infallible;
455
456    fn respond(
457        &mut self,
458        _challenge: ClientAuthenticationChallenge,
459    ) -> ClientAuthenticationFuture<'_, ClientAuthenticationResponse, Self::Error> {
460        Box::pin(async { Ok(ClientAuthenticationResponse::Verified) })
461    }
462
463    fn authenticated(self) -> Pin<Box<dyn Future<Output = Result<Self::Evidence, Self::Error>>>> {
464        Box::pin(async { Ok(()) })
465    }
466}
467
468impl fmt::Debug for TrustClientAuthentication {
469    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
470        formatter.write_str("Trust")
471    }
472}
473
474/// Application-defined destination supplied to the connector.
475#[derive(Clone, Eq, PartialEq)]
476pub struct ConnectTarget {
477    name: String,
478    metadata: BTreeMap<String, String>,
479}
480
481impl fmt::Debug for ConnectTarget {
482    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
483        formatter.write_str("ConnectTarget(<redacted>)")
484    }
485}
486
487impl ConnectTarget {
488    /// Creates a named destination.
489    #[must_use]
490    pub fn new(name: impl Into<String>) -> Self {
491        Self {
492            name: name.into(),
493            metadata: BTreeMap::new(),
494        }
495    }
496
497    /// Returns its application-defined name or address.
498    #[must_use]
499    pub fn name(&self) -> &str {
500        &self.name
501    }
502
503    /// Adds routing metadata retained in the connection context.
504    #[must_use]
505    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
506        self.metadata.insert(key.into(), value.into());
507        self
508    }
509
510    /// Returns application-owned routing metadata.
511    #[must_use]
512    pub const fn metadata(&self) -> &BTreeMap<String, String> {
513        &self.metadata
514    }
515}
516
517/// Structured startup fields and extension parameters.
518#[derive(Clone, Default, Eq, PartialEq)]
519pub struct StartupParameters {
520    user: Option<String>,
521    database: Option<String>,
522    extensions: BTreeMap<String, String>,
523}
524
525impl fmt::Debug for StartupParameters {
526    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
527        formatter
528            .debug_struct("StartupParameters")
529            .field("user", &self.user.as_ref().map(|_| "<redacted>"))
530            .field("database", &self.database.as_ref().map(|_| "<redacted>"))
531            .field("extensions", &"<redacted>")
532            .finish()
533    }
534}
535
536impl StartupParameters {
537    /// Creates parameters containing a PostgreSQL user.
538    #[must_use]
539    pub fn new(user: impl Into<String>) -> Self {
540        Self {
541            user: Some(user.into()),
542            ..Self::default()
543        }
544    }
545
546    /// Returns the configured user, when present.
547    #[must_use]
548    pub fn user(&self) -> Option<&str> {
549        self.user.as_deref()
550    }
551
552    /// Returns the configured database, when present.
553    #[must_use]
554    pub fn database_name(&self) -> Option<&str> {
555        self.database.as_deref()
556    }
557
558    pub(crate) fn from_wire(message: &StartupMessage) -> io::Result<Self> {
559        let mut parameters = Self::default();
560        for (name, value) in &message.parameters {
561            let name = std::str::from_utf8(name).map_err(|_| {
562                io::Error::new(
563                    io::ErrorKind::InvalidData,
564                    "startup parameter name is not UTF-8",
565                )
566            })?;
567            let value = std::str::from_utf8(value).map_err(|_| {
568                io::Error::new(
569                    io::ErrorKind::InvalidData,
570                    "startup parameter value is not UTF-8",
571                )
572            })?;
573            match name {
574                "user" => parameters.user = Some(value.to_owned()),
575                "database" => parameters.database = Some(value.to_owned()),
576                _ => {
577                    parameters
578                        .extensions
579                        .insert(name.to_owned(), value.to_owned());
580                }
581            }
582        }
583        Ok(parameters)
584    }
585
586    /// Overrides the database field.
587    #[must_use]
588    pub fn database(mut self, database: impl Into<String>) -> Self {
589        self.database = Some(database.into());
590        self
591    }
592
593    /// Adds a non-standard startup extension parameter.
594    ///
595    /// # Errors
596    ///
597    /// Returns an error when `name` is a structured standard field.
598    pub fn extension(
599        mut self,
600        name: impl Into<String>,
601        value: impl Into<String>,
602    ) -> Result<Self, StartupParameterError> {
603        let name = name.into();
604        if matches!(name.as_str(), "user" | "database") {
605            return Err(StartupParameterError::ReservedExtension(name));
606        }
607        self.extensions.insert(name, value.into());
608        Ok(self)
609    }
610
611    fn merged_with(mut self, overrides: Self) -> Self {
612        if overrides.user.is_some() {
613            self.user = overrides.user;
614        }
615        if overrides.database.is_some() {
616            self.database = overrides.database;
617        }
618        self.extensions.extend(overrides.extensions);
619        self
620    }
621
622    fn into_message(self) -> Result<StartupMessage, StartupParameterError> {
623        let user = self.user.ok_or(StartupParameterError::MissingUser)?;
624        let mut parameters = self
625            .extensions
626            .into_iter()
627            .map(|(key, value)| (Bytes::from(key), Bytes::from(value)))
628            .collect::<BTreeMap<_, _>>();
629        parameters.insert(Bytes::from_static(b"user"), Bytes::from(user));
630        if let Some(database) = self.database {
631            parameters.insert(Bytes::from_static(b"database"), Bytes::from(database));
632        }
633        Ok(StartupMessage {
634            version: ProtocolVersion::V3_2,
635            parameters,
636        })
637    }
638}
639
640/// Invalid structured startup configuration.
641#[derive(Clone, Eq, PartialEq)]
642pub enum StartupParameterError {
643    /// A standard structured field was supplied through the extension map.
644    ReservedExtension(String),
645    /// Neither reusable defaults nor per-call overrides supplied a user.
646    MissingUser,
647}
648
649impl fmt::Debug for StartupParameterError {
650    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
651        formatter.write_str(match self {
652            Self::ReservedExtension(_) => "ReservedExtension(<redacted>)",
653            Self::MissingUser => "MissingUser",
654        })
655    }
656}
657
658impl fmt::Display for StartupParameterError {
659    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
660        match self {
661            Self::ReservedExtension(name) => {
662                write!(formatter, "startup extension name `{name}` is reserved")
663            }
664            Self::MissingUser => formatter.write_str("startup user is required"),
665        }
666    }
667}
668
669impl std::error::Error for StartupParameterError {}
670
671/// Conservative protocol allocation limits.
672#[derive(Clone, Copy, Debug, Eq, PartialEq)]
673pub struct ProtocolLimits {
674    max_frame_len: usize,
675}
676
677impl Default for ProtocolLimits {
678    fn default() -> Self {
679        Self {
680            max_frame_len: 1024 * 1024,
681        }
682    }
683}
684
685impl ProtocolLimits {
686    /// Sets the maximum complete tagged frame size.
687    ///
688    /// # Errors
689    ///
690    /// Returns an error outside PostgreSQL's tagged-frame range.
691    pub fn max_frame_len(mut self, limit: usize) -> Result<Self, ProtocolLimitError> {
692        if !(5..=i32::MAX as usize + 1).contains(&limit) {
693            return Err(ProtocolLimitError);
694        }
695        self.max_frame_len = limit;
696        Ok(self)
697    }
698
699    /// Explicitly selects the largest tagged frame PostgreSQL can encode.
700    #[must_use]
701    pub fn without_frame_limit(mut self) -> Self {
702        self.max_frame_len = i32::MAX as usize + 1;
703        self
704    }
705}
706
707/// Invalid protocol limit configuration.
708#[derive(Clone, Copy, Debug, Eq, PartialEq)]
709pub struct ProtocolLimitError;
710
711impl fmt::Display for ProtocolLimitError {
712    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
713        formatter.write_str("frame limit is outside PostgreSQL's tagged-frame range")
714    }
715}
716
717impl std::error::Error for ProtocolLimitError {}
718
719/// Immutable facts retained by a client-role connection.
720#[derive(Clone, Debug, Eq, PartialEq)]
721pub struct ClientConnectionContext<Evidence = ()> {
722    target: ConnectTarget,
723    tls: Option<ClientTlsStatus>,
724    identity: Option<Evidence>,
725    backend_key: Option<crate::demux::CancelKey>,
726}
727
728/// Progressively discovered transport security for a client connection.
729#[derive(Clone, Copy, Debug, Eq, PartialEq)]
730pub enum ClientTlsStatus {
731    /// The connection uses explicitly selected plaintext.
732    Plaintext,
733    /// The connection is protected by TLS.
734    Encrypted,
735}
736
737impl<Evidence> ClientConnectionContext<Evidence> {
738    /// Returns the destination used for this connection.
739    #[must_use]
740    pub const fn target(&self) -> &ConnectTarget {
741        &self.target
742    }
743
744    /// Returns transport security once negotiation has completed.
745    #[must_use]
746    pub const fn tls(&self) -> Option<ClientTlsStatus> {
747        self.tls
748    }
749
750    /// Returns application-defined evidence from the authentication policy.
751    ///
752    /// # Panics
753    ///
754    /// Panics when called from middleware before authentication has completed.
755    #[must_use]
756    pub const fn identity(&self) -> &Evidence {
757        match &self.identity {
758            Some(identity) => identity,
759            None => panic!("identity is not known before authentication"),
760        }
761    }
762
763    /// Returns evidence only after authentication has enriched the context.
764    #[must_use]
765    pub const fn identity_if_known(&self) -> Option<&Evidence> {
766        self.identity.as_ref()
767    }
768
769    /// Returns the upstream cancellation key captured during startup readiness.
770    #[must_use]
771    pub const fn backend_key(&self) -> Option<&crate::demux::CancelKey> {
772        self.backend_key.as_ref()
773    }
774}
775
776/// Identity middleware handler.
777#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
778pub struct IdentityHandler;
779impl<C> crate::MiddlewareFactory<C> for IdentityHandler {
780    type Handler = Self;
781    fn create(&self, _: &C) -> Self {
782        *self
783    }
784}
785impl<S, C> crate::ClientMiddleware<S, C> for IdentityHandler {}
786
787/// Immutable facts available when a client middleware handler is created.
788pub struct ClientInitialContext {
789    target: ConnectTarget,
790}
791impl ClientInitialContext {
792    /// Destination selected for this connection.
793    #[must_use]
794    pub const fn target(&self) -> &ConnectTarget {
795        &self.target
796    }
797}
798
799/// Reusable client-role component.
800pub struct Client<
801    Connector = (),
802    Tls = ClientTlsPolicy,
803    Authentication = TrustClientAuthentication,
804    Middleware = IdentityHandler,
805> {
806    connector: Connector,
807    tls: Tls,
808    authentication: Authentication,
809    defaults: StartupParameters,
810    limits: ProtocolLimits,
811    middleware: Middleware,
812}
813
814impl<Connector: Clone, Tls: Clone, Authentication: Clone, Middleware: Clone> Clone
815    for Client<Connector, Tls, Authentication, Middleware>
816{
817    fn clone(&self) -> Self {
818        Self {
819            connector: self.connector.clone(),
820            tls: self.tls.clone(),
821            authentication: self.authentication.clone(),
822            defaults: self.defaults.clone(),
823            limits: self.limits,
824            middleware: self.middleware.clone(),
825        }
826    }
827}
828
829impl Client<()> {
830    /// Starts client-role configuration.
831    #[must_use]
832    pub fn builder() -> ClientBuilder {
833        ClientBuilder::default()
834    }
835}
836
837impl<Connector, Tls: fmt::Debug, Authentication: fmt::Debug, Middleware> fmt::Debug
838    for Client<Connector, Tls, Authentication, Middleware>
839{
840    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
841        formatter
842            .debug_struct("Client")
843            .field("tls", &self.tls)
844            .field("authentication", &self.authentication)
845            .finish_non_exhaustive()
846    }
847}
848
849/// Ordinary generic builder for a reusable client-role component.
850pub struct ClientBuilder<
851    Connector = (),
852    Tls = (),
853    Authentication = (),
854    Middleware = IdentityHandler,
855> {
856    connector: Option<Connector>,
857    tls: Option<Tls>,
858    authentication: Option<Authentication>,
859    defaults: StartupParameters,
860    limits: ProtocolLimits,
861    middleware: Middleware,
862}
863
864impl Default for ClientBuilder<()> {
865    fn default() -> Self {
866        Self {
867            connector: None,
868            tls: None,
869            authentication: None,
870            defaults: StartupParameters::default(),
871            limits: ProtocolLimits::default(),
872            middleware: IdentityHandler,
873        }
874    }
875}
876
877impl ClientBuilder<()> {
878    /// Configures the reusable application-supplied connector.
879    #[must_use]
880    pub fn connector<Next, Work, Transport, Error>(
881        self,
882        connector: Next,
883    ) -> ClientBuilder<Next, (), ()>
884    where
885        Next: Fn(&ConnectTarget) -> Work,
886        Work: Future<Output = Result<Transport, Error>>,
887    {
888        ClientBuilder {
889            connector: Some(connector),
890            tls: self.tls,
891            authentication: self.authentication,
892            defaults: self.defaults,
893            limits: self.limits,
894            middleware: self.middleware,
895        }
896    }
897}
898
899impl<Connector, Tls, Authentication, Middleware>
900    ClientBuilder<Connector, Tls, Authentication, Middleware>
901{
902    /// Selects an explicit TLS policy.
903    #[must_use]
904    pub fn tls<Next>(
905        self,
906        tls: Next,
907    ) -> ClientBuilder<Connector, Next, Authentication, Middleware> {
908        ClientBuilder {
909            connector: self.connector,
910            tls: Some(tls),
911            authentication: self.authentication,
912            defaults: self.defaults,
913            limits: self.limits,
914            middleware: self.middleware,
915        }
916    }
917
918    /// Selects explicit trust authentication.
919    #[must_use]
920    pub fn authentication<Next>(
921        self,
922        authentication: Next,
923    ) -> ClientBuilder<Connector, Tls, Next, Middleware> {
924        ClientBuilder {
925            connector: self.connector,
926            tls: self.tls,
927            authentication: Some(authentication),
928            defaults: self.defaults,
929            limits: self.limits,
930            middleware: self.middleware,
931        }
932    }
933
934    /// Sets reusable startup defaults which per-call values override explicitly.
935    #[must_use]
936    pub fn startup_parameters(mut self, defaults: StartupParameters) -> Self {
937        self.defaults = defaults;
938        self
939    }
940
941    /// Replaces conservative protocol limits with an explicit policy.
942    #[must_use]
943    pub fn protocol_limits(mut self, limits: ProtocolLimits) -> Self {
944        self.limits = limits;
945        self
946    }
947
948    /// Appends a middleware factory. Stages run in declaration order.
949    #[must_use]
950    pub fn middleware<Next>(
951        self,
952        factory: Next,
953    ) -> ClientBuilder<Connector, Tls, Authentication, crate::MiddlewareChain<Middleware, Next>>
954    {
955        ClientBuilder {
956            connector: self.connector,
957            tls: self.tls,
958            authentication: self.authentication,
959            defaults: self.defaults,
960            limits: self.limits,
961            middleware: crate::MiddlewareChain(self.middleware, factory),
962        }
963    }
964
965    /// Validates and creates the reusable component.
966    ///
967    /// # Errors
968    ///
969    /// Returns the first missing mandatory configuration category.
970    pub fn build(self) -> Result<Client<Connector, Tls, Authentication, Middleware>, BuildError> {
971        Ok(Client {
972            connector: self.connector.ok_or(BuildError::MissingConnector)?,
973            tls: self.tls.ok_or(BuildError::MissingTls)?,
974            authentication: self
975                .authentication
976                .ok_or(BuildError::MissingAuthentication)?,
977            defaults: self.defaults,
978            limits: self.limits,
979            middleware: self.middleware,
980        })
981    }
982}
983
984/// Failure while establishing a client-role connection, distinct from [`BuildError`].
985#[derive(Debug)]
986pub enum ConnectError<ConnectorError, TlsError = Infallible, AuthenticationError = Infallible> {
987    /// The application-supplied connector failed.
988    Connector(ConnectorError),
989    /// The reloadable TLS provider or handshake failed.
990    Tls(TlsError),
991    /// The application authentication policy failed.
992    Authentication(AuthenticationError),
993    /// Structured startup values were invalid before network establishment.
994    Startup(StartupParameterError),
995    /// PostgreSQL framing, startup, authentication, or readiness failed.
996    Protocol(io::Error),
997}
998
999/// Failure while sending a one-shot PostgreSQL cancellation packet.
1000#[derive(Debug)]
1001pub enum CancelError<ConnectorError> {
1002    /// The configured connector could not open the cancellation transport.
1003    Connector(ConnectorError),
1004    /// The key could not be encoded or the raw packet could not be written.
1005    Protocol(io::Error),
1006}
1007
1008impl<E: fmt::Display> fmt::Display for CancelError<E> {
1009    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1010        match self {
1011            Self::Connector(error) => error.fmt(formatter),
1012            Self::Protocol(error) => error.fmt(formatter),
1013        }
1014    }
1015}
1016impl<E: std::error::Error + 'static> std::error::Error for CancelError<E> {
1017    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1018        match self {
1019            Self::Connector(error) => Some(error),
1020            Self::Protocol(error) => Some(error),
1021        }
1022    }
1023}
1024
1025impl<ConnectorError: fmt::Display, TlsError: fmt::Display, AuthenticationError: fmt::Display>
1026    fmt::Display for ConnectError<ConnectorError, TlsError, AuthenticationError>
1027{
1028    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1029        match self {
1030            Self::Connector(error) => error.fmt(f),
1031            Self::Tls(error) => error.fmt(f),
1032            Self::Authentication(error) => error.fmt(f),
1033            Self::Startup(error) => error.fmt(f),
1034            Self::Protocol(error) => error.fmt(f),
1035        }
1036    }
1037}
1038impl<ConnectorError, TlsError, AuthenticationError> std::error::Error
1039    for ConnectError<ConnectorError, TlsError, AuthenticationError>
1040where
1041    ConnectorError: std::error::Error + 'static,
1042    TlsError: std::error::Error + 'static,
1043    AuthenticationError: std::error::Error + 'static,
1044{
1045    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1046        match self {
1047            Self::Connector(error) => Some(error),
1048            Self::Tls(error) => Some(error),
1049            Self::Authentication(error) => Some(error),
1050            Self::Startup(error) => Some(error),
1051            Self::Protocol(error) => Some(error),
1052        }
1053    }
1054}
1055
1056/// Evidence that a connection has not performed a state-changing operation.
1057#[derive(Debug)]
1058pub enum ConnectionClean {}
1059
1060/// Evidence that an operation may have changed session-local state.
1061#[derive(Debug)]
1062pub enum ConnectionChanged {}
1063
1064/// Operational client-role connection.
1065pub struct ClientConnection<
1066    Transport,
1067    State,
1068    Cleanliness = ConnectionClean,
1069    Evidence = (),
1070    Handler = IdentityHandler,
1071> {
1072    core: ClientConnectionCore<Transport, Cleanliness, Evidence, Handler>,
1073    state: State,
1074}
1075
1076pub(crate) struct ClientConnectionCore<Transport, Cleanliness, Evidence, Handler> {
1077    connection: Conn<Buffered<Transport, Backend>, Ready, Cleanliness>,
1078    handler: Handler,
1079    context: ClientConnectionContext<Evidence>,
1080}
1081
1082impl<Transport, State, Cleanliness, Evidence, Handler>
1083    ClientConnection<Transport, State, Cleanliness, Evidence, Handler>
1084{
1085    /// Returns immutable connection facts.
1086    #[must_use]
1087    pub const fn context(&self) -> &ClientConnectionContext<Evidence> {
1088        &self.core.context
1089    }
1090
1091    /// Returns the caller-owned connection state.
1092    ///
1093    #[must_use]
1094    pub const fn state(&self) -> &State {
1095        &self.state
1096    }
1097
1098    /// Receives one backend message at the operational inspection boundary.
1099    ///
1100    /// # Errors
1101    ///
1102    /// Returns a transport, decoding, or configured frame-limit error.
1103    ///
1104    pub async fn receive_wire(&mut self) -> io::Result<crate::codec::BackendMessage>
1105    where
1106        Transport: AsyncRead + Unpin,
1107        Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1108    {
1109        let message = self.core.receive_wire_raw().await?;
1110        Ok(self.core.intercept_backend(&mut self.state, message))
1111    }
1112
1113    /// Recovers every owned connection part deliberately.
1114    ///
1115    pub fn into_parts(self) -> (Transport, State, Handler, ClientConnectionContext<Evidence>) {
1116        (
1117            self.core.connection.into_transport().into_inner(),
1118            self.state,
1119            self.core.handler,
1120            self.core.context,
1121        )
1122    }
1123}
1124
1125impl<Transport, Cleanliness, Evidence, Handler>
1126    ClientConnectionCore<Transport, Cleanliness, Evidence, Handler>
1127{
1128    pub(crate) const fn context(&self) -> &ClientConnectionContext<Evidence> {
1129        &self.context
1130    }
1131    pub(crate) async fn receive_wire_raw(&mut self) -> io::Result<crate::codec::BackendMessage>
1132    where
1133        Transport: AsyncRead + Unpin,
1134    {
1135        self.connection.receive_backend_wire().await
1136    }
1137
1138    pub(crate) fn intercept_backend<State>(
1139        &mut self,
1140        state: &mut State,
1141        message: crate::codec::BackendMessage,
1142    ) -> crate::codec::BackendMessage
1143    where
1144        Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1145    {
1146        self.handler.backend(&self.context, state, message)
1147    }
1148
1149    pub(crate) fn intercept_frontend<State>(
1150        &mut self,
1151        state: &mut State,
1152        message: FrontendMessage,
1153    ) -> FrontendMessage
1154    where
1155        Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1156    {
1157        self.handler.frontend(&self.context, state, message)
1158    }
1159
1160    pub(crate) async fn send_wire_raw(&mut self, message: FrontendMessage) -> io::Result<()>
1161    where
1162        Transport: AsyncWrite + Unpin,
1163    {
1164        self.connection.push_frame(message.to_frame()?)?;
1165        self.connection.flush().await
1166    }
1167
1168    pub(crate) fn into_parts(self) -> (Transport, Handler, ClientConnectionContext<Evidence>) {
1169        (
1170            self.connection.into_transport().into_inner(),
1171            self.handler,
1172            self.context,
1173        )
1174    }
1175}
1176
1177fn intercept_auth_response<State, Evidence, Handler>(
1178    handler: &mut Handler,
1179    context: &ClientConnectionContext<Evidence>,
1180    state: &mut State,
1181    response: Bytes,
1182) -> Result<Bytes, AuthenticationDriveError<Infallible>>
1183where
1184    Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1185{
1186    match handler.frontend(
1187        context,
1188        state,
1189        crate::codec::FrontendMessage::PasswordResponse(response),
1190    ) {
1191        crate::codec::FrontendMessage::PasswordResponse(response) => Ok(response),
1192        _ => Err(AuthenticationDriveError::InvalidResponse),
1193    }
1194}
1195
1196async fn complete_password<Transport, Policy, State, Handler>(
1197    connection: Conn<Buffered<Transport, Backend>, crate::auth::PasswordResponse>,
1198    challenge: ClientAuthenticationChallenge,
1199    policy: &mut Policy,
1200    context: &ClientConnectionContext<Policy::Evidence>,
1201    state: &mut State,
1202    handler: &mut Handler,
1203) -> Result<
1204    Conn<Buffered<Transport, Backend>, crate::auth::AwaitingStartupReady>,
1205    AuthenticationDriveError<Policy::Error>,
1206>
1207where
1208    Transport: AsyncRead + AsyncWrite + Unpin,
1209    Policy: ClientAuthenticationSession,
1210    Handler: crate::ClientMiddleware<State, ClientConnectionContext<Policy::Evidence>>,
1211{
1212    let response = policy
1213        .respond(challenge)
1214        .await
1215        .map_err(AuthenticationDriveError::Policy)?;
1216    let ClientAuthenticationResponse::Password(password) = response else {
1217        let _ = connection.into_transport();
1218        return Err(AuthenticationDriveError::InvalidResponse);
1219    };
1220    let password = intercept_auth_response(handler, context, state, password)
1221        .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1222    let (mut awaiting, frame) = connection
1223        .password(&password)
1224        .map_err(AuthenticationDriveError::Protocol)?;
1225    awaiting
1226        .push_frame(frame)
1227        .map_err(AuthenticationDriveError::Protocol)?;
1228    awaiting
1229        .flush()
1230        .await
1231        .map_err(AuthenticationDriveError::Protocol)?;
1232    let message = awaiting
1233        .receive_backend_wire()
1234        .await
1235        .map_err(AuthenticationDriveError::Protocol)?;
1236    let message = handler.backend(context, state, message);
1237    match awaiting.offer(message) {
1238        Ok(crate::auth::AuthCompletion::Ok(connection)) => Ok(connection),
1239        Ok(crate::auth::AuthCompletion::Error { conn, .. }) => {
1240            let _ = conn.into_transport();
1241            Err(AuthenticationDriveError::Rejected)
1242        }
1243        Err((conn, _)) => {
1244            let _ = conn.into_transport();
1245            Err(AuthenticationDriveError::Protocol(io::Error::new(
1246                io::ErrorKind::InvalidData,
1247                "illegal authentication completion",
1248            )))
1249        }
1250    }
1251}
1252
1253async fn complete_sasl<Transport, Policy, State, Handler>(
1254    connection: Conn<Buffered<Transport, Backend>, crate::auth::SaslInitial>,
1255    mechanisms: Vec<Bytes>,
1256    policy: &mut Policy,
1257    context: &ClientConnectionContext<Policy::Evidence>,
1258    state: &mut State,
1259    handler: &mut Handler,
1260) -> Result<
1261    Conn<Buffered<Transport, Backend>, crate::auth::AwaitingStartupReady>,
1262    AuthenticationDriveError<Policy::Error>,
1263>
1264where
1265    Transport: AsyncRead + AsyncWrite + Unpin,
1266    Policy: ClientAuthenticationSession,
1267    Handler: crate::ClientMiddleware<State, ClientConnectionContext<Policy::Evidence>>,
1268{
1269    let response = policy
1270        .respond(ClientAuthenticationChallenge::Sasl(mechanisms))
1271        .await
1272        .map_err(AuthenticationDriveError::Policy)?;
1273    let ClientAuthenticationResponse::SaslInitial {
1274        mechanism,
1275        response,
1276    } = response
1277    else {
1278        let _ = connection.into_transport();
1279        return Err(AuthenticationDriveError::InvalidResponse);
1280    };
1281    let response = intercept_auth_response(handler, context, state, response)
1282        .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1283    let (mut sasl, frame) = connection
1284        .sasl(&mechanism, &response)
1285        .map_err(AuthenticationDriveError::Protocol)?;
1286    sasl.push_frame(frame)
1287        .map_err(AuthenticationDriveError::Protocol)?;
1288    sasl.flush()
1289        .await
1290        .map_err(AuthenticationDriveError::Protocol)?;
1291    loop {
1292        let message = sasl
1293            .receive_backend_wire()
1294            .await
1295            .map_err(AuthenticationDriveError::Protocol)?;
1296        let message = handler.backend(context, state, message);
1297        match sasl.offer_backend(message) {
1298            Ok(crate::auth::SaslEvent::Continue { conn, challenge }) => {
1299                let response = policy
1300                    .respond(ClientAuthenticationChallenge::SaslContinue(challenge))
1301                    .await
1302                    .map_err(AuthenticationDriveError::Policy)?;
1303                let ClientAuthenticationResponse::Sasl(response) = response else {
1304                    let _ = conn.into_transport();
1305                    return Err(AuthenticationDriveError::InvalidResponse);
1306                };
1307                let response = intercept_auth_response(handler, context, state, response)
1308                    .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1309                let (mut next, frame) = conn.respond(response);
1310                next.push_frame(frame)
1311                    .map_err(AuthenticationDriveError::Protocol)?;
1312                next.flush()
1313                    .await
1314                    .map_err(AuthenticationDriveError::Protocol)?;
1315                sasl = next;
1316            }
1317            Ok(crate::auth::SaslEvent::Final { conn, server_final }) => {
1318                let response = policy
1319                    .respond(ClientAuthenticationChallenge::SaslFinal(server_final))
1320                    .await
1321                    .map_err(AuthenticationDriveError::Policy)?;
1322                if response != ClientAuthenticationResponse::Verified {
1323                    let _ = conn.into_transport();
1324                    return Err(AuthenticationDriveError::InvalidResponse);
1325                }
1326                let mut awaiting = conn.verified();
1327                let message = awaiting
1328                    .receive_backend_wire()
1329                    .await
1330                    .map_err(AuthenticationDriveError::Protocol)?;
1331                let message = handler.backend(context, state, message);
1332                return match awaiting.offer(message) {
1333                    Ok(crate::auth::AuthCompletion::Ok(connection)) => Ok(connection),
1334                    Ok(crate::auth::AuthCompletion::Error { conn, .. }) => {
1335                        let _ = conn.into_transport();
1336                        Err(AuthenticationDriveError::Rejected)
1337                    }
1338                    Err((conn, _)) => {
1339                        let _ = conn.into_transport();
1340                        Err(AuthenticationDriveError::Protocol(io::Error::new(
1341                            io::ErrorKind::InvalidData,
1342                            "illegal SASL authentication completion",
1343                        )))
1344                    }
1345                };
1346            }
1347            Ok(crate::auth::SaslEvent::Error { conn, .. }) => {
1348                let _ = conn.into_transport();
1349                return Err(AuthenticationDriveError::Rejected);
1350            }
1351            Err((conn, _)) => {
1352                let _ = conn.into_transport();
1353                return Err(AuthenticationDriveError::Protocol(io::Error::new(
1354                    io::ErrorKind::InvalidData,
1355                    "illegal SASL authentication message",
1356                )));
1357            }
1358        }
1359    }
1360}
1361
1362async fn complete_token<Transport, Policy, State, Handler>(
1363    connection: Conn<Buffered<Transport, Backend>, crate::auth::TokenResponse>,
1364    challenge: ClientAuthenticationChallenge,
1365    policy: &mut Policy,
1366    context: &ClientConnectionContext<Policy::Evidence>,
1367    state: &mut State,
1368    handler: &mut Handler,
1369) -> Result<
1370    Conn<Buffered<Transport, Backend>, crate::auth::AwaitingStartupReady>,
1371    AuthenticationDriveError<Policy::Error>,
1372>
1373where
1374    Transport: AsyncRead + AsyncWrite + Unpin,
1375    Policy: ClientAuthenticationSession,
1376    Handler: crate::ClientMiddleware<State, ClientConnectionContext<Policy::Evidence>>,
1377{
1378    let response = policy
1379        .respond(challenge)
1380        .await
1381        .map_err(AuthenticationDriveError::Policy)?;
1382    let ClientAuthenticationResponse::Token(token) = response else {
1383        let _ = connection.into_transport();
1384        return Err(AuthenticationDriveError::InvalidResponse);
1385    };
1386    let token = intercept_auth_response(handler, context, state, token)
1387        .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1388    let (mut waiting, frame) = connection.respond(token);
1389    waiting
1390        .push_frame(frame)
1391        .map_err(AuthenticationDriveError::Protocol)?;
1392    waiting
1393        .flush()
1394        .await
1395        .map_err(AuthenticationDriveError::Protocol)?;
1396    loop {
1397        let message = waiting
1398            .receive_backend_wire()
1399            .await
1400            .map_err(AuthenticationDriveError::Protocol)?;
1401        let message = handler.backend(context, state, message);
1402        match waiting.offer(message) {
1403            Ok(crate::auth::TokenAuthEvent::Continue { conn, token }) => {
1404                let response = policy
1405                    .respond(ClientAuthenticationChallenge::TokenContinue(token))
1406                    .await
1407                    .map_err(AuthenticationDriveError::Policy)?;
1408                let ClientAuthenticationResponse::Token(token) = response else {
1409                    let _ = conn.into_transport();
1410                    return Err(AuthenticationDriveError::InvalidResponse);
1411                };
1412                let token = intercept_auth_response(handler, context, state, token)
1413                    .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1414                let (mut next, frame) = conn.respond(token);
1415                next.push_frame(frame)
1416                    .map_err(AuthenticationDriveError::Protocol)?;
1417                next.flush()
1418                    .await
1419                    .map_err(AuthenticationDriveError::Protocol)?;
1420                waiting = next;
1421            }
1422            Ok(crate::auth::TokenAuthEvent::Ok(connection)) => return Ok(connection),
1423            Ok(crate::auth::TokenAuthEvent::Error { conn, .. }) => {
1424                let _ = conn.into_transport();
1425                return Err(AuthenticationDriveError::Rejected);
1426            }
1427            Err((conn, _)) => {
1428                let _ = conn.into_transport();
1429                return Err(AuthenticationDriveError::Protocol(io::Error::new(
1430                    io::ErrorKind::InvalidData,
1431                    "illegal token authentication message",
1432                )));
1433            }
1434        }
1435    }
1436}
1437
1438enum SessionEstablishError<AuthenticationError> {
1439    Authentication(ClientAuthenticationError<AuthenticationError>),
1440    Protocol(io::Error),
1441}
1442
1443fn replace_session_item(
1444    item: SessionItem,
1445    replacement: crate::codec::BackendMessage,
1446) -> Option<SessionItem> {
1447    match (item, replacement) {
1448        (SessionItem::Message(_), message) => Some(SessionItem::Message(message)),
1449        (
1450            SessionItem::ReadyForQuery {
1451                parameters_changed, ..
1452            },
1453            crate::codec::BackendMessage::ReadyForQuery(status),
1454        ) => Some(SessionItem::ReadyForQuery {
1455            status,
1456            parameters_changed,
1457        }),
1458        (
1459            SessionItem::CommandComplete {
1460                command, notices, ..
1461            },
1462            crate::codec::BackendMessage::CommandComplete(tag),
1463        ) => Some(SessionItem::CommandComplete {
1464            tag,
1465            command,
1466            notices,
1467        }),
1468        _ => None,
1469    }
1470}
1471
1472#[allow(clippy::too_many_arguments, clippy::too_many_lines)] // Typestate plus middleware lifecycle.
1473async fn establish_client_session<Transport, Authentication, State, Handler>(
1474    transport: ClientTransport<Transport>,
1475    startup: &StartupMessage,
1476    target: &ConnectTarget,
1477    authentication_policy: &Authentication,
1478    max_frame_len: usize,
1479    context: &ClientConnectionContext<Authentication::Evidence>,
1480    state: &mut State,
1481    handler: &mut Handler,
1482) -> Result<
1483    (
1484        Conn<Buffered<ClientTransport<Transport>, Backend>, Ready>,
1485        Authentication::Evidence,
1486        Option<crate::demux::CancelKey>,
1487    ),
1488    SessionEstablishError<Authentication::Error>,
1489>
1490where
1491    Transport: AsyncRead + AsyncWrite + Unpin,
1492    Authentication: ClientAuthentication,
1493    Handler: crate::ClientMiddleware<State, ClientConnectionContext<Authentication::Evidence>>,
1494{
1495    let mut policy = authentication_policy.begin(target).await.map_err(|error| {
1496        SessionEstablishError::Authentication(ClientAuthenticationError::Policy(error))
1497    })?;
1498    let buffered = Buffered::with_max_frame_len(transport, max_frame_len)
1499        .map_err(SessionEstablishError::Protocol)?;
1500    let (mut startup_connection, packet) = Conn::new(buffered)
1501        .startup(startup)
1502        .map_err(SessionEstablishError::Protocol)?;
1503    startup_connection.push_startup_packet(&packet);
1504    let mut authentication = startup_connection.authentication();
1505    if let Err(error) = authentication.flush().await {
1506        let _ = authentication.into_transport();
1507        return Err(SessionEstablishError::Protocol(error));
1508    }
1509    let awaiting_ready = loop {
1510        let message = match authentication.receive_backend_wire().await {
1511            Ok(message) => message,
1512            Err(error) => {
1513                let _ = authentication.into_transport();
1514                return Err(SessionEstablishError::Protocol(error));
1515            }
1516        };
1517        let message = handler.backend(context, state, message);
1518        match authentication.offer_backend(message) {
1519            Ok(AuthEvent::Authentication(AuthOffer::Ok(connection))) => break connection,
1520            Ok(AuthEvent::Negotiate { conn, .. }) => authentication = conn,
1521            Ok(AuthEvent::Authentication(AuthOffer::Cleartext(connection))) => {
1522                break complete_password(
1523                    connection,
1524                    ClientAuthenticationChallenge::CleartextPassword,
1525                    &mut policy,
1526                    context,
1527                    state,
1528                    handler,
1529                )
1530                .await
1531                .map_err(session_authentication_error)?;
1532            }
1533            Ok(AuthEvent::Authentication(AuthOffer::Md5 { conn, salt })) => {
1534                break complete_password(
1535                    conn,
1536                    ClientAuthenticationChallenge::Md5Password(salt),
1537                    &mut policy,
1538                    context,
1539                    state,
1540                    handler,
1541                )
1542                .await
1543                .map_err(session_authentication_error)?;
1544            }
1545            Ok(AuthEvent::Authentication(AuthOffer::Sasl { conn, mechanisms })) => {
1546                break complete_sasl(conn, mechanisms, &mut policy, context, state, handler)
1547                    .await
1548                    .map_err(session_authentication_error)?;
1549            }
1550            Ok(AuthEvent::Authentication(AuthOffer::Gss(conn))) => {
1551                break complete_token(
1552                    conn,
1553                    ClientAuthenticationChallenge::Gss,
1554                    &mut policy,
1555                    context,
1556                    state,
1557                    handler,
1558                )
1559                .await
1560                .map_err(session_authentication_error)?;
1561            }
1562            Ok(AuthEvent::Authentication(AuthOffer::Sspi(conn))) => {
1563                break complete_token(
1564                    conn,
1565                    ClientAuthenticationChallenge::Sspi,
1566                    &mut policy,
1567                    context,
1568                    state,
1569                    handler,
1570                )
1571                .await
1572                .map_err(session_authentication_error)?;
1573            }
1574            Ok(AuthEvent::Authentication(AuthOffer::KerberosV5(conn))) => {
1575                break complete_token(
1576                    conn,
1577                    ClientAuthenticationChallenge::KerberosV5,
1578                    &mut policy,
1579                    context,
1580                    state,
1581                    handler,
1582                )
1583                .await
1584                .map_err(session_authentication_error)?;
1585            }
1586            Ok(AuthEvent::Error { conn, .. }) => {
1587                let _ = conn.into_transport();
1588                return Err(SessionEstablishError::Authentication(
1589                    ClientAuthenticationError::Rejected,
1590                ));
1591            }
1592            Err((conn, _, source)) => {
1593                authentication = conn;
1594                if let Some(source) = source {
1595                    let _ = authentication.into_transport();
1596                    return Err(SessionEstablishError::Protocol(source));
1597                }
1598            }
1599        }
1600    };
1601    let mut awaiting_ready = awaiting_ready;
1602    let mut backend_key = None;
1603    let ready = loop {
1604        let item = match awaiting_ready.receive().await {
1605            Ok(item) => item,
1606            Err(error) => {
1607                let _ = awaiting_ready.into_transport();
1608                return Err(SessionEstablishError::Protocol(error));
1609            }
1610        };
1611        if let crate::codec::BackendMessage::BackendKeyData {
1612            process_id,
1613            secret_key,
1614        } = item.clone().into_backend_message()
1615        {
1616            backend_key = Some(crate::demux::CancelKey {
1617                process_id,
1618                secret_key,
1619            });
1620        }
1621        let replacement = handler.backend(context, state, item.clone().into_backend_message());
1622        let Some(item) = replace_session_item(item, replacement) else {
1623            let _ = awaiting_ready.into_transport();
1624            return Err(SessionEstablishError::Protocol(io::Error::new(
1625                io::ErrorKind::InvalidData,
1626                "middleware replacement during startup readiness is not phase-compatible",
1627            )));
1628        };
1629        match awaiting_ready.offer_ready(item) {
1630            Ok(ready) => break ready,
1631            Err((connection, SessionItem::Message(_))) => awaiting_ready = connection,
1632            Err((connection, _)) => {
1633                let _ = connection.into_transport();
1634                return Err(SessionEstablishError::Protocol(io::Error::new(
1635                    io::ErrorKind::InvalidData,
1636                    "startup did not reach an idle operational phase",
1637                )));
1638            }
1639        }
1640    };
1641    let identity = policy.authenticated().await.map_err(|error| {
1642        SessionEstablishError::Authentication(ClientAuthenticationError::Policy(error))
1643    })?;
1644    Ok((ready, identity, backend_key))
1645}
1646
1647fn session_authentication_error<Error>(
1648    error: AuthenticationDriveError<Error>,
1649) -> SessionEstablishError<Error> {
1650    match error {
1651        AuthenticationDriveError::Policy(error) => {
1652            SessionEstablishError::Authentication(ClientAuthenticationError::Policy(error))
1653        }
1654        AuthenticationDriveError::Rejected => {
1655            SessionEstablishError::Authentication(ClientAuthenticationError::Rejected)
1656        }
1657        AuthenticationDriveError::InvalidResponse => {
1658            SessionEstablishError::Authentication(ClientAuthenticationError::InvalidResponse)
1659        }
1660        AuthenticationDriveError::Protocol(error) => SessionEstablishError::Protocol(error),
1661    }
1662}
1663
1664fn map_session_error<ConnectorError, TlsError, AuthenticationError>(
1665    error: SessionEstablishError<AuthenticationError>,
1666) -> ConnectError<ConnectorError, TlsError, ClientAuthenticationError<AuthenticationError>> {
1667    match error {
1668        SessionEstablishError::Authentication(error) => ConnectError::Authentication(error),
1669        SessionEstablishError::Protocol(error) => ConnectError::Protocol(error),
1670    }
1671}
1672
1673/// Failure while executing an operational client-role action.
1674#[derive(Debug)]
1675pub struct QueryError(io::Error);
1676
1677impl fmt::Display for QueryError {
1678    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1679        self.0.fmt(formatter)
1680    }
1681}
1682
1683impl std::error::Error for QueryError {
1684    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1685        Some(&self.0)
1686    }
1687}
1688
1689impl<Transport, State, Cleanliness, Evidence, Handler>
1690    ClientConnection<Transport, State, Cleanliness, Evidence, Handler>
1691where
1692    Transport: AsyncRead + AsyncWrite + Unpin,
1693    Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1694{
1695    /// Executes a simple query through typed completion back to operational readiness.
1696    ///
1697    /// The returned connection is conservatively marked dirty because arbitrary
1698    /// SQL may retain session-local state.
1699    ///
1700    /// # Errors
1701    ///
1702    /// Returns an error for invalid query text, I/O or framing failure, an
1703    /// illegal peer transition, COPY entry, or a backend error response.
1704    ///
1705    #[allow(clippy::too_many_lines)]
1706    pub async fn simple_query(
1707        self,
1708        query: &[u8],
1709    ) -> Result<
1710        (
1711            ClientConnection<Transport, State, ConnectionChanged, Evidence, Handler>,
1712            Vec<crate::codec::BackendMessage>,
1713        ),
1714        QueryError,
1715    > {
1716        let Self {
1717            core:
1718                ClientConnectionCore {
1719                    connection,
1720                    mut handler,
1721                    context,
1722                },
1723            mut state,
1724        } = self;
1725        let outbound = handler.frontend(
1726            &context,
1727            &mut state,
1728            crate::codec::FrontendMessage::Query(Bytes::copy_from_slice(query)),
1729        );
1730        let crate::codec::FrontendMessage::Query(query) = outbound else {
1731            let _ = connection.into_transport();
1732            return Err(QueryError(io::Error::new(
1733                io::ErrorKind::InvalidData,
1734                "middleware replaced Query with an incompatible message",
1735            )));
1736        };
1737        let (mut query_connection, frame) = connection.push_query(&query).map_err(QueryError)?;
1738        if let Err(error) = query_connection.push_frame(frame) {
1739            let _ = query_connection.into_transport();
1740            return Err(QueryError(error));
1741        }
1742        if let Err(error) = query_connection.flush().await {
1743            let _ = query_connection.into_transport();
1744            return Err(QueryError(error));
1745        }
1746        let mut messages = Vec::new();
1747        loop {
1748            let item = match query_connection.receive().await {
1749                Ok(item) => item,
1750                Err(error) => {
1751                    let _ = query_connection.into_transport();
1752                    return Err(QueryError(error));
1753                }
1754            };
1755            let observed =
1756                handler.backend(&context, &mut state, item.clone().into_backend_message());
1757            let Some(item) = replace_session_item(item, observed.clone()) else {
1758                let _ = query_connection.into_transport();
1759                return Err(QueryError(io::Error::new(
1760                    io::ErrorKind::InvalidData,
1761                    "middleware replacement during simple query is not phase-compatible",
1762                )));
1763            };
1764            messages.push(observed);
1765            match query_connection.offer(item) {
1766                Ok(SimpleTransition::Continue(connection, _)) => query_connection = connection,
1767                Ok(SimpleTransition::Ready(
1768                    ReadyState::Clean(connection)
1769                    | ReadyState::Dirty {
1770                        conn: connection, ..
1771                    },
1772                )) => {
1773                    return Ok((
1774                        ClientConnection {
1775                            core: ClientConnectionCore {
1776                                connection: connection.transition(),
1777                                handler,
1778                                context,
1779                            },
1780                            state,
1781                        },
1782                        messages,
1783                    ));
1784                }
1785                Ok(SimpleTransition::Error(connection, _)) => {
1786                    let _ = connection.into_transport();
1787                    return Err(QueryError(io::Error::other(
1788                        "backend rejected simple query",
1789                    )));
1790                }
1791                Ok(SimpleTransition::CopyIn(connection, _)) => {
1792                    let _ = connection.into_transport();
1793                    return Err(QueryError(io::Error::other("simple query entered COPY IN")));
1794                }
1795                Ok(SimpleTransition::CopyOut(connection, _)) => {
1796                    let _ = connection.into_transport();
1797                    return Err(QueryError(io::Error::other(
1798                        "simple query entered COPY OUT",
1799                    )));
1800                }
1801                Ok(SimpleTransition::CopyBoth(connection, _)) => {
1802                    let _ = connection.into_transport();
1803                    return Err(QueryError(io::Error::other(
1804                        "simple query entered COPY BOTH",
1805                    )));
1806                }
1807                Err((connection, _)) => {
1808                    let _ = connection.into_transport();
1809                    return Err(QueryError(io::Error::new(
1810                        io::ErrorKind::InvalidData,
1811                        "illegal simple-query response",
1812                    )));
1813                }
1814            }
1815        }
1816    }
1817}
1818
1819impl<Connector, Tls, Authentication, Middleware, Work, Transport, Error>
1820    Client<Connector, Tls, Authentication, Middleware>
1821where
1822    Connector: Fn(&ConnectTarget) -> Work,
1823    Work: Future<Output = Result<Transport, Error>>,
1824    Transport: AsyncRead + AsyncWrite + Unpin,
1825    Authentication: ClientAuthentication,
1826    Tls: ClientTlsConfiguration,
1827    Middleware: crate::MiddlewareFactory<ClientInitialContext>,
1828{
1829    /// Establishes transport, startup, and configured authentication before returning.
1830    ///
1831    /// Per-call startup values explicitly override component defaults.
1832    ///
1833    /// # Errors
1834    ///
1835    /// Returns a connection-time error for connector, startup, framing,
1836    /// authentication, or readiness failures.
1837    pub async fn connect<State>(
1838        &self,
1839        target: ConnectTarget,
1840        overrides: StartupParameters,
1841        mut state: State,
1842    ) -> Result<
1843        ClientConnection<
1844            ClientTransport<Transport>,
1845            State,
1846            ConnectionClean,
1847            Authentication::Evidence,
1848            <Middleware as crate::MiddlewareFactory<ClientInitialContext>>::Handler,
1849        >,
1850        ConnectError<
1851            Error,
1852            ClientTlsError<<Tls::Provider as ClientTlsProvider>::Error>,
1853            ClientAuthenticationError<Authentication::Error>,
1854        >,
1855    >
1856    where
1857        <Middleware as crate::MiddlewareFactory<ClientInitialContext>>::Handler:
1858            crate::ClientMiddleware<State, ClientConnectionContext<Authentication::Evidence>>,
1859    {
1860        let core = self.connect_core(target, overrides, &mut state).await?;
1861        Ok(ClientConnection {
1862            core: ClientConnectionCore {
1863                connection: core.connection.transition(),
1864                handler: core.handler,
1865                context: core.context,
1866            },
1867            state,
1868        })
1869    }
1870
1871    /// Establishes a client role while borrowing facade-owned state, ensuring
1872    /// the state remains recoverable on every connection failure.
1873    pub(crate) async fn connect_core<State>(
1874        &self,
1875        target: ConnectTarget,
1876        overrides: StartupParameters,
1877        state: &mut State,
1878    ) -> Result<
1879        ClientConnectionCore<
1880            ClientTransport<Transport>,
1881            Pristine,
1882            Authentication::Evidence,
1883            <Middleware as crate::MiddlewareFactory<ClientInitialContext>>::Handler,
1884        >,
1885        ConnectError<
1886            Error,
1887            ClientTlsError<<Tls::Provider as ClientTlsProvider>::Error>,
1888            ClientAuthenticationError<Authentication::Error>,
1889        >,
1890    >
1891    where
1892        <Middleware as crate::MiddlewareFactory<ClientInitialContext>>::Handler:
1893            crate::ClientMiddleware<State, ClientConnectionContext<Authentication::Evidence>>,
1894    {
1895        let mut handler = self.middleware.create(&ClientInitialContext {
1896            target: target.clone(),
1897        });
1898        let startup = self
1899            .defaults
1900            .clone()
1901            .merged_with(overrides)
1902            .into_message()
1903            .map_err(ConnectError::Startup)?;
1904        let mut context = ClientConnectionContext {
1905            target: target.clone(),
1906            tls: None,
1907            identity: None,
1908            backend_key: None,
1909        };
1910        let transport = (self.connector)(&target)
1911            .await
1912            .map_err(ConnectError::Connector)?;
1913        let configured_tls = self.tls.configured();
1914        let transport = match configured_tls {
1915            None => {
1916                context.tls = Some(ClientTlsStatus::Plaintext);
1917                ClientTransport::Plain(transport)
1918            }
1919            Some((mode, provider)) => negotiate_client_tls(
1920                transport,
1921                mode,
1922                provider,
1923                &target,
1924                &mut context,
1925                state,
1926                &mut handler,
1927            )
1928            .await
1929            .map_err(ConnectError::Tls)?,
1930        };
1931        let first_startup = handler.startup(&context, state, startup.clone());
1932        let first = establish_client_session(
1933            transport,
1934            &first_startup,
1935            &target,
1936            &self.authentication,
1937            self.limits.max_frame_len,
1938            &context,
1939            state,
1940            &mut handler,
1941        )
1942        .await;
1943        let retry_provider = match configured_tls {
1944            Some((crate::pre_startup::SslMode::Allow, provider)) if first.is_err() => {
1945                Some(provider)
1946            }
1947            _ => None,
1948        };
1949        let (ready, identity, backend_key) = if let Some(provider) = retry_provider {
1950            let transport = (self.connector)(&target)
1951                .await
1952                .map_err(ConnectError::Connector)?;
1953            // The retry is a new transport attempt within the same logical
1954            // connection; do not expose the failed plaintext attempt as its
1955            // current transport fact.
1956            context.tls = None;
1957            let transport = negotiate_client_tls(
1958                transport,
1959                crate::pre_startup::SslMode::Require,
1960                provider,
1961                &target,
1962                &mut context,
1963                state,
1964                &mut handler,
1965            )
1966            .await
1967            .map_err(ConnectError::Tls)?;
1968            let retry_startup = handler.startup(&context, state, startup);
1969            establish_client_session(
1970                transport,
1971                &retry_startup,
1972                &target,
1973                &self.authentication,
1974                self.limits.max_frame_len,
1975                &context,
1976                state,
1977                &mut handler,
1978            )
1979            .await
1980            .map_err(map_session_error)?
1981        } else {
1982            first.map_err(map_session_error)?
1983        };
1984        Ok(ClientConnectionCore {
1985            connection: ready,
1986            handler,
1987            context: {
1988                context.identity = Some(identity);
1989                context.backend_key = backend_key;
1990                context
1991            },
1992        })
1993    }
1994}
1995
1996impl<Connector, Tls, Authentication, Middleware, Work, Transport, Error>
1997    Client<Connector, Tls, Authentication, Middleware>
1998where
1999    Connector: Fn(&ConnectTarget) -> Work,
2000    Work: Future<Output = Result<Transport, Error>>,
2001    Transport: AsyncWrite + Unpin,
2002{
2003    /// Opens a fresh transport and writes exactly one raw cancellation packet.
2004    ///
2005    /// Cancellation deliberately performs neither TLS negotiation nor startup
2006    /// authentication, as required by PostgreSQL's out-of-band protocol.
2007    ///
2008    /// # Errors
2009    ///
2010    /// Returns the connector's typed error or a cancellation encode/write error.
2011    pub async fn cancel(
2012        &self,
2013        target: &ConnectTarget,
2014        key: &crate::demux::CancelKey,
2015    ) -> Result<(), CancelError<Error>> {
2016        let mut transport = (self.connector)(target)
2017            .await
2018            .map_err(CancelError::Connector)?;
2019        let packet = crate::pre_startup::PreStartupMessage::CancelRequest {
2020            process_id: key.process_id,
2021            secret_key: key.secret_key.clone(),
2022        }
2023        .to_packet()
2024        .map_err(CancelError::Protocol)?;
2025        tokio::io::AsyncWriteExt::write_all(&mut transport, &packet)
2026            .await
2027            .map_err(CancelError::Protocol)?;
2028        tokio::io::AsyncWriteExt::shutdown(&mut transport)
2029            .await
2030            .map_err(CancelError::Protocol)
2031    }
2032}
pg_proto/client_component.rs

pg_proto/
client_component.rs
