1#[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 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 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 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 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 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 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 }
256 Err(e) => {
257 if gss_enc_mode == GssEncMode::Require {
258 return Err(e);
259 }
260 tracing::debug!(
262 host = %host,
263 port = %port,
264 error = %e,
265 "gssenc_prefer_fallthrough"
266 );
267 }
268 }
269 }
270
271 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 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 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 tcp_stream.write_all(&GSSENC_REQUEST).await?;
378 tcp_stream.flush().await?;
379
380 let mut response = [0u8; 1];
384 tcp_stream.read_exact(&mut response).await?;
385
386 match response[0] {
387 b'G' => {
388 let mut peek_buf = [0u8; 1];
391 match tcp_stream.try_read(&mut peek_buf) {
392 Ok(0) => {} Ok(_n) => {
394 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 }
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 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 ¶ms.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 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 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), sql_buf: BytesMut::with_capacity(512),
578 params_buf: Vec::with_capacity(16), 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 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 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 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 tcp_stream.write_all(&SSL_REQUEST).await?;
766
767 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 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 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 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 tcp_stream.write_all(&SSL_REQUEST).await?;
982
983 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 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 #[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}