From b2c96047d01f4be7935f421081e764e2e3cfbcb5 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Wed, 25 Oct 2023 08:30:47 +0100 Subject: [PATCH] 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 | 60 ++++++++++++++++++++++++++++--- proxy/src/auth/backend/classic.rs | 42 +++++----------------- proxy/src/auth/backend/hacks.rs | 25 +++---------- 3 files changed, 69 insertions(+), 58 deletions(-) diff --git a/proxy/src/auth/backend.rs b/proxy/src/auth/backend.rs index aa19acf65a..850d3a1837 100644 --- a/proxy/src/auth/backend.rs +++ b/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 { @@ -117,22 +120,27 @@ impl<'a, T, E> BackendType<'a, Result> { } } +pub enum ComputeCredentials { + Password(Vec), + 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, allow_cleartext: bool, config: &'static AuthenticationConfig, -) -> auth::Result> { +) -> auth::Result> { // 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, + allow_cleartext: bool, + config: &'static AuthenticationConfig, +) -> auth::Result> { + 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 { diff --git a/proxy/src/auth/backend/classic.rs b/proxy/src/auth/backend/classic.rs index af628efc5e..e7013a4cac 100644 --- a/proxy/src/auth/backend/classic.rs +++ b/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, config: &'static AuthenticationConfig, -) -> auth::Result> { +) -> auth::Result> { 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, + )), }) } diff --git a/proxy/src/auth/backend/hacks.rs b/proxy/src/auth/backend/hacks.rs index dcc93ec04c..53958bf337 100644 --- a/proxy/src/auth/backend/hacks.rs +++ b/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, -) -> auth::Result> { +) -> auth::Result> { 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, -) -> auth::Result> { +) -> auth::Result> { 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), }) }