From bd283876dd20c237d37adad146b3164fd72850ff Mon Sep 17 00:00:00 2001 From: Thom Wright Date: Fri, 26 Sep 2025 17:50:45 +0100 Subject: [PATCH 1/3] Cache PG salted password and client key --- sqlx-postgres/src/connection/mod.rs | 2 + sqlx-postgres/src/connection/sasl.rs | 133 +++++++++++++++++++++++++-- sqlx-postgres/src/message/mod.rs | 2 +- sqlx-postgres/src/options/mod.rs | 9 +- 4 files changed, 133 insertions(+), 13 deletions(-) diff --git a/sqlx-postgres/src/connection/mod.rs b/sqlx-postgres/src/connection/mod.rs index d5db20ad05..00069e14f3 100644 --- a/sqlx-postgres/src/connection/mod.rs +++ b/sqlx-postgres/src/connection/mod.rs @@ -23,6 +23,8 @@ use sqlx_core::sql_str::SqlSafeStr; pub use self::stream::PgStream; +pub use sasl::ClientKeyCache; + #[cfg(feature = "offline")] mod describe; mod establish; diff --git a/sqlx-postgres/src/connection/sasl.rs b/sqlx-postgres/src/connection/sasl.rs index 94fdfc689f..5190eb1ef3 100644 --- a/sqlx-postgres/src/connection/sasl.rs +++ b/sqlx-postgres/src/connection/sasl.rs @@ -1,6 +1,12 @@ +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + use crate::connection::stream::PgStream; use crate::error::Error; -use crate::message::{Authentication, AuthenticationSasl, SaslInitialResponse, SaslResponse}; +use crate::message::{ + Authentication, AuthenticationSasl, AuthenticationSaslContinue, SaslInitialResponse, + SaslResponse, +}; use crate::rt; use crate::PgConnectOptions; use hmac::{Hmac, Mac}; @@ -16,6 +22,100 @@ const USERNAME_ATTR: &str = "n"; const CLIENT_PROOF_ATTR: &str = "p"; const NONCE_ATTR: &str = "r"; +/// A single-entry cache for the client key derived from the password. +/// +/// Salting the password and deriving the client key can be expensive, so this cache stores the most +/// recently used client key along with the parameters used to derive it. +/// +/// According to [RFC-7677](https://datatracker.ietf.org/doc/html/rfc7677): +/// +/// > This computational cost can be avoided by caching the ClientKey (assuming the Salt and hash +/// > iteration-count is stable). +#[derive(Debug, Clone)] +pub struct ClientKeyCache { + inner: Arc>>, +} + +#[derive(Debug, PartialEq, Eq)] +struct CacheKey { + host: String, + port: u16, + socket: Option, + database: Option, + username: String, + password: String, + salt: Vec, + iterations: u32, +} + +impl From<(&PgConnectOptions, &AuthenticationSaslContinue)> for CacheKey { + fn from((options, cont): (&PgConnectOptions, &AuthenticationSaslContinue)) -> Self { + CacheKey { + host: options.host.clone(), + port: options.port, + socket: options.socket.clone(), + database: options.database.clone(), + username: options.username.clone(), + password: options.password.clone().unwrap_or_default(), + salt: cont.salt.clone(), + iterations: cont.iterations, + } + } +} + +#[derive(Debug)] +struct CacheInner { + cache_key: CacheKey, + salted_password: [u8; 32], + client_key: Hmac, +} + +impl ClientKeyCache { + pub fn new() -> Self { + ClientKeyCache { + inner: Arc::new(Mutex::new(None)), + } + } + + fn get( + &self, + options: &PgConnectOptions, + cont: &AuthenticationSaslContinue, + ) -> Option<([u8; 32], Hmac)> { + let key = CacheKey::from((options, cont)); + + self.inner + .lock() + .expect("BUG: panicked while holding a lock") + .as_ref() + .and_then(|inner| { + if inner.cache_key == key { + Some((inner.salted_password, inner.client_key.clone())) + } else { + None + } + }) + } + + fn set( + &self, + options: &PgConnectOptions, + cont: &AuthenticationSaslContinue, + salted_password: [u8; 32], + client_key: Hmac, + ) { + let mut inner = self + .inner + .lock() + .expect("BUG: panicked while holding a lock"); + *inner = Some(CacheInner { + cache_key: CacheKey::from((options, cont)), + salted_password, + client_key, + }); + } +} + pub(crate) async fn authenticate( stream: &mut PgStream, options: &PgConnectOptions, @@ -86,16 +186,29 @@ pub(crate) async fn authenticate( } }; - // SaltedPassword := Hi(Normalize(password), salt, i) - let salted_password = hi( - options.password.as_deref().unwrap_or_default(), - &cont.salt, - cont.iterations, - ) - .await?; + let (salted_password, mut mac) = { + if let Some(cached) = options.sasl_client_key_cache.get(options, &cont) { + cached + } else { + // SaltedPassword := Hi(Normalize(password), salt, i) + let salted_password = hi( + options.password.as_deref().unwrap_or_default(), + &cont.salt, + cont.iterations, + ) + .await?; + + // ClientKey := HMAC(SaltedPassword, "Client Key") + let mac = Hmac::::new_from_slice(&salted_password).map_err(Error::protocol)?; + + options + .sasl_client_key_cache + .set(options, &cont, salted_password, mac.clone()); + + (salted_password, mac) + } + }; - // ClientKey := HMAC(SaltedPassword, "Client Key") - let mut mac = Hmac::::new_from_slice(&salted_password).map_err(Error::protocol)?; mac.update(b"Client Key"); let client_key = mac.finalize().into_bytes(); diff --git a/sqlx-postgres/src/message/mod.rs b/sqlx-postgres/src/message/mod.rs index e62f9bebb3..e7648f4419 100644 --- a/sqlx-postgres/src/message/mod.rs +++ b/sqlx-postgres/src/message/mod.rs @@ -30,7 +30,7 @@ mod startup; mod sync; mod terminate; -pub use authentication::{Authentication, AuthenticationSasl}; +pub use authentication::{Authentication, AuthenticationSasl, AuthenticationSaslContinue}; pub use backend_key_data::BackendKeyData; pub use bind::Bind; pub use close::Close; diff --git a/sqlx-postgres/src/options/mod.rs b/sqlx-postgres/src/options/mod.rs index 21e6628cae..b7ff288f0b 100644 --- a/sqlx-postgres/src/options/mod.rs +++ b/sqlx-postgres/src/options/mod.rs @@ -5,7 +5,10 @@ use std::path::{Path, PathBuf}; pub use ssl_mode::PgSslMode; -use crate::{connection::LogSettings, net::tls::CertificateInput}; +use crate::{ + connection::{ClientKeyCache, LogSettings}, + net::tls::CertificateInput, +}; mod connect; mod parse; @@ -30,6 +33,7 @@ pub struct PgConnectOptions { pub(crate) log_settings: LogSettings, pub(crate) extra_float_digits: Option>, pub(crate) options: Option, + pub(crate) sasl_client_key_cache: ClientKeyCache, } impl Default for PgConnectOptions { @@ -97,6 +101,7 @@ impl PgConnectOptions { extra_float_digits: Some("2".into()), log_settings: Default::default(), options: var("PGOPTIONS").ok(), + sasl_client_key_cache: ClientKeyCache::new(), } } @@ -274,7 +279,7 @@ impl PgConnectOptions { /// -----BEGIN CERTIFICATE----- /// /// -----END CERTIFICATE-----"; - /// + /// /// let options = PgConnectOptions::new() /// // Providing a CA certificate with less than VerifyCa is pointless /// .ssl_mode(PgSslMode::VerifyCa) From 41b8688224f78dfaea1c52c3a3e70a20077d416c Mon Sep 17 00:00:00 2001 From: Thom Wright Date: Thu, 26 Mar 2026 15:17:31 +0000 Subject: [PATCH 2/3] Replace key with just salt and iterations And invalidate on password change. --- Cargo.lock | 4 +-- sqlx-postgres/src/connection/sasl.rs | 48 +++++++--------------------- sqlx-postgres/src/options/mod.rs | 2 ++ 3 files changed, 15 insertions(+), 39 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 236039f0ff..581c3eab76 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1377,9 +1377,9 @@ dependencies = [ [[package]] name = "flume" -version = "0.11.1" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da0e4dd2a88388a1f4ccc7c9ce104604dab68d9f408dc34cd45823d5a9069095" +checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" dependencies = [ "futures-core", "futures-sink", diff --git a/sqlx-postgres/src/connection/sasl.rs b/sqlx-postgres/src/connection/sasl.rs index 5190eb1ef3..1c382be4f6 100644 --- a/sqlx-postgres/src/connection/sasl.rs +++ b/sqlx-postgres/src/connection/sasl.rs @@ -1,4 +1,3 @@ -use std::path::PathBuf; use std::sync::{Arc, Mutex}; use crate::connection::stream::PgStream; @@ -27,6 +26,10 @@ const NONCE_ATTR: &str = "r"; /// Salting the password and deriving the client key can be expensive, so this cache stores the most /// recently used client key along with the parameters used to derive it. /// +/// The cache is keyed on the salt and iteration count, which are the server-provided parameters +/// that affect the HMAC result. The password is not included in the cache key because it can only +/// change via `&mut self` on `PgConnectOptions`, which replaces the cache instance. +/// /// According to [RFC-7677](https://datatracker.ietf.org/doc/html/rfc7677): /// /// > This computational cost can be avoided by caching the ClientKey (assuming the Salt and hash @@ -36,36 +39,10 @@ pub struct ClientKeyCache { inner: Arc>>, } -#[derive(Debug, PartialEq, Eq)] -struct CacheKey { - host: String, - port: u16, - socket: Option, - database: Option, - username: String, - password: String, - salt: Vec, - iterations: u32, -} - -impl From<(&PgConnectOptions, &AuthenticationSaslContinue)> for CacheKey { - fn from((options, cont): (&PgConnectOptions, &AuthenticationSaslContinue)) -> Self { - CacheKey { - host: options.host.clone(), - port: options.port, - socket: options.socket.clone(), - database: options.database.clone(), - username: options.username.clone(), - password: options.password.clone().unwrap_or_default(), - salt: cont.salt.clone(), - iterations: cont.iterations, - } - } -} - #[derive(Debug)] struct CacheInner { - cache_key: CacheKey, + salt: Vec, + iterations: u32, salted_password: [u8; 32], client_key: Hmac, } @@ -79,17 +56,14 @@ impl ClientKeyCache { fn get( &self, - options: &PgConnectOptions, cont: &AuthenticationSaslContinue, ) -> Option<([u8; 32], Hmac)> { - let key = CacheKey::from((options, cont)); - self.inner .lock() .expect("BUG: panicked while holding a lock") .as_ref() .and_then(|inner| { - if inner.cache_key == key { + if inner.salt == cont.salt && inner.iterations == cont.iterations { Some((inner.salted_password, inner.client_key.clone())) } else { None @@ -99,7 +73,6 @@ impl ClientKeyCache { fn set( &self, - options: &PgConnectOptions, cont: &AuthenticationSaslContinue, salted_password: [u8; 32], client_key: Hmac, @@ -109,7 +82,8 @@ impl ClientKeyCache { .lock() .expect("BUG: panicked while holding a lock"); *inner = Some(CacheInner { - cache_key: CacheKey::from((options, cont)), + salt: cont.salt.clone(), + iterations: cont.iterations, salted_password, client_key, }); @@ -187,7 +161,7 @@ pub(crate) async fn authenticate( }; let (salted_password, mut mac) = { - if let Some(cached) = options.sasl_client_key_cache.get(options, &cont) { + if let Some(cached) = options.sasl_client_key_cache.get(&cont) { cached } else { // SaltedPassword := Hi(Normalize(password), salt, i) @@ -203,7 +177,7 @@ pub(crate) async fn authenticate( options .sasl_client_key_cache - .set(options, &cont, salted_password, mac.clone()); + .set(&cont, salted_password, mac.clone()); (salted_password, mac) } diff --git a/sqlx-postgres/src/options/mod.rs b/sqlx-postgres/src/options/mod.rs index b7ff288f0b..0efd2e032c 100644 --- a/sqlx-postgres/src/options/mod.rs +++ b/sqlx-postgres/src/options/mod.rs @@ -193,6 +193,8 @@ impl PgConnectOptions { /// ``` pub fn password(mut self, password: &str) -> Self { self.password = Some(password.to_owned()); + // Invalidate the cached SASL client key, since it was derived from the old password. + self.sasl_client_key_cache = ClientKeyCache::new(); self } From 7b46dd692533882579a3a80a3df84766feb4fdc8 Mon Sep 17 00:00:00 2001 From: Thom Wright Date: Thu, 26 Mar 2026 15:46:51 +0000 Subject: [PATCH 3/3] Use async RwLock for SASL client key cache Prevents redundant concurrent hi() computations during pool startup. --- Cargo.lock | 1 + sqlx-postgres/Cargo.toml | 1 + sqlx-postgres/src/connection/sasl.rs | 102 +++++++++++++-------------- 3 files changed, 52 insertions(+), 52 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 581c3eab76..ccbfa4787d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3938,6 +3938,7 @@ dependencies = [ name = "sqlx-postgres" version = "0.9.0-alpha.1" dependencies = [ + "async-lock", "atoi", "base64 0.22.1", "bigdecimal", diff --git a/sqlx-postgres/Cargo.toml b/sqlx-postgres/Cargo.toml index 2943049f0b..85b6bdc70e 100644 --- a/sqlx-postgres/Cargo.toml +++ b/sqlx-postgres/Cargo.toml @@ -52,6 +52,7 @@ time = { workspace = true, optional = true } uuid = { workspace = true, optional = true } # Misc +async-lock = "3.4" atoi = "2.0" base64 = { version = "0.22.0", default-features = false, features = ["std"] } bitflags = { version = "2", default-features = false } diff --git a/sqlx-postgres/src/connection/sasl.rs b/sqlx-postgres/src/connection/sasl.rs index 1c382be4f6..d7e3e0d921 100644 --- a/sqlx-postgres/src/connection/sasl.rs +++ b/sqlx-postgres/src/connection/sasl.rs @@ -1,4 +1,6 @@ -use std::sync::{Arc, Mutex}; +use std::sync::Arc; + +use async_lock::{RwLock, RwLockUpgradableReadGuard}; use crate::connection::stream::PgStream; use crate::error::Error; @@ -30,63 +32,77 @@ const NONCE_ATTR: &str = "r"; /// that affect the HMAC result. The password is not included in the cache key because it can only /// change via `&mut self` on `PgConnectOptions`, which replaces the cache instance. /// +/// An async `RwLock` is used so that only one caller computes the key at a time; subsequent callers +/// wait and then read the cached result. +/// /// According to [RFC-7677](https://datatracker.ietf.org/doc/html/rfc7677): /// /// > This computational cost can be avoided by caching the ClientKey (assuming the Salt and hash /// > iteration-count is stable). #[derive(Debug, Clone)] pub struct ClientKeyCache { - inner: Arc>>, + inner: Arc>>, } #[derive(Debug)] -struct CacheInner { +struct CacheEntry { + // Keys salt: Vec, iterations: u32, + + // Values salted_password: [u8; 32], client_key: Hmac, } +impl CacheEntry { + fn matches(&self, cont: &AuthenticationSaslContinue) -> bool { + self.salt == cont.salt && self.iterations == cont.iterations + } +} + impl ClientKeyCache { pub fn new() -> Self { ClientKeyCache { - inner: Arc::new(Mutex::new(None)), + inner: Arc::new(RwLock::new(None)), } } - fn get( + /// Returns the cached salted password and client key HMAC if the cache matches the given + /// salt and iteration count. Otherwise, computes and caches them. + async fn get_or_compute( &self, + password: &str, cont: &AuthenticationSaslContinue, - ) -> Option<([u8; 32], Hmac)> { - self.inner - .lock() - .expect("BUG: panicked while holding a lock") - .as_ref() - .and_then(|inner| { - if inner.salt == cont.salt && inner.iterations == cont.iterations { - Some((inner.salted_password, inner.client_key.clone())) - } else { - None - } - }) - } + ) -> Result<([u8; 32], Hmac), Error> { + let guard = self.inner.upgradable_read().await; - fn set( - &self, - cont: &AuthenticationSaslContinue, - salted_password: [u8; 32], - client_key: Hmac, - ) { - let mut inner = self - .inner - .lock() - .expect("BUG: panicked while holding a lock"); - *inner = Some(CacheInner { + if let Some(entry) = guard.as_ref().filter(|e| e.matches(cont)) { + return Ok((entry.salted_password, entry.client_key.clone())); + } + + let mut guard = RwLockUpgradableReadGuard::upgrade(guard).await; + + // Re-check after acquiring the write lock, in case another caller populated the cache. + if let Some(entry) = guard.as_ref().filter(|e| e.matches(cont)) { + return Ok((entry.salted_password, entry.client_key.clone())); + } + + // SaltedPassword := Hi(Normalize(password), salt, i) + let salted_password = hi(password, &cont.salt, cont.iterations).await?; + + // ClientKey := HMAC(SaltedPassword, "Client Key") + let client_key = + Hmac::::new_from_slice(&salted_password).map_err(Error::protocol)?; + + *guard = Some(CacheEntry { salt: cont.salt.clone(), iterations: cont.iterations, salted_password, - client_key, + client_key: client_key.clone(), }); + + Ok((salted_password, client_key)) } } @@ -160,28 +176,10 @@ pub(crate) async fn authenticate( } }; - let (salted_password, mut mac) = { - if let Some(cached) = options.sasl_client_key_cache.get(&cont) { - cached - } else { - // SaltedPassword := Hi(Normalize(password), salt, i) - let salted_password = hi( - options.password.as_deref().unwrap_or_default(), - &cont.salt, - cont.iterations, - ) - .await?; - - // ClientKey := HMAC(SaltedPassword, "Client Key") - let mac = Hmac::::new_from_slice(&salted_password).map_err(Error::protocol)?; - - options - .sasl_client_key_cache - .set(&cont, salted_password, mac.clone()); - - (salted_password, mac) - } - }; + let (salted_password, mut mac) = options + .sasl_client_key_cache + .get_or_compute(options.password.as_deref().unwrap_or_default(), &cont) + .await?; mac.update(b"Client Key");