Add auth passthough (auth_query)

Adds a feature that allows setting auth passthrough for md5 auth.

It adds 3 new (general and pool) config parameters:

- `auth_query`: An string containing a query that will be executed on boot
to obtain the hash of a given user. This query have to use a placeholder `$1`,
so pgcat can replace it with the user its trying to fetch the hash from.
- `auth_query_user`: The user to use for connecting to the server and executing the
auth_query.
- `auth_query_password`: The password to use for connecting to the server and executing the
auth_query.

The configuration can be done either on the general config (so pools share them) or in a per-pool basis.

The behavior is, at boot time, when validating server connections, a hash is fetched per server
and stored in the pool. When new server connections are created, and no cleartext password is specified,
the obtained hash is used for creating them, if the hash could not be obtained for whatever reason, it retries
it.

When client authentication is tried, it uses cleartext passwords if specified, it not, it checks whether
we have query_auth set up, if so, it tries to use the obtained hash for making client auth. If there is no
hash (we could not obtain one when validating the connection), a new fetch is tried.

Once we have a hash, we authenticate using it against whathever the client has sent us, if there is a failure
we refetch the hash and retry auth (so password changes can be done).

The idea with this 'retrial' mechanism is to make it fault tolerant, so if for whatever reason hash could not be
obtained during connection validation, or the password has change, we can still connect later.

Author: magec

.circleci/config.yml

@@ -46,6 +46,14 @@ jobs:
           POSTGRES_PASSWORD: postgres
           POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
 
+      - image: postgres:14
+        command: ["postgres", "-p", "10432", "-c", "shared_preload_libraries=pg_stat_statements"]
+        environment:
+          POSTGRES_USER: postgres
+          POSTGRES_DB: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_INITDB_ARGS: --auth-local=md5 --auth-host=md5 --auth=md5
+
     # Add steps to the job
     # See: https://circleci.com/docs/2.0/configuration-reference/#steps
     steps:

.circleci/run_tests.sh

@@ -19,6 +19,7 @@ PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/q
 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 7432 -U postgres -f tests/sharding/query_routing_setup.sql
 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 8432 -U postgres -f tests/sharding/query_routing_setup.sql
 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 9432 -U postgres -f tests/sharding/query_routing_setup.sql
+PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 10432 -U postgres -f tests/sharding/query_routing_setup.sql
 
 PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard0 -i
 PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard1 -i

.rustfmt.toml

@@ -0,0 +1,2 @@
+edition = "2021"
+hard_tabs = false

Cargo.lock

@@ -4,7 +4,6 @@ version = "0.6.0-alpha1"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 tokio = { version = "1", features = ["full"] }
 bytes = "1"
@@ -37,6 +36,8 @@ exitcode = "1.1.2"
 futures = "0.3"
 socket2 = { version = "0.4.7", features = ["all"] }
 nix = "0.26.2"
+postgres-protocol = "0.6.4"
+fallible-iterator = "0.2"
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemallocator = "0.5.0"

Cargo.toml

@@ -56,6 +56,13 @@ services:
       POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
       PGPORT: 9432
     command: ["postgres", "-p", "9432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
+  pg5:
+    <<: *common-definition-pg
+    environment:
+      <<: *common-env-pg
+      POSTGRES_INITDB_ARGS: --auth-local=md5 --auth-host=md5 --auth=md5
+      PGPORT: 10432
+    command: ["postgres", "-p", "5432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
 
   toxiproxy:
     build: .
@@ -69,6 +76,7 @@ services:
       - pg2
       - pg3
       - pg4
+      - pg5
 
   pgcat-shell:
     stdin_open: true

dev/docker-compose.yaml

@@ -0,0 +1,107 @@
+use crate::errors::Error;
+use crate::server::Server;
+use log::debug;
+
+#[derive(Clone, Debug)]
+pub struct AuthPassthrough {
+    password: String,
+    query: String,
+    user: String,
+}
+
+impl AuthPassthrough {
+    /// Initializes an AuthPassthrough.
+    pub fn new(query: &str, user: &str, password: &str) -> Self {
+        AuthPassthrough {
+            password: password.to_string(),
+            query: query.to_string(),
+            user: user.to_string(),
+        }
+    }
+
+    /// Returns an AuthPassthrough given the pool configuration.
+    /// If any of required values is not set, None is returned.
+    pub fn from_pool_config(pool_config: &crate::config::Pool) -> Option<Self> {
+        if pool_config.is_auth_query_configured() {
+            return Some(AuthPassthrough::new(
+                pool_config.auth_query.as_ref().unwrap(),
+                pool_config.auth_query_user.as_ref().unwrap(),
+                pool_config.auth_query_password.as_ref().unwrap(),
+            ));
+        }
+
+        None
+    }
+
+    /// Returns an AuthPassthrough given the pool settings.
+    /// If any of required values is not set, None is returned.
+    pub fn from_pool_settings(pool_settings: &crate::pool::PoolSettings) -> Option<Self> {
+        let pool_config = crate::config::Pool {
+            auth_query: pool_settings.auth_query.clone(),
+            auth_query_password: pool_settings.auth_query_password.clone(),
+            auth_query_user: pool_settings.auth_query_user.clone(),
+            ..Default::default()
+        };
+
+        AuthPassthrough::from_pool_config(&pool_config)
+    }
+
+    /// Connects to server and executes auth_query for the specified address.
+    /// If the response is a row with two columns containing the username set in the address.
+    /// and its MD5 hash, the MD5 hash returned.
+    ///
+    /// Note that the query is executed, changing $1 with the name of the user
+    /// this is so we only hold in memory (and transfer) the least amount of 'sensitive' data.
+    /// Also, it is compatible with pgbouncer.
+    ///
+    /// # Arguments
+    ///
+    /// * `address` - An Address of the server we want to connect to. The username for the hash will be obtained from this value.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use pgcat::auth_passthrough::AuthPassthrough;
+    /// use pgcat::config::Address;
+    /// let auth_passthrough = AuthPassthrough::new("SELECT * FROM public.user_lookup('$1');", "postgres", "postgres");
+    /// auth_passthrough.fetch_hash(&Address::default());
+    /// ```
+    ///
+    pub async fn fetch_hash(&self, address: &crate::config::Address) -> Result<String, Error> {
+        let auth_user = crate::config::User {
+            username: self.user.clone(),
+            password: Some(self.password.clone()),
+            pool_size: 1,
+            statement_timeout: 0,
+        };
+
+        let user = &address.username;
+
+        debug!("Connecting to server to obtain auth hashes.");
+        let auth_query = self.query.replace("$1", user);
+        match Server::exec_simple_query(address, &auth_user, &auth_query).await {
+            Ok(password_data) => {
+                if password_data.len() == 2 && password_data.first().unwrap() == user {
+		    if let Some(stripped_hash) = password_data.last().unwrap().to_string().strip_prefix("md5") {
+			Ok(stripped_hash.to_string())
+		    }
+		    else {
+			Err(Error::AuthPassthroughError(
+			    "Obtained hash from auth_query does not seem to be in md5 format.".to_string(),
+			))
+		    }
+                } else {
+                    Err(Error::AuthPassthroughError(
+                        "Data obtained from query does not follow the scheme 'user','hash'."
+                            .to_string(),
+                    ))
+                 }
+            }
+            Err(err) => {
+		Err(Error::AuthPassthroughError(
+		    format!("Error trying to obtain password from auth_query, ignoring hash for user '{}'. Error: {:?}",
+			    user, err)))
+            }
+	}
+    }
+}

src/auth_passthrough.rs

@@ -13,14 +13,13 @@ use tokio::sync::mpsc::Sender;
 
 use crate::admin::{generate_server_info_for_admin, handle_admin};
 use crate::config::{get_config, Address, PoolMode};
-use crate::constants::*;
-
 use crate::messages::*;
 use crate::pool::{get_pool, ClientServerMap, ConnectionPool};
 use crate::query_router::{Command, QueryRouter};
 use crate::server::Server;
 use crate::stats::{get_reporter, Reporter};
 use crate::tls::Tls;
+use crate::{auth_passthrough::AuthPassthrough, constants::*};
 
 use tokio_rustls::server::TlsStream;
 
@@ -362,6 +361,20 @@ pub async fn startup_tls(
     }
 }
 
+async fn refetch_auth_hash(pool: &ConnectionPool) -> Result<String, Error> {
+    let address = pool.address(0, 0);
+    if let Some(apt) = AuthPassthrough::from_pool_settings(&pool.settings) {
+        let hash = apt.fetch_hash(address).await?;
+
+        return Ok(hash);
+    }
+
+    Err(Error::ClientError(format!(
+        "Could not obtain hash for {{ username: {:?}, database: {:?} }}. Auth passthrough not enabled.",
+        address.username, address.database
+    )))
+}
+
 impl<S, T> Client<S, T>
 where
     S: tokio::io::AsyncRead + std::marker::Unpin,
@@ -495,14 +508,68 @@ where
                 }
             };
 
-            // Compare server and client hashes.
-            let password_hash = md5_hash_password(username, &pool.settings.user.password, &salt);
+            // Obtain the hash to compare, we give preference to that written in cleartext in config
+            // if there is nothing set in cleartext and auth passthrough (auth_query) is configured, we use the hash obtained
+            // when the pool was created. If there is no hash there, we try to fetch it one more time.
+            let password_hash = if let Some(password) = &pool.settings.user.password {
+                Some(md5_hash_password(username, password, &salt))
+            } else {
+                if !get_config().is_auth_query_configured() {
+                    return Err(Error::ClientError(format!("Client auth not possible, no cleartext password set for username: {:?} in config and auth passthrough (query_auth) is not set up.", username)));
+                }
 
-            if password_hash != password_response {
-                warn!("Invalid password {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}", username, pool_name, application_name);
-                wrong_password(&mut write, username).await?;
+                let mut hash = (*pool.auth_hash.read()).clone();
 
-                return Err(Error::ClientError(format!("Invalid password {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}", username, pool_name, application_name)));
+                if hash.is_none() {
+                    warn!("Query auth configured but no hash password found for pool {}. Will try to refetch it.", pool_name);
+                    match refetch_auth_hash(&pool).await {
+                        Ok(fetched_hash) => {
+                            warn!("Password for {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}, obtained. Updating.", username, pool_name, application_name);
+                            {
+                                let mut pool_auth_hash = pool.auth_hash.write();
+                                *pool_auth_hash = Some(fetched_hash.clone());
+                            }
+
+                            hash = Some(fetched_hash);
+                        }
+                        Err(err) => {
+                            return Err(
+				Error::ClientError(
+				    format!("No cleartext password set, and no auth passthrough could not obtain the hash from server for {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}, the error was: {:?}",
+					    username,
+					    pool_name,
+					    application_name,
+					    err)
+				)
+			    );
+                        }
+                    }
+                };
+
+                Some(md5_hash_second_pass(&hash.unwrap(), &salt))
+            };
+
+            // Once we have the resulting hash, we compare with what the client gave us.
+            // If they do not match and auth query is set up, we try to refetch the hash one more time
+            // to see if the password has changed since the pool was created.
+            //
+            // @TODO: we could end up fetching again the same password twice (see above).
+            if password_hash.unwrap() != password_response {
+                warn!("Invalid password {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}, will try to refetch it.", username, pool_name, application_name);
+                let fetched_hash = refetch_auth_hash(&pool).await?;
+                let new_password_hash = md5_hash_second_pass(&fetched_hash, &salt);
+
+                // Ok password changed in server an auth is possible.
+                if new_password_hash == password_response {
+                    warn!("Password for {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}, changed in server. Updating.", username, pool_name, application_name);
+                    {
+                        let mut pool_auth_hash = pool.auth_hash.write();
+                        *pool_auth_hash = Some(fetched_hash);
+                    }
+                } else {
+                    wrong_password(&mut write, username).await?;
+                    return Err(Error::ClientError(format!("Invalid password {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}", username, pool_name, application_name)));
+                }
             }
 
             let transaction_mode = pool.settings.pool_mode == PoolMode::Transaction;