From ef2e809564fa37213a40710015a7ce39231e9c7f Mon Sep 17 00:00:00 2001 From: Thom Wright Date: Wed, 3 Sep 2025 19:12:45 +0100 Subject: [PATCH 1/2] Yield while salting password for PG SASL To prevent spenting too much time doing synchronous work on the event loop, we yield every few iterations of the hmac-sha256. --- sqlx-postgres/src/connection/sasl.rs | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/sqlx-postgres/src/connection/sasl.rs b/sqlx-postgres/src/connection/sasl.rs index 729cc1fcc5..3e58fbee03 100644 --- a/sqlx-postgres/src/connection/sasl.rs +++ b/sqlx-postgres/src/connection/sasl.rs @@ -1,6 +1,7 @@ use crate::connection::stream::PgStream; use crate::error::Error; use crate::message::{Authentication, AuthenticationSasl, SaslInitialResponse, SaslResponse}; +use crate::rt; use crate::PgConnectOptions; use hmac::{Hmac, Mac}; use rand::Rng; @@ -90,7 +91,8 @@ pub(crate) async fn authenticate( options.password.as_deref().unwrap_or_default(), &cont.salt, cont.iterations, - )?; + ) + .await?; // ClientKey := HMAC(SaltedPassword, "Client Key") let mut mac = Hmac::::new_from_slice(&salted_password).map_err(Error::protocol)?; @@ -187,7 +189,7 @@ fn gen_nonce() -> String { } // Hi(str, salt, i): -fn hi<'a>(s: &'a str, salt: &'a [u8], iter_count: u32) -> Result<[u8; 32], Error> { +async fn hi<'a>(s: &'a str, salt: &'a [u8], iter_count: u32) -> Result<[u8; 32], Error> { let mut mac = Hmac::::new_from_slice(s.as_bytes()).map_err(Error::protocol)?; mac.update(salt); @@ -196,10 +198,18 @@ fn hi<'a>(s: &'a str, salt: &'a [u8], iter_count: u32) -> Result<[u8; 32], Error let mut u = mac.finalize_reset().into_bytes(); let mut hi = u; - for _ in 1..iter_count { + for i in 1..iter_count { mac.update(u.as_slice()); u = mac.finalize_reset().into_bytes(); hi = hi.iter().zip(u.iter()).map(|(&a, &b)| a ^ b).collect(); + + // For large iteration counts, this process can take a long time and block the event loop. + // It was measured as taking ~50ms for 4096 iterations (the default) on a developer machine. + // If we want to yield every 10-100us (as generally advised for tokio), then we can yield + // every 5 iterations which should be every ~50us. + if i % 5 == 0 { + rt::yield_now().await; + } } Ok(hi.into()) @@ -207,19 +217,21 @@ fn hi<'a>(s: &'a str, salt: &'a [u8], iter_count: u32) -> Result<[u8; 32], Error #[cfg(all(test, not(debug_assertions)))] #[bench] -fn bench_sasl_hi(b: &mut test::Bencher) { +async fn bench_sasl_hi(b: &mut test::Bencher) { use test::black_box; + let runtime = tokio::runtime::Runtime::new().unwrap(); + let mut rng = rand::thread_rng(); let nonce: Vec = std::iter::repeat(()) .map(|()| rng.sample(rand::distributions::Alphanumeric)) .take(64) .collect(); - b.iter(|| { - let _ = hi( + b.to_async(&runtime).iter(|| { + hi( test::black_box("secret_password"), test::black_box(&nonce), test::black_box(4096), - ); + ) }); } From 2d68935e9b5a7cdc33e246bff785adad4f1d3990 Mon Sep 17 00:00:00 2001 From: Thom Wright Date: Thu, 4 Sep 2025 10:27:35 +0100 Subject: [PATCH 2/2] Remove unused bench --- sqlx-postgres/src/connection/sasl.rs | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/sqlx-postgres/src/connection/sasl.rs b/sqlx-postgres/src/connection/sasl.rs index 3e58fbee03..94fdfc689f 100644 --- a/sqlx-postgres/src/connection/sasl.rs +++ b/sqlx-postgres/src/connection/sasl.rs @@ -214,24 +214,3 @@ async fn hi<'a>(s: &'a str, salt: &'a [u8], iter_count: u32) -> Result<[u8; 32], Ok(hi.into()) } - -#[cfg(all(test, not(debug_assertions)))] -#[bench] -async fn bench_sasl_hi(b: &mut test::Bencher) { - use test::black_box; - - let runtime = tokio::runtime::Runtime::new().unwrap(); - - let mut rng = rand::thread_rng(); - let nonce: Vec = std::iter::repeat(()) - .map(|()| rng.sample(rand::distributions::Alphanumeric)) - .take(64) - .collect(); - b.to_async(&runtime).iter(|| { - hi( - test::black_box("secret_password"), - test::black_box(&nonce), - test::black_box(4096), - ) - }); -}