SCRAM for server connections

pgdog/Cargo.toml

@@ -30,3 +30,4 @@ toml = "0.8"
 pgdog-plugin = { path = "../pgdog-plugin", version = "0.1.0" }
 tokio-util = { version = "0.7", features = ["rt"] }
 fnv = "1"
+scram = "0.6"

pgdog/src/auth/scram/client.rs

@@ -0,0 +1,68 @@
+//! SCRAM-SHA-256 client.
+
+use super::Error;
+
+use scram::{
+    client::{ClientFinal, ServerFinal, ServerFirst},
+    ScramClient,
+};
+
+enum State<'a> {
+    Initial(ScramClient<'a>),
+    First(ServerFirst<'a>),
+    Final(ClientFinal),
+    ServerFinal(ServerFinal),
+}
+
+/// SASL SCRAM client.
+pub struct Client<'a> {
+    state: Option<State<'a>>,
+}
+
+impl<'a> Client<'a> {
+    /// Create new SCRAM client.
+    pub fn new(user: &'a str, password: &'a str) -> Self {
+        Self {
+            state: Some(State::Initial(ScramClient::new(user, password, None))),
+        }
+    }
+
+    /// Client first message.
+    pub fn first(&mut self) -> Result<String, Error> {
+        let (scram, client_first) = match self.state.take() {
+            Some(State::Initial(scram)) => scram.client_first(),
+            _ => return Err(Error::OutOfOrder),
+        };
+        self.state = Some(State::First(scram));
+        Ok(client_first)
+    }
+
+    /// Handle server first message.
+    pub fn server_first(&mut self, message: &str) -> Result<(), Error> {
+        let scram = match self.state.take() {
+            Some(State::First(scram)) => scram.handle_server_first(message)?,
+            _ => return Err(Error::OutOfOrder),
+        };
+        self.state = Some(State::Final(scram));
+        Ok(())
+    }
+
+    /// Client last message.
+    pub fn last(&mut self) -> Result<String, Error> {
+        let (scram, client_final) = match self.state.take() {
+            Some(State::Final(scram)) => scram.client_final(),
+            _ => return Err(Error::OutOfOrder),
+        };
+        self.state = Some(State::ServerFinal(scram));
+        Ok(client_final)
+    }
+
+    /// Verify server last message.
+    pub fn server_last(&mut self, message: &str) -> Result<(), Error> {
+        match self.state.take() {
+            Some(State::ServerFinal(scram)) => scram.handle_server_final(message)?,
+            _ => return Err(Error::OutOfOrder),
+        };
+        Ok(())
+    }
+}

pgdog/src/auth/scram/error.rs

@@ -0,0 +1,11 @@
+//! SCRAM errors.
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("out of order auth")]
+    OutOfOrder,
+
+    #[error("invalid server first message")]
+    InvalidServerFirst(#[from] scram::Error),
+}

pgdog/src/auth/scram/mod.rs

@@ -1 +1,7 @@
 //! SCRAM-SHA-256 authentication.
+pub mod client;
+pub mod error;
+pub mod state;
+
+pub use client::Client;
+pub use error::Error;

pgdog/src/auth/scram/state.rs

@@ -0,0 +1 @@
+//! SCRAM-SHA-256 state.

pgdog/src/backend/error.rs

@@ -44,6 +44,9 @@ pub enum Error {
 
     #[error("no cluster connected")]
     NoCluster,
+
+    #[error("scram auth failed")]
+    ScramAuth(#[from] crate::auth::scram::Error),
 }
 
 impl Error {

pgdog/src/backend/server.rs

@@ -11,17 +11,13 @@ use tokio::{
 use tracing::{debug, info, trace};
 
 use super::{pool::Address, Error};
-use crate::net::{
-    messages::{hello::SslReply, FromBytes, Protocol, Startup, ToBytes},
-    parameter::Parameters,
-    tls::connector,
-    Parameter, Stream,
-};
+use crate::net::{parameter::Parameters, tls::connector, Parameter, Stream};
 use crate::state::State;
 use crate::{
+    auth::scram::Client,
     net::messages::{
-        Authentication, BackendKeyData, ErrorResponse, Message, ParameterStatus, Query,
-        ReadyForQuery, Terminate,
+        hello::SslReply, Authentication, BackendKeyData, ErrorResponse, FromBytes, Message,
+        ParameterStatus, Password, Protocol, Query, ReadyForQuery, Startup, Terminate, ToBytes,
     },
     stats::ConnStats,
 };
@@ -71,6 +67,7 @@ impl Server {
         stream.flush().await?;
 
         // Perform authentication.
+        let mut scram = Client::new(&addr.user, &addr.password);
         loop {
             let message = stream.read().await?;
 
@@ -84,6 +81,20 @@ impl Server {
 
                     match auth {
                         Authentication::Ok => break,
+                        Authentication::AuthenticationSASL(_) => {
+                            let initial = Password::sasl_initial(&scram.first()?);
+                            stream.send_flush(initial).await?;
+                        }
+                        Authentication::AuthenticationSASLContinue(data) => {
+                            scram.server_first(&data)?;
+                            let response = Password::SASLResponse {
+                                response: scram.last()?,
+                            };
+                            stream.send_flush(response).await?;
+                        }
+                        Authentication::AuthenticationSASLFinal(data) => {
+                            scram.server_last(&data)?;
+                        }
                     }
                 }
 

pgdog/src/config/mod.rs

@@ -28,6 +28,7 @@ pub fn load(config: &PathBuf, users: &PathBuf) -> Result<ConfigAndUsers, Error>
     Ok(config)
 }
 
+/// pgdog.toml and users.toml.
 #[derive(Debug, Clone, Default)]
 pub struct ConfigAndUsers {
     /// pgdog.toml

pgdog/src/net/messages/auth/mod.rs

@@ -1,14 +1,25 @@
 //! Authentication messages.
 
+use crate::net::c_string_buf;
+
 use super::{code, prelude::*};
 
 use super::FromBytes;
 
+pub mod password;
+pub use password::Password;
+
 /// Authentication messages.
 #[derive(Debug)]
 pub enum Authentication {
     /// AuthenticationOk (F)
     Ok,
+    /// AuthenticationSASL (B)
+    AuthenticationSASL(String),
+    /// AuthenticationSASLContinue (B)
+    AuthenticationSASLContinue(String),
+    /// AuthenticationSASLFinal (B)
+    AuthenticationSASLFinal(String),
 }
 
 impl FromBytes for Authentication {
@@ -21,6 +32,18 @@ impl FromBytes for Authentication {
 
         match status {
             0 => Ok(Authentication::Ok),
+            10 => {
+                let mechanism = c_string_buf(&mut bytes);
+                Ok(Authentication::AuthenticationSASL(mechanism))
+            }
+            11 => {
+                let data = c_string_buf(&mut bytes);
+                Ok(Authentication::AuthenticationSASLContinue(data))
+            }
+            12 => {
+                let data = c_string_buf(&mut bytes);
+                Ok(Authentication::AuthenticationSASLFinal(data))
+            }
             status => Err(Error::UnsupportedAuthentication(status)),
         }
     }
@@ -42,6 +65,27 @@ impl ToBytes for Authentication {
 
                 Ok(payload.freeze())
             }
+
+            Authentication::AuthenticationSASL(mechanism) => {
+                payload.put_i32(10);
+                payload.put_string(&mechanism);
+
+                Ok(payload.freeze())
+            }
+
+            Authentication::AuthenticationSASLContinue(data) => {
+                payload.put_i32(11);
+                payload.put_string(&data);
+
+                Ok(payload.freeze())
+            }
+
+            Authentication::AuthenticationSASLFinal(data) => {
+                payload.put_i32(12);
+                payload.put_string(&data);
+
+                Ok(payload.freeze())
+            }
         }
     }
 }

pgdog/src/net/messages/auth/password.rs

@@ -0,0 +1,74 @@
+//! Password messages.
+
+use crate::net::c_string_buf;
+
+use super::super::code;
+use super::super::prelude::*;
+
+/// Password message.
+#[derive(Debug)]
+pub enum Password {
+    /// SASLInitialResponse (F)
+    SASLInitialResponse { name: String, response: String },
+    /// SASLResponse (F)
+    SASLResponse { response: String },
+}
+
+impl Password {
+    /// Create new SASL initial response.
+    pub fn sasl_initial(response: &str) -> Self {
+        Self::SASLInitialResponse {
+            name: "SCRAM-SHA-256".to_string(),
+            response: response.to_owned(),
+        }
+    }
+}
+
+impl FromBytes for Password {
+    fn from_bytes(mut bytes: Bytes) -> Result<Self, Error> {
+        code!(bytes, 'p');
+        let _len = bytes.get_i32();
+        let content = c_string_buf(&mut bytes);
+
+        if bytes.has_remaining() {
+            let len = bytes.get_i32();
+            let response = if len >= 0 {
+                c_string_buf(&mut bytes)
+            } else {
+                String::new()
+            };
+
+            Ok(Self::SASLInitialResponse {
+                name: content,
+                response,
+            })
+        } else {
+            Ok(Password::SASLResponse { response: content })
+        }
+    }
+}
+
+impl ToBytes for Password {
+    fn to_bytes(&self) -> Result<Bytes, Error> {
+        let mut payload = Payload::named(self.code());
+        match self {
+            Password::SASLInitialResponse { name, response } => {
+                payload.put_string(name);
+                payload.put_i32(response.len() as i32);
+                payload.put(Bytes::copy_from_slice(response.as_bytes()));
+            }
+
+            Password::SASLResponse { response } => {
+                payload.put(Bytes::copy_from_slice(response.as_bytes()));
+            }
+        }
+
+        Ok(payload.freeze())
+    }
+}
+
+impl Protocol for Password {
+    fn code(&self) -> char {
+        'p'
+    }
+}

pgdog/src/net/messages/mod.rs

@@ -15,7 +15,7 @@ pub mod rfq;
 pub mod row_description;
 pub mod terminate;
 
-pub use auth::Authentication;
+pub use auth::{Authentication, Password};
 pub use backend_key::BackendKeyData;
 pub use bind::Bind;
 pub use data_row::{DataRow, ToDataRowColumn};