From 021ab8365f1c9d86d29cecc467039389e8e71b5b Mon Sep 17 00:00:00 2001
From: Dmitry Ivanov <ivadmi5@gmail.com>
Date: Wed, 22 Feb 2023 20:37:18 +0300
Subject: [PATCH] [proxy] Refactoring in the classic auth backend

---
 proxy/src/auth/backend.rs         |  4 +--
 proxy/src/auth/backend/classic.rs | 48 +++++++++++++++++++------------
 2 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/proxy/src/auth/backend.rs b/proxy/src/auth/backend.rs
index b8599adaeb..3c20f6d10e 100644
--- a/proxy/src/auth/backend.rs
+++ b/proxy/src/auth/backend.rs
@@ -59,8 +59,8 @@ impl std::fmt::Display for BackendType<'_, ()> {
     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         use BackendType::*;
         match self {
-            Console(endpoint, _) => fmt.debug_tuple("Console").field(&endpoint.url()).finish(),
-            Postgres(endpoint, _) => fmt.debug_tuple("Postgres").field(&endpoint.url()).finish(),
+            Console(api, _) => fmt.debug_tuple("Console").field(&api.url()).finish(),
+            Postgres(api, _) => fmt.debug_tuple("Postgres").field(&api.url()).finish(),
             Link(url) => fmt.debug_tuple("Link").field(&url.as_str()).finish(),
         }
     }
diff --git a/proxy/src/auth/backend/classic.rs b/proxy/src/auth/backend/classic.rs
index 6753e7ed7f..51d73ffcad 100644
--- a/proxy/src/auth/backend/classic.rs
+++ b/proxy/src/auth/backend/classic.rs
@@ -7,8 +7,36 @@ use crate::{
     stream::PqStream,
 };
 use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_postgres::config::AuthKeys;
 use tracing::info;
 
+async fn do_scram(
+    secret: scram::ServerSecret,
+    creds: &ClientCredentials<'_>,
+    client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+) -> auth::Result<compute::ScramKeys> {
+    let outcome = AuthFlow::new(client)
+        .begin(auth::Scram(&secret))
+        .await?
+        .authenticate()
+        .await?;
+
+    let client_key = match outcome {
+        sasl::Outcome::Success(key) => key,
+        sasl::Outcome::Failure(reason) => {
+            info!("auth backend failed with an error: {reason}");
+            return Err(auth::AuthError::auth_failed(creds.user));
+        }
+    };
+
+    let keys = compute::ScramKeys {
+        client_key: client_key.as_bytes(),
+        server_key: secret.server_key.as_bytes(),
+    };
+
+    Ok(keys)
+}
+
 pub(super) async fn authenticate(
     api: &impl console::Api,
     extra: &ConsoleReqExtra<'_>,
@@ -24,7 +52,6 @@ pub(super) async fn authenticate(
         AuthInfo::Scram(scram::ServerSecret::mock(creds.user, rand::random()))
     });
 
-    let flow = AuthFlow::new(client);
     let scram_keys = match info {
         AuthInfo::Md5(_) => {
             info!("auth endpoint chooses MD5");
@@ -32,27 +59,12 @@ pub(super) async fn authenticate(
         }
         AuthInfo::Scram(secret) => {
             info!("auth endpoint chooses SCRAM");
-            let scram = auth::Scram(&secret);
-            let client_key = match flow.begin(scram).await?.authenticate().await? {
-                sasl::Outcome::Success(key) => key,
-                sasl::Outcome::Failure(reason) => {
-                    info!("auth backend failed with an error: {reason}");
-                    return Err(auth::AuthError::auth_failed(creds.user));
-                }
-            };
-
-            Some(compute::ScramKeys {
-                client_key: client_key.as_bytes(),
-                server_key: secret.server_key.as_bytes(),
-            })
+            do_scram(secret, creds, client).await?
         }
     };
 
     let mut node = api.wake_compute(extra, creds).await?;
-    if let Some(keys) = scram_keys {
-        use tokio_postgres::config::AuthKeys;
-        node.config.auth_keys(AuthKeys::ScramSha256(keys));
-    }
+    node.config.auth_keys(AuthKeys::ScramSha256(scram_keys));
 
     Ok(AuthSuccess {
         reported_auth_ok: false,