postgres: Allow configuration of SSL negotiation procedure

Author: danielfrankcom

sqlx-postgres/src/connection/tls.rs

@@ -3,7 +3,7 @@ use crate::net::tls::{self, TlsConfig};
 use crate::net::{Socket, SocketIntoBox, WithSocket};
 
 use crate::message::SslRequest;
-use crate::{PgConnectOptions, PgSslMode};
+use crate::{PgConnectOptions, PgSslMode, PgSslNegotiation};
 
 pub struct MaybeUpgradeTls<'a>(pub &'a PgConnectOptions);
 
@@ -19,28 +19,52 @@ async fn maybe_upgrade<S: Socket>(
     mut socket: S,
     options: &PgConnectOptions,
 ) -> Result<Box<dyn Socket>, Error> {
-    // https://www.postgresql.org/docs/12/libpq-ssl.html#LIBPQ-SSL-SSLMODE-STATEMENTS
-    match options.ssl_mode {
-        // FIXME: Implement ALLOW
-        PgSslMode::Allow | PgSslMode::Disable => return Ok(Box::new(socket)),
-
-        PgSslMode::Prefer => {
-            if !tls::available() {
-                return Ok(Box::new(socket));
-            }
-
-            // try upgrade, but its okay if we fail
-            if !request_upgrade(&mut socket, options).await? {
-                return Ok(Box::new(socket));
+    // https://www.postgresql.org/docs/17/libpq-connect.html#LIBPQ-CONNECT-SSLNEGOTIATION
+    match options.ssl_negotiation {
+        PgSslNegotiation::Postgres => {
+            // https://www.postgresql.org/docs/12/libpq-ssl.html#LIBPQ-SSL-SSLMODE-STATEMENTS
+            match options.ssl_mode {
+                // FIXME: Implement ALLOW
+                PgSslMode::Allow | PgSslMode::Disable => return Ok(Box::new(socket)),
+
+                PgSslMode::Prefer => {
+                    if !tls::available() {
+                        return Ok(Box::new(socket));
+                    }
+
+                    // try upgrade, but its okay if we fail
+                    if !request_upgrade(&mut socket, options).await? {
+                        return Ok(Box::new(socket));
+                    }
+                }
+
+                PgSslMode::Require | PgSslMode::VerifyFull | PgSslMode::VerifyCa => {
+                    tls::error_if_unavailable()?;
+
+                    if !request_upgrade(&mut socket, options).await? {
+                        // upgrade failed, die
+                        return Err(Error::Tls("server does not support TLS".into()));
+                    }
+                }
             }
         }
-
-        PgSslMode::Require | PgSslMode::VerifyFull | PgSslMode::VerifyCa => {
-            tls::error_if_unavailable()?;
-
-            if !request_upgrade(&mut socket, options).await? {
-                // upgrade failed, die
-                return Err(Error::Tls("server does not support TLS".into()));
+        PgSslNegotiation::Direct => {
+            // Direct TLS negotiation without PostgreSQL handshake
+            match options.ssl_mode {
+                PgSslMode::Disable | PgSslMode::Allow | PgSslMode::Prefer => {
+                    return Err(Error::Tls(
+                        format!(
+                            "SSL mode {:?} is incompatible with direct SSL negotiation",
+                            options.ssl_mode
+                        )
+                        .into(),
+                    ));
+                }
+
+                PgSslMode::Require | PgSslMode::VerifyFull | PgSslMode::VerifyCa => {
+                    tls::error_if_unavailable()?;
+                    // No need to request an upgrade. We go straight to TLS handshake.
+                }
             }
         }
     }

sqlx-postgres/src/lib.rs

@@ -56,7 +56,7 @@ pub use database::Postgres;
 pub use error::{PgDatabaseError, PgErrorPosition};
 pub use listener::{PgListener, PgNotification};
 pub use message::PgSeverity;
-pub use options::{PgConnectOptions, PgSslMode};
+pub use options::{PgConnectOptions, PgSslMode, PgSslNegotiation};
 pub use query_result::PgQueryResult;
 pub use row::PgRow;
 pub use statement::PgStatement;

sqlx-postgres/src/options/doc.md

@@ -33,6 +33,7 @@ if a parameter is not passed in via URL, it is populated by reading
 | `port`             | `PGPORT`             | `5432`                                                      |
 | `dbname`           | `PGDATABASE`         | Unset; defaults to the username server-side.                |
 | `sslmode`          | `PGSSLMODE`          | `prefer`. See [`PgSslMode`] for details.                    |
+| `sslnegotiation`   | `PGSSLNEGOTIATION`   | `postgres`. See [`PgSslNegotiation`] for details.           |
 | `sslrootcert`      | `PGSSLROOTCERT`      | Unset. See [Note: SSL](#note-ssl).                          |
 | `sslcert`          | `PGSSLCERT`          | Unset. See [Note: SSL](#note-ssl).                          |
 | `sslkey`           | `PGSSLKEY`           | Unset. See [Note: SSL](#note-ssl).                          |

sqlx-postgres/src/options/mod.rs

@@ -4,13 +4,15 @@ use std::fmt::{self, Display, Write};
 use std::path::{Path, PathBuf};
 
 pub use ssl_mode::PgSslMode;
+pub use ssl_negotiation::PgSslNegotiation;
 
 use crate::{connection::LogSettings, net::tls::CertificateInput};
 
 mod connect;
 mod parse;
 mod pgpass;
 mod ssl_mode;
+mod ssl_negotiation;
 
 #[doc = include_str!("doc.md")]
 #[derive(Debug, Clone)]
@@ -22,6 +24,7 @@ pub struct PgConnectOptions {
     pub(crate) password: Option<String>,
     pub(crate) database: Option<String>,
     pub(crate) ssl_mode: PgSslMode,
+    pub(crate) ssl_negotiation: PgSslNegotiation,
     pub(crate) ssl_root_cert: Option<CertificateInput>,
     pub(crate) ssl_client_cert: Option<CertificateInput>,
     pub(crate) ssl_client_key: Option<CertificateInput>,
@@ -85,6 +88,10 @@ impl PgConnectOptions {
                 .ok()
                 .and_then(|v| v.parse().ok())
                 .unwrap_or_default(),
+            ssl_negotiation: var("PGSSLNEGOTIATION")
+                .ok()
+                .and_then(|v| v.parse().ok())
+                .unwrap_or_default(),
             statement_cache_capacity: 100,
             application_name: var("PGAPPNAME").ok(),
             extra_float_digits: Some("2".into()),
@@ -218,6 +225,26 @@ impl PgConnectOptions {
         self
     }
 
+    /// Sets the protocol with which the secure SSL TCP/IP connection will be negotiated with
+    /// the server.
+    ///
+    /// By default, the protocol is [`Postgres`](PgSslNegotiation::Postgres), and the client will
+    /// first check whether the server supports SSL, and fallback to a non-SSL connection if not.
+    ///
+    /// Ignored for Unix domain socket communication.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # use sqlx_postgres::{PgSslNegotiation, PgConnectOptions};
+    /// let options = PgConnectOptions::new()
+    ///     .ssl_negotiation(PgSslNegotiation::Postgres);
+    /// ```
+    pub fn ssl_negotiation(mut self, procedure: PgSslNegotiation) -> Self {
+        self.ssl_negotiation = procedure;
+        self
+    }
+
     /// Sets the name of a file containing SSL certificate authority (CA) certificate(s).
     /// If the file exists, the server's certificate will be verified to be signed by
     /// one of these authorities.
@@ -546,6 +573,19 @@ impl PgConnectOptions {
         self.ssl_mode
     }
 
+    /// Get the SSL negotiation protocol.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # use sqlx_postgres::{PgConnectOptions, PgSslNegotiation};
+    /// let options = PgConnectOptions::new();
+    /// assert!(matches!(options.get_ssl_negotiation(), PgSslNegotiation::Postgres));
+    /// ```
+    pub fn get_ssl_negotiation(&self) -> PgSslNegotiation {
+        self.ssl_negotiation
+    }
+
     /// Get the application name.
     ///
     /// # Example

sqlx-postgres/src/options/parse.rs

@@ -1,5 +1,5 @@
 use crate::error::Error;
-use crate::{PgConnectOptions, PgSslMode};
+use crate::{PgConnectOptions, PgSslMode, PgSslNegotiation};
 use sqlx_core::percent_encoding::{percent_decode_str, utf8_percent_encode, NON_ALPHANUMERIC};
 use sqlx_core::Url;
 use std::net::IpAddr;
@@ -53,6 +53,10 @@ impl PgConnectOptions {
                     options = options.ssl_mode(value.parse().map_err(Error::config)?);
                 }
 
+                "sslnegotiation" | "ssl-negotiation" => {
+                    options = options.ssl_negotiation(value.parse().map_err(Error::config)?);
+                }
+
                 "sslrootcert" | "ssl-root-cert" | "ssl-ca" => {
                     options = options.ssl_root_cert(&*value);
                 }
@@ -146,6 +150,13 @@ impl PgConnectOptions {
         };
         url.query_pairs_mut().append_pair("sslmode", ssl_mode);
 
+        let ssl_negotiation = match self.ssl_negotiation {
+            PgSslNegotiation::Postgres => "postgres",
+            PgSslNegotiation::Direct => "direct",
+        };
+        url.query_pairs_mut()
+            .append_pair("sslnegotiation", ssl_negotiation);
+
         if let Some(ssl_root_cert) = &self.ssl_root_cert {
             url.query_pairs_mut()
                 .append_pair("sslrootcert", &ssl_root_cert.to_string());
@@ -266,6 +277,22 @@ fn it_parses_password_with_non_ascii_chars_correctly() {
     assert_eq!(Some("p@ssw0rd".into()), opts.password);
 }
 
+#[test]
+fn it_parses_sslmode_correctly_from_parameter() {
+    let url = "postgres://?sslmode=verify-full";
+    let opts = PgConnectOptions::from_str(url).unwrap();
+
+    assert!(matches!(opts.ssl_mode, PgSslMode::VerifyFull));
+}
+
+#[test]
+fn it_parses_sslnegotiation_correctly_from_parameter() {
+    let url = "postgres://?sslnegotiation=direct";
+    let opts = PgConnectOptions::from_str(url).unwrap();
+
+    assert!(matches!(opts.ssl_negotiation, PgSslNegotiation::Direct));
+}
+
 #[test]
 fn it_parses_socket_correctly_percent_encoded() {
     let url = "postgres://%2Fvar%2Flib%2Fpostgres/database";
@@ -310,7 +337,7 @@ fn it_returns_the_parsed_url_when_socket() {
 
     let mut expected_url = Url::parse(url).unwrap();
     // PgConnectOptions defaults
-    let query_string = "sslmode=prefer&statement-cache-capacity=100";
+    let query_string = "sslmode=prefer&sslnegotiation=postgres&statement-cache-capacity=100";
     let port = 5432;
     expected_url.set_query(Some(query_string));
     let _ = expected_url.set_port(Some(port));
@@ -325,7 +352,7 @@ fn it_returns_the_parsed_url_when_host() {
 
     let mut expected_url = Url::parse(url).unwrap();
     // PgConnectOptions defaults
-    let query_string = "sslmode=prefer&statement-cache-capacity=100";
+    let query_string = "sslmode=prefer&sslnegotiation=postgres&statement-cache-capacity=100";
     expected_url.set_query(Some(query_string));
 
     assert_eq!(expected_url, opts.build_url());

sqlx-postgres/src/options/ssl_negotiation.rs

@@ -0,0 +1,35 @@
+use crate::error::Error;
+use std::str::FromStr;
+
+/// Options for controlling the connection establishment procedure for PostgreSQL SSL connections.
+///
+/// It is used by the [`sslnegotiation`](super::PgConnectOptions::ssl_negotiation) method.
+#[derive(Debug, Clone, Copy, Default)]
+pub enum PgSslNegotiation {
+    /// The client first asks the server if SSL is supported.
+    ///
+    /// This is the default if no other mode is specified.
+    #[default]
+    Postgres,
+
+    /// The client starts the standard SSL handshake directly after establishing the TCP/IP
+    /// connection.
+    Direct,
+}
+
+impl FromStr for PgSslNegotiation {
+    type Err = Error;
+
+    fn from_str(s: &str) -> Result<Self, Error> {
+        Ok(match &*s.to_ascii_lowercase() {
+            "postgres" => PgSslNegotiation::Postgres,
+            "direct" => PgSslNegotiation::Direct,
+
+            _ => {
+                return Err(Error::Configuration(
+                    format!("unknown value {s:?} for `ssl_negotiation`").into(),
+                ));
+            }
+        })
+    }
+}