Allow passing precomputed SCRAM keys via Config

funbringer · funbringer · commit a27a40657dc1 · 2022-12-16T20:25:53.000+03:00
According to https://datatracker.ietf.org/doc/html/rfc5802#section-3, SCRAM protocol explicitly allows client to use a `ClientKey` & `ServerKey` pair instead of a password to perform authentication. This is also useful for proxy implementations which would like to leverage `rust-postgres`. This patch adds the ability to do that.
diff --git a/postgres-protocol/src/authentication/sasl.rs b/postgres-protocol/src/authentication/sasl.rs
@@ -96,14 +96,32 @@ impl ChannelBinding {
     }
 }
 
+/// A pair of keys for the SCRAM-SHA-256 mechanism.
+/// See <https://datatracker.ietf.org/doc/html/rfc5802#section-3> for details.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct ScramKeys<const N: usize> {
+    /// Used by server to authenticate client.
+    pub client_key: [u8; N],
+    /// Used by client to verify server's signature.
+    pub server_key: [u8; N],
+}
+
+/// Password or keys which were derived from it.
+enum Credentials<const N: usize> {
+    /// A regular password as a vector of bytes.
+    Password(Vec<u8>),
+    /// A precomputed pair of keys.
+    Keys(Box<ScramKeys<N>>),
+}
+
 enum State {
     Update {
         nonce: String,
-        password: Vec<u8>,
+        password: Credentials<32>,
         channel_binding: ChannelBinding,
     },
     Finish {
-        salted_password: [u8; 32],
+        server_key: [u8; 32],
         auth_message: String,
     },
     Done,
@@ -129,30 +147,43 @@ pub struct ScramSha256 {
     state: State,
 }
 
+fn nonce() -> String {
+    // rand 0.5's ThreadRng is cryptographically secure
+    let mut rng = rand::thread_rng();
+    (0..NONCE_LENGTH)
+        .map(|_| {
+            let mut v = rng.gen_range(0x21u8..0x7e);
+            if v == 0x2c {
+                v = 0x7e
+            }
+            v as char
+        })
+        .collect()
+}
+
 impl ScramSha256 {
     /// Constructs a new instance which will use the provided password for authentication.
     pub fn new(password: &[u8], channel_binding: ChannelBinding) -> ScramSha256 {
-        // rand 0.5's ThreadRng is cryptographically secure
-        let mut rng = rand::thread_rng();
-        let nonce = (0..NONCE_LENGTH)
-            .map(|_| {
-                let mut v = rng.gen_range(0x21u8..0x7e);
-                if v == 0x2c {
-                    v = 0x7e
-                }
-                v as char
-            })
-            .collect::<String>();
+        let password = Credentials::Password(normalize(password));
+        ScramSha256::new_inner(password, channel_binding, nonce())
+    }
 
-        ScramSha256::new_inner(password, channel_binding, nonce)
+    /// Constructs a new instance which will use the provided key pair for authentication.
+    pub fn new_with_keys(keys: ScramKeys<32>, channel_binding: ChannelBinding) -> ScramSha256 {
+        let password = Credentials::Keys(keys.into());
+        ScramSha256::new_inner(password, channel_binding, nonce())
     }
 
-    fn new_inner(password: &[u8], channel_binding: ChannelBinding, nonce: String) -> ScramSha256 {
+    fn new_inner(
+        password: Credentials<32>,
+        channel_binding: ChannelBinding,
+        nonce: String,
+    ) -> ScramSha256 {
         ScramSha256 {
             message: format!("{}n=,r={}", channel_binding.gs2_header(), nonce),
             state: State::Update {
                 nonce,
-                password: normalize(password),
+                password,
                 channel_binding,
             },
         }
@@ -189,20 +220,32 @@ impl ScramSha256 {
             return Err(io::Error::new(io::ErrorKind::InvalidInput, "invalid nonce"));
         }
 
-        let salt = match base64::decode(parsed.salt) {
-            Ok(salt) => salt,
-            Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
-        };
+        let (client_key, server_key) = match password {
+            Credentials::Password(password) => {
+                let salt = match base64::decode(parsed.salt) {
+                    Ok(salt) => salt,
+                    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);
 
-        let mut hmac = Hmac::<Sha256>::new_from_slice(&salted_password)
-            .expect("HMAC is able to accept all key sizes");
-        hmac.update(b"Client Key");
-        let client_key = hmac.finalize().into_bytes();
+                let make_key = |name| {
+                    let mut hmac = Hmac::<Sha256>::new_from_slice(&salted_password)
+                        .expect("HMAC is able to accept all key sizes");
+                    hmac.update(name);
+
+                    let mut key = [0u8; 32];
+                    key.copy_from_slice(hmac.finalize().into_bytes().as_slice());
+                    key
+                };
+
+                (make_key(b"Client Key"), make_key(b"Server Key"))
+            }
+            Credentials::Keys(keys) => (keys.client_key, keys.server_key),
+        };
 
         let mut hash = Sha256::default();
-        hash.update(client_key.as_slice());
+        hash.update(client_key);
         let stored_key = hash.finalize_fixed();
 
         let mut cbind_input = vec![];
@@ -225,10 +268,10 @@ impl ScramSha256 {
             *proof ^= signature;
         }
 
-        write!(&mut self.message, ",p={}", base64::encode(&*client_proof)).unwrap();
+        write!(&mut self.message, ",p={}", base64::encode(client_proof)).unwrap();
 
         self.state = State::Finish {
-            salted_password,
+            server_key,
             auth_message,
         };
         Ok(())
@@ -239,11 +282,11 @@ impl ScramSha256 {
     /// This should be called when the backend sends an `AuthenticationSASLFinal` message.
     /// Authentication has only succeeded if this method returns `Ok(())`.
     pub fn finish(&mut self, message: &[u8]) -> io::Result<()> {
-        let (salted_password, auth_message) = match mem::replace(&mut self.state, State::Done) {
+        let (server_key, auth_message) = match mem::replace(&mut self.state, State::Done) {
             State::Finish {
-                salted_password,
+                server_key,
                 auth_message,
-            } => (salted_password, auth_message),
+            } => (server_key, auth_message),
             _ => return Err(io::Error::new(io::ErrorKind::Other, "invalid SCRAM state")),
         };
 
@@ -267,11 +310,6 @@ impl ScramSha256 {
             Err(e) => return Err(io::Error::new(io::ErrorKind::InvalidInput, e)),
         };
 
-        let mut hmac = Hmac::<Sha256>::new_from_slice(&salted_password)
-            .expect("HMAC is able to accept all key sizes");
-        hmac.update(b"Server Key");
-        let server_key = hmac.finalize().into_bytes();
-
         let mut hmac = Hmac::<Sha256>::new_from_slice(&server_key)
             .expect("HMAC is able to accept all key sizes");
         hmac.update(auth_message.as_bytes());
@@ -458,7 +496,7 @@ mod test {
         let server_final = "v=U+ppxD5XUKtradnv8e2MkeupiA8FU87Sg8CXzXHDAzw=";
 
         let mut scram = ScramSha256::new_inner(
-            password.as_bytes(),
+            Credentials::Password(normalize(password.as_bytes())),
             ChannelBinding::unsupported(),
             nonce.to_string(),
         );
diff --git a/postgres/src/config.rs b/postgres/src/config.rs
@@ -12,7 +12,9 @@ use std::sync::Arc;
 use std::time::Duration;
 use tokio::runtime;
 #[doc(inline)]
-pub use tokio_postgres::config::{ChannelBinding, Host, SslMode, TargetSessionAttrs};
+pub use tokio_postgres::config::{
+    AuthKeys, ChannelBinding, Host, ScramKeys, SslMode, TargetSessionAttrs,
+};
 use tokio_postgres::error::DbError;
 use tokio_postgres::tls::{MakeTlsConnect, TlsConnect};
 use tokio_postgres::{Error, Socket};
@@ -149,6 +151,20 @@ impl Config {
         self.config.get_password()
     }
 
+    /// Sets precomputed protocol-specific keys to authenticate with.
+    /// When set, this option will override `password`.
+    /// See [`AuthKeys`] for more information.
+    pub fn auth_keys(&mut self, keys: AuthKeys) -> &mut Config {
+        self.config.auth_keys(keys);
+        self
+    }
+
+    /// Gets precomputed protocol-specific keys to authenticate with.
+    /// if one has been configured with the `auth_keys` method.
+    pub fn get_auth_keys(&self) -> Option<AuthKeys> {
+        self.config.get_auth_keys()
+    }
+
     /// Sets the name of the database to connect to.
     ///
     /// Defaults to the user.
diff --git a/tokio-postgres/src/config.rs b/tokio-postgres/src/config.rs
@@ -23,6 +23,8 @@ use std::time::Duration;
 use std::{error, fmt, iter, mem};
 use tokio::io::{AsyncRead, AsyncWrite};
 
+pub use postgres_protocol::authentication::sasl::ScramKeys;
+
 /// Properties required of a session.
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 #[non_exhaustive]
@@ -79,6 +81,13 @@ pub enum Host {
     Unix(PathBuf),
 }
 
+/// Precomputed keys which may override password during auth.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum AuthKeys {
+    /// A `ClientKey` & `ServerKey` pair for `SCRAM-SHA-256`.
+    ScramSha256(ScramKeys<32>),
+}
+
 /// Connection configuration.
 ///
 /// Configuration can be parsed from libpq-style connection strings. These strings come in two formats:
@@ -163,6 +172,7 @@ pub enum Host {
 pub struct Config {
     pub(crate) user: Option<String>,
     pub(crate) password: Option<Vec<u8>>,
+    pub(crate) auth_keys: Option<Box<AuthKeys>>,
     pub(crate) dbname: Option<String>,
     pub(crate) options: Option<String>,
     pub(crate) application_name: Option<String>,
@@ -194,6 +204,7 @@ impl Config {
         Config {
             user: None,
             password: None,
+            auth_keys: None,
             dbname: None,
             options: None,
             application_name: None,
@@ -238,6 +249,20 @@ impl Config {
         self.password.as_deref()
     }
 
+    /// Sets precomputed protocol-specific keys to authenticate with.
+    /// When set, this option will override `password`.
+    /// See [`AuthKeys`] for more information.
+    pub fn auth_keys(&mut self, keys: AuthKeys) -> &mut Config {
+        self.auth_keys = Some(Box::new(keys));
+        self
+    }
+
+    /// Gets precomputed protocol-specific keys to authenticate with.
+    /// if one has been configured with the `auth_keys` method.
+    pub fn get_auth_keys(&self) -> Option<AuthKeys> {
+        self.auth_keys.as_deref().copied()
+    }
+
     /// Sets the name of the database to connect to.
     ///
     /// Defaults to the user.
diff --git a/tokio-postgres/src/connect_raw.rs b/tokio-postgres/src/connect_raw.rs
@@ -1,5 +1,5 @@
 use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
-use crate::config::{self, Config, ReplicationMode};
+use crate::config::{self, AuthKeys, Config, ReplicationMode};
 use crate::connect_tls::connect_tls;
 use crate::maybe_tls_stream::MaybeTlsStream;
 use crate::tls::{TlsConnect, TlsStream};
@@ -234,11 +234,6 @@ where
     S: AsyncRead + AsyncWrite + Unpin,
     T: TlsStream + Unpin,
 {
-    let password = config
-        .password
-        .as_ref()
-        .ok_or_else(|| Error::config("password missing".into()))?;
-
     let mut has_scram = false;
     let mut has_scram_plus = false;
     let mut mechanisms = body.mechanisms();
@@ -276,7 +271,13 @@ where
         can_skip_channel_binding(config)?;
     }
 
-    let mut scram = ScramSha256::new(password, channel_binding);
+    let mut scram = if let Some(AuthKeys::ScramSha256(keys)) = config.get_auth_keys() {
+        ScramSha256::new_with_keys(keys, channel_binding)
+    } else if let Some(password) = config.get_password() {
+        ScramSha256::new(password, channel_binding)
+    } else {
+        return Err(Error::config("password or auth keys missing".into()));
+    };
 
     let mut buf = BytesMut::new();
     frontend::sasl_initial_response(mechanism, scram.message(), &mut buf).map_err(Error::encode)?;
