diff --git a/Cargo.lock b/Cargo.lock index 08453120c7..588afb7e1e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4478,6 +4478,8 @@ dependencies = [ "bytes", "fallible-iterator", "postgres-protocol", + "serde", + "serde_json", ] [[package]] @@ -8022,6 +8024,7 @@ dependencies = [ "num-traits", "once_cell", "parquet", + "postgres-types", "prettyplease", "proc-macro2", "prost", @@ -8048,6 +8051,7 @@ dependencies = [ "time", "time-macros", "tokio", + "tokio-postgres", "tokio-rustls 0.26.0", "tokio-stream", "tokio-util", diff --git a/proxy/Cargo.toml b/proxy/Cargo.toml index f362a45035..2a90afe197 100644 --- a/proxy/Cargo.toml +++ b/proxy/Cargo.toml @@ -121,4 +121,4 @@ rcgen.workspace = true rstest.workspace = true walkdir.workspace = true rand_distr = "0.4" -tokio-postgres.workspace = true +tokio-postgres = { workspace = true, features = ["with-serde_json-1"] } diff --git a/proxy/src/control_plane/client/cplane_proxy_v1.rs b/proxy/src/control_plane/client/cplane_proxy_v1.rs index ece03156d1..65a9a51af6 100644 --- a/proxy/src/control_plane/client/cplane_proxy_v1.rs +++ b/proxy/src/control_plane/client/cplane_proxy_v1.rs @@ -153,9 +153,22 @@ impl NeonControlPlaneClient { .proxy .allowed_ips_number .observe(allowed_ips.len() as f64); + + let auth_rules = body + .jwks + .into_iter() + .map(|jwks| AuthRule { + id: jwks.id, + jwks_url: jwks.jwks_url, + audience: jwks.jwt_audience, + role_names: jwks.role_names, + }) + .collect(); + Ok(AuthInfo { secret, allowed_ips, + auth_rules, project_id: body.project_id, }) } diff --git a/proxy/src/control_plane/client/mock.rs b/proxy/src/control_plane/client/mock.rs index 5f8bda0f35..0b821559e8 100644 --- a/proxy/src/control_plane/client/mock.rs +++ b/proxy/src/control_plane/client/mock.rs @@ -5,8 +5,9 @@ use std::sync::Arc; use futures::TryFutureExt; use thiserror::Error; +use tokio_postgres::types::Json; use tokio_postgres::Client; -use tracing::{error, info, info_span, warn, Instrument}; +use tracing::{error, info, warn}; use crate::auth::backend::jwt::AuthRule; use crate::auth::backend::ComputeUserInfo; @@ -17,7 +18,7 @@ use crate::control_plane::client::{CachedAllowedIps, CachedRoleSecret}; use crate::control_plane::errors::{ ControlPlaneError, GetAuthInfoError, GetEndpointJwksError, WakeComputeError, }; -use crate::control_plane::messages::MetricsAuxInfo; +use crate::control_plane::messages::{JwksSettings, MetricsAuxInfo}; use crate::control_plane::{AuthInfo, AuthSecret, CachedNodeInfo, NodeInfo}; use crate::error::io_error; use crate::intern::RoleNameInt; @@ -65,61 +66,70 @@ impl MockControlPlane { &self, user_info: &ComputeUserInfo, ) -> Result { - let (secret, allowed_ips) = async { - // Perhaps we could persist this connection, but then we'd have to - // write more code for reopening it if it got closed, which doesn't - // seem worth it. - let (client, connection) = - tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls).await?; + // Perhaps we could persist this connection, but then we'd have to + // write more code for reopening it if it got closed, which doesn't + // seem worth it. + let (client, connection) = + tokio_postgres::connect(self.endpoint.as_str(), tokio_postgres::NoTls).await?; - tokio::spawn(connection); + tokio::spawn(connection); - let secret = if let Some(entry) = get_execute_postgres_query( - &client, - "select rolpassword from pg_catalog.pg_authid where rolname = $1", - &[&&*user_info.user], - "rolpassword", - ) - .await? - { - info!("got a secret: {entry}"); // safe since it's not a prod scenario - let secret = scram::ServerSecret::parse(&entry).map(AuthSecret::Scram); - secret.or_else(|| parse_md5(&entry).map(AuthSecret::Md5)) - } else { - warn!("user '{}' does not exist", user_info.user); - None - }; + let secret = if let Some(entry) = get_execute_postgres_query( + &client, + "select rolpassword from pg_catalog.pg_authid where rolname = $1", + &[&&*user_info.user], + "rolpassword", + ) + .await? + { + info!("got a secret: {entry}"); // safe since it's not a prod scenario + let secret = scram::ServerSecret::parse(&entry).map(AuthSecret::Scram); + secret.or_else(|| parse_md5(&entry).map(AuthSecret::Md5)) + } else { + warn!("user '{}' does not exist", user_info.user); + None + }; - let allowed_ips = if self.ip_allowlist_check_enabled { - match get_execute_postgres_query( - &client, - "select allowed_ips from neon_control_plane.endpoints where endpoint_id = $1", - &[&user_info.endpoint.as_str()], - "allowed_ips", - ) - .await? - { - Some(s) => { - info!("got allowed_ips: {s}"); - s.split(',') - .map(|s| { - IpPattern::from_str(s).expect("mocked ip pattern should be correct") - }) - .collect() - } - None => vec![], + let (allowed_ips, auth_rules) = if self.ip_allowlist_check_enabled { + let row = client.query_opt("select allowed_ips, jwks from neon_control_plane.endpoints where endpoint_id = $1", &[&user_info.endpoint.as_str()]).await?; + match row { + Some(row) => { + let allowed_ips: String = row + .try_get("allowed_ips") + .map_err(MockApiError::PasswordNotSet)?; + let jwks: Json> = + row.try_get("jwks").map_err(MockApiError::PasswordNotSet)?; + + info!("got allowed_ips: {allowed_ips}"); + let allowed_ips = allowed_ips + .split(',') + .map(|s| { + IpPattern::from_str(s).expect("mocked ip pattern should be correct") + }) + .collect(); + + let auth_rules = jwks + .0 + .into_iter() + .map(|jwks| AuthRule { + id: jwks.id, + jwks_url: jwks.jwks_url, + audience: jwks.jwt_audience, + role_names: jwks.role_names, + }) + .collect(); + + (allowed_ips, auth_rules) } - } else { - vec![] - }; + None => (vec![], vec![]), + } + } else { + (vec![], vec![]) + }; - Ok((secret, allowed_ips)) - } - .inspect_err(|e: &GetAuthInfoError| tracing::error!("{e}")) - .instrument(info_span!("postgres", url = self.endpoint.as_str())) - .await?; Ok(AuthInfo { secret, + auth_rules, allowed_ips, project_id: None, }) @@ -203,7 +213,7 @@ async fn get_execute_postgres_query( } impl super::ControlPlaneApi for MockControlPlane { - #[tracing::instrument(skip_all)] + #[tracing::instrument(skip_all, fields(url = self.endpoint.as_str()))] async fn get_role_secret( &self, _ctx: &RequestContext, @@ -214,6 +224,7 @@ impl super::ControlPlaneApi for MockControlPlane { )) } + #[tracing::instrument(skip_all, fields(url = self.endpoint.as_str()))] async fn get_allowed_ips_and_secret( &self, _ctx: &RequestContext, @@ -227,6 +238,7 @@ impl super::ControlPlaneApi for MockControlPlane { )) } + #[tracing::instrument(skip_all, fields(url = self.endpoint.as_str()))] async fn get_endpoint_jwks( &self, _ctx: &RequestContext, @@ -235,7 +247,7 @@ impl super::ControlPlaneApi for MockControlPlane { self.do_get_endpoint_jwks(endpoint).await } - #[tracing::instrument(skip_all)] + #[tracing::instrument(skip_all, fields(url = self.endpoint.as_str()))] async fn wake_compute( &self, _ctx: &RequestContext, diff --git a/proxy/src/control_plane/messages.rs b/proxy/src/control_plane/messages.rs index d068614b24..3215918294 100644 --- a/proxy/src/control_plane/messages.rs +++ b/proxy/src/control_plane/messages.rs @@ -229,6 +229,7 @@ pub(crate) struct GetEndpointAccessControl { pub(crate) allowed_ips: Option>, pub(crate) project_id: Option, pub(crate) allowed_vpc_endpoint_ids: Option>, + pub(crate) jwks: Vec, } /// Response which holds compute node's `host:port` pair. diff --git a/proxy/src/control_plane/mod.rs b/proxy/src/control_plane/mod.rs index 1dca26d686..ec62b1cc4d 100644 --- a/proxy/src/control_plane/mod.rs +++ b/proxy/src/control_plane/mod.rs @@ -54,6 +54,7 @@ pub(crate) struct AuthInfo { pub(crate) allowed_ips: Vec, /// Project ID. This is used for cache invalidation. pub(crate) project_id: Option, + pub(crate) auth_rules: Vec, } /// Info for establishing a connection to a compute node. diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index e22e452a52..9d5a50b890 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -3645,7 +3645,7 @@ def static_proxy( vanilla_pg.safe_psql("create user proxy with login superuser password 'password'") vanilla_pg.safe_psql("CREATE SCHEMA IF NOT EXISTS neon_control_plane") vanilla_pg.safe_psql( - "CREATE TABLE neon_control_plane.endpoints (endpoint_id VARCHAR(255) PRIMARY KEY, allowed_ips VARCHAR(255))" + "CREATE TABLE neon_control_plane.endpoints (endpoint_id VARCHAR(255) PRIMARY KEY, allowed_ips VARCHAR(255), jwks jsonb default '[]'::jsonb)" ) proxy_port = port_distributor.get_port() diff --git a/workspace_hack/Cargo.toml b/workspace_hack/Cargo.toml index 0ffeeead18..ae83aa0e25 100644 --- a/workspace_hack/Cargo.toml +++ b/workspace_hack/Cargo.toml @@ -65,6 +65,7 @@ num-rational = { version = "0.4", default-features = false, features = ["num-big num-traits = { version = "0.2", features = ["i128", "libm"] } once_cell = { version = "1" } parquet = { version = "53", default-features = false, features = ["zstd"] } +postgres-types = { git = "https://github.com/neondatabase/rust-postgres.git", branch = "neon", default-features = false, features = ["with-serde_json-1"] } prost = { version = "0.13", features = ["no-recursion-limit", "prost-derive"] } rand = { version = "0.8", features = ["small_rng"] } regex = { version = "1" } @@ -86,6 +87,7 @@ tikv-jemalloc-ctl = { version = "0.6", features = ["stats", "use_std"] } tikv-jemalloc-sys = { version = "0.6", features = ["profiling", "stats", "unprefixed_malloc_on_supported_platforms"] } time = { version = "0.3", features = ["macros", "serde-well-known"] } tokio = { version = "1", features = ["full", "test-util"] } +tokio-postgres = { git = "https://github.com/neondatabase/rust-postgres.git", branch = "neon", features = ["with-serde_json-1"] } tokio-rustls = { version = "0.26", default-features = false, features = ["logging", "ring", "tls12"] } tokio-stream = { version = "0.1", features = ["net"] } tokio-util = { version = "0.7", features = ["codec", "compat", "io", "rt"] }