Change auth_query config to be per pool instead of global.

Given that we can connect to different username/databases/servers using
connection pools, it makes sense that `auth_query` should be configured in a
per pool basis.

This change implements that and also leaves the global parameters so pool
configuration can inherit global ones.

Note that now, `auth_query_database` is dropped in favor of the pool configured
database.

Author: magec

.circleci/pgcat.toml

@@ -53,7 +53,6 @@ admin_password = "admin_pass"
 # auth_query = "SELECT * FROM public.user_lookup('$1');"
 # auth_query_user = "md5_auth_user"
 # auth_query_password = "secret"
-# auth_query_database = "postgres"
 
 # pool
 # configs are structured as pool.<pool_name>

.circleci/query_auth_test.sh

@@ -12,16 +12,18 @@ export LOCAL_IP=$(hostname -i)
 #  auth_query = "SELECT * FROM public.user_lookup('$1');"
 #  auth_query_user = "md5_auth_user"
 #  auth_query_password = "secret"
-#  auth_query_database = "postgres"
 #  ...
 
 # Before (sets up auth_query in postgres and pgcat)
 PGDATABASE=postgres PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_setup.sql
+PGDATABASE=shard0 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_setup_function.sql
+PGDATABASE=shard1 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_setup_function.sql
+PGDATABASE=shard2 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_setup_function.sql
 sed -i 's/^# auth_query/auth_query/' .circleci/pgcat.toml
 
 # TEST_WRONG_AUTH_QUERY BEGIN
 #    When auth_query fails...
-PGDATABASE=postgres \
+PGDATABASE=shard0 \
     PGPASSWORD=postgres \
     psql -e -h 127.0.0.1 -p 5432 -U postgres -c "REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;"
 
@@ -33,7 +35,7 @@ echo "When query_auth_config is wrong, we fall back to passwords set in cleartex
 psql -U sharding_user -h 127.0.0.1 -p 6432 -c 'SELECT 1'
 
 # After
-PGDATABASE=postgres \
+PGDATABASE=shard0 \
     PGPASSWORD=postgres \
     psql -e -h 127.0.0.1 -p 5432 -U postgres -c "GRANT EXECUTE ON FUNCTION public.user_lookup(text) TO md5_auth_user;"
 # TEST_WRONG_AUTH_QUERY END
@@ -80,6 +82,9 @@ PGPASSWORD=another_sharding_password psql -U sharding_user -h "${LOCAL_IP}" -p 6
 # TEST_PASSWORD_CHANGE END
 
 # After
+PGDATABASE=shard0 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_teardown_function.sql
+PGDATABASE=shard1 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_teardown_function.sql
+PGDATABASE=shard2 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_teardown_function.sql
 PGDATABASE=postgres PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_teardown.sql
 sed -i 's/^auth_query/# auth_query/' .circleci/pgcat.toml
 sed -i 's/^# password =/password =/' .circleci/pgcat.toml

src/auth_passthrough.rs

@@ -7,14 +7,12 @@ use bytes::BytesMut;
 use fallible_iterator::FallibleIterator;
 use parking_lot::Mutex;
 
+use crate::pool::ClientServerMap;
 use log::{debug, trace, warn};
 use postgres_protocol::message;
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use crate::config::get_config;
-use crate::pool::ClientServerMap;
-
 pub struct AuthPassthrough {
     password: String,
     query: String,
@@ -33,23 +31,35 @@ impl AuthPassthrough {
         }
     }
 
-    /// Returns an AuthPassthrough given the configuration.
+    /// Returns an AuthPassthrough given the pool configuration.
     /// If any of required values is not set, None is returned.
-    pub fn from_config() -> Option<Self> {
-        let config = get_config();
-
-        if config.is_auth_query_configured() {
+    pub fn from_pool_config(pool_config: &crate::config::Pool, database: &String) -> Option<Self> {
+        if pool_config.is_auth_query_configured() {
             return Some(AuthPassthrough::new(
-                config.general.auth_query.as_ref().unwrap(),
-                config.general.auth_query_user.as_ref().unwrap(),
-                config.general.auth_query_password.as_ref().unwrap(),
-                config.general.auth_query_database.as_ref().unwrap(),
+                pool_config.auth_query.as_ref().unwrap(),
+                pool_config.auth_query_user.as_ref().unwrap(),
+                pool_config.auth_query_password.as_ref().unwrap(),
+                database,
             ));
         }
 
         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,
+        database: &String,
+    ) -> Option<Self> {
+        let mut pool_config = crate::config::Pool::default();
+        pool_config.auth_query = pool_settings.auth_query.clone();
+        pool_config.auth_query_password = pool_settings.auth_query_password.clone();
+        pool_config.auth_query_user = pool_settings.auth_query_user.clone();
+
+        AuthPassthrough::from_pool_config(&pool_config, database)
+    }
+
     /// 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.
@@ -220,7 +230,7 @@ async fn recv_data(server: &mut Server) -> Result<BytesMut, Error> {
 }
 
 async fn send_auth_query(server: &mut Server, query: &str) -> Result<(), Error> {
-    return match server.send(simple_query(query)).await {
+    return match server.send(&simple_query(query)).await {
         Ok(()) => Ok(()),
         Err(err) => {
             let message = format!(

src/client.rs

@@ -507,7 +507,7 @@ where
                 }
             }
             if password_hash.is_none() {
-                warn!("Clien auth is not possible, you either have not set a valid auth_query or a password for {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}", username, pool_name, application_name);
+                warn!("Client auth is not possible, you either have not set a valid auth_query or a password for {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}", username, pool_name, application_name);
                 wrong_password(&mut write, username).await?;
 
                 return Err(Error::ClientError(format!("Invalid client auth {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}", username, pool_name, application_name)));

src/config.rs

@@ -1,5 +1,6 @@
 /// Parse the configuration file.
 use arc_swap::ArcSwap;
+use hmac::digest::typenum::private::IsNotEqualPrivate;
 use log::{error, info};
 use once_cell::sync::Lazy;
 use serde_derive::{Deserialize, Serialize};
@@ -193,7 +194,6 @@ pub struct General {
     pub auth_query: Option<String>,
     pub auth_query_user: Option<String>,
     pub auth_query_password: Option<String>,
-    pub auth_query_database: Option<String>,
 }
 
 impl General {
@@ -258,7 +258,6 @@ impl Default for General {
             auth_query: None,
             auth_query_user: None,
             auth_query_password: None,
-            auth_query_database: None,
         }
     }
 }
@@ -329,9 +328,19 @@ pub struct Pool {
 
     pub shards: BTreeMap<String, Shard>,
     pub users: BTreeMap<String, User>,
+
+    pub auth_query: Option<String>,
+    pub auth_query_user: Option<String>,
+    pub auth_query_password: Option<String>,
 }
 
 impl Pool {
+    pub fn is_auth_query_configured(&self) -> bool {
+        self.auth_query_password.is_some()
+            && self.auth_query_user.is_some()
+            && self.auth_query_password.is_some()
+    }
+
     pub fn default_pool_mode() -> PoolMode {
         PoolMode::Transaction
     }
@@ -390,6 +399,9 @@ impl Default for Pool {
             automatic_sharding_key: None,
             connect_timeout: None,
             idle_timeout: None,
+            auth_query: None,
+            auth_query_user: None,
+            auth_query_password: None,
         }
     }
 }
@@ -481,6 +493,12 @@ pub struct Config {
 }
 
 impl Config {
+    pub fn is_auth_query_configured(&self) -> bool {
+        self.pools
+            .iter()
+            .any(|(_name, pool)| pool.is_auth_query_configured())
+    }
+
     pub fn default_path() -> String {
         String::from("pgcat.toml")
     }
@@ -494,6 +512,22 @@ impl Config {
             self.general.auth_query_password = Some(val);
         }
     }
+
+    pub fn fill_up_auth_query_config(&mut self) {
+        for (_name, pool) in self.pools.iter_mut() {
+            if pool.auth_query.is_none() {
+                pool.auth_query = self.general.auth_query.clone();
+            }
+
+            if pool.auth_query_user.is_none() {
+                pool.auth_query_user = self.general.auth_query_user.clone();
+            }
+
+            if pool.auth_query_password.is_none() {
+                pool.auth_query_password = self.general.auth_query_password.clone();
+            }
+        }
+    }
 }
 
 impl Default for Config {
@@ -698,21 +732,13 @@ impl Config {
         }
     }
 
-    pub fn is_auth_query_configured(&self) -> bool {
-        self.general.auth_query_password.is_some()
-            && self.general.auth_query_user.is_some()
-            && self.general.auth_query_password.is_some()
-            && self.general.auth_query_database.is_some()
-    }
-
     pub fn validate(&mut self) -> Result<(), Error> {
         // Validation for auth_query feature
         if self.general.auth_query.is_some()
             && (self.general.auth_query_user.is_none()
-                || self.general.auth_query_password.is_none()
-                || self.general.auth_query_database.is_none())
+                || self.general.auth_query_password.is_none())
         {
-            error!("If auth_query is specified, you need to provide a value for `auth_query_user`, `auth_query_password` and `auth_query_database`");
+            error!("If auth_query is specified, you need to provide a value for `auth_query_user`, `auth_query_password`");
             return Err(Error::BadConfig);
         }
 
@@ -801,6 +827,7 @@ pub async fn parse(path: &str) -> Result<(), Error> {
     };
 
     config.overwrite_passwords();
+    config.fill_up_auth_query_config();
     config.validate()?;
 
     config.path = path.to_string();
@@ -916,7 +943,6 @@ mod test {
         );
         assert_eq!(get_config().pools["simple_db"].users["0"].pool_size, 5);
         assert_eq!(get_config().general.auth_query, None);
-        assert_eq!(get_config().general.auth_query_database, None);
         assert_eq!(get_config().general.auth_query_user, None);
         assert_eq!(get_config().general.auth_query_password, None);
     }

src/pool.rs

@@ -92,6 +92,11 @@ pub struct PoolSettings {
 
     // Health check delay
     pub healthcheck_delay: u64,
+
+    // Auth query parameters
+    pub auth_query: Option<String>,
+    pub auth_query_user: Option<String>,
+    pub auth_query_password: Option<String>,
 }
 
 impl Default for PoolSettings {
@@ -108,6 +113,9 @@ impl Default for PoolSettings {
             automatic_sharding_key: None,
             healthcheck_delay: General::default_healthcheck_delay(),
             healthcheck_timeout: General::default_healthcheck_timeout(),
+            auth_query: None,
+            auth_query_user: None,
+            auth_query_password: None,
         }
     }
 }
@@ -144,15 +152,19 @@ pub struct ConnectionPool {
 
 impl ConnectionPool {
     async fn auth_hash_has_changed(&mut self) -> bool {
-        if let Some(apt) = AuthPassthrough::from_config() {
-            if let Some(shard) = self.addresses.first() {
-                if let Some(address) = shard.first() {
+        if let Some(shard) = self.addresses.first() {
+            if let Some(address) = shard.first() {
+                if let Some(apt) =
+                    AuthPassthrough::from_pool_settings(&self.settings, &address.database)
+                {
                     match apt.fetch_hash(address).await {
 			Ok(md5) => {
 			    if let Some(auth_hash) = self.auth_hash.as_ref() {
 				if auth_hash != &md5 {
 				    info!("Password hash changed for user: {:?}. Pool will be reloaded.", address.username);
 				    return true
+				} else {
+				    debug!("Password hash has not changed for user: {:?}.", address.username);
 				}
 			    }
 			},
@@ -167,7 +179,6 @@ impl ConnectionPool {
     /// Construct the connection pool from the configuration.
     pub async fn from_config(client_server_map: ClientServerMap) -> Result<(), Error> {
         let config = get_config();
-        let auth_passthrough = AuthPassthrough::from_config();
 
         let mut new_pools = HashMap::new();
         let mut address_id = 0;
@@ -247,14 +258,16 @@ impl ConnectionPool {
                         }
 
                         // We assume every server in the pool share user/passwords
-                        // TODO: Make it server dependant.
                         let mut auth_hash = None;
+                        let auth_passthrough =
+                            AuthPassthrough::from_pool_config(pool_config, &shard.database);
                         if let Some(apt) = auth_passthrough.as_ref() {
                             match apt.fetch_hash(&address).await {
 				Ok(ok) => {
 				    if let Some(pool_auth_hash_value) = pool_auth_hash {
 					if ok != pool_auth_hash_value {
-					    warn!("Hash is not the same across servers of the same pool, client auth will be done using last obtained hash.");
+					    warn!("Hash is not the same across shards of the same pool, client auth will \
+						   be done using last obtained hash. Server: {}:{}, Database: {}", server.host, server.port, shard.database);
 					}
 				    }
 				    pool_auth_hash = Some(ok.clone());
@@ -302,6 +315,12 @@ impl ConnectionPool {
                 }
 
                 assert_eq!(shards.len(), addresses.len());
+                if pool_auth_hash.is_some() {
+                    info!(
+                        "Auth hash obtained from query_auth for pool {{ name: {}, user: {} }}",
+                        pool_name, user.username
+                    );
+                }
 
                 let mut pool = ConnectionPool {
                     databases: shards,
@@ -328,6 +347,9 @@ impl ConnectionPool {
                         automatic_sharding_key: pool_config.automatic_sharding_key.clone(),
                         healthcheck_delay: config.general.healthcheck_delay,
                         healthcheck_timeout: config.general.healthcheck_timeout,
+                        auth_query: pool_config.auth_query.clone(),
+                        auth_query_user: pool_config.auth_query_user.clone(),
+                        auth_query_password: pool_config.auth_query_password.clone(),
                     },
                 };
 

tests/sharding/query_auth_setup.sql

@@ -4,15 +4,3 @@
 
 ALTER ROLE sharding_user ENCRYPTED PASSWORD 'md5fa9d23e5a874c61a91bf37e1e4a9c86e'; --- sharding_user
 CREATE ROLE md5_auth_user ENCRYPTED PASSWORD 'md54ab2c5d00339c4b2a4e921d2dc4edec7' LOGIN; --- secret
-
-CREATE OR REPLACE FUNCTION public.user_lookup(in i_username text, out uname text, out phash text)
-RETURNS record AS $$
-BEGIN
-    SELECT usename, passwd FROM pg_catalog.pg_shadow
-    WHERE usename = i_username INTO uname, phash;
-    RETURN;
-END;
-$$ LANGUAGE plpgsql SECURITY DEFINER;
-
-REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;
-GRANT EXECUTE ON FUNCTION public.user_lookup(text) TO md5_auth_user;

tests/sharding/query_auth_setup_function.sql

@@ -0,0 +1,11 @@
+CREATE OR REPLACE FUNCTION public.user_lookup(in i_username text, out uname text, out phash text)
+RETURNS record AS $$
+BEGIN
+    SELECT usename, passwd FROM pg_catalog.pg_shadow
+    WHERE usename = i_username INTO uname, phash;
+    RETURN;
+END;
+$$ LANGUAGE plpgsql SECURITY DEFINER;
+
+REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;
+GRANT EXECUTE ON FUNCTION public.user_lookup(text) TO md5_auth_user;

tests/sharding/query_auth_teardown.sql

@@ -3,7 +3,4 @@
 ---   - Drops auth query function and user.
 
 ALTER ROLE sharding_user ENCRYPTED PASSWORD 'sharding_user' LOGIN;
-
-REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;
-DROP FUNCTION public.user_lookup(in i_username text, out uname text, out phash text);
 DROP ROLE md5_auth_user;

tests/sharding/query_auth_teardown_function.sql

@@ -0,0 +1,2 @@
+REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;
+DROP FUNCTION public.user_lookup(in i_username text, out uname text, out phash text);