make password hashing async aware (#29)

* make password hashing async aware

Author: conradludgate

.github/workflows/ci.yml

@@ -57,7 +57,7 @@ jobs:
       - run: docker compose up -d
       - uses: sfackler/actions/rustup@master
         with:
-          version: 1.67.0
+          version: 1.75.0
       - run: echo "::set-output name=version::$(rustc --version)"
         id: rust-version
       - uses: actions/cache@v1

postgres-protocol/Cargo.toml

@@ -20,3 +20,7 @@ memchr = "2.0"
 rand = "0.8"
 sha2 = "0.10"
 stringprep = "0.1"
+tokio = { version = "1.0", features = ["rt"] }
+
+[dev-dependencies]
+tokio = { version = "1.0", features = ["full"] }

postgres-protocol/src/authentication/sasl.rs

@@ -9,6 +9,7 @@ use std::io;
 use std::iter;
 use std::mem;
 use std::str;
+use tokio::task::yield_now;
 
 const NONCE_LENGTH: usize = 24;
 
@@ -32,7 +33,7 @@ fn normalize(pass: &[u8]) -> Vec<u8> {
     }
 }
 
-pub(crate) fn hi(str: &[u8], salt: &[u8], i: u32) -> [u8; 32] {
+pub(crate) async fn hi(str: &[u8], salt: &[u8], iterations: u32) -> [u8; 32] {
     let mut hmac =
         Hmac::<Sha256>::new_from_slice(str).expect("HMAC is able to accept all key sizes");
     hmac.update(salt);
@@ -41,14 +42,19 @@ pub(crate) fn hi(str: &[u8], salt: &[u8], i: u32) -> [u8; 32] {
 
     let mut hi = prev;
 
-    for _ in 1..i {
+    for i in 1..iterations {
         let mut hmac = Hmac::<Sha256>::new_from_slice(str).expect("already checked above");
         hmac.update(&prev);
         prev = hmac.finalize().into_bytes();
 
         for (hi, prev) in hi.iter_mut().zip(prev) {
             *hi ^= prev;
         }
+        // yield every ~250us
+        // hopefully reduces tail latencies
+        if i % 1024 == 0 {
+            yield_now().await
+        }
     }
 
     hi.into()
@@ -200,7 +206,7 @@ impl ScramSha256 {
     /// Updates the state machine with the response from the backend.
     ///
     /// This should be called when an `AuthenticationSASLContinue` message is received.
-    pub fn update(&mut self, message: &[u8]) -> io::Result<()> {
+    pub async fn update(&mut self, message: &[u8]) -> io::Result<()> {
         let (client_nonce, password, channel_binding) =
             match mem::replace(&mut self.state, State::Done) {
                 State::Update {
@@ -227,7 +233,7 @@ impl ScramSha256 {
                     Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
                 };
 
-                let salted_password = hi(&password, &salt, parsed.iteration_count);
+                let salted_password = hi(&password, &salt, parsed.iteration_count).await;
 
                 let make_key = |name| {
                     let mut hmac = Hmac::<Sha256>::new_from_slice(&salted_password)
@@ -481,8 +487,8 @@ mod test {
     }
 
     // recorded auth exchange from psql
-    #[test]
-    fn exchange() {
+    #[tokio::test]
+    async fn exchange() {
         let password = "foobar";
         let nonce = "9IZ2O01zb9IgiIZ1WJ/zgpJB";
 
@@ -502,7 +508,7 @@ mod test {
         );
         assert_eq!(str::from_utf8(scram.message()).unwrap(), client_first);
 
-        scram.update(server_first.as_bytes()).unwrap();
+        scram.update(server_first.as_bytes()).await.unwrap();
         assert_eq!(str::from_utf8(scram.message()).unwrap(), client_final);
 
         scram.finish(server_final.as_bytes()).unwrap();

postgres-protocol/src/password/mod.rs

@@ -24,16 +24,19 @@ const SCRAM_DEFAULT_SALT_LEN: usize = 16;
 ///
 /// The client may assume the returned string doesn't contain any
 /// special characters that would require escaping in an SQL command.
-pub fn scram_sha_256(password: &[u8]) -> String {
+pub async fn scram_sha_256(password: &[u8]) -> String {
     let mut salt: [u8; SCRAM_DEFAULT_SALT_LEN] = [0; SCRAM_DEFAULT_SALT_LEN];
     let mut rng = rand::thread_rng();
     rng.fill_bytes(&mut salt);
-    scram_sha_256_salt(password, salt)
+    scram_sha_256_salt(password, salt).await
 }
 
 // Internal implementation of scram_sha_256 with a caller-provided
 // salt. This is useful for testing.
-pub(crate) fn scram_sha_256_salt(password: &[u8], salt: [u8; SCRAM_DEFAULT_SALT_LEN]) -> String {
+pub(crate) async fn scram_sha_256_salt(
+    password: &[u8],
+    salt: [u8; SCRAM_DEFAULT_SALT_LEN],
+) -> String {
     // Prepare the password, per [RFC
     // 4013](https://tools.ietf.org/html/rfc4013), if possible.
     //
@@ -58,7 +61,7 @@ pub(crate) fn scram_sha_256_salt(password: &[u8], salt: [u8; SCRAM_DEFAULT_SALT_
     };
 
     // salt password
-    let salted_password = sasl::hi(&prepared, &salt, SCRAM_DEFAULT_ITERATIONS);
+    let salted_password = sasl::hi(&prepared, &salt, SCRAM_DEFAULT_ITERATIONS).await;
 
     // client key
     let mut hmac = Hmac::<Sha256>::new_from_slice(&salted_password)

postgres-protocol/src/password/test.rs

@@ -1,11 +1,11 @@
 use crate::password;
 
-#[test]
-fn test_encrypt_scram_sha_256() {
+#[tokio::test]
+async fn test_encrypt_scram_sha_256() {
     // Specify the salt to make the test deterministic. Any bytes will do.
     let salt: [u8; 16] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
     assert_eq!(
-        password::scram_sha_256_salt(b"secret", salt),
+        password::scram_sha_256_salt(b"secret", salt).await,
         "SCRAM-SHA-256$4096:AQIDBAUGBwgJCgsMDQ4PEA==$8rrDg00OqaiWXJ7p+sCgHEIaBSHY89ZJl3mfIsf32oY=:05L1f+yZbiN8O0AnO40Og85NNRhvzTS57naKRWCcsIA="
     );
 }

postgres-types/src/chrono_04.rs

@@ -111,7 +111,7 @@ impl<'a> FromSql<'a> for NaiveDate {
         let jd = types::date_from_sql(raw)?;
         base()
             .date()
-            .checked_add_signed(Duration::days(i64::from(jd)))
+            .checked_add_signed(Duration::try_days(i64::from(jd)).unwrap())
             .ok_or_else(|| "value too large to decode".into())
     }
 

tokio-postgres/src/connect_raw.rs

@@ -300,6 +300,7 @@ where
 
     scram
         .update(body.data())
+        .await
         .map_err(|e| Error::authentication(e.into()))?;
 
     let mut buf = BytesMut::new();

tokio-postgres/tests/test/main.rs

@@ -486,17 +486,17 @@ async fn simple_query() {
     }
     match &messages[2] {
         SimpleQueryMessage::Row(row) => {
-            assert_eq!(row.columns().get(0).map(|c| c.name()), Some("id"));
-            assert_eq!(row.columns().get(1).map(|c| c.name()), Some("name"));
+            assert_eq!(row.columns()[0].name(), "id");
+            assert_eq!(row.columns()[1].name(), "name");
             assert_eq!(row.get(0), Some("1"));
             assert_eq!(row.get(1), Some("steven"));
         }
         _ => panic!("unexpected message"),
     }
     match &messages[3] {
         SimpleQueryMessage::Row(row) => {
-            assert_eq!(row.columns().get(0).map(|c| c.name()), Some("id"));
-            assert_eq!(row.columns().get(1).map(|c| c.name()), Some("name"));
+            assert_eq!(row.columns()[0].name(), "id");
+            assert_eq!(row.columns()[1].name(), "name");
             assert_eq!(row.get(0), Some("2"));
             assert_eq!(row.get(1), Some("joe"));
         }