From 76515cdae3cebef79523efd55442c4a6bf32658d Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Tue, 17 Sep 2024 14:32:59 +0100
Subject: [PATCH] split out auth info from conn info, return the jwt as the
 auth keys

---
 proxy/src/auth/backend.rs             |  2 ++
 proxy/src/console/provider.rs         |  1 +
 proxy/src/serverless/backend.rs       | 11 ++++++-----
 proxy/src/serverless/sql_over_http.rs |  4 ++--
 4 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/proxy/src/auth/backend.rs b/proxy/src/auth/backend.rs
index 24ac5d99b0..afa0af336b 100644
--- a/proxy/src/auth/backend.rs
+++ b/proxy/src/auth/backend.rs
@@ -171,6 +171,8 @@ impl ComputeUserInfo {
 }
 
 pub(crate) enum ComputeCredentialKeys {
+    #[expect(dead_code, reason = "WIP")]
+    Jwt(String),
     #[cfg(any(test, feature = "testing"))]
     Password(Vec<u8>),
     AuthKeys(AuthKeys),
diff --git a/proxy/src/console/provider.rs b/proxy/src/console/provider.rs
index 95097f2de9..5fbcf6c4e2 100644
--- a/proxy/src/console/provider.rs
+++ b/proxy/src/console/provider.rs
@@ -309,6 +309,7 @@ impl NodeInfo {
             #[cfg(any(test, feature = "testing"))]
             ComputeCredentialKeys::Password(password) => self.config.password(password),
             ComputeCredentialKeys::AuthKeys(auth_keys) => self.config.auth_keys(*auth_keys),
+            ComputeCredentialKeys::Jwt(_) => panic!("unsupported keys"),
             ComputeCredentialKeys::None => &mut self.config,
         };
     }
diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs
index 65bfdd5fc8..995b7a7cda 100644
--- a/proxy/src/serverless/backend.rs
+++ b/proxy/src/serverless/backend.rs
@@ -1,7 +1,7 @@
 use std::{sync::Arc, time::Duration};
 
 use async_trait::async_trait;
-use tracing::{field::display, info};
+use tracing::{debug, field::display, info};
 
 use crate::{
     auth::{
@@ -105,7 +105,7 @@ impl PoolingBackend {
         ctx: &RequestMonitoring,
         config: &AuthenticationConfig,
         user_info: &ComputeUserInfo,
-        jwt: &str,
+        jwt: String,
     ) -> Result<ComputeCredentials, AuthError> {
         match &self.config.auth_backend {
             crate::auth::Backend::Console(console, ()) => {
@@ -116,13 +116,13 @@ impl PoolingBackend {
                         user_info.endpoint.clone(),
                         &user_info.user,
                         &**console,
-                        jwt,
+                        &jwt,
                     )
                     .await
                     .map_err(|e| AuthError::auth_failed(e.to_string()))?;
                 Ok(ComputeCredentials {
                     info: user_info.clone(),
-                    keys: crate::auth::backend::ComputeCredentialKeys::None,
+                    keys: crate::auth::backend::ComputeCredentialKeys::Jwt(jwt),
                 })
             }
             crate::auth::Backend::Web(_, ()) => Err(AuthError::auth_failed(
@@ -136,12 +136,13 @@ impl PoolingBackend {
                         user_info.endpoint.clone(),
                         &user_info.user,
                         &StaticAuthRules,
-                        jwt,
+                        &jwt,
                     )
                     .await
                     .map_err(|e| AuthError::auth_failed(e.to_string()))?;
                 Ok(ComputeCredentials {
                     info: user_info.clone(),
+                    // todo: rewrite JWT signature with key shared somehow between local proxy and postgres
                     keys: crate::auth::backend::ComputeCredentialKeys::None,
                 })
             }
diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs
index 1dee34671d..97e280d252 100644
--- a/proxy/src/serverless/sql_over_http.rs
+++ b/proxy/src/serverless/sql_over_http.rs
@@ -563,14 +563,14 @@ async fn handle_inner(
 
     let authenticate_and_connect = Box::pin(
         async {
-            let keys = match &conn_info.auth {
+            let keys = match conn_info.auth {
                 AuthData::Password(pw) => {
                     backend
                         .authenticate_with_password(
                             ctx,
                             &config.authentication_config,
                             &conn_info.conn_info.user_info,
-                            pw,
+                            &pw,
                         )
                         .await?
                 }