From 08c7f933a307dc9852b86fa5de6ea9f457972817 Mon Sep 17 00:00:00 2001 From: Conrad Ludgate Date: Wed, 28 Aug 2024 14:30:26 +0100 Subject: [PATCH] add support for console backend jwt --- proxy/src/auth/backend.rs | 3 ++- proxy/src/auth/backend/local.rs | 4 +--- proxy/src/bin/local_proxy.rs | 6 +++++- proxy/src/bin/proxy.rs | 2 ++ proxy/src/config.rs | 6 +++++- proxy/src/serverless/backend.rs | 23 +++++++++++++++++++---- proxy/src/serverless/sql_over_http.rs | 7 ++++++- 7 files changed, 40 insertions(+), 11 deletions(-) diff --git a/proxy/src/auth/backend.rs b/proxy/src/auth/backend.rs index 5dbfa5cc09..24ac5d99b0 100644 --- a/proxy/src/auth/backend.rs +++ b/proxy/src/auth/backend.rs @@ -565,7 +565,7 @@ mod tests { stream::{PqStream, Stream}, }; - use super::{auth_quirks, AuthRateLimiter}; + use super::{auth_quirks, jwt::JwkCache, AuthRateLimiter}; struct Auth { ips: Vec, @@ -611,6 +611,7 @@ mod tests { } static CONFIG: Lazy = Lazy::new(|| AuthenticationConfig { + jwks_cache: JwkCache::default(), thread_pool: ThreadPool::new(1), scram_protocol_timeout: std::time::Duration::from_secs(5), rate_limiter_enabled: true, diff --git a/proxy/src/auth/backend/local.rs b/proxy/src/auth/backend/local.rs index 2ab53f2c6a..f56b0a0a6d 100644 --- a/proxy/src/auth/backend/local.rs +++ b/proxy/src/auth/backend/local.rs @@ -14,17 +14,15 @@ use crate::{ EndpointId, }; -use super::jwt::{AuthRule, FetchAuthRules, JwkCache}; +use super::jwt::{AuthRule, FetchAuthRules}; pub struct LocalBackend { - pub(crate) jwks_cache: JwkCache, pub(crate) node_info: NodeInfo, } impl LocalBackend { pub fn new(postgres_addr: SocketAddr) -> Self { LocalBackend { - jwks_cache: JwkCache::default(), node_info: NodeInfo { config: { let mut cfg = ConnCfg::new(); diff --git a/proxy/src/bin/local_proxy.rs b/proxy/src/bin/local_proxy.rs index 1b3f465686..06505d29d6 100644 --- a/proxy/src/bin/local_proxy.rs +++ b/proxy/src/bin/local_proxy.rs @@ -6,7 +6,10 @@ use compute_api::spec::LocalProxySpec; use dashmap::DashMap; use futures::future::Either; use proxy::{ - auth::backend::local::{LocalBackend, JWKS_ROLE_MAP}, + auth::backend::{ + jwt::JwkCache, + local::{LocalBackend, JWKS_ROLE_MAP}, + }, cancellation::CancellationHandlerMain, config::{self, AuthenticationConfig, HttpConfig, ProxyConfig, RetryConfig}, console::{ @@ -267,6 +270,7 @@ fn build_config(args: &LocalProxyCliArgs) -> anyhow::Result<&'static ProxyConfig allow_self_signed_compute: false, http_config, authentication_config: AuthenticationConfig { + jwks_cache: JwkCache::default(), thread_pool: ThreadPool::new(0), scram_protocol_timeout: Duration::from_secs(10), rate_limiter_enabled: false, diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index 141005788d..3cca754290 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -8,6 +8,7 @@ use aws_config::web_identity_token::WebIdentityTokenCredentialsProvider; use aws_config::Region; use futures::future::Either; use proxy::auth; +use proxy::auth::backend::jwt::JwkCache; use proxy::auth::backend::AuthRateLimiter; use proxy::auth::backend::MaybeOwned; use proxy::cancellation::CancelMap; @@ -689,6 +690,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> { max_response_size_bytes: args.sql_over_http.sql_over_http_max_response_size_bytes, }; let authentication_config = AuthenticationConfig { + jwks_cache: JwkCache::default(), thread_pool, scram_protocol_timeout: args.scram_protocol_timeout, rate_limiter_enabled: args.auth_rate_limit_enabled, diff --git a/proxy/src/config.rs b/proxy/src/config.rs index 373e4cf650..3b044e2f5d 100644 --- a/proxy/src/config.rs +++ b/proxy/src/config.rs @@ -1,5 +1,8 @@ use crate::{ - auth::{self, backend::AuthRateLimiter}, + auth::{ + self, + backend::{jwt::JwkCache, AuthRateLimiter}, + }, console::locks::ApiLocks, rate_limiter::{RateBucketInfo, RateLimitAlgorithm, RateLimiterConfig}, scram::threadpool::ThreadPool, @@ -67,6 +70,7 @@ pub struct AuthenticationConfig { pub rate_limiter: AuthRateLimiter, pub rate_limit_ip_subnet: u8, pub ip_allowlist_check_enabled: bool, + pub jwks_cache: JwkCache, } impl TlsConfig { diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index 607eb0caf6..65bfdd5fc8 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -103,18 +103,33 @@ impl PoolingBackend { pub(crate) async fn authenticate_with_jwt( &self, ctx: &RequestMonitoring, + config: &AuthenticationConfig, user_info: &ComputeUserInfo, jwt: &str, ) -> Result { match &self.config.auth_backend { - crate::auth::Backend::Console(_, ()) => { - Err(AuthError::auth_failed("JWT login is not yet supported")) + crate::auth::Backend::Console(console, ()) => { + config + .jwks_cache + .check_jwt( + ctx, + user_info.endpoint.clone(), + &user_info.user, + &**console, + jwt, + ) + .await + .map_err(|e| AuthError::auth_failed(e.to_string()))?; + Ok(ComputeCredentials { + info: user_info.clone(), + keys: crate::auth::backend::ComputeCredentialKeys::None, + }) } crate::auth::Backend::Web(_, ()) => Err(AuthError::auth_failed( "JWT login over web auth proxy is not supported", )), - crate::auth::Backend::Local(cache) => { - cache + crate::auth::Backend::Local(_) => { + config .jwks_cache .check_jwt( ctx, diff --git a/proxy/src/serverless/sql_over_http.rs b/proxy/src/serverless/sql_over_http.rs index 7c78439a0a..1dee34671d 100644 --- a/proxy/src/serverless/sql_over_http.rs +++ b/proxy/src/serverless/sql_over_http.rs @@ -576,7 +576,12 @@ async fn handle_inner( } AuthData::Jwt(jwt) => { backend - .authenticate_with_jwt(ctx, &conn_info.conn_info.user_info, jwt) + .authenticate_with_jwt( + ctx, + &config.authentication_config, + &conn_info.conn_info.user_info, + jwt, + ) .await? } };