From 8b1ffa1718fc8444e12ab482e2187733b6377cf1 Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Thu, 29 May 2025 15:46:40 +0100
Subject: [PATCH] simplify cplane authentication

---
 proxy/src/auth/backend/mod.rs | 66 +++++++++++------------------------
 proxy/src/proxy/mod.rs        | 25 +++++++------
 2 files changed, 33 insertions(+), 58 deletions(-)

diff --git a/proxy/src/auth/backend/mod.rs b/proxy/src/auth/backend/mod.rs
index a211d0b65a..fe6886e31e 100644
--- a/proxy/src/auth/backend/mod.rs
+++ b/proxy/src/auth/backend/mod.rs
@@ -135,16 +135,6 @@ impl<'a, T> Backend<'a, T> {
         }
     }
 }
-impl<'a, T, E> Backend<'a, Result<T, E>> {
-    /// Very similar to [`std::option::Option::transpose`].
-    /// This is most useful for error handling.
-    pub(crate) fn transpose(self) -> Result<Backend<'a, T>, E> {
-        match self {
-            Self::ControlPlane(c, x) => x.map(|x| Backend::ControlPlane(c, x)),
-            Self::Local(l) => Ok(Backend::Local(l)),
-        }
-    }
-}
 
 pub(crate) struct ComputeCredentials {
     pub(crate) info: ComputeUserInfo,
@@ -415,53 +405,39 @@ async fn authenticate_with_secret(
     classic::authenticate(ctx, info, client, config, secret).await
 }
 
-impl<'a> Backend<'a, ComputeUserInfoMaybeEndpoint> {
-    /// Get username from the credentials.
-    pub(crate) fn get_user(&self) -> &str {
-        match self {
-            Self::ControlPlane(_, user_info) => &user_info.user,
-            Self::Local(_) => "local",
-        }
-    }
-
+impl ControlPlaneClient {
     /// Authenticate the client via the requested backend, possibly using credentials.
     #[tracing::instrument(fields(allow_cleartext = allow_cleartext), skip_all)]
     pub(crate) async fn authenticate(
-        self,
+        &self,
         ctx: &RequestContext,
         client: &mut stream::PqStream<Stream<impl AsyncRead + AsyncWrite + Unpin>>,
+        user_info: ComputeUserInfoMaybeEndpoint,
         allow_cleartext: bool,
         config: &'static AuthenticationConfig,
         endpoint_rate_limiter: Arc<EndpointRateLimiter>,
-    ) -> auth::Result<Backend<'a, ComputeCredentials>> {
-        let res = match self {
-            Self::ControlPlane(api, user_info) => {
-                debug!(
-                    user = &*user_info.user,
-                    project = user_info.endpoint(),
-                    "performing authentication using the console"
-                );
+    ) -> auth::Result<ComputeCredentials> {
+        debug!(
+            user = &*user_info.user,
+            project = user_info.endpoint(),
+            "performing authentication using the console"
+        );
 
-                let credentials = auth_quirks(
-                    ctx,
-                    &*api,
-                    user_info,
-                    client,
-                    allow_cleartext,
-                    config,
-                    endpoint_rate_limiter,
-                )
-                .await?;
-                Ok(Backend::ControlPlane(api, credentials))
-            }
-            Self::Local(_) => {
-                return Err(auth::AuthError::bad_auth_method("invalid for local proxy"));
-            }
-        };
+        let credentials = auth_quirks(
+            ctx,
+            self,
+            user_info,
+            client,
+            allow_cleartext,
+            config,
+            endpoint_rate_limiter,
+        )
+        .await?;
 
         // TODO: replace with some metric
         info!("user successfully authenticated");
-        res
+
+        Ok(credentials)
     }
 }
 
diff --git a/proxy/src/proxy/mod.rs b/proxy/src/proxy/mod.rs
index d177a92802..1f0db188f2 100644
--- a/proxy/src/proxy/mod.rs
+++ b/proxy/src/proxy/mod.rs
@@ -275,6 +275,11 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     conn_gauge: NumClientConnectionsGuard<'static>,
     tracker: TaskTrackerToken,
 ) -> Result<Option<ProxyPassthrough<S>>, ClientRequestError> {
+    let cplane = match auth_backend {
+        auth::Backend::ControlPlane(cplane, ()) => &**cplane,
+        auth::Backend::Local(_) => unreachable!("local proxy does not run tcp proxy service"),
+    };
+
     debug!(
         protocol = %ctx.protocol(),
         "handling interactive connection from client"
@@ -346,21 +351,18 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     let common_names = tls.map(|tls| &tls.common_names);
 
     // Extract credentials which we're going to use for auth.
-    let result = auth_backend
-        .as_ref()
-        .map(|()| auth::ComputeUserInfoMaybeEndpoint::parse(ctx, &params, hostname, common_names))
-        .transpose();
-
+    let result = auth::ComputeUserInfoMaybeEndpoint::parse(ctx, &params, hostname, common_names);
     let user_info = match result {
         Ok(user_info) => user_info,
         Err(e) => stream.throw_error(e, Some(ctx)).await?,
     };
 
-    let user = user_info.get_user().to_owned();
-    let user_info = match user_info
+    let user = user_info.user.clone();
+    let compute_creds = match cplane
         .authenticate(
             ctx,
             &mut stream,
+            user_info,
             mode.allow_cleartext(),
             &config.authentication_config,
             endpoint_rate_limiter,
@@ -380,10 +382,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
         }
     };
 
-    let compute_user_info = match &user_info {
-        auth::Backend::ControlPlane(_, info) => &info.info,
-        auth::Backend::Local(_) => unreachable!("local proxy does not run tcp proxy service"),
-    };
+    let compute_user_info = compute_creds.info.clone();
     let params_compat = compute_user_info
         .options
         .get(NeonOptions::PARAMS_COMPAT)
@@ -392,12 +391,12 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin + Send>(
     let mut node = connect_to_compute(
         ctx,
         &TcpMechanism {
-            user_info: compute_user_info.clone(),
+            user_info: compute_user_info,
             params_compat,
             params: &params,
             locks: &config.connect_compute_locks,
         },
-        &user_info,
+        &auth::Backend::ControlPlane(auth::backend::MaybeOwned::Borrowed(cplane), compute_creds),
         config.wake_compute_retry_config,
         &config.connect_to_compute,
     )