Skip to content

Commit 3d6dddf

Browse files
magecmagec
committed
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.
1 parent 304093c commit 3d6dddf

File tree

10 files changed

+105
-50
lines changed

10 files changed

+105
-50
lines changed

‎.circleci/pgcat.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ admin_password = "admin_pass"
5353
# auth_query = "SELECT * FROM public.user_lookup('$1');"
5454
# auth_query_user = "md5_auth_user"
5555
# auth_query_password = "secret"
56-
# auth_query_database = "postgres"
5756

5857
# pool
5958
# configs are structured as pool.<pool_name>

‎.circleci/query_auth_test.sh

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,18 @@ export LOCAL_IP=$(hostname -i)
1212
# auth_query = "SELECT * FROM public.user_lookup('$1');"
1313
# auth_query_user = "md5_auth_user"
1414
# auth_query_password = "secret"
15-
# auth_query_database = "postgres"
1615
# ...
1716

1817
# Before (sets up auth_query in postgres and pgcat)
1918
PGDATABASE=postgres PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_setup.sql
19+
PGDATABASE=shard0 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_setup_function.sql
20+
PGDATABASE=shard1 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_setup_function.sql
21+
PGDATABASE=shard2 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_setup_function.sql
2022
sed -i 's/^# auth_query/auth_query/' .circleci/pgcat.toml
2123

2224
# TEST_WRONG_AUTH_QUERY BEGIN
2325
# When auth_query fails...
24-
PGDATABASE=postgres \
26+
PGDATABASE=shard0 \
2527
PGPASSWORD=postgres \
2628
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;"
2729

@@ -33,7 +35,7 @@ echo "When query_auth_config is wrong, we fall back to passwords set in cleartex
3335
psql -U sharding_user -h 127.0.0.1 -p 6432 -c 'SELECT 1'
3436

3537
# After
36-
PGDATABASE=postgres \
38+
PGDATABASE=shard0 \
3739
PGPASSWORD=postgres \
3840
psql -e -h 127.0.0.1 -p 5432 -U postgres -c "GRANT EXECUTE ON FUNCTION public.user_lookup(text) TO md5_auth_user;"
3941
# TEST_WRONG_AUTH_QUERY END
@@ -80,6 +82,9 @@ PGPASSWORD=another_sharding_password psql -U sharding_user -h "${LOCAL_IP}" -p 6
8082
# TEST_PASSWORD_CHANGE END
8183

8284
# After
85+
PGDATABASE=shard0 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_teardown_function.sql
86+
PGDATABASE=shard1 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_teardown_function.sql
87+
PGDATABASE=shard2 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_teardown_function.sql
8388
PGDATABASE=postgres PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_auth_teardown.sql
8489
sed -i 's/^auth_query/# auth_query/' .circleci/pgcat.toml
8590
sed -i 's/^# password =/password =/' .circleci/pgcat.toml

‎src/auth_passthrough.rs

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,14 +7,12 @@ use bytes::BytesMut;
77
use fallible_iterator::FallibleIterator;
88
use parking_lot::Mutex;
99

10+
use crate::pool::ClientServerMap;
1011
use log::{debug, trace, warn};
1112
use postgres_protocol::message;
1213
use std::collections::HashMap;
1314
use std::sync::Arc;
1415

15-
use crate::config::get_config;
16-
use crate::pool::ClientServerMap;
17-
1816
pub struct AuthPassthrough {
1917
password: String,
2018
query: String,
@@ -33,23 +31,35 @@ impl AuthPassthrough {
3331
}
3432
}
3533

36-
/// Returns an AuthPassthrough given the configuration.
34+
/// Returns an AuthPassthrough given the pool configuration.
3735
/// If any of required values is not set, None is returned.
38-
pub fn from_config() -> Option<Self> {
39-
let config = get_config();
40-
41-
if config.is_auth_query_configured() {
36+
pub fn from_pool_config(pool_config: &crate::config::Pool, database: &String) -> Option<Self> {
37+
if pool_config.is_auth_query_configured() {
4238
return Some(AuthPassthrough::new(
43-
config.general.auth_query.as_ref().unwrap(),
44-
config.general.auth_query_user.as_ref().unwrap(),
45-
config.general.auth_query_password.as_ref().unwrap(),
46-
config.general.auth_query_database.as_ref().unwrap(),
39+
pool_config.auth_query.as_ref().unwrap(),
40+
pool_config.auth_query_user.as_ref().unwrap(),
41+
pool_config.auth_query_password.as_ref().unwrap(),
42+
database,
4743
));
4844
}
4945

5046
None
5147
}
5248

49+
/// Returns an AuthPassthrough given the pool settings.
50+
/// If any of required values is not set, None is returned.
51+
pub fn from_pool_settings(
52+
pool_settings: &crate::pool::PoolSettings,
53+
database: &String,
54+
) -> Option<Self> {
55+
let mut pool_config = crate::config::Pool::default();
56+
pool_config.auth_query = pool_settings.auth_query.clone();
57+
pool_config.auth_query_password = pool_settings.auth_query_password.clone();
58+
pool_config.auth_query_user = pool_settings.auth_query_user.clone();
59+
60+
AuthPassthrough::from_pool_config(&pool_config, database)
61+
}
62+
5363
/// Connects to server and executes auth_query for the specified address.
5464
/// If the response is a row with two columns containing the username set in the address.
5565
/// and its MD5 hash, the MD5 hash returned.

‎src/client.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -507,7 +507,7 @@ where
507507
}
508508
}
509509
if password_hash.is_none() {
510-
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);
510+
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);
511511
wrong_password(&mut write, username).await?;
512512

513513
return Err(Error::ClientError(format!("Invalid client auth {{ username: {:?}, pool_name: {:?}, application_name: {:?} }}", username, pool_name, application_name)));

‎src/config.rs

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
/// Parse the configuration file.
22
use arc_swap::ArcSwap;
3+
use hmac::digest::typenum::private::IsNotEqualPrivate;
34
use log::{error, info};
45
use once_cell::sync::Lazy;
56
use serde_derive::{Deserialize, Serialize};
@@ -193,7 +194,6 @@ pub struct General {
193194
pub auth_query: Option<String>,
194195
pub auth_query_user: Option<String>,
195196
pub auth_query_password: Option<String>,
196-
pub auth_query_database: Option<String>,
197197
}
198198

199199
impl General {
@@ -258,7 +258,6 @@ impl Default for General {
258258
auth_query: None,
259259
auth_query_user: None,
260260
auth_query_password: None,
261-
auth_query_database: None,
262261
}
263262
}
264263
}
@@ -308,9 +307,19 @@ pub struct Pool {
308307

309308
pub shards: BTreeMap<String, Shard>,
310309
pub users: BTreeMap<String, User>,
310+
311+
pub auth_query: Option<String>,
312+
pub auth_query_user: Option<String>,
313+
pub auth_query_password: Option<String>,
311314
}
312315

313316
impl Pool {
317+
pub fn is_auth_query_configured(&self) -> bool {
318+
self.auth_query_password.is_some()
319+
&& self.auth_query_user.is_some()
320+
&& self.auth_query_password.is_some()
321+
}
322+
314323
pub fn default_pool_mode() -> PoolMode {
315324
PoolMode::Transaction
316325
}
@@ -364,6 +373,9 @@ impl Default for Pool {
364373
automatic_sharding_key: None,
365374
connect_timeout: None,
366375
idle_timeout: None,
376+
auth_query: None,
377+
auth_query_user: None,
378+
auth_query_password: None,
367379
}
368380
}
369381
}
@@ -455,6 +467,12 @@ pub struct Config {
455467
}
456468

457469
impl Config {
470+
pub fn is_auth_query_configured(&self) -> bool {
471+
self.pools
472+
.iter()
473+
.any(|(_name, pool)| pool.is_auth_query_configured())
474+
}
475+
458476
pub fn default_path() -> String {
459477
String::from("pgcat.toml")
460478
}
@@ -468,6 +486,22 @@ impl Config {
468486
self.general.auth_query_password = Some(val);
469487
}
470488
}
489+
490+
pub fn fill_up_auth_query_config(&mut self) {
491+
for (_name, pool) in self.pools.iter_mut() {
492+
if pool.auth_query.is_none() {
493+
pool.auth_query = self.general.auth_query.clone();
494+
}
495+
496+
if pool.auth_query_user.is_none() {
497+
pool.auth_query_user = self.general.auth_query_user.clone();
498+
}
499+
500+
if pool.auth_query_password.is_none() {
501+
pool.auth_query_password = self.general.auth_query_password.clone();
502+
}
503+
}
504+
}
471505
}
472506

473507
impl Default for Config {
@@ -664,21 +698,13 @@ impl Config {
664698
}
665699
}
666700

667-
pub fn is_auth_query_configured(&self) -> bool {
668-
self.general.auth_query_password.is_some()
669-
&& self.general.auth_query_user.is_some()
670-
&& self.general.auth_query_password.is_some()
671-
&& self.general.auth_query_database.is_some()
672-
}
673-
674701
pub fn validate(&mut self) -> Result<(), Error> {
675702
// Validation for auth_query feature
676703
if self.general.auth_query.is_some()
677704
&& (self.general.auth_query_user.is_none()
678-
|| self.general.auth_query_password.is_none()
679-
|| self.general.auth_query_database.is_none())
705+
|| self.general.auth_query_password.is_none())
680706
{
681-
error!("If auth_query is specified, you need to provide a value for `auth_query_user`, `auth_query_password` and `auth_query_database`");
707+
error!("If auth_query is specified, you need to provide a value for `auth_query_user`, `auth_query_password`");
682708
return Err(Error::BadConfig);
683709
}
684710

@@ -767,6 +793,7 @@ pub async fn parse(path: &str) -> Result<(), Error> {
767793
};
768794

769795
config.overwrite_passwords();
796+
config.fill_up_auth_query_config();
770797
config.validate()?;
771798

772799
config.path = path.to_string();
@@ -882,7 +909,6 @@ mod test {
882909
);
883910
assert_eq!(get_config().pools["simple_db"].users["0"].pool_size, 5);
884911
assert_eq!(get_config().general.auth_query, None);
885-
assert_eq!(get_config().general.auth_query_database, None);
886912
assert_eq!(get_config().general.auth_query_user, None);
887913
assert_eq!(get_config().general.auth_query_password, None);
888914
}

‎src/pool.rs

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,11 @@ pub struct PoolSettings {
8989

9090
// Health check delay
9191
pub healthcheck_delay: u64,
92+
93+
// Auth query parameters
94+
pub auth_query: Option<String>,
95+
pub auth_query_user: Option<String>,
96+
pub auth_query_password: Option<String>,
9297
}
9398

9499
impl Default for PoolSettings {
@@ -104,6 +109,9 @@ impl Default for PoolSettings {
104109
automatic_sharding_key: None,
105110
healthcheck_delay: General::default_healthcheck_delay(),
106111
healthcheck_timeout: General::default_healthcheck_timeout(),
112+
auth_query: None,
113+
auth_query_user: None,
114+
auth_query_password: None,
107115
}
108116
}
109117
}
@@ -140,15 +148,19 @@ pub struct ConnectionPool {
140148

141149
impl ConnectionPool {
142150
async fn auth_hash_has_changed(&mut self) -> bool {
143-
if let Some(apt) = AuthPassthrough::from_config() {
144-
if let Some(shard) = self.addresses.first() {
145-
if let Some(address) = shard.first() {
151+
if let Some(shard) = self.addresses.first() {
152+
if let Some(address) = shard.first() {
153+
if let Some(apt) =
154+
AuthPassthrough::from_pool_settings(&self.settings, &address.database)
155+
{
146156
match apt.fetch_hash(address).await {
147157
Ok(md5) => {
148158
if let Some(auth_hash) = self.auth_hash.as_ref() {
149159
if auth_hash != &md5 {
150160
info!("Password hash changed for user: {:?}. Pool will be reloaded.", address.username);
151161
return true
162+
} else {
163+
debug!("Password hash has not changed for user: {:?}.", address.username);
152164
}
153165
}
154166
},
@@ -163,7 +175,6 @@ impl ConnectionPool {
163175
/// Construct the connection pool from the configuration.
164176
pub async fn from_config(client_server_map: ClientServerMap) -> Result<(), Error> {
165177
let config = get_config();
166-
let auth_passthrough = AuthPassthrough::from_config();
167178

168179
let mut new_pools = HashMap::new();
169180
let mut address_id = 0;
@@ -245,12 +256,15 @@ impl ConnectionPool {
245256
// We assume every server in the pool share user/passwords
246257
// TODO: Make it server dependant.
247258
let mut auth_hash = None;
259+
let auth_passthrough =
260+
AuthPassthrough::from_pool_config(pool_config, &shard.database);
248261
if let Some(apt) = auth_passthrough.as_ref() {
249262
match apt.fetch_hash(&address).await {
250263
Ok(ok) => {
251264
if let Some(pool_auth_hash_value) = pool_auth_hash {
252265
if ok != pool_auth_hash_value {
253-
warn!("Hash is not the same across servers of the same pool, client auth will be done using last obtained hash.");
266+
warn!("Hash is not the same across shards of the same pool, client auth will \
267+
be done using last obtained hash. Server: {}:{}, Database: {}", server.host, server.port, shard.database);
254268
}
255269
}
256270
pool_auth_hash = Some(ok.clone());
@@ -323,6 +337,9 @@ impl ConnectionPool {
323337
automatic_sharding_key: pool_config.automatic_sharding_key.clone(),
324338
healthcheck_delay: config.general.healthcheck_delay,
325339
healthcheck_timeout: config.general.healthcheck_timeout,
340+
auth_query: pool_config.auth_query.clone(),
341+
auth_query_user: pool_config.auth_query_user.clone(),
342+
auth_query_password: pool_config.auth_query_password.clone(),
326343
},
327344
};
328345

‎tests/sharding/query_auth_setup.sql

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,3 @@
44

55
ALTER ROLE sharding_user ENCRYPTED PASSWORD 'md5fa9d23e5a874c61a91bf37e1e4a9c86e'; --- sharding_user
66
CREATE ROLE md5_auth_user ENCRYPTED PASSWORD 'md54ab2c5d00339c4b2a4e921d2dc4edec7' LOGIN; --- secret
7-
8-
CREATE OR REPLACE FUNCTION public.user_lookup(in i_username text, out uname text, out phash text)
9-
RETURNS record AS $$
10-
BEGIN
11-
SELECT usename, passwd FROM pg_catalog.pg_shadow
12-
WHERE usename = i_username INTO uname, phash;
13-
RETURN;
14-
END;
15-
$$ LANGUAGE plpgsql SECURITY DEFINER;
16-
17-
REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;
18-
GRANT EXECUTE ON FUNCTION public.user_lookup(text) TO md5_auth_user;

‎tests/sharding/query_auth_setup_function.sql

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
CREATE OR REPLACE FUNCTION public.user_lookup(in i_username text, out uname text, out phash text)
2+
RETURNS record AS $$
3+
BEGIN
4+
SELECT usename, passwd FROM pg_catalog.pg_shadow
5+
WHERE usename = i_username INTO uname, phash;
6+
RETURN;
7+
END;
8+
$$ LANGUAGE plpgsql SECURITY DEFINER;
9+
10+
REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;
11+
GRANT EXECUTE ON FUNCTION public.user_lookup(text) TO md5_auth_user;

‎tests/sharding/query_auth_teardown.sql

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,4 @@
33
--- - Drops auth query function and user.
44

55
ALTER ROLE sharding_user ENCRYPTED PASSWORD 'sharding_user' LOGIN;
6-
7-
REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;
8-
DROP FUNCTION public.user_lookup(in i_username text, out uname text, out phash text);
96
DROP ROLE md5_auth_user;

‎tests/sharding/query_auth_teardown_function.sql

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
REVOKE ALL ON FUNCTION public.user_lookup(text) FROM public, md5_auth_user;
2+
DROP FUNCTION public.user_lookup(in i_username text, out uname text, out phash text);

0 commit comments

Comments
 (0)