From 112c0bdae8a6d471dc753b3812171abb4a901061 Mon Sep 17 00:00:00 2001
From: Lev Kokotov <lev.kokotov@gmail.com>
Date: Thu, 30 Mar 2023 15:19:52 -0700
Subject: [PATCH] Rebased

---
 dev/docker-compose.yaml |  2 +-
 src/auth.rs             | 30 ++++++++++++++++++++++--------
 src/auth_passthrough.rs | 26 ++++++++++++++------------
 src/client.rs           |  1 -
 4 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/dev/docker-compose.yaml b/dev/docker-compose.yaml
index 71704bc..7684553 100644
--- a/dev/docker-compose.yaml
+++ b/dev/docker-compose.yaml
@@ -64,7 +64,7 @@ services:
       <<: *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"]
+    command: ["postgres", "-p", "10432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
 
   toxiproxy:
     build: .
diff --git a/src/auth.rs b/src/auth.rs
index cd6fae3..7b5bf21 100644
--- a/src/auth.rs
+++ b/src/auth.rs
@@ -24,11 +24,15 @@ async fn refetch_auth_hash<S>(
 where
     S: tokio::io::AsyncWrite + std::marker::Unpin + std::marker::Send,
 {
-    let address = pool.address(0, 0);
-    if let Some(apt) = AuthPassthrough::from_pool_settings(&pool.settings) {
-        let hash = apt.fetch_hash(address).await?;
+    let config = get_config();
 
-        return Ok(hash);
+    if config.is_auth_query_configured() {
+        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);
+        }
     }
 
     error_response(
@@ -42,7 +46,7 @@ where
 
     Err(Error::ClientError(format!(
         "Could not obtain hash for {{ username: {:?}, database: {:?} }}. Auth passthrough not enabled.",
-        address.username, address.database
+        pool_name, username
     )))
 }
 
@@ -331,10 +335,19 @@ impl Md5 {
                             let hash = (*pool.auth_hash.read()).clone();
 
                             let hash = match hash {
-                                Some(hash) => hash.to_string(),
+                                Some(hash) => hash.clone(),
                                 None => {
-                                    refetch_auth_hash(&pool, write, &self.username, &self.pool_name)
-                                        .await?
+                                    let hash = refetch_auth_hash(
+                                        &pool,
+                                        write,
+                                        &self.username,
+                                        &self.pool_name,
+                                    )
+                                    .await?;
+
+                                    (*pool.auth_hash.write()) = Some(hash.clone());
+
+                                    hash
                                 }
                             };
 
@@ -350,6 +363,7 @@ impl Md5 {
                                     &self.pool_name,
                                 )
                                 .await?;
+
                                 let our_hash = md5_hash_second_pass(&hash, &self.salt);
 
                                 if our_hash != password_hash {
diff --git a/src/auth_passthrough.rs b/src/auth_passthrough.rs
index 33a9fb1..9cf64b7 100644
--- a/src/auth_passthrough.rs
+++ b/src/auth_passthrough.rs
@@ -79,18 +79,19 @@ impl AuthPassthrough {
         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(),
-			))
-		    }
+                    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'."
@@ -98,11 +99,12 @@ impl AuthPassthrough {
                     ))
                  }
             }
+
             Err(err) => {
-		Err(Error::AuthPassthroughError(
-		    format!("Error trying to obtain password from auth_query, ignoring hash for user '{}'. Error: {:?}",
-			    user, err)))
+                Err(Error::AuthPassthroughError(
+                    format!("Error trying to obtain password from auth_query, ignoring hash for user '{}'. Error: {:?}",
+                        user, err)))
+                    }
             }
-	}
     }
 }
diff --git a/src/client.rs b/src/client.rs
index 5ef5fe4..c494b1a 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -12,7 +12,6 @@ use tokio::sync::broadcast::Receiver;
 use tokio::sync::mpsc::Sender;
 
 use crate::admin::{generate_server_info_for_admin, handle_admin};
-use crate::auth_passthrough::AuthPassthrough;
 use crate::config::{get_config, get_idle_client_in_transaction_timeout, Address, PoolMode};
 use crate::constants::*;
 use crate::messages::*;