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.

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 regex::Regex;
@@ -202,7 +203,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 {
@@ -285,7 +285,6 @@ impl Default for General {
             auth_query: None,
             auth_query_user: None,
             auth_query_password: None,
-            auth_query_database: None,
         }
     }
 }
@@ -360,6 +359,10 @@ 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>,
     // Note, don't put simple fields below these configs. There's a compatability issue with TOML that makes it
     // incompatible to have simple fields in TOML after complex objects. See
     // https://users.rust-lang.org/t/why-toml-to-string-get-error-valueaftertable/85903
@@ -372,6 +375,12 @@ impl Pool {
         s.finish()
     }
 
+    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
     }
@@ -445,6 +454,9 @@ impl Default for Pool {
             sharding_key_regex: None,
             shard_id_regex: None,
             regex_search_limit: Some(1000),
+            auth_query: None,
+            auth_query_user: None,
+            auth_query_password: None,
         }
     }
 }
@@ -536,6 +548,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")
     }
@@ -549,6 +567,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 {
@@ -754,21 +788,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);
         }
 
@@ -857,6 +883,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();
@@ -972,7 +999,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

@@ -113,6 +113,11 @@ pub struct PoolSettings {
 
     // Limit how much of each query is searched for a potential shard regex match
     pub regex_search_limit: usize,
+
+    // Auth query parameters
+    pub auth_query: Option<String>,
+    pub auth_query_user: Option<String>,
+    pub auth_query_password: Option<String>,
 }
 
 impl Default for PoolSettings {
@@ -133,6 +138,9 @@ impl Default for PoolSettings {
             sharding_key_regex: None,
             shard_id_regex: None,
             regex_search_limit: 1000,
+            auth_query: None,
+            auth_query_user: None,
+            auth_query_password: None,
         }
     }
 }
@@ -182,15 +190,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);
 				}
 			    }
 			},
@@ -205,7 +217,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;
@@ -283,14 +294,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());
@@ -338,6 +351,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 pool = ConnectionPool {
                     databases: shards,
@@ -375,6 +394,9 @@ impl ConnectionPool {
                             .clone()
                             .map(|regex| Regex::new(regex.as_str()).unwrap()),
                         regex_search_limit: pool_config.regex_search_limit.unwrap_or(1000),
+                        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(),
                     },
                     validated: Arc::new(AtomicBool::new(false)),
                     paused: Arc::new(AtomicBool::new(false)),

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);