rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-10-25 13:13:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: introduce jwks settings in cplane response

Browse Source
This commit is contained in:
Conrad Ludgate
2025-01-14 11:26:15 +00:00
parent a039f8381f
commit d4eedb4069
8 changed files with 87 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+4
View File
@@ -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",
proxy/Cargo.toml
+1 -1
View File
@@ -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"] }
proxy/src/control_plane/client/cplane_proxy_v1.rs
+13
View File
@@ -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,
})
}
proxy/src/control_plane/client/mock.rs
+64 -52
View File
@@ -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<AuthInfo, GetAuthInfoError> {
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<Vec<JwksSettings>> =
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,
proxy/src/control_plane/messages.rs
+1
View File
@@ -229,6 +229,7 @@ pub(crate) struct GetEndpointAccessControl {
pub(crate) allowed_ips: Option<Vec<IpPattern>>,
pub(crate) project_id: Option<ProjectIdInt>,
pub(crate) allowed_vpc_endpoint_ids: Option<Vec<EndpointIdInt>>,
pub(crate) jwks: Vec<JwksSettings>,
}
/// Response which holds compute node's `host:port` pair.
proxy/src/control_plane/mod.rs
+1
View File
@@ -54,6 +54,7 @@ pub(crate) struct AuthInfo {
pub(crate) allowed_ips: Vec<IpPattern>,
/// Project ID. This is used for cache invalidation.
pub(crate) project_id: Option<ProjectIdInt>,
pub(crate) auth_rules: Vec<AuthRule>,
}
/// Info for establishing a connection to a compute node.
test_runner/fixtures/neon_fixtures.py
+1 -1
View File
@@ -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()
workspace_hack/Cargo.toml
+2
View File
@@ -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"] }