move wake compute after the auth quirks logic (#5642)

## Problem

https://github.com/neondatabase/neon/issues/5568#issuecomment-1777015606

## Summary of changes

Make the auth_quirks_creds return the authentication information, and
push the wake_compute loop to after, inside `auth_quirks`

proxy/src/auth/backend.rs

@@ -3,7 +3,9 @@ mod hacks;
 mod link;
 
 pub use link::LinkAuthError;
+use tokio_postgres::config::AuthKeys;
 
+use crate::proxy::{handle_try_wake, retry_after};
 use crate::{
     auth::{self, ClientCredentials},
     config::AuthenticationConfig,
@@ -16,8 +18,9 @@ use crate::{
 };
 use futures::TryFutureExt;
 use std::borrow::Cow;
+use std::ops::ControlFlow;
 use tokio::io::{AsyncRead, AsyncWrite};
-use tracing::info;
+use tracing::{error, info, warn};
 
 /// A product of successful authentication.
 pub struct AuthSuccess<T> {
@@ -117,22 +120,27 @@ impl<'a, T, E> BackendType<'a, Result<T, E>> {
     }
 }
 
+pub enum ComputeCredentials {
+    Password(Vec<u8>),
+    AuthKeys(AuthKeys),
+}
+
 /// True to its name, this function encapsulates our current auth trade-offs.
 /// Here, we choose the appropriate auth flow based on circumstances.
-async fn auth_quirks(
+async fn auth_quirks_creds(
     api: &impl console::Api,
     extra: &ConsoleReqExtra<'_>,
     creds: &mut ClientCredentials<'_>,
     client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     allow_cleartext: bool,
     config: &'static AuthenticationConfig,
-) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
+) -> auth::Result<AuthSuccess<ComputeCredentials>> {
     // If there's no project so far, that entails that client doesn't
     // support SNI or other means of passing the endpoint (project) name.
     // We now expect to see a very specific payload in the place of password.
     if creds.project.is_none() {
         // Password will be checked by the compute node later.
-        return hacks::password_hack(api, extra, creds, client).await;
+        return hacks::password_hack(creds, client).await;
     }
 
     // Password hack should set the project name.
@@ -143,13 +151,55 @@ async fn auth_quirks(
     // Currently, we use it for websocket connections (latency).
     if allow_cleartext {
         // Password will be checked by the compute node later.
-        return hacks::cleartext_hack(api, extra, creds, client).await;
+        return hacks::cleartext_hack(client).await;
     }
 
     // Finally, proceed with the main auth flow (SCRAM-based).
     classic::authenticate(api, extra, creds, client, config).await
 }
 
+/// True to its name, this function encapsulates our current auth trade-offs.
+/// Here, we choose the appropriate auth flow based on circumstances.
+async fn auth_quirks(
+    api: &impl console::Api,
+    extra: &ConsoleReqExtra<'_>,
+    creds: &mut ClientCredentials<'_>,
+    client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    allow_cleartext: bool,
+    config: &'static AuthenticationConfig,
+) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
+    let auth_stuff = auth_quirks_creds(api, extra, creds, client, allow_cleartext, config).await?;
+
+    let mut num_retries = 0;
+    let mut node = loop {
+        let wake_res = api.wake_compute(extra, creds).await;
+        match handle_try_wake(wake_res, num_retries) {
+            Err(e) => {
+                error!(error = ?e, num_retries, retriable = false, "couldn't wake compute node");
+                return Err(e.into());
+            }
+            Ok(ControlFlow::Continue(e)) => {
+                warn!(error = ?e, num_retries, retriable = true, "couldn't wake compute node");
+            }
+            Ok(ControlFlow::Break(n)) => break n,
+        }
+
+        let wait_duration = retry_after(num_retries);
+        num_retries += 1;
+        tokio::time::sleep(wait_duration).await;
+    };
+
+    match auth_stuff.value {
+        ComputeCredentials::Password(password) => node.config.password(password),
+        ComputeCredentials::AuthKeys(auth_keys) => node.config.auth_keys(auth_keys),
+    };
+
+    Ok(AuthSuccess {
+        reported_auth_ok: auth_stuff.reported_auth_ok,
+        value: node,
+    })
+}
+
 impl BackendType<'_, ClientCredentials<'_>> {
     /// Get compute endpoint name from the credentials.
     pub fn get_endpoint(&self) -> Option<String> {

proxy/src/auth/backend/classic.rs

@@ -1,17 +1,14 @@
-use std::ops::ControlFlow;
-
-use super::AuthSuccess;
+use super::{AuthSuccess, ComputeCredentials};
 use crate::{
     auth::{self, AuthFlow, ClientCredentials},
     compute,
     config::AuthenticationConfig,
-    console::{self, AuthInfo, CachedNodeInfo, ConsoleReqExtra},
-    proxy::{handle_try_wake, retry_after},
+    console::{self, AuthInfo, ConsoleReqExtra},
     sasl, scram,
     stream::PqStream,
 };
 use tokio::io::{AsyncRead, AsyncWrite};
-use tracing::{error, info, warn};
+use tracing::{info, warn};
 
 pub(super) async fn authenticate(
     api: &impl console::Api,
@@ -19,7 +16,7 @@ pub(super) async fn authenticate(
     creds: &ClientCredentials<'_>,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     config: &'static AuthenticationConfig,
-) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
+) -> auth::Result<AuthSuccess<ComputeCredentials>> {
     info!("fetching user's authentication info");
     let info = api.get_auth_info(extra, creds).await?.unwrap_or_else(|| {
         // If we don't have an authentication secret, we mock one to
@@ -66,38 +63,17 @@ pub(super) async fn authenticate(
                 }
             };
 
-            Some(compute::ScramKeys {
+            compute::ScramKeys {
                 client_key: client_key.as_bytes(),
                 server_key: secret.server_key.as_bytes(),
-            })
+            }
         }
     };
 
-    let mut num_retries = 0;
-    let mut node = loop {
-        let wake_res = api.wake_compute(extra, creds).await;
-        match handle_try_wake(wake_res, num_retries) {
-            Err(e) => {
-                error!(error = ?e, num_retries, retriable = false, "couldn't wake compute node");
-                return Err(e.into());
-            }
-            Ok(ControlFlow::Continue(e)) => {
-                warn!(error = ?e, num_retries, retriable = true, "couldn't wake compute node");
-            }
-            Ok(ControlFlow::Break(n)) => break n,
-        }
-
-        let wait_duration = retry_after(num_retries);
-        num_retries += 1;
-        tokio::time::sleep(wait_duration).await;
-    };
-    if let Some(keys) = scram_keys {
-        use tokio_postgres::config::AuthKeys;
-        node.config.auth_keys(AuthKeys::ScramSha256(keys));
-    }
-
     Ok(AuthSuccess {
         reported_auth_ok: false,
-        value: node,
+        value: ComputeCredentials::AuthKeys(tokio_postgres::config::AuthKeys::ScramSha256(
+            scram_keys,
+        )),
     })
 }

proxy/src/auth/backend/hacks.rs

@@ -1,10 +1,6 @@
-use super::AuthSuccess;
+use super::{AuthSuccess, ComputeCredentials};
 use crate::{
     auth::{self, AuthFlow, ClientCredentials},
-    console::{
-        self,
-        provider::{CachedNodeInfo, ConsoleReqExtra},
-    },
     stream,
 };
 use tokio::io::{AsyncRead, AsyncWrite};
@@ -15,11 +11,8 @@ use tracing::{info, warn};
 /// These properties are benefical for serverless JS workers, so we
 /// use this mechanism for websocket connections.
 pub async fn cleartext_hack(
-    api: &impl console::Api,
-    extra: &ConsoleReqExtra<'_>,
-    creds: &mut ClientCredentials<'_>,
     client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
+) -> auth::Result<AuthSuccess<ComputeCredentials>> {
     warn!("cleartext auth flow override is enabled, proceeding");
     let password = AuthFlow::new(client)
         .begin(auth::CleartextPassword)
@@ -27,24 +20,19 @@ pub async fn cleartext_hack(
         .authenticate()
         .await?;
 
-    let mut node = api.wake_compute(extra, creds).await?;
-    node.config.password(password);
-
     // Report tentative success; compute node will check the password anyway.
     Ok(AuthSuccess {
         reported_auth_ok: false,
-        value: node,
+        value: ComputeCredentials::Password(password),
     })
 }
 
 /// Workaround for clients which don't provide an endpoint (project) name.
 /// Very similar to [`cleartext_hack`], but there's a specific password format.
 pub async fn password_hack(
-    api: &impl console::Api,
-    extra: &ConsoleReqExtra<'_>,
     creds: &mut ClientCredentials<'_>,
     client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
-) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
+) -> auth::Result<AuthSuccess<ComputeCredentials>> {
     warn!("project not specified, resorting to the password hack auth flow");
     let payload = AuthFlow::new(client)
         .begin(auth::PasswordHack)
@@ -55,12 +43,9 @@ pub async fn password_hack(
     info!(project = &payload.endpoint, "received missing parameter");
     creds.project = Some(payload.endpoint);
 
-    let mut node = api.wake_compute(extra, creds).await?;
-    node.config.password(payload.password);
-
     // Report tentative success; compute node will check the password anyway.
     Ok(AuthSuccess {
         reported_auth_ok: false,
-        value: node,
+        value: ComputeCredentials::Password(payload.password),
     })
 }