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

pg_proto/
client_component.rs
