diff --git a/proxy/src/auth/backend/mod.rs b/proxy/src/auth/backend/mod.rs index e331c27fbb..fe8dd021af 100644 --- a/proxy/src/auth/backend/mod.rs +++ b/proxy/src/auth/backend/mod.rs @@ -20,8 +20,7 @@ use crate::auth::credentials::check_peer_addr_is_in_list; use crate::auth::{validate_password_and_exchange, AuthError}; use crate::cache::Cached; use crate::context::RequestMonitoring; -use crate::control_plane::errors::GetAuthInfoError; -use crate::control_plane::provider::{CachedRoleSecret, ControlPlaneBackend}; +use crate::control_plane::provider::ControlPlaneBackend; use crate::control_plane::AuthSecret; use crate::intern::EndpointIdInt; use crate::metrics::Metrics; @@ -32,47 +31,21 @@ use crate::stream::Stream; use crate::{ auth::{self, ComputeUserInfoMaybeEndpoint}, config::AuthenticationConfig, - control_plane::{ - self, - provider::{CachedAllowedIps, CachedNodeInfo}, - Api, - }, + control_plane::{self, provider::CachedNodeInfo, Api}, stream, }; use crate::{scram, EndpointCacheKey, EndpointId, RoleName}; -/// Alternative to [`std::borrow::Cow`] but doesn't need `T: ToOwned` as we don't need that functionality -pub enum MaybeOwned<'a, T> { - Owned(T), - Borrowed(&'a T), -} - -impl std::ops::Deref for MaybeOwned<'_, T> { - type Target = T; - - fn deref(&self) -> &Self::Target { - match self { - MaybeOwned::Owned(t) => t, - MaybeOwned::Borrowed(t) => t, - } - } -} - -/// This type serves two purposes: -/// -/// * When `T` is `()`, it's just a regular auth backend selector -/// which we use in [`crate::config::ProxyConfig`]. -/// -/// * However, when we substitute `T` with [`ComputeUserInfoMaybeEndpoint`], -/// this helps us provide the credentials only to those auth -/// backends which require them for the authentication process. -pub enum Backend<'a, T> { +pub enum Backend<'a> { /// Cloud API (V2). - ControlPlane(MaybeOwned<'a, ControlPlaneBackend>, T), + ControlPlane(&'a ControlPlaneBackend), /// Local proxy uses configured auth credentials and does not wake compute - Local(MaybeOwned<'a, LocalBackend>), + Local(&'a LocalBackend), } +#[cfg(test)] +use crate::control_plane::provider::{CachedAllowedIps, CachedRoleSecret}; + #[cfg(test)] pub(crate) trait TestBackend: Send + Sync + 'static { fn wake_compute(&self) -> Result; @@ -89,56 +62,20 @@ impl Clone for Box { } } -impl std::fmt::Display for Backend<'_, ()> { +impl std::fmt::Display for ControlPlaneBackend { fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Self::ControlPlane(api, ()) => match &**api { - ControlPlaneBackend::Management(endpoint) => fmt - .debug_tuple("ControlPlane::Management") - .field(&endpoint.url()) - .finish(), - #[cfg(any(test, feature = "testing"))] - ControlPlaneBackend::PostgresMock(endpoint) => fmt - .debug_tuple("ControlPlane::PostgresMock") - .field(&endpoint.url()) - .finish(), - #[cfg(test)] - ControlPlaneBackend::Test(_) => fmt.debug_tuple("ControlPlane::Test").finish(), - }, - Self::Local(_) => fmt.debug_tuple("Local").finish(), - } - } -} - -impl Backend<'_, T> { - /// Very similar to [`std::option::Option::as_ref`]. - /// This helps us pass structured config to async tasks. - pub(crate) fn as_ref(&self) -> Backend<'_, &T> { - match self { - Self::ControlPlane(c, x) => Backend::ControlPlane(MaybeOwned::Borrowed(c), x), - Self::Local(l) => Backend::Local(MaybeOwned::Borrowed(l)), - } - } -} - -impl<'a, T> Backend<'a, T> { - /// Very similar to [`std::option::Option::map`]. - /// Maps [`Backend`] to [`Backend`] by applying - /// a function to a contained value. - pub(crate) fn map(self, f: impl FnOnce(T) -> R) -> Backend<'a, R> { - match self { - Self::ControlPlane(c, x) => Backend::ControlPlane(c, f(x)), - Self::Local(l) => Backend::Local(l), - } - } -} -impl<'a, T, E> Backend<'a, Result> { - /// Very similar to [`std::option::Option::transpose`]. - /// This is most useful for error handling. - pub(crate) fn transpose(self) -> Result, E> { - match self { - Self::ControlPlane(c, x) => x.map(|x| Backend::ControlPlane(c, x)), - Self::Local(l) => Ok(Backend::Local(l)), + ControlPlaneBackend::Management(endpoint) => fmt + .debug_tuple("ControlPlane::Management") + .field(&endpoint.url()) + .finish(), + #[cfg(any(test, feature = "testing"))] + ControlPlaneBackend::PostgresMock(endpoint) => fmt + .debug_tuple("ControlPlane::PostgresMock") + .field(&endpoint.url()) + .finish(), + #[cfg(test)] + ControlPlaneBackend::Test(_) => fmt.debug_tuple("ControlPlane::Test").finish(), } } } @@ -403,96 +340,79 @@ 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", - } - } - - /// Authenticate the client via the requested backend, possibly using credentials. +impl ControlPlaneBackend { #[tracing::instrument(fields(allow_cleartext = allow_cleartext), skip_all)] pub(crate) async fn authenticate( - self, + &self, ctx: &RequestMonitoring, + user_info: ComputeUserInfoMaybeEndpoint, client: &mut stream::PqStream>, allow_cleartext: bool, config: &'static AuthenticationConfig, endpoint_rate_limiter: Arc, - ) -> auth::Result> { - let res = match self { - Self::ControlPlane(api, user_info) => { - info!( - user = &*user_info.user, - project = user_info.endpoint(), - "performing authentication using the console" - ); + ) -> auth::Result { + info!( + 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?; - 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?; info!("user successfully authenticated"); - Ok(res) + Ok(ControlPlaneComputeBackend { + api: self, + creds: credentials, + }) + } + + pub(crate) fn attach_to_credentials( + &self, + creds: ComputeCredentials, + ) -> ControlPlaneComputeBackend { + ControlPlaneComputeBackend { api: self, creds } } } -impl Backend<'_, ComputeUserInfo> { - pub(crate) async fn get_role_secret( - &self, - ctx: &RequestMonitoring, - ) -> Result { - match self { - Self::ControlPlane(api, user_info) => api.get_role_secret(ctx, user_info).await, - Self::Local(_) => Ok(Cached::new_uncached(None)), - } - } - - pub(crate) async fn get_allowed_ips_and_secret( - &self, - ctx: &RequestMonitoring, - ) -> Result<(CachedAllowedIps, Option), GetAuthInfoError> { - match self { - Self::ControlPlane(api, user_info) => { - api.get_allowed_ips_and_secret(ctx, user_info).await - } - Self::Local(_) => Ok((Cached::new_uncached(Arc::new(vec![])), None)), - } - } +pub struct ControlPlaneComputeBackend<'a> { + pub(crate) api: &'a ControlPlaneBackend, + pub(crate) creds: ComputeCredentials, } #[async_trait::async_trait] -impl ComputeConnectBackend for Backend<'_, ComputeCredentials> { +impl ComputeConnectBackend for ControlPlaneComputeBackend<'_> { async fn wake_compute( &self, ctx: &RequestMonitoring, ) -> Result { - match self { - Self::ControlPlane(api, creds) => api.wake_compute(ctx, &creds.info).await, - Self::Local(local) => Ok(Cached::new_uncached(local.node_info.clone())), - } + self.api.wake_compute(ctx, &self.creds.info).await } fn get_keys(&self) -> &ComputeCredentialKeys { - match self { - Self::ControlPlane(_, creds) => &creds.keys, - Self::Local(_) => &ComputeCredentialKeys::None, - } + &self.creds.keys + } +} + +#[async_trait::async_trait] +impl ComputeConnectBackend for LocalBackend { + async fn wake_compute( + &self, + _ctx: &RequestMonitoring, + ) -> Result { + Ok(Cached::new_uncached(self.node_info.clone())) + } + + fn get_keys(&self) -> &ComputeCredentialKeys { + &ComputeCredentialKeys::None } } diff --git a/proxy/src/bin/local_proxy.rs b/proxy/src/bin/local_proxy.rs index c92ebbc51f..0733c7c610 100644 --- a/proxy/src/bin/local_proxy.rs +++ b/proxy/src/bin/local_proxy.rs @@ -197,7 +197,7 @@ async fn main() -> anyhow::Result<()> { let task = serverless::task_main( config, - auth_backend, + auth::Backend::Local(auth_backend), http_listener, shutdown.clone(), Arc::new(CancellationHandlerMain::new( @@ -289,12 +289,8 @@ fn build_config(args: &LocalProxyCliArgs) -> anyhow::Result<&'static ProxyConfig } /// auth::Backend is created at proxy startup, and lives forever. -fn build_auth_backend( - args: &LocalProxyCliArgs, -) -> anyhow::Result<&'static auth::Backend<'static, ()>> { - let auth_backend = proxy::auth::Backend::Local(proxy::auth::backend::MaybeOwned::Owned( - LocalBackend::new(args.compute), - )); +fn build_auth_backend(args: &LocalProxyCliArgs) -> anyhow::Result<&'static LocalBackend> { + let auth_backend = LocalBackend::new(args.compute); Ok(Box::leak(Box::new(auth_backend))) } diff --git a/proxy/src/bin/proxy.rs b/proxy/src/bin/proxy.rs index 3c0e66dec3..f7b4466204 100644 --- a/proxy/src/bin/proxy.rs +++ b/proxy/src/bin/proxy.rs @@ -11,7 +11,6 @@ use proxy::auth; use proxy::auth::backend::jwt::JwkCache; use proxy::auth::backend::AuthRateLimiter; use proxy::auth::backend::ConsoleRedirectBackend; -use proxy::auth::backend::MaybeOwned; use proxy::cancellation::CancelMap; use proxy::cancellation::CancellationHandler; use proxy::config::remote_storage_from_toml; @@ -22,6 +21,7 @@ use proxy::config::ProjectInfoCacheOptions; use proxy::config::ProxyProtocolV2; use proxy::context::parquet::ParquetUploadArgs; use proxy::control_plane; +use proxy::control_plane::provider::ControlPlaneBackend; use proxy::http; use proxy::http::health_server::AppMetrics; use proxy::metrics::Metrics; @@ -480,7 +480,7 @@ async fn main() -> anyhow::Result<()> { if let Some(serverless_listener) = serverless_listener { client_tasks.spawn(serverless::task_main( config, - auth_backend, + auth::Backend::ControlPlane(auth_backend), serverless_listener, cancellation_token.clone(), cancellation_handler.clone(), @@ -528,40 +528,38 @@ async fn main() -> anyhow::Result<()> { )); } - if let Either::Left(auth::Backend::ControlPlane(api, _)) = &auth_backend { - if let proxy::control_plane::provider::ControlPlaneBackend::Management(api) = &**api { - match (redis_notifications_client, regional_redis_client.clone()) { - (None, None) => {} - (client1, client2) => { - let cache = api.caches.project_info.clone(); - if let Some(client) = client1 { - maintenance_tasks.spawn(notifications::task_main( - client, - cache.clone(), - cancel_map.clone(), - args.region.clone(), - )); - } - if let Some(client) = client2 { - maintenance_tasks.spawn(notifications::task_main( - client, - cache.clone(), - cancel_map.clone(), - args.region.clone(), - )); - } - maintenance_tasks.spawn(async move { cache.clone().gc_worker().await }); + if let Either::Left(ControlPlaneBackend::Management(api)) = &auth_backend { + match (redis_notifications_client, regional_redis_client.clone()) { + (None, None) => {} + (client1, client2) => { + let cache = api.caches.project_info.clone(); + if let Some(client) = client1 { + maintenance_tasks.spawn(notifications::task_main( + client, + cache.clone(), + cancel_map.clone(), + args.region.clone(), + )); } + if let Some(client) = client2 { + maintenance_tasks.spawn(notifications::task_main( + client, + cache.clone(), + cancel_map.clone(), + args.region.clone(), + )); + } + maintenance_tasks.spawn(async move { cache.clone().gc_worker().await }); } - if let Some(regional_redis_client) = regional_redis_client { - let cache = api.caches.endpoints_cache.clone(); - let con = regional_redis_client; - let span = tracing::info_span!("endpoints_cache"); - maintenance_tasks.spawn( - async move { cache.do_read(con, cancellation_token.clone()).await } - .instrument(span), - ); - } + } + if let Some(regional_redis_client) = regional_redis_client { + let cache = api.caches.endpoints_cache.clone(); + let con = regional_redis_client; + let span = tracing::info_span!("endpoints_cache"); + maintenance_tasks.spawn( + async move { cache.do_read(con, cancellation_token.clone()).await } + .instrument(span), + ); } } @@ -707,7 +705,7 @@ fn build_config(args: &ProxyCliArgs) -> anyhow::Result<&'static ProxyConfig> { /// auth::Backend is created at proxy startup, and lives forever. fn build_auth_backend( args: &ProxyCliArgs, -) -> anyhow::Result, &'static ConsoleRedirectBackend>> { +) -> anyhow::Result> { match &args.auth_backend { AuthBackendType::Console => { let wake_compute_cache_config: CacheOptions = args.wake_compute_cache.parse()?; @@ -757,8 +755,7 @@ fn build_auth_backend( locks, wake_compute_endpoint_rate_limiter, ); - let api = control_plane::provider::ControlPlaneBackend::Management(api); - let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ()); + let auth_backend = control_plane::provider::ControlPlaneBackend::Management(api); let config = Box::leak(Box::new(auth_backend)); @@ -769,9 +766,7 @@ fn build_auth_backend( AuthBackendType::Postgres => { let url = args.auth_endpoint.parse()?; let api = control_plane::provider::mock::Api::new(url, !args.is_private_access_proxy); - let api = control_plane::provider::ControlPlaneBackend::PostgresMock(api); - - let auth_backend = auth::Backend::ControlPlane(MaybeOwned::Owned(api), ()); + let auth_backend = control_plane::provider::ControlPlaneBackend::PostgresMock(api); let config = Box::leak(Box::new(auth_backend)); diff --git a/proxy/src/proxy/mod.rs b/proxy/src/proxy/mod.rs index 01eb5cf251..6ec7b6f22d 100644 --- a/proxy/src/proxy/mod.rs +++ b/proxy/src/proxy/mod.rs @@ -11,6 +11,7 @@ pub use copy_bidirectional::copy_bidirectional_client_compute; pub use copy_bidirectional::ErrorSource; use crate::config::ProxyProtocolV2; +use crate::control_plane::provider::ControlPlaneBackend; use crate::{ auth, cancellation::{self, CancellationHandlerMain, CancellationHandlerMainInternal}, @@ -61,7 +62,7 @@ pub async fn run_until_cancelled( pub async fn task_main( config: &'static ProxyConfig, - auth_backend: &'static auth::Backend<'static, ()>, + auth_backend: &'static ControlPlaneBackend, listener: tokio::net::TcpListener, cancellation_token: CancellationToken, cancellation_handler: Arc, @@ -248,7 +249,7 @@ impl ReportableError for ClientRequestError { #[allow(clippy::too_many_arguments)] pub(crate) async fn handle_client( config: &'static ProxyConfig, - auth_backend: &'static auth::Backend<'static, ()>, + auth_backend: &'static ControlPlaneBackend, ctx: &RequestMonitoring, cancellation_handler: Arc, stream: S, @@ -289,20 +290,17 @@ pub(crate) async fn handle_client( 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, ¶ms, hostname, common_names)) - .transpose(); - + let result = auth::ComputeUserInfoMaybeEndpoint::parse(ctx, ¶ms, hostname, common_names); let user_info = match result { Ok(user_info) => user_info, Err(e) => stream.throw_error(e).await?, }; - let user = user_info.get_user().to_owned(); - let user_info = match user_info + let user = user_info.user.clone(); + let user_info = match auth_backend .authenticate( ctx, + user_info, &mut stream, mode.allow_cleartext(), &config.authentication_config, diff --git a/proxy/src/proxy/tests/mod.rs b/proxy/src/proxy/tests/mod.rs index 58fb36dba7..2e2caf0416 100644 --- a/proxy/src/proxy/tests/mod.rs +++ b/proxy/src/proxy/tests/mod.rs @@ -8,7 +8,7 @@ use super::connect_compute::ConnectMechanism; use super::retry::CouldRetry; use super::*; use crate::auth::backend::{ - ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, MaybeOwned, TestBackend, + ComputeCredentialKeys, ComputeCredentials, ComputeUserInfo, TestBackend, }; use crate::config::{CertResolver, RetryConfig}; use crate::control_plane::messages::{ControlPlaneError, Details, MetricsAuxInfo, Status}; @@ -20,6 +20,7 @@ use crate::error::ErrorKind; use crate::{sasl, scram, BranchId, EndpointId, ProjectId}; use anyhow::{bail, Context}; use async_trait::async_trait; +use auth::backend::ControlPlaneComputeBackend; use http::StatusCode; use retry::{retry_after, ShouldRetryWakeCompute}; use rstest::rstest; @@ -552,19 +553,21 @@ fn helper_create_cached_node_info(cache: &'static NodeInfoCache) -> CachedNodeIn fn helper_create_connect_info( mechanism: &TestConnectMechanism, -) -> auth::Backend<'static, ComputeCredentials> { - let user_info = auth::Backend::ControlPlane( - MaybeOwned::Owned(ControlPlaneBackend::Test(Box::new(mechanism.clone()))), - ComputeCredentials { - info: ComputeUserInfo { - endpoint: "endpoint".into(), - user: "user".into(), - options: NeonOptions::parse_options_raw(""), - }, - keys: ComputeCredentialKeys::Password("password".into()), +) -> ControlPlaneComputeBackend<'static> { + let api = Box::leak(Box::new(ControlPlaneBackend::Test(Box::new( + mechanism.clone(), + )))); + + let creds = ComputeCredentials { + info: ComputeUserInfo { + endpoint: "endpoint".into(), + user: "user".into(), + options: NeonOptions::parse_options_raw(""), }, - ); - user_info + keys: ComputeCredentialKeys::Password("password".into()), + }; + + ControlPlaneComputeBackend { api, creds } } #[tokio::test] diff --git a/proxy/src/serverless/backend.rs b/proxy/src/serverless/backend.rs index 0f2d842eee..02e5686920 100644 --- a/proxy/src/serverless/backend.rs +++ b/proxy/src/serverless/backend.rs @@ -17,7 +17,7 @@ use crate::{ errors::{GetAuthInfoError, WakeComputeError}, locks::ApiLocks, provider::ApiLockError, - CachedNodeInfo, + Api, CachedNodeInfo, }, error::{ErrorKind, ReportableError, UserFacingError}, intern::EndpointIdInt, @@ -38,7 +38,7 @@ pub(crate) struct PoolingBackend { pub(crate) http_conn_pool: Arc, pub(crate) pool: Arc>, pub(crate) config: &'static ProxyConfig, - pub(crate) auth_backend: &'static crate::auth::Backend<'static, ()>, + pub(crate) auth_backend: crate::auth::Backend<'static>, pub(crate) endpoint_rate_limiter: Arc, } @@ -49,9 +49,16 @@ impl PoolingBackend { user_info: &ComputeUserInfo, password: &[u8], ) -> Result { - let user_info = user_info.clone(); - let backend = self.auth_backend.as_ref().map(|()| user_info.clone()); - let (allowed_ips, maybe_secret) = backend.get_allowed_ips_and_secret(ctx).await?; + let cplane = match &self.auth_backend { + crate::auth::Backend::ControlPlane(cplane) => cplane, + crate::auth::Backend::Local(_local) => { + return Err(AuthError::bad_auth_method( + "password authentication not supported by local_proxy", + )) + } + }; + + let (allowed_ips, maybe_secret) = cplane.get_allowed_ips_and_secret(ctx, user_info).await?; if self.config.authentication_config.ip_allowlist_check_enabled && !check_peer_addr_is_in_list(&ctx.peer_addr(), &allowed_ips) { @@ -65,7 +72,7 @@ impl PoolingBackend { } let cached_secret = match maybe_secret { Some(secret) => secret, - None => backend.get_role_secret(ctx).await?, + None => cplane.get_role_secret(ctx, user_info).await?, }; let secret = match cached_secret.value.clone() { @@ -100,7 +107,7 @@ impl PoolingBackend { } }; res.map(|key| ComputeCredentials { - info: user_info, + info: user_info.clone(), keys: key, }) } @@ -112,7 +119,7 @@ impl PoolingBackend { jwt: String, ) -> Result<(), AuthError> { match &self.auth_backend { - crate::auth::Backend::ControlPlane(console, ()) => { + crate::auth::Backend::ControlPlane(console) => { self.config .authentication_config .jwks_cache @@ -173,21 +180,41 @@ impl PoolingBackend { let conn_id = uuid::Uuid::new_v4(); tracing::Span::current().record("conn_id", display(conn_id)); info!(%conn_id, "pool: opening a new connection '{conn_info}'"); - let backend = self.auth_backend.as_ref().map(|()| keys); - crate::proxy::connect_compute::connect_to_compute( - ctx, - &TokioMechanism { - conn_id, - conn_info, - pool: self.pool.clone(), - locks: &self.config.connect_compute_locks, - }, - &backend, - false, // do not allow self signed compute for http flow - self.config.wake_compute_retry_config, - self.config.connect_to_compute_retry_config, - ) - .await + + match &self.auth_backend { + crate::auth::Backend::ControlPlane(cplane) => { + crate::proxy::connect_compute::connect_to_compute( + ctx, + &TokioMechanism { + conn_id, + conn_info, + pool: self.pool.clone(), + locks: &self.config.connect_compute_locks, + }, + &cplane.attach_to_credentials(keys), + false, // do not allow self signed compute for http flow + self.config.wake_compute_retry_config, + self.config.connect_to_compute_retry_config, + ) + .await + } + crate::auth::Backend::Local(local_proxy) => { + crate::proxy::connect_compute::connect_to_compute( + ctx, + &TokioMechanism { + conn_id, + conn_info, + pool: self.pool.clone(), + locks: &self.config.connect_compute_locks, + }, + &**local_proxy, + false, // do not allow self signed compute for http flow + self.config.wake_compute_retry_config, + self.config.connect_to_compute_retry_config, + ) + .await + } + } } // Wake up the destination if needed @@ -197,6 +224,13 @@ impl PoolingBackend { ctx: &RequestMonitoring, conn_info: ConnInfo, ) -> Result { + let cplane = match &self.auth_backend { + crate::auth::Backend::Local(_) => { + panic!("connect to local_proxy should not be called if we are already local_proxy") + } + crate::auth::Backend::ControlPlane(cplane) => cplane, + }; + info!("pool: looking for an existing connection"); if let Some(client) = self.http_conn_pool.get(ctx, &conn_info) { return Ok(client); @@ -205,7 +239,8 @@ impl PoolingBackend { let conn_id = uuid::Uuid::new_v4(); tracing::Span::current().record("conn_id", display(conn_id)); info!(%conn_id, "pool: opening a new connection '{conn_info}'"); - let backend = self.auth_backend.as_ref().map(|()| ComputeCredentials { + + let backend = cplane.attach_to_credentials(ComputeCredentials { info: conn_info.user_info.clone(), keys: crate::auth::backend::ComputeCredentialKeys::None, }); diff --git a/proxy/src/serverless/mod.rs b/proxy/src/serverless/mod.rs index d3ce56882c..ecc6a2fb7d 100644 --- a/proxy/src/serverless/mod.rs +++ b/proxy/src/serverless/mod.rs @@ -32,6 +32,7 @@ use tokio::time::timeout; use tokio_rustls::TlsAcceptor; use tokio_util::task::TaskTracker; +use crate::auth::Backend; use crate::cancellation::CancellationHandlerMain; use crate::config::ProxyConfig; use crate::context::RequestMonitoring; @@ -54,7 +55,7 @@ pub(crate) const SERVERLESS_DRIVER_SNI: &str = "api"; pub async fn task_main( config: &'static ProxyConfig, - auth_backend: &'static crate::auth::Backend<'static, ()>, + auth_backend: crate::auth::Backend<'static>, ws_listener: TcpListener, cancellation_token: CancellationToken, cancellation_handler: Arc, @@ -379,6 +380,10 @@ async fn request_handler( if config.http_config.accept_websockets && framed_websockets::upgrade::is_upgrade_request(&request) { + let Backend::ControlPlane(auth_backend) = backend.auth_backend else { + return json_response(StatusCode::BAD_REQUEST, "query is not supported"); + }; + let ctx = RequestMonitoring::new( session_id, peer_addr, @@ -396,7 +401,7 @@ async fn request_handler( async move { if let Err(e) = websocket::serve_websocket( config, - backend.auth_backend, + auth_backend, ctx, websocket, cancellation_handler, diff --git a/proxy/src/serverless/websocket.rs b/proxy/src/serverless/websocket.rs index 0c04943855..0d2d2ea1a5 100644 --- a/proxy/src/serverless/websocket.rs +++ b/proxy/src/serverless/websocket.rs @@ -1,3 +1,4 @@ +use crate::control_plane::provider::ControlPlaneBackend; use crate::proxy::ErrorSource; use crate::{ cancellation::CancellationHandlerMain, @@ -129,7 +130,7 @@ impl AsyncBufRead for WebSocketRw { pub(crate) async fn serve_websocket( config: &'static ProxyConfig, - auth_backend: &'static crate::auth::Backend<'static, ()>, + auth_backend: &'static ControlPlaneBackend, ctx: RequestMonitoring, websocket: OnUpgrade, cancellation_handler: Arc,