Set up timeout for scram protocol execution (#5551)

## Problem
Context:
https://github.com/neondatabase/neon/issues/5511#issuecomment-1759649679

Some of out scram protocol execution timed out only after 17 minutes. 
## Summary of changes
Make timeout for scram execution meaningful and configurable.

proxy/src/auth/backend.rs

@@ -6,6 +6,7 @@ pub use link::LinkAuthError;
 
 use crate::{
     auth::{self, ClientCredentials},
+    config::AuthenticationConfig,
     console::{
         self,
         provider::{CachedNodeInfo, ConsoleReqExtra},
@@ -124,6 +125,7 @@ async fn auth_quirks(
     creds: &mut ClientCredentials<'_>,
     client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
     allow_cleartext: bool,
+    config: &'static AuthenticationConfig,
 ) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
     // If there's no project so far, that entails that client doesn't
     // support SNI or other means of passing the endpoint (project) name.
@@ -145,7 +147,7 @@ async fn auth_quirks(
     }
 
     // Finally, proceed with the main auth flow (SCRAM-based).
-    classic::authenticate(api, extra, creds, client).await
+    classic::authenticate(api, extra, creds, client, config).await
 }
 
 impl BackendType<'_, ClientCredentials<'_>> {
@@ -180,6 +182,7 @@ impl BackendType<'_, ClientCredentials<'_>> {
         extra: &ConsoleReqExtra<'_>,
         client: &mut stream::PqStream<impl AsyncRead + AsyncWrite + Unpin>,
         allow_cleartext: bool,
+        config: &'static AuthenticationConfig,
     ) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
         use BackendType::*;
 
@@ -192,7 +195,7 @@ impl BackendType<'_, ClientCredentials<'_>> {
                 );
 
                 let api = api.as_ref();
-                auth_quirks(api, extra, creds, client, allow_cleartext).await?
+                auth_quirks(api, extra, creds, client, allow_cleartext, config).await?
             }
             Postgres(api, creds) => {
                 info!(
@@ -202,7 +205,7 @@ impl BackendType<'_, ClientCredentials<'_>> {
                 );
 
                 let api = api.as_ref();
-                auth_quirks(api, extra, creds, client, allow_cleartext).await?
+                auth_quirks(api, extra, creds, client, allow_cleartext, config).await?
             }
             // NOTE: this auth backend doesn't use client credentials.
             Link(url) => {

proxy/src/auth/backend/classic.rs

@@ -4,6 +4,7 @@ use super::AuthSuccess;
 use crate::{
     auth::{self, AuthFlow, ClientCredentials},
     compute,
+    config::AuthenticationConfig,
     console::{self, AuthInfo, CachedNodeInfo, ConsoleReqExtra},
     proxy::{handle_try_wake, retry_after},
     sasl, scram,
@@ -17,6 +18,7 @@ pub(super) async fn authenticate(
     extra: &ConsoleReqExtra<'_>,
     creds: &ClientCredentials<'_>,
     client: &mut PqStream<impl AsyncRead + AsyncWrite + Unpin>,
+    config: &'static AuthenticationConfig,
 ) -> auth::Result<AuthSuccess<CachedNodeInfo>> {
     info!("fetching user's authentication info");
     let info = api.get_auth_info(extra, creds).await?.unwrap_or_else(|| {
@@ -42,7 +44,16 @@ pub(super) async fn authenticate(
                 error
             })?;
 
-            let auth_outcome = auth_flow.authenticate().await.map_err(|error| {
+            let auth_outcome = tokio::time::timeout(
+                config.scram_protocol_timeout,
+                auth_flow.authenticate(),
+            )
+            .await
+            .map_err(|error| {
+                warn!("error processing scram messages error = authentication timed out, execution time exeeded {} seconds", config.scram_protocol_timeout.as_secs());
+                auth::io::Error::new(auth::io::ErrorKind::TimedOut, error)
+            })?
+            .map_err(|error| {
                 warn!(?error, "error processing scram messages");
                 error
             })?;

proxy/src/bin/proxy.rs

@@ -1,5 +1,6 @@
 use futures::future::Either;
 use proxy::auth;
+use proxy::config::AuthenticationConfig;
 use proxy::config::HttpConfig;
 use proxy::console;
 use proxy::http;
@@ -83,7 +84,9 @@ struct ProxyCliArgs {
     /// timeout for http connections
     #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
     sql_over_http_timeout: tokio::time::Duration,
-
+    /// timeout for scram authentication protocol
+    #[clap(long, default_value = "15s", value_parser = humantime::parse_duration)]
+    scram_protocol_timeout: tokio::time::Duration,
     /// Require that all incoming requests have a Proxy Protocol V2 packet **and** have an IP address associated.
     #[clap(long, default_value_t = false, value_parser = clap::builder::BoolishValueParser::new(), action = clap::ArgAction::Set)]
     require_client_ip: bool,
@@ -231,12 +234,16 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
     let http_config = HttpConfig {
         sql_over_http_timeout: args.sql_over_http_timeout,
     };
+    let authentication_config = AuthenticationConfig {
+        scram_protocol_timeout: args.scram_protocol_timeout,
+    };
     let config = Box::leak(Box::new(ProxyConfig {
         tls_config,
         auth_backend,
         metric_collection,
         allow_self_signed_compute: args.allow_self_signed_compute,
         http_config,
+        authentication_config,
         require_client_ip: args.require_client_ip,
     }));
 

proxy/src/config.rs

@@ -14,6 +14,7 @@ pub struct ProxyConfig {
     pub metric_collection: Option<MetricCollectionConfig>,
     pub allow_self_signed_compute: bool,
     pub http_config: HttpConfig,
+    pub authentication_config: AuthenticationConfig,
     pub require_client_ip: bool,
 }
 
@@ -32,6 +33,10 @@ pub struct HttpConfig {
     pub sql_over_http_timeout: tokio::time::Duration,
 }
 
+pub struct AuthenticationConfig {
+    pub scram_protocol_timeout: tokio::time::Duration,
+}
+
 impl TlsConfig {
     pub fn to_server_config(&self) -> Arc<rustls::ServerConfig> {
         self.config.clone()

proxy/src/proxy.rs

@@ -5,7 +5,7 @@ use crate::{
     auth::{self, backend::AuthSuccess},
     cancellation::{self, CancelMap},
     compute::{self, PostgresConnection},
-    config::{ProxyConfig, TlsConfig},
+    config::{AuthenticationConfig, ProxyConfig, TlsConfig},
     console::{self, errors::WakeComputeError, messages::MetricsAuxInfo, Api},
     http::StatusCode,
     metrics::{Ids, USAGE_METRICS},
@@ -340,7 +340,7 @@ pub async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
         mode.allow_self_signed_compute(config),
     );
     cancel_map
-        .with_session(|session| client.connect_to_db(session, mode))
+        .with_session(|session| client.connect_to_db(session, mode, &config.authentication_config))
         .await
 }
 
@@ -818,6 +818,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> Client<'_, S> {
         self,
         session: cancellation::Session<'_>,
         mode: ClientMode,
+        config: &'static AuthenticationConfig,
     ) -> anyhow::Result<()> {
         let Self {
             mut stream,
@@ -835,7 +836,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin> Client<'_, S> {
         let latency_timer = LatencyTimer::new(mode.protocol_label());
 
         let auth_result = match creds
-            .authenticate(&extra, &mut stream, mode.allow_cleartext())
+            .authenticate(&extra, &mut stream, mode.allow_cleartext(), config)
             .await
         {
             Ok(auth_result) => auth_result,