From 1e90e792d663b921260797a61d8b0849a5a956b8 Mon Sep 17 00:00:00 2001
From: Folke Behrens <folke@neon.tech>
Date: Wed, 2 Oct 2024 12:10:56 +0200
Subject: [PATCH] proxy: Add timeout to webauth confirmation wait (#9227)

```shell
$ cargo run -p proxy --bin proxy -- --auth-backend=web --webauth-confirmation-timeout=5s
```

```
$ psql -h localhost -p 4432
NOTICE:  Welcome to Neon!
Authenticate by visiting within 5s:
    http://localhost:3000/psql_session/e946900c8a9bc6e9


psql: error: connection to server at "localhost" (::1), port 4432 failed: Connection refused
	Is the server running on that host and accepting TCP/IP connections?
connection to server at "localhost" (127.0.0.1), port 4432 failed: ERROR:  Disconnected due to inactivity after 5s.
```
---
 proxy/src/auth.rs             | 9 +++++++++
 proxy/src/auth/backend.rs     | 1 +
 proxy/src/auth/backend/web.rs | 7 ++++++-
 proxy/src/bin/local_proxy.rs  | 1 +
 proxy/src/bin/proxy.rs        | 5 +++++
 proxy/src/config.rs           | 1 +
 6 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/proxy/src/auth.rs b/proxy/src/auth.rs
index 7c408f817c..13639af3aa 100644
--- a/proxy/src/auth.rs
+++ b/proxy/src/auth.rs
@@ -73,6 +73,9 @@ pub(crate) enum AuthErrorImpl {
 
     #[error("Authentication timed out")]
     UserTimeout(Elapsed),
+
+    #[error("Disconnected due to inactivity after {0}.")]
+    ConfirmationTimeout(humantime::Duration),
 }
 
 #[derive(Debug, Error)]
@@ -103,6 +106,10 @@ impl AuthError {
     pub(crate) fn user_timeout(elapsed: Elapsed) -> Self {
         AuthErrorImpl::UserTimeout(elapsed).into()
     }
+
+    pub(crate) fn confirmation_timeout(timeout: humantime::Duration) -> Self {
+        AuthErrorImpl::ConfirmationTimeout(timeout).into()
+    }
 }
 
 impl<E: Into<AuthErrorImpl>> From<E> for AuthError {
@@ -125,6 +132,7 @@ impl UserFacingError for AuthError {
             AuthErrorImpl::IpAddressNotAllowed(_) => self.to_string(),
             AuthErrorImpl::TooManyConnections => self.to_string(),
             AuthErrorImpl::UserTimeout(_) => self.to_string(),
+            AuthErrorImpl::ConfirmationTimeout(_) => self.to_string(),
         }
     }
 }
@@ -143,6 +151,7 @@ impl ReportableError for AuthError {
             AuthErrorImpl::IpAddressNotAllowed(_) => crate::error::ErrorKind::User,
             AuthErrorImpl::TooManyConnections => crate::error::ErrorKind::RateLimit,
             AuthErrorImpl::UserTimeout(_) => crate::error::ErrorKind::User,
+            AuthErrorImpl::ConfirmationTimeout(_) => crate::error::ErrorKind::User,
         }
     }
 }
diff --git a/proxy/src/auth/backend.rs b/proxy/src/auth/backend.rs
index 52ddfd90fb..0eeed27fb2 100644
--- a/proxy/src/auth/backend.rs
+++ b/proxy/src/auth/backend.rs
@@ -620,6 +620,7 @@ mod tests {
         ip_allowlist_check_enabled: true,
         is_auth_broker: false,
         accept_jwts: false,
+        webauth_confirmation_timeout: std::time::Duration::from_secs(5),
     });
 
     async fn read_message(r: &mut (impl AsyncRead + Unpin), b: &mut BytesMut) -> PgMessage {
diff --git a/proxy/src/auth/backend/web.rs b/proxy/src/auth/backend/web.rs
index 05f437355e..45710d244d 100644
--- a/proxy/src/auth/backend/web.rs
+++ b/proxy/src/auth/backend/web.rs
@@ -89,7 +89,12 @@ pub(super) async fn authenticate(
 
     // Wait for web console response (see `mgmt`).
     info!(parent: &span, "waiting for console's reply...");
-    let db_info = waiter.await.map_err(WebAuthError::from)?;
+    let db_info = tokio::time::timeout(auth_config.webauth_confirmation_timeout, waiter)
+        .await
+        .map_err(|_elapsed| {
+            auth::AuthError::confirmation_timeout(auth_config.webauth_confirmation_timeout.into())
+        })?
+        .map_err(WebAuthError::from)?;
 
     if auth_config.ip_allowlist_check_enabled {
         if let Some(allowed_ips) = &db_info.allowed_ips {
diff --git a/proxy/src/bin/local_proxy.rs b/proxy/src/bin/local_proxy.rs
index 49887576c7..b18810adbe 100644
--- a/proxy/src/bin/local_proxy.rs
+++ b/proxy/src/bin/local_proxy.rs
@@ -279,6 +279,7 @@ fn build_config(args: &LocalProxyCliArgs) -> anyhow::Result<&'static ProxyConfig
             ip_allowlist_check_enabled: true,
             is_auth_broker: false,
             accept_jwts: true,
+            webauth_confirmation_timeout: Duration::ZERO,
         },
         proxy_protocol_v2: config::ProxyProtocolV2::Rejected,
         handshake_timeout: Duration::from_secs(10),
diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs
index fa4fb264f2..0585902c3b 100644
--- a/proxy/src/bin/proxy.rs
+++ b/proxy/src/bin/proxy.rs
@@ -236,6 +236,10 @@ struct ProxyCliArgs {
     // TODO(conradludgate): switch default to rejected or required once we've updated all deployments
     #[clap(value_enum, long, default_value_t = ProxyProtocolV2::Supported)]
     proxy_protocol_v2: ProxyProtocolV2,
+
+    /// Time the proxy waits for the webauth session to be confirmed by the control plane.
+    #[clap(long, default_value = "2m", value_parser = humantime::parse_duration)]
+    webauth_confirmation_timeout: std::time::Duration,
 }
 
 #[derive(clap::Args, Clone, Copy, Debug)]
@@ -719,6 +723,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> {
         ip_allowlist_check_enabled: !args.is_private_access_proxy,
         is_auth_broker: args.is_auth_broker,
         accept_jwts: args.is_auth_broker,
+        webauth_confirmation_timeout: args.webauth_confirmation_timeout,
     };
 
     let config = Box::leak(Box::new(ProxyConfig {
diff --git a/proxy/src/config.rs b/proxy/src/config.rs
index 7d86ef4348..e0d666adf7 100644
--- a/proxy/src/config.rs
+++ b/proxy/src/config.rs
@@ -84,6 +84,7 @@ pub struct AuthenticationConfig {
     pub jwks_cache: JwkCache,
     pub is_auth_broker: bool,
     pub accept_jwts: bool,
+    pub webauth_confirmation_timeout: tokio::time::Duration,
 }
 
 impl TlsConfig {