Skip to main content

qail_pg/driver/connection/
connect.rs

1//! Connection establishment — connect_*, TLS, mTLS, Unix socket.
2
3#[cfg(all(target_os = "linux", feature = "native-io-uring"))]
4use super::helpers::should_try_uring_plain;
5use super::helpers::{
6    connect_backend_for_stream, plain_connect_attempt_backend, record_connect_attempt,
7    record_connect_result,
8};
9use super::types::{
10    BUFFER_CAPACITY, CONNECT_BACKEND_TOKIO, CONNECT_TRANSPORT_GSSENC, CONNECT_TRANSPORT_MTLS,
11    CONNECT_TRANSPORT_PLAIN, CONNECT_TRANSPORT_TLS, ConnectParams, DEFAULT_CONNECT_TIMEOUT,
12    GSSENC_REQUEST, GssEncNegotiationResult, PgConnection, SSL_REQUEST, STMT_CACHE_CAPACITY,
13    StatementCache, TlsConfig, has_logical_replication_startup_mode,
14};
15use crate::driver::stream::PgStream;
16use crate::driver::{AuthSettings, ConnectOptions, GssEncMode, PgError, PgResult, TlsMode};
17use crate::protocol::PROTOCOL_VERSION_3_0;
18use crate::protocol::wire::FrontendMessage;
19use bytes::BytesMut;
20use std::collections::{HashMap, VecDeque};
21use std::sync::Arc;
22use std::time::Instant;
23use tokio::io::AsyncWriteExt;
24use tokio::net::TcpStream;
25
26#[inline]
27fn protocol_version_from_minor(minor: u16) -> i32 {
28    ((3i32) << 16) | i32::from(minor)
29}
30
31fn socket_addr(host: &str, port: u16) -> String {
32    if host.contains(':') && !host.starts_with('[') {
33        format!("[{}]:{}", host, port)
34    } else {
35        format!("{}:{}", host, port)
36    }
37}
38
39fn is_explicit_protocol_version_rejection(err: &PgError) -> bool {
40    let msg = match err {
41        PgError::Connection(msg) | PgError::Protocol(msg) | PgError::Auth(msg) => msg,
42        PgError::Query(msg) => msg,
43        PgError::QueryServer(server) => &server.message,
44        _ => return false,
45    };
46
47    let lower = msg.to_ascii_lowercase();
48    lower.contains("unsupported frontend protocol")
49        || lower.contains("frontend protocol") && lower.contains("unsupported")
50        || lower.contains("protocol version") && lower.contains("not support")
51}
52
53impl PgConnection {
54    /// Connect to PostgreSQL server without authentication (trust mode).
55    ///
56    /// # Arguments
57    ///
58    /// * `host` — PostgreSQL server hostname or IP.
59    /// * `port` — TCP port (typically 5432).
60    /// * `user` — PostgreSQL role name.
61    /// * `database` — Target database name.
62    pub async fn connect(host: &str, port: u16, user: &str, database: &str) -> PgResult<Self> {
63        Self::connect_with_password(host, port, user, database, None).await
64    }
65
66    /// Connect to PostgreSQL server with optional password authentication.
67    /// Includes a default 10-second timeout covering TCP connect + handshake.
68    ///
69    /// Startup requests protocol 3.2 by default and performs a one-shot retry
70    /// with protocol 3.0 only when startup fails due to explicit
71    /// protocol-version rejection from the server.
72    pub async fn connect_with_password(
73        host: &str,
74        port: u16,
75        user: &str,
76        database: &str,
77        password: Option<&str>,
78    ) -> PgResult<Self> {
79        Self::connect_with_password_and_auth(
80            host,
81            port,
82            user,
83            database,
84            password,
85            AuthSettings::default(),
86        )
87        .await
88    }
89
90    /// Connect to PostgreSQL with explicit enterprise options.
91    ///
92    /// Negotiation preface order follows libpq:
93    ///   1. If gss_enc_mode != Disable → try GSSENCRequest on fresh TCP
94    ///   2. If GSSENC rejected/unavailable and tls_mode != Disable → try SSLRequest
95    ///   3. If both rejected/unavailable → plain StartupMessage
96    ///
97    /// The StartupMessage protocol version behavior is the same as
98    /// `connect_with_password`: request protocol 3.2 first, then retry once
99    /// with 3.0 only on explicit protocol-version rejection.
100    pub async fn connect_with_options(
101        host: &str,
102        port: u16,
103        user: &str,
104        database: &str,
105        password: Option<&str>,
106        options: ConnectOptions,
107    ) -> PgResult<Self> {
108        let ConnectOptions {
109            tls_mode,
110            gss_enc_mode,
111            tls_ca_cert_pem,
112            mtls,
113            gss_token_provider,
114            auth,
115            io_uring,
116            startup_params,
117        } = options;
118
119        if mtls.is_some() && matches!(tls_mode, TlsMode::Disable) {
120            return Err(PgError::Connection(
121                "Invalid connect options: mTLS requires tls_mode=Prefer or Require".to_string(),
122            ));
123        }
124
125        // Enforce gss_enc_mode policy before mTLS early-return.
126        // GSSENC and mTLS are both transport-level encryption; using
127        // both simultaneously is not supported by the PostgreSQL protocol.
128        if gss_enc_mode == GssEncMode::Require && mtls.is_some() {
129            return Err(PgError::Connection(
130                "gssencmode=require is incompatible with mTLS — both provide \
131                 transport encryption; use one or the other"
132                    .to_string(),
133            ));
134        }
135
136        if let Some(mtls_config) = mtls {
137            // gss_enc_mode is Disable or Prefer here (Require rejected above).
138            // mTLS already provides transport encryption; skip GSSENC.
139            return Self::connect_mtls_with_password_and_auth_and_gss(
140                ConnectParams {
141                    host,
142                    port,
143                    user,
144                    database,
145                    password,
146                    auth_settings: auth,
147                    gss_token_provider,
148                    io_uring,
149                    protocol_minor: Self::default_protocol_minor(),
150                    startup_params: startup_params.clone(),
151                },
152                mtls_config,
153            )
154            .await;
155        }
156
157        // ── Phase 1: Try GSSENC if requested ──────────────────────────
158        if gss_enc_mode != GssEncMode::Disable {
159            match Self::try_gssenc_request(host, port).await {
160                Ok(GssEncNegotiationResult::Accepted(tcp_stream)) => {
161                    let connect_started = Instant::now();
162                    record_connect_attempt(CONNECT_TRANSPORT_GSSENC, CONNECT_BACKEND_TOKIO);
163                    #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
164                    {
165                        let default_minor = Self::default_protocol_minor();
166                        let gss_params = ConnectParams {
167                            host,
168                            port,
169                            user,
170                            database,
171                            password,
172                            auth_settings: auth,
173                            gss_token_provider: gss_token_provider.clone(),
174                            io_uring,
175                            protocol_minor: default_minor,
176                            startup_params: startup_params.clone(),
177                        };
178                        let mut result = Self::connect_gssenc_accepted_with_timeout(
179                            tcp_stream,
180                            gss_params.clone(),
181                        )
182                        .await;
183                        if let Err(err) = &result
184                            && default_minor > 0
185                            && is_explicit_protocol_version_rejection(err)
186                        {
187                            let downgrade_minor = (PROTOCOL_VERSION_3_0 & 0xFFFF) as u16;
188                            let retry_stream = match Self::try_gssenc_request(host, port).await {
189                                Ok(GssEncNegotiationResult::Accepted(stream)) => stream,
190                                Ok(GssEncNegotiationResult::Rejected) => {
191                                    return Err(PgError::Connection(
192                                        "Protocol downgrade retry failed: server rejected GSSENCRequest"
193                                            .to_string(),
194                                    ));
195                                }
196                                Ok(GssEncNegotiationResult::ServerError) => {
197                                    return Err(PgError::Connection(
198                                        "Protocol downgrade retry failed: server returned error to GSSENCRequest"
199                                            .to_string(),
200                                    ));
201                                }
202                                Err(e) => {
203                                    return Err(e);
204                                }
205                            };
206                            let mut retry_params = gss_params;
207                            retry_params.protocol_minor = downgrade_minor;
208                            result = Self::connect_gssenc_accepted_with_timeout(
209                                retry_stream,
210                                retry_params,
211                            )
212                            .await;
213                        }
214                        record_connect_result(
215                            CONNECT_TRANSPORT_GSSENC,
216                            CONNECT_BACKEND_TOKIO,
217                            &result,
218                            connect_started.elapsed(),
219                        );
220                        return result;
221                    }
222                    #[cfg(not(all(feature = "enterprise-gssapi", target_os = "linux")))]
223                    {
224                        let _ = tcp_stream;
225                        let err = PgError::Connection(
226                            "Server accepted GSSENCRequest but GSSAPI encryption requires \
227                             feature enterprise-gssapi on Linux"
228                                .to_string(),
229                        );
230                        metrics::histogram!(
231                            "qail_pg_connect_duration_seconds",
232                            "transport" => CONNECT_TRANSPORT_GSSENC,
233                            "backend" => CONNECT_BACKEND_TOKIO,
234                            "outcome" => "error"
235                        )
236                        .record(connect_started.elapsed().as_secs_f64());
237                        metrics::counter!(
238                            "qail_pg_connect_failure_total",
239                            "transport" => CONNECT_TRANSPORT_GSSENC,
240                            "backend" => CONNECT_BACKEND_TOKIO,
241                            "error_kind" => super::helpers::connect_error_kind(&err)
242                        )
243                        .increment(1);
244                        return Err(err);
245                    }
246                }
247                Ok(GssEncNegotiationResult::Rejected)
248                | Ok(GssEncNegotiationResult::ServerError) => {
249                    if gss_enc_mode == GssEncMode::Require {
250                        return Err(PgError::Connection(
251                            "gssencmode=require but server rejected GSSENCRequest".to_string(),
252                        ));
253                    }
254                    // gss_enc_mode == Prefer — fall through to TLS / plain
255                }
256                Err(e) => {
257                    if gss_enc_mode == GssEncMode::Require {
258                        return Err(e);
259                    }
260                    // gss_enc_mode == Prefer — connection error, fall through
261                    tracing::debug!(
262                        host = %host,
263                        port = %port,
264                        error = %e,
265                        "gssenc_prefer_fallthrough"
266                    );
267                }
268            }
269        }
270
271        // ── Phase 2: TLS / plain per sslmode ──────────────────────────
272        match tls_mode {
273            TlsMode::Disable => {
274                Self::connect_with_password_and_auth_and_gss(ConnectParams {
275                    host,
276                    port,
277                    user,
278                    database,
279                    password,
280                    auth_settings: auth,
281                    gss_token_provider,
282                    io_uring,
283                    protocol_minor: Self::default_protocol_minor(),
284                    startup_params: startup_params.clone(),
285                })
286                .await
287            }
288            TlsMode::Require => {
289                Self::connect_tls_with_auth_and_gss(
290                    ConnectParams {
291                        host,
292                        port,
293                        user,
294                        database,
295                        password,
296                        auth_settings: auth,
297                        gss_token_provider,
298                        io_uring,
299                        protocol_minor: Self::default_protocol_minor(),
300                        startup_params: startup_params.clone(),
301                    },
302                    tls_ca_cert_pem.as_deref(),
303                )
304                .await
305            }
306            TlsMode::Prefer => {
307                match Self::connect_tls_with_auth_and_gss(
308                    ConnectParams {
309                        host,
310                        port,
311                        user,
312                        database,
313                        password,
314                        auth_settings: auth,
315                        gss_token_provider: gss_token_provider.clone(),
316                        io_uring,
317                        protocol_minor: Self::default_protocol_minor(),
318                        startup_params: startup_params.clone(),
319                    },
320                    tls_ca_cert_pem.as_deref(),
321                )
322                .await
323                {
324                    Ok(conn) => Ok(conn),
325                    Err(PgError::Connection(msg))
326                        if msg.contains("Server does not support TLS") =>
327                    {
328                        Self::connect_with_password_and_auth_and_gss(ConnectParams {
329                            host,
330                            port,
331                            user,
332                            database,
333                            password,
334                            auth_settings: auth,
335                            gss_token_provider,
336                            io_uring,
337                            protocol_minor: Self::default_protocol_minor(),
338                            startup_params: startup_params.clone(),
339                        })
340                        .await
341                    }
342                    Err(e) => Err(e),
343                }
344            }
345        }
346    }
347
348    /// Attempt GSSAPI session encryption negotiation.
349    ///
350    /// Opens a fresh TCP connection, sends GSSENCRequest (80877104),
351    /// reads exactly one byte (CVE-2021-23222 safe), and returns
352    /// the result.  The entire operation is bounded by
353    /// `DEFAULT_CONNECT_TIMEOUT`.
354    async fn try_gssenc_request(host: &str, port: u16) -> PgResult<GssEncNegotiationResult> {
355        tokio::time::timeout(
356            DEFAULT_CONNECT_TIMEOUT,
357            Self::try_gssenc_request_inner(host, port),
358        )
359        .await
360        .map_err(|_| {
361            PgError::Connection(format!(
362                "GSSENCRequest timeout after {:?}",
363                DEFAULT_CONNECT_TIMEOUT
364            ))
365        })?
366    }
367
368    /// Inner GSSENCRequest logic without timeout wrapper.
369    async fn try_gssenc_request_inner(host: &str, port: u16) -> PgResult<GssEncNegotiationResult> {
370        use tokio::io::AsyncReadExt;
371
372        let addr = socket_addr(host, port);
373        let mut tcp_stream = TcpStream::connect(&addr).await?;
374        tcp_stream.set_nodelay(true)?;
375
376        // Send the 8-byte GSSENCRequest.
377        tcp_stream.write_all(&GSSENC_REQUEST).await?;
378        tcp_stream.flush().await?;
379
380        // CVE-2021-23222: Read exactly one byte.  The server must
381        // respond with a single 'G' or 'N'.  Any additional bytes
382        // in the buffer indicate a buffer-stuffing attack.
383        let mut response = [0u8; 1];
384        tcp_stream.read_exact(&mut response).await?;
385
386        match response[0] {
387            b'G' => {
388                // CVE-2021-23222 check: verify no extra bytes are buffered.
389                // Use a non-blocking peek to detect leftover data.
390                let mut peek_buf = [0u8; 1];
391                match tcp_stream.try_read(&mut peek_buf) {
392                    Ok(0) => {} // EOF — fine (shouldn't happen yet but harmless)
393                    Ok(_n) => {
394                        // Extra bytes after 'G' — possible buffer-stuffing.
395                        return Err(PgError::Connection(
396                            "Protocol violation: extra bytes after GSSENCRequest 'G' response \
397                             (possible CVE-2021-23222 buffer-stuffing attack)"
398                                .to_string(),
399                        ));
400                    }
401                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
402                        // No extra data — this is the expected path.
403                    }
404                    Err(e) => {
405                        return Err(PgError::Io(e));
406                    }
407                }
408                Ok(GssEncNegotiationResult::Accepted(tcp_stream))
409            }
410            b'N' => Ok(GssEncNegotiationResult::Rejected),
411            b'E' => {
412                // Server sent an ErrorMessage.  Per CVE-2024-10977 we
413                // must NOT display this to users since the server has
414                // not been authenticated.  Log at trace only.
415                tracing::trace!(
416                    host = %host,
417                    port = %port,
418                    "gssenc_request_server_error (suppressed per CVE-2024-10977)"
419                );
420                Ok(GssEncNegotiationResult::ServerError)
421            }
422            other => Err(PgError::Connection(format!(
423                "Unexpected response to GSSENCRequest: 0x{:02X} \
424                     (expected 'G'=0x47 or 'N'=0x4E)",
425                other
426            ))),
427        }
428    }
429
430    #[cfg(all(feature = "enterprise-gssapi", target_os = "linux"))]
431    async fn connect_gssenc_accepted_with_timeout(
432        tcp_stream: TcpStream,
433        params: ConnectParams<'_>,
434    ) -> PgResult<Self> {
435        let gssenc_fut = async {
436            let gss_stream = super::super::gss::gssenc_handshake(tcp_stream, params.host)
437                .await
438                .map_err(PgError::Auth)?;
439            let mut conn = Self {
440                stream: PgStream::GssEnc(gss_stream),
441                buffer: BytesMut::with_capacity(BUFFER_CAPACITY),
442                write_buf: BytesMut::with_capacity(BUFFER_CAPACITY),
443                sql_buf: BytesMut::with_capacity(512),
444                params_buf: Vec::with_capacity(16),
445                prepared_statements: HashMap::new(),
446                stmt_cache: StatementCache::new(STMT_CACHE_CAPACITY),
447                column_info_cache: HashMap::new(),
448                process_id: 0,
449                cancel_key_bytes: Vec::new(),
450                requested_protocol_minor: params.protocol_minor,
451                negotiated_protocol_minor: params.protocol_minor,
452                notifications: VecDeque::new(),
453                replication_stream_active: false,
454                replication_mode_enabled: has_logical_replication_startup_mode(
455                    &params.startup_params,
456                ),
457                last_replication_wal_end: None,
458                io_desynced: false,
459                pending_statement_closes: Vec::new(),
460                draining_statement_closes: false,
461            };
462            conn.send(FrontendMessage::Startup {
463                user: params.user.to_string(),
464                database: params.database.to_string(),
465                protocol_version: protocol_version_from_minor(params.protocol_minor),
466                startup_params: params.startup_params.clone(),
467            })
468            .await?;
469            conn.handle_startup(
470                params.user,
471                params.password,
472                params.auth_settings,
473                params.gss_token_provider,
474            )
475            .await?;
476            Ok(conn)
477        };
478        tokio::time::timeout(DEFAULT_CONNECT_TIMEOUT, gssenc_fut)
479            .await
480            .map_err(|_| {
481                PgError::Connection(format!(
482                    "GSSENC connection timeout after {:?} (handshake + auth)",
483                    DEFAULT_CONNECT_TIMEOUT
484                ))
485            })?
486    }
487
488    /// Connect to PostgreSQL server with optional password authentication and auth policy.
489    pub async fn connect_with_password_and_auth(
490        host: &str,
491        port: u16,
492        user: &str,
493        database: &str,
494        password: Option<&str>,
495        auth_settings: AuthSettings,
496    ) -> PgResult<Self> {
497        Self::connect_with_password_and_auth_and_gss(ConnectParams {
498            host,
499            port,
500            user,
501            database,
502            password,
503            auth_settings,
504            gss_token_provider: None,
505            io_uring: false,
506            protocol_minor: Self::default_protocol_minor(),
507            startup_params: Vec::new(),
508        })
509        .await
510    }
511
512    async fn connect_with_password_and_auth_and_gss(params: ConnectParams<'_>) -> PgResult<Self> {
513        let first = Self::connect_with_password_and_auth_and_gss_once(params.clone()).await;
514        if let Err(err) = &first
515            && params.protocol_minor > 0
516            && is_explicit_protocol_version_rejection(err)
517        {
518            let mut downgraded = params;
519            downgraded.protocol_minor = (PROTOCOL_VERSION_3_0 & 0xFFFF) as u16;
520            return Self::connect_with_password_and_auth_and_gss_once(downgraded).await;
521        }
522        first
523    }
524
525    async fn connect_with_password_and_auth_and_gss_once(
526        params: ConnectParams<'_>,
527    ) -> PgResult<Self> {
528        let connect_started = Instant::now();
529        let attempt_backend = plain_connect_attempt_backend(params.io_uring);
530        record_connect_attempt(CONNECT_TRANSPORT_PLAIN, attempt_backend);
531        let result = tokio::time::timeout(
532            DEFAULT_CONNECT_TIMEOUT,
533            Self::connect_with_password_inner(params),
534        )
535        .await
536        .map_err(|_| {
537            PgError::Connection(format!(
538                "Connection timeout after {:?} (TCP connect + handshake)",
539                DEFAULT_CONNECT_TIMEOUT
540            ))
541        })?;
542        let backend = result
543            .as_ref()
544            .map(|conn| connect_backend_for_stream(&conn.stream))
545            .unwrap_or(attempt_backend);
546        record_connect_result(
547            CONNECT_TRANSPORT_PLAIN,
548            backend,
549            &result,
550            connect_started.elapsed(),
551        );
552        result
553    }
554
555    /// Inner connection logic without timeout wrapper.
556    async fn connect_with_password_inner(params: ConnectParams<'_>) -> PgResult<Self> {
557        let ConnectParams {
558            host,
559            port,
560            user,
561            database,
562            password,
563            auth_settings,
564            gss_token_provider,
565            io_uring,
566            protocol_minor,
567            startup_params,
568        } = params;
569        let replication_mode_enabled = has_logical_replication_startup_mode(&startup_params);
570        let addr = socket_addr(host, port);
571        let stream = Self::connect_plain_stream(&addr, io_uring).await?;
572
573        let mut conn = Self {
574            stream,
575            buffer: BytesMut::with_capacity(BUFFER_CAPACITY),
576            write_buf: BytesMut::with_capacity(BUFFER_CAPACITY), // 64KB write buffer
577            sql_buf: BytesMut::with_capacity(512),
578            params_buf: Vec::with_capacity(16), // SQL encoding buffer
579            prepared_statements: HashMap::new(),
580            stmt_cache: StatementCache::new(STMT_CACHE_CAPACITY),
581            column_info_cache: HashMap::new(),
582            process_id: 0,
583            cancel_key_bytes: Vec::new(),
584            requested_protocol_minor: protocol_minor,
585            negotiated_protocol_minor: protocol_minor,
586            notifications: VecDeque::new(),
587            replication_stream_active: false,
588            replication_mode_enabled,
589            last_replication_wal_end: None,
590            io_desynced: false,
591            pending_statement_closes: Vec::new(),
592            draining_statement_closes: false,
593        };
594
595        conn.send(FrontendMessage::Startup {
596            user: user.to_string(),
597            database: database.to_string(),
598            protocol_version: protocol_version_from_minor(protocol_minor),
599            startup_params,
600        })
601        .await?;
602
603        conn.handle_startup(user, password, auth_settings, gss_token_provider)
604            .await?;
605
606        Ok(conn)
607    }
608
609    async fn connect_plain_stream(addr: &str, io_uring: bool) -> PgResult<PgStream> {
610        let tcp_stream = TcpStream::connect(addr).await?;
611        tcp_stream.set_nodelay(true)?;
612
613        #[cfg(all(target_os = "linux", feature = "native-io-uring"))]
614        {
615            if should_try_uring_plain(io_uring) {
616                let std_stream = tcp_stream.into_std()?;
617                let fallback_std = std_stream.try_clone()?;
618                match super::super::uring::UringTcpStream::from_std(std_stream) {
619                    Ok(uring_stream) => {
620                        tracing::info!(
621                            addr = %addr,
622                            "qail-pg: using io_uring plain TCP transport"
623                        );
624                        return Ok(PgStream::Uring(uring_stream));
625                    }
626                    Err(e) => {
627                        tracing::warn!(
628                            addr = %addr,
629                            error = %e,
630                            "qail-pg: io_uring stream conversion failed; falling back to tokio TCP"
631                        );
632                        fallback_std.set_nonblocking(true)?;
633                        let fallback = TcpStream::from_std(fallback_std)?;
634                        return Ok(PgStream::Tcp(fallback));
635                    }
636                }
637            }
638        }
639        #[cfg(not(all(target_os = "linux", feature = "native-io-uring")))]
640        {
641            let _ = io_uring;
642        }
643
644        Ok(PgStream::Tcp(tcp_stream))
645    }
646
647    /// Connect to PostgreSQL server with TLS encryption.
648    /// Includes a default 10-second timeout covering TCP connect + TLS + handshake.
649    pub async fn connect_tls(
650        host: &str,
651        port: u16,
652        user: &str,
653        database: &str,
654        password: Option<&str>,
655    ) -> PgResult<Self> {
656        Self::connect_tls_with_auth(
657            host,
658            port,
659            user,
660            database,
661            password,
662            AuthSettings::default(),
663            None,
664        )
665        .await
666    }
667
668    /// Connect to PostgreSQL over TLS with explicit auth policy and optional custom CA bundle.
669    pub async fn connect_tls_with_auth(
670        host: &str,
671        port: u16,
672        user: &str,
673        database: &str,
674        password: Option<&str>,
675        auth_settings: AuthSettings,
676        ca_cert_pem: Option<&[u8]>,
677    ) -> PgResult<Self> {
678        Self::connect_tls_with_auth_and_gss(
679            ConnectParams {
680                host,
681                port,
682                user,
683                database,
684                password,
685                auth_settings,
686                gss_token_provider: None,
687                io_uring: false,
688                protocol_minor: Self::default_protocol_minor(),
689                startup_params: Vec::new(),
690            },
691            ca_cert_pem,
692        )
693        .await
694    }
695
696    async fn connect_tls_with_auth_and_gss(
697        params: ConnectParams<'_>,
698        ca_cert_pem: Option<&[u8]>,
699    ) -> PgResult<Self> {
700        let first = Self::connect_tls_with_auth_and_gss_once(params.clone(), ca_cert_pem).await;
701        if let Err(err) = &first
702            && params.protocol_minor > 0
703            && is_explicit_protocol_version_rejection(err)
704        {
705            let mut downgraded = params;
706            downgraded.protocol_minor = (PROTOCOL_VERSION_3_0 & 0xFFFF) as u16;
707            return Self::connect_tls_with_auth_and_gss_once(downgraded, ca_cert_pem).await;
708        }
709        first
710    }
711
712    async fn connect_tls_with_auth_and_gss_once(
713        params: ConnectParams<'_>,
714        ca_cert_pem: Option<&[u8]>,
715    ) -> PgResult<Self> {
716        let connect_started = Instant::now();
717        record_connect_attempt(CONNECT_TRANSPORT_TLS, CONNECT_BACKEND_TOKIO);
718        let result = tokio::time::timeout(
719            DEFAULT_CONNECT_TIMEOUT,
720            Self::connect_tls_inner(params, ca_cert_pem),
721        )
722        .await
723        .map_err(|_| {
724            PgError::Connection(format!(
725                "TLS connection timeout after {:?}",
726                DEFAULT_CONNECT_TIMEOUT
727            ))
728        })?;
729        record_connect_result(
730            CONNECT_TRANSPORT_TLS,
731            CONNECT_BACKEND_TOKIO,
732            &result,
733            connect_started.elapsed(),
734        );
735        result
736    }
737
738    /// Inner TLS connection logic without timeout wrapper.
739    async fn connect_tls_inner(
740        params: ConnectParams<'_>,
741        ca_cert_pem: Option<&[u8]>,
742    ) -> PgResult<Self> {
743        let ConnectParams {
744            host,
745            port,
746            user,
747            database,
748            password,
749            auth_settings,
750            gss_token_provider,
751            io_uring: _,
752            protocol_minor,
753            startup_params,
754        } = params;
755        let replication_mode_enabled = has_logical_replication_startup_mode(&startup_params);
756        use tokio::io::AsyncReadExt;
757        use tokio_rustls::TlsConnector;
758        use tokio_rustls::rustls::ClientConfig;
759        use tokio_rustls::rustls::pki_types::{CertificateDer, ServerName, pem::PemObject};
760
761        let addr = socket_addr(host, port);
762        let mut tcp_stream = TcpStream::connect(&addr).await?;
763
764        // Send SSLRequest
765        tcp_stream.write_all(&SSL_REQUEST).await?;
766
767        // Read response
768        let mut response = [0u8; 1];
769        tcp_stream.read_exact(&mut response).await?;
770
771        if response[0] != b'S' {
772            return Err(PgError::Connection(
773                "Server does not support TLS".to_string(),
774            ));
775        }
776
777        let mut root_cert_store = tokio_rustls::rustls::RootCertStore::empty();
778
779        if let Some(ca_pem) = ca_cert_pem {
780            let certs = CertificateDer::pem_slice_iter(ca_pem)
781                .collect::<Result<Vec<_>, _>>()
782                .map_err(|e| PgError::Connection(format!("Invalid CA certificate PEM: {}", e)))?;
783            if certs.is_empty() {
784                return Err(PgError::Connection(
785                    "No CA certificates found in provided PEM".to_string(),
786                ));
787            }
788            for cert in certs {
789                let _ = root_cert_store.add(cert);
790            }
791        } else {
792            let certs = rustls_native_certs::load_native_certs();
793            for cert in certs.certs {
794                let _ = root_cert_store.add(cert);
795            }
796        }
797
798        let config = ClientConfig::builder()
799            .with_root_certificates(root_cert_store)
800            .with_no_client_auth();
801
802        let connector = TlsConnector::from(Arc::new(config));
803        let server_name = ServerName::try_from(host.to_string())
804            .map_err(|_| PgError::Connection("Invalid hostname for TLS".to_string()))?;
805
806        let tls_stream = connector
807            .connect(server_name, tcp_stream)
808            .await
809            .map_err(|e| PgError::Connection(format!("TLS handshake failed: {}", e)))?;
810
811        let mut conn = Self {
812            stream: PgStream::Tls(Box::new(tls_stream)),
813            buffer: BytesMut::with_capacity(BUFFER_CAPACITY),
814            write_buf: BytesMut::with_capacity(BUFFER_CAPACITY),
815            sql_buf: BytesMut::with_capacity(512),
816            params_buf: Vec::with_capacity(16),
817            prepared_statements: HashMap::new(),
818            stmt_cache: StatementCache::new(STMT_CACHE_CAPACITY),
819            column_info_cache: HashMap::new(),
820            process_id: 0,
821            cancel_key_bytes: Vec::new(),
822            requested_protocol_minor: protocol_minor,
823            negotiated_protocol_minor: protocol_minor,
824            notifications: VecDeque::new(),
825            replication_stream_active: false,
826            replication_mode_enabled,
827            last_replication_wal_end: None,
828            io_desynced: false,
829            pending_statement_closes: Vec::new(),
830            draining_statement_closes: false,
831        };
832
833        conn.send(FrontendMessage::Startup {
834            user: user.to_string(),
835            database: database.to_string(),
836            protocol_version: protocol_version_from_minor(protocol_minor),
837            startup_params,
838        })
839        .await?;
840
841        conn.handle_startup(user, password, auth_settings, gss_token_provider)
842            .await?;
843
844        Ok(conn)
845    }
846
847    /// Connect with mutual TLS (client certificate authentication).
848    /// # Arguments
849    /// * `host` - PostgreSQL server hostname
850    /// * `port` - PostgreSQL server port
851    /// * `user` - Database user
852    /// * `database` - Database name
853    /// * `config` - TLS configuration with client cert/key
854    /// # Example
855    /// ```ignore
856    /// let config = TlsConfig {
857    ///     client_cert_pem: include_bytes!("client.crt").to_vec(),
858    ///     client_key_pem: include_bytes!("client.key").to_vec(),
859    ///     ca_cert_pem: Some(include_bytes!("ca.crt").to_vec()),
860    /// };
861    /// let conn = PgConnection::connect_mtls("localhost", 5432, "user", "db", config).await?;
862    /// ```
863    pub async fn connect_mtls(
864        host: &str,
865        port: u16,
866        user: &str,
867        database: &str,
868        config: TlsConfig,
869    ) -> PgResult<Self> {
870        Self::connect_mtls_with_password_and_auth(
871            host,
872            port,
873            user,
874            database,
875            None,
876            config,
877            AuthSettings::default(),
878        )
879        .await
880    }
881
882    /// Connect with mutual TLS and optional password fallback.
883    pub async fn connect_mtls_with_password_and_auth(
884        host: &str,
885        port: u16,
886        user: &str,
887        database: &str,
888        password: Option<&str>,
889        config: TlsConfig,
890        auth_settings: AuthSettings,
891    ) -> PgResult<Self> {
892        Self::connect_mtls_with_password_and_auth_and_gss(
893            ConnectParams {
894                host,
895                port,
896                user,
897                database,
898                password,
899                auth_settings,
900                gss_token_provider: None,
901                io_uring: false,
902                protocol_minor: Self::default_protocol_minor(),
903                startup_params: Vec::new(),
904            },
905            config,
906        )
907        .await
908    }
909
910    async fn connect_mtls_with_password_and_auth_and_gss(
911        params: ConnectParams<'_>,
912        config: TlsConfig,
913    ) -> PgResult<Self> {
914        let first =
915            Self::connect_mtls_with_password_and_auth_and_gss_once(params.clone(), config.clone())
916                .await;
917        if let Err(err) = &first
918            && params.protocol_minor > 0
919            && is_explicit_protocol_version_rejection(err)
920        {
921            let mut downgraded = params;
922            downgraded.protocol_minor = (PROTOCOL_VERSION_3_0 & 0xFFFF) as u16;
923            return Self::connect_mtls_with_password_and_auth_and_gss_once(downgraded, config)
924                .await;
925        }
926        first
927    }
928
929    async fn connect_mtls_with_password_and_auth_and_gss_once(
930        params: ConnectParams<'_>,
931        config: TlsConfig,
932    ) -> PgResult<Self> {
933        let connect_started = Instant::now();
934        record_connect_attempt(CONNECT_TRANSPORT_MTLS, CONNECT_BACKEND_TOKIO);
935        let result = tokio::time::timeout(
936            DEFAULT_CONNECT_TIMEOUT,
937            Self::connect_mtls_inner(params, config),
938        )
939        .await
940        .map_err(|_| {
941            PgError::Connection(format!(
942                "mTLS connection timeout after {:?}",
943                DEFAULT_CONNECT_TIMEOUT
944            ))
945        })?;
946        record_connect_result(
947            CONNECT_TRANSPORT_MTLS,
948            CONNECT_BACKEND_TOKIO,
949            &result,
950            connect_started.elapsed(),
951        );
952        result
953    }
954
955    /// Inner mTLS connection logic without timeout wrapper.
956    async fn connect_mtls_inner(params: ConnectParams<'_>, config: TlsConfig) -> PgResult<Self> {
957        let ConnectParams {
958            host,
959            port,
960            user,
961            database,
962            password,
963            auth_settings,
964            gss_token_provider,
965            io_uring: _,
966            protocol_minor,
967            startup_params,
968        } = params;
969        let replication_mode_enabled = has_logical_replication_startup_mode(&startup_params);
970        use tokio::io::AsyncReadExt;
971        use tokio_rustls::TlsConnector;
972        use tokio_rustls::rustls::{
973            ClientConfig,
974            pki_types::{CertificateDer, PrivateKeyDer, ServerName, pem::PemObject},
975        };
976
977        let addr = socket_addr(host, port);
978        let mut tcp_stream = TcpStream::connect(&addr).await?;
979
980        // Send SSLRequest
981        tcp_stream.write_all(&SSL_REQUEST).await?;
982
983        // Read response
984        let mut response = [0u8; 1];
985        tcp_stream.read_exact(&mut response).await?;
986
987        if response[0] != b'S' {
988            return Err(PgError::Connection(
989                "Server does not support TLS".to_string(),
990            ));
991        }
992
993        let mut root_cert_store = tokio_rustls::rustls::RootCertStore::empty();
994
995        if let Some(ca_pem) = &config.ca_cert_pem {
996            let certs = CertificateDer::pem_slice_iter(ca_pem)
997                .collect::<Result<Vec<_>, _>>()
998                .map_err(|e| PgError::Connection(format!("Invalid CA certificate PEM: {}", e)))?;
999            if certs.is_empty() {
1000                return Err(PgError::Connection(
1001                    "No CA certificates found in provided PEM".to_string(),
1002                ));
1003            }
1004            for cert in certs {
1005                let _ = root_cert_store.add(cert);
1006            }
1007        } else {
1008            // Use system certs
1009            let certs = rustls_native_certs::load_native_certs();
1010            for cert in certs.certs {
1011                let _ = root_cert_store.add(cert);
1012            }
1013        }
1014
1015        let client_certs: Vec<CertificateDer<'static>> =
1016            CertificateDer::pem_slice_iter(&config.client_cert_pem)
1017                .collect::<Result<Vec<_>, _>>()
1018                .map_err(|e| PgError::Connection(format!("Invalid client cert PEM: {}", e)))?;
1019        if client_certs.is_empty() {
1020            return Err(PgError::Connection(
1021                "No client certificates found in PEM".to_string(),
1022            ));
1023        }
1024
1025        let client_key = PrivateKeyDer::from_pem_slice(&config.client_key_pem)
1026            .map_err(|e| PgError::Connection(format!("Invalid client key PEM: {}", e)))?;
1027
1028        let tls_config = ClientConfig::builder()
1029            .with_root_certificates(root_cert_store)
1030            .with_client_auth_cert(client_certs, client_key)
1031            .map_err(|e| PgError::Connection(format!("Invalid client cert/key: {}", e)))?;
1032
1033        let connector = TlsConnector::from(Arc::new(tls_config));
1034        let server_name = ServerName::try_from(host.to_string())
1035            .map_err(|_| PgError::Connection("Invalid hostname for TLS".to_string()))?;
1036
1037        let tls_stream = connector
1038            .connect(server_name, tcp_stream)
1039            .await
1040            .map_err(|e| PgError::Connection(format!("mTLS handshake failed: {}", e)))?;
1041
1042        let mut conn = Self {
1043            stream: PgStream::Tls(Box::new(tls_stream)),
1044            buffer: BytesMut::with_capacity(BUFFER_CAPACITY),
1045            write_buf: BytesMut::with_capacity(BUFFER_CAPACITY),
1046            sql_buf: BytesMut::with_capacity(512),
1047            params_buf: Vec::with_capacity(16),
1048            prepared_statements: HashMap::new(),
1049            stmt_cache: StatementCache::new(STMT_CACHE_CAPACITY),
1050            column_info_cache: HashMap::new(),
1051            process_id: 0,
1052            cancel_key_bytes: Vec::new(),
1053            requested_protocol_minor: protocol_minor,
1054            negotiated_protocol_minor: protocol_minor,
1055            notifications: VecDeque::new(),
1056            replication_stream_active: false,
1057            replication_mode_enabled,
1058            last_replication_wal_end: None,
1059            io_desynced: false,
1060            pending_statement_closes: Vec::new(),
1061            draining_statement_closes: false,
1062        };
1063
1064        conn.send(FrontendMessage::Startup {
1065            user: user.to_string(),
1066            database: database.to_string(),
1067            protocol_version: protocol_version_from_minor(protocol_minor),
1068            startup_params,
1069        })
1070        .await?;
1071
1072        conn.handle_startup(user, password, auth_settings, gss_token_provider)
1073            .await?;
1074
1075        Ok(conn)
1076    }
1077
1078    /// Connect to PostgreSQL server via Unix domain socket.
1079    #[cfg(unix)]
1080    pub async fn connect_unix(
1081        socket_path: &str,
1082        user: &str,
1083        database: &str,
1084        password: Option<&str>,
1085    ) -> PgResult<Self> {
1086        let default_minor = Self::default_protocol_minor();
1087        let first =
1088            Self::connect_unix_with_protocol(socket_path, user, database, password, default_minor)
1089                .await;
1090        if let Err(err) = &first
1091            && default_minor > 0
1092            && is_explicit_protocol_version_rejection(err)
1093        {
1094            let downgrade_minor = (PROTOCOL_VERSION_3_0 & 0xFFFF) as u16;
1095            return Self::connect_unix_with_protocol(
1096                socket_path,
1097                user,
1098                database,
1099                password,
1100                downgrade_minor,
1101            )
1102            .await;
1103        }
1104        first
1105    }
1106
1107    #[cfg(unix)]
1108    async fn connect_unix_with_protocol(
1109        socket_path: &str,
1110        user: &str,
1111        database: &str,
1112        password: Option<&str>,
1113        protocol_minor: u16,
1114    ) -> PgResult<Self> {
1115        use tokio::net::UnixStream;
1116
1117        let unix_stream = UnixStream::connect(socket_path).await?;
1118
1119        let mut conn = Self {
1120            stream: PgStream::Unix(unix_stream),
1121            buffer: BytesMut::with_capacity(BUFFER_CAPACITY),
1122            write_buf: BytesMut::with_capacity(BUFFER_CAPACITY),
1123            sql_buf: BytesMut::with_capacity(512),
1124            params_buf: Vec::with_capacity(16),
1125            prepared_statements: HashMap::new(),
1126            stmt_cache: StatementCache::new(STMT_CACHE_CAPACITY),
1127            column_info_cache: HashMap::new(),
1128            process_id: 0,
1129            cancel_key_bytes: Vec::new(),
1130            requested_protocol_minor: protocol_minor,
1131            negotiated_protocol_minor: protocol_minor,
1132            notifications: VecDeque::new(),
1133            replication_stream_active: false,
1134            replication_mode_enabled: false,
1135            last_replication_wal_end: None,
1136            io_desynced: false,
1137            pending_statement_closes: Vec::new(),
1138            draining_statement_closes: false,
1139        };
1140
1141        conn.send(FrontendMessage::Startup {
1142            user: user.to_string(),
1143            database: database.to_string(),
1144            protocol_version: protocol_version_from_minor(protocol_minor),
1145            startup_params: Vec::new(),
1146        })
1147        .await?;
1148
1149        conn.handle_startup(user, password, AuthSettings::default(), None)
1150            .await?;
1151
1152        Ok(conn)
1153    }
1154}
1155
1156#[cfg(test)]
1157mod tests {
1158    use super::{is_explicit_protocol_version_rejection, protocol_version_from_minor, socket_addr};
1159    use crate::driver::PgError;
1160
1161    #[test]
1162    fn protocol_version_from_minor_encodes_major_3() {
1163        assert_eq!(protocol_version_from_minor(2), 196610);
1164        assert_eq!(protocol_version_from_minor(0), 196608);
1165    }
1166
1167    #[test]
1168    fn socket_addr_brackets_ipv6_hosts() {
1169        assert_eq!(socket_addr("127.0.0.1", 5432), "127.0.0.1:5432");
1170        assert_eq!(socket_addr("::1", 5432), "[::1]:5432");
1171        assert_eq!(socket_addr("[::1]", 5432), "[::1]:5432");
1172    }
1173
1174    #[test]
1175    fn explicit_protocol_rejection_detection_is_case_insensitive() {
1176        let err = PgError::Connection("Unsupported frontend protocol 3.2".to_string());
1177        assert!(is_explicit_protocol_version_rejection(&err));
1178
1179        let err = PgError::Protocol("server: Protocol VERSION not supported".to_string());
1180        assert!(is_explicit_protocol_version_rejection(&err));
1181    }
1182
1183    #[test]
1184    fn explicit_protocol_rejection_does_not_match_unrelated_errors() {
1185        let err = PgError::Connection("connection reset by peer".to_string());
1186        assert!(!is_explicit_protocol_version_rejection(&err));
1187    }
1188}