rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-01-17 02:12:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

validate jwt during auth_quirks

Browse Source
This commit is contained in:
Conrad Ludgate
2025-01-14 11:59:45 +00:00
parent d4eedb4069
commit 97d0147ed9
12 changed files with 288 additions and 206 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
proxy/src/auth/backend/console_redirect.rs
View File

@@ -94,7 +94,7 @@ impl BackendIpAllowlist for ConsoleRedirectBackend {
self.api
.get_allowed_ips_and_secret(ctx, user_info)
.await
.map(|(ips, _)| ips.as_ref().clone())
.map(|(ips, _)| ips.0.clone())
.map_err(|e| e.into())
}
}

12
proxy/src/auth/backend/jwt.rs
View File

@@ -44,6 +44,18 @@ pub(crate) trait FetchAuthRules: Clone + Send + Sync + 'static {
) -> impl Future<Output = Result<Vec<AuthRule>, FetchAuthRulesError>> + Send;
}
#[derive(Clone)]
pub(crate) struct StaticAuthRules(pub Vec<AuthRule>);
impl FetchAuthRules for StaticAuthRules {
async fn fetch_auth_rules(
&self,
_ctx: &RequestContext,
_endpoint: EndpointId,
) -> Result<Vec<AuthRule>, FetchAuthRulesError> {
Ok(self.0.clone())
}
}
#[derive(Error, Debug)]
pub(crate) enum FetchAuthRulesError {
#[error(transparent)]

64
proxy/src/auth/backend/mod.rs
View File

@@ -10,6 +10,7 @@ use std::sync::Arc;
pub use console_redirect::ConsoleRedirectBackend;
pub(crate) use console_redirect::ConsoleRedirectError;
use ipnet::{Ipv4Net, Ipv6Net};
use jwt::{JwkCache, StaticAuthRules};
use local::LocalBackend;
use postgres_client::config::AuthKeys;
use tokio::io::{AsyncRead, AsyncWrite};
@@ -259,6 +260,7 @@ pub(crate) trait BackendIpAllowlist {
/// Here, we choose the appropriate auth flow based on circumstances.
///
/// All authentication flows will emit an AuthenticationOk message if successful.
#[allow(clippy::too_many_arguments)]
async fn auth_quirks(
ctx: &RequestContext,
api: &impl control_plane::ControlPlaneApi,
@@ -267,6 +269,7 @@ async fn auth_quirks(
allow_cleartext: bool,
config: &'static AuthenticationConfig,
endpoint_rate_limiter: Arc<EndpointRateLimiter>,
jwks_cache: Arc<JwkCache>,
) -> auth::Result<(ComputeCredentials, Option<Vec<IpPattern>>)> {
// If there's no project so far, that entails that client doesn't
// support SNI or other means of passing the endpoint (project) name.
@@ -282,11 +285,53 @@ async fn auth_quirks(
};
debug!("fetching user's authentication info");
let (allowed_ips, maybe_secret) = api.get_allowed_ips_and_secret(ctx, &info).await?;
let (x, maybe_secret) = api.get_allowed_ips_and_secret(ctx, &info).await?;
let (allowed_ips, auth_rules) = &**x;
// we expect a jwt in the options field
if !auth_rules.is_empty() {
match info.options.get("jwt") {
Some(jwt) => {
let creds = jwks_cache
.check_jwt(
ctx,
info.endpoint.clone(),
&info.user,
&StaticAuthRules(auth_rules.clone()),
&jwt,
)
.await?;
let token = match creds {
ComputeCredentialKeys::JwtPayload(payload) => {
serde_json::from_slice::<serde_json::Value>(&payload)
.expect("jwt payload is valid json")
}
_ => unreachable!(),
};
// the token has a required IP claim.
if let Some(expected_ip) = token.get("ip") {
// todo: don't panic here, obviously.
let expected_ip: IpAddr = expected_ip
.as_str()
.expect("jwt should not have an invalid IP claim")
.parse()
.expect("jwt should not have an invalid IP claim");
if ctx.peer_addr() != expected_ip {
return Err(auth::AuthError::ip_address_not_allowed(ctx.peer_addr()));
}
}
}
None => {
return Err(AuthError::bad_auth_method("needs jwt"));
}
}
}
// check allowed list
if config.ip_allowlist_check_enabled
&& !check_peer_addr_is_in_list(&ctx.peer_addr(), &allowed_ips)
&& !check_peer_addr_is_in_list(&ctx.peer_addr(), allowed_ips)
{
return Err(auth::AuthError::ip_address_not_allowed(ctx.peer_addr()));
}
@@ -326,7 +371,7 @@ async fn auth_quirks(
)
.await
{
Ok(keys) => Ok((keys, Some(allowed_ips.as_ref().clone()))),
Ok(keys) => Ok((keys, Some(allowed_ips.clone()))),
Err(e) => {
if e.is_password_failed() {
// The password could have been changed, so we invalidate the cache.
@@ -396,6 +441,7 @@ impl<'a> Backend<'a, ComputeUserInfoMaybeEndpoint> {
allow_cleartext: bool,
config: &'static AuthenticationConfig,
endpoint_rate_limiter: Arc<EndpointRateLimiter>,
jwks_cache: Arc<JwkCache>,
) -> auth::Result<(Backend<'a, ComputeCredentials>, Option<Vec<IpPattern>>)> {
let res = match self {
Self::ControlPlane(api, user_info) => {
@@ -413,6 +459,7 @@ impl<'a> Backend<'a, ComputeUserInfoMaybeEndpoint> {
allow_cleartext,
config,
endpoint_rate_limiter,
jwks_cache,
)
.await?;
Ok((Backend::ControlPlane(api, credentials), ip_allowlist))
@@ -447,7 +494,7 @@ impl Backend<'_, ComputeUserInfo> {
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)),
Self::Local(_) => Ok((Cached::new_uncached(Arc::new((vec![], vec![]))), None)),
}
}
}
@@ -461,11 +508,11 @@ impl BackendIpAllowlist for Backend<'_, ()> {
) -> auth::Result<Vec<auth::IpPattern>> {
let auth_data = match self {
Self::ControlPlane(api, ()) => api.get_allowed_ips_and_secret(ctx, user_info).await,
Self::Local(_) => Ok((Cached::new_uncached(Arc::new(vec![])), None)),
Self::Local(_) => Ok((Cached::new_uncached(Arc::new((vec![], vec![]))), None)),
};
auth_data
.map(|(ips, _)| ips.as_ref().clone())
.map(|(ips, _)| ips.0.clone())
.map_err(|e| e.into())
}
}
@@ -543,7 +590,7 @@ mod tests {
control_plane::errors::GetAuthInfoError,
> {
Ok((
CachedAllowedIps::new_uncached(Arc::new(self.ips.clone())),
CachedAllowedIps::new_uncached(Arc::new((self.ips.clone(), vec![]))),
Some(CachedRoleSecret::new_uncached(Some(self.secret.clone()))),
))
}
@@ -703,6 +750,7 @@ mod tests {
false,
&CONFIG,
endpoint_rate_limiter,
Arc::new(JwkCache::default()),
)
.await
.unwrap();
@@ -758,6 +806,7 @@ mod tests {
true,
&CONFIG,
endpoint_rate_limiter,
Arc::new(JwkCache::default()),
)
.await
.unwrap();
@@ -811,6 +860,7 @@ mod tests {
true,
&CONFIG,
endpoint_rate_limiter,
Arc::new(JwkCache::default()),
)
.await
.unwrap();

383
proxy/src/cache/project_info.rs vendored
View File

@@ -1,3 +1,5 @@
#![allow(clippy::type_complexity)]
use std::collections::HashSet;
use std::convert::Infallible;
use std::sync::atomic::AtomicU64;
@@ -13,6 +15,7 @@ use tokio::time::Instant;
use tracing::{debug, info};
use super::{Cache, Cached};
use crate::auth::backend::jwt::AuthRule;
use crate::auth::IpPattern;
use crate::config::ProjectInfoCacheOptions;
use crate::control_plane::AuthSecret;
@@ -50,7 +53,7 @@ impl<T> From<T> for Entry<T> {
#[derive(Default)]
struct EndpointInfo {
secret: std::collections::HashMap<RoleNameInt, Entry<Option<AuthSecret>>>,
allowed_ips: Option<Entry<Arc<Vec<IpPattern>>>>,
allowed_ips: Option<Entry<Arc<(Vec<IpPattern>, Vec<AuthRule>)>>>,
}
impl EndpointInfo {
@@ -81,7 +84,7 @@ impl EndpointInfo {
&self,
valid_since: Instant,
ignore_cache_since: Option<Instant>,
) -> Option<(Arc<Vec<IpPattern>>, bool)> {
) -> Option<(Arc<(Vec<IpPattern>, Vec<AuthRule>)>, bool)> {
if let Some(allowed_ips) = &self.allowed_ips {
if valid_since < allowed_ips.created_at {
return Some((
@@ -211,7 +214,7 @@ impl ProjectInfoCacheImpl {
pub(crate) fn get_allowed_ips(
&self,
endpoint_id: &EndpointId,
) -> Option<Cached<&Self, Arc<Vec<IpPattern>>>> {
) -> Option<Cached<&Self, Arc<(Vec<IpPattern>, Vec<AuthRule>)>>> {
let endpoint_id = EndpointIdInt::get(endpoint_id)?;
let (valid_since, ignore_cache_since) = self.get_cache_times();
let endpoint_info = self.cache.get(&endpoint_id)?;
@@ -247,7 +250,7 @@ impl ProjectInfoCacheImpl {
&self,
project_id: ProjectIdInt,
endpoint_id: EndpointIdInt,
allowed_ips: Arc<Vec<IpPattern>>,
allowed_ips: Arc<(Vec<IpPattern>, Vec<AuthRule>)>,
) {
if self.cache.len() >= self.config.size {
// If there are too many entries, wait until the next gc cycle.
@@ -364,209 +367,209 @@ impl Cache for ProjectInfoCacheImpl {
}
}
#[cfg(test)]
#[expect(clippy::unwrap_used)]
mod tests {
use super::*;
use crate::scram::ServerSecret;
use crate::types::ProjectId;
// #[cfg(test)]
// #[expect(clippy::unwrap_used)]
// mod tests {
// use super::*;
// use crate::scram::ServerSecret;
// use crate::types::ProjectId;
#[tokio::test]
async fn test_project_info_cache_settings() {
tokio::time::pause();
let cache = ProjectInfoCacheImpl::new(ProjectInfoCacheOptions {
size: 2,
max_roles: 2,
ttl: Duration::from_secs(1),
gc_interval: Duration::from_secs(600),
});
let project_id: ProjectId = "project".into();
let endpoint_id: EndpointId = "endpoint".into();
let user1: RoleName = "user1".into();
let user2: RoleName = "user2".into();
let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
let secret2 = None;
let allowed_ips = Arc::new(vec![
"127.0.0.1".parse().unwrap(),
"127.0.0.2".parse().unwrap(),
]);
cache.insert_role_secret(
(&project_id).into(),
(&endpoint_id).into(),
(&user1).into(),
secret1.clone(),
);
cache.insert_role_secret(
(&project_id).into(),
(&endpoint_id).into(),
(&user2).into(),
secret2.clone(),
);
cache.insert_allowed_ips(
(&project_id).into(),
(&endpoint_id).into(),
allowed_ips.clone(),
);
// #[tokio::test]
// async fn test_project_info_cache_settings() {
// tokio::time::pause();
// let cache = ProjectInfoCacheImpl::new(ProjectInfoCacheOptions {
// size: 2,
// max_roles: 2,
// ttl: Duration::from_secs(1),
// gc_interval: Duration::from_secs(600),
// });
// let project_id: ProjectId = "project".into();
// let endpoint_id: EndpointId = "endpoint".into();
// let user1: RoleName = "user1".into();
// let user2: RoleName = "user2".into();
// let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
// let secret2 = None;
// let allowed_ips = Arc::new(vec![
// "127.0.0.1".parse().unwrap(),
// "127.0.0.2".parse().unwrap(),
// ]);
// cache.insert_role_secret(
// (&project_id).into(),
// (&endpoint_id).into(),
// (&user1).into(),
// secret1.clone(),
// );
// cache.insert_role_secret(
// (&project_id).into(),
// (&endpoint_id).into(),
// (&user2).into(),
// secret2.clone(),
// );
// cache.insert_allowed_ips(
// (&project_id).into(),
// (&endpoint_id).into(),
// allowed_ips.clone(),
// );
let cached = cache.get_role_secret(&endpoint_id, &user1).unwrap();
assert!(cached.cached());
assert_eq!(cached.value, secret1);
let cached = cache.get_role_secret(&endpoint_id, &user2).unwrap();
assert!(cached.cached());
assert_eq!(cached.value, secret2);
// let cached = cache.get_role_secret(&endpoint_id, &user1).unwrap();
// assert!(cached.cached());
// assert_eq!(cached.value, secret1);
// let cached = cache.get_role_secret(&endpoint_id, &user2).unwrap();
// assert!(cached.cached());
// assert_eq!(cached.value, secret2);
// Shouldn't add more than 2 roles.
let user3: RoleName = "user3".into();
let secret3 = Some(AuthSecret::Scram(ServerSecret::mock([3; 32])));
cache.insert_role_secret(
(&project_id).into(),
(&endpoint_id).into(),
(&user3).into(),
secret3.clone(),
);
assert!(cache.get_role_secret(&endpoint_id, &user3).is_none());
// // Shouldn't add more than 2 roles.
// let user3: RoleName = "user3".into();
// let secret3 = Some(AuthSecret::Scram(ServerSecret::mock([3; 32])));
// cache.insert_role_secret(
// (&project_id).into(),
// (&endpoint_id).into(),
// (&user3).into(),
// secret3.clone(),
// );
// assert!(cache.get_role_secret(&endpoint_id, &user3).is_none());
let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
assert!(cached.cached());
assert_eq!(cached.value, allowed_ips);
// let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
// assert!(cached.cached());
// assert_eq!(cached.value, allowed_ips);
tokio::time::advance(Duration::from_secs(2)).await;
let cached = cache.get_role_secret(&endpoint_id, &user1);
assert!(cached.is_none());
let cached = cache.get_role_secret(&endpoint_id, &user2);
assert!(cached.is_none());
let cached = cache.get_allowed_ips(&endpoint_id);
assert!(cached.is_none());
}
// tokio::time::advance(Duration::from_secs(2)).await;
// let cached = cache.get_role_secret(&endpoint_id, &user1);
// assert!(cached.is_none());
// let cached = cache.get_role_secret(&endpoint_id, &user2);
// assert!(cached.is_none());
// let cached = cache.get_allowed_ips(&endpoint_id);
// assert!(cached.is_none());
// }
#[tokio::test]
async fn test_project_info_cache_invalidations() {
tokio::time::pause();
let cache = Arc::new(ProjectInfoCacheImpl::new(ProjectInfoCacheOptions {
size: 2,
max_roles: 2,
ttl: Duration::from_secs(1),
gc_interval: Duration::from_secs(600),
}));
cache.clone().increment_active_listeners().await;
tokio::time::advance(Duration::from_secs(2)).await;
// #[tokio::test]
// async fn test_project_info_cache_invalidations() {
// tokio::time::pause();
// let cache = Arc::new(ProjectInfoCacheImpl::new(ProjectInfoCacheOptions {
// size: 2,
// max_roles: 2,
// ttl: Duration::from_secs(1),
// gc_interval: Duration::from_secs(600),
// }));
// cache.clone().increment_active_listeners().await;
// tokio::time::advance(Duration::from_secs(2)).await;
let project_id: ProjectId = "project".into();
let endpoint_id: EndpointId = "endpoint".into();
let user1: RoleName = "user1".into();
let user2: RoleName = "user2".into();
let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
let secret2 = Some(AuthSecret::Scram(ServerSecret::mock([2; 32])));
let allowed_ips = Arc::new(vec![
"127.0.0.1".parse().unwrap(),
"127.0.0.2".parse().unwrap(),
]);
cache.insert_role_secret(
(&project_id).into(),
(&endpoint_id).into(),
(&user1).into(),
secret1.clone(),
);
cache.insert_role_secret(
(&project_id).into(),
(&endpoint_id).into(),
(&user2).into(),
secret2.clone(),
);
cache.insert_allowed_ips(
(&project_id).into(),
(&endpoint_id).into(),
allowed_ips.clone(),
);
// let project_id: ProjectId = "project".into();
// let endpoint_id: EndpointId = "endpoint".into();
// let user1: RoleName = "user1".into();
// let user2: RoleName = "user2".into();
// let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
// let secret2 = Some(AuthSecret::Scram(ServerSecret::mock([2; 32])));
// let allowed_ips = Arc::new(vec![
// "127.0.0.1".parse().unwrap(),
// "127.0.0.2".parse().unwrap(),
// ]);
// cache.insert_role_secret(
// (&project_id).into(),
// (&endpoint_id).into(),
// (&user1).into(),
// secret1.clone(),
// );
// cache.insert_role_secret(
// (&project_id).into(),
// (&endpoint_id).into(),
// (&user2).into(),
// secret2.clone(),
// );
// cache.insert_allowed_ips(
// (&project_id).into(),
// (&endpoint_id).into(),
// allowed_ips.clone(),
// );
tokio::time::advance(Duration::from_secs(2)).await;
// Nothing should be invalidated.
// tokio::time::advance(Duration::from_secs(2)).await;
// // Nothing should be invalidated.
let cached = cache.get_role_secret(&endpoint_id, &user1).unwrap();
// TTL is disabled, so it should be impossible to invalidate this value.
assert!(!cached.cached());
assert_eq!(cached.value, secret1);
// let cached = cache.get_role_secret(&endpoint_id, &user1).unwrap();
// // TTL is disabled, so it should be impossible to invalidate this value.
// assert!(!cached.cached());
// assert_eq!(cached.value, secret1);
cached.invalidate(); // Shouldn't do anything.
let cached = cache.get_role_secret(&endpoint_id, &user1).unwrap();
assert_eq!(cached.value, secret1);
// cached.invalidate(); // Shouldn't do anything.
// let cached = cache.get_role_secret(&endpoint_id, &user1).unwrap();
// assert_eq!(cached.value, secret1);
let cached = cache.get_role_secret(&endpoint_id, &user2).unwrap();
assert!(!cached.cached());
assert_eq!(cached.value, secret2);
// let cached = cache.get_role_secret(&endpoint_id, &user2).unwrap();
// assert!(!cached.cached());
// assert_eq!(cached.value, secret2);
// The only way to invalidate this value is to invalidate via the api.
cache.invalidate_role_secret_for_project((&project_id).into(), (&user2).into());
assert!(cache.get_role_secret(&endpoint_id, &user2).is_none());
// // The only way to invalidate this value is to invalidate via the api.
// cache.invalidate_role_secret_for_project((&project_id).into(), (&user2).into());
// assert!(cache.get_role_secret(&endpoint_id, &user2).is_none());
let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
assert!(!cached.cached());
assert_eq!(cached.value, allowed_ips);
}
// let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
// assert!(!cached.cached());
// assert_eq!(cached.value, allowed_ips);
// }
#[tokio::test]
async fn test_increment_active_listeners_invalidate_added_before() {
tokio::time::pause();
let cache = Arc::new(ProjectInfoCacheImpl::new(ProjectInfoCacheOptions {
size: 2,
max_roles: 2,
ttl: Duration::from_secs(1),
gc_interval: Duration::from_secs(600),
}));
// #[tokio::test]
// async fn test_increment_active_listeners_invalidate_added_before() {
// tokio::time::pause();
// let cache = Arc::new(ProjectInfoCacheImpl::new(ProjectInfoCacheOptions {
// size: 2,
// max_roles: 2,
// ttl: Duration::from_secs(1),
// gc_interval: Duration::from_secs(600),
// }));
let project_id: ProjectId = "project".into();
let endpoint_id: EndpointId = "endpoint".into();
let user1: RoleName = "user1".into();
let user2: RoleName = "user2".into();
let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
let secret2 = Some(AuthSecret::Scram(ServerSecret::mock([2; 32])));
let allowed_ips = Arc::new(vec![
"127.0.0.1".parse().unwrap(),
"127.0.0.2".parse().unwrap(),
]);
cache.insert_role_secret(
(&project_id).into(),
(&endpoint_id).into(),
(&user1).into(),
secret1.clone(),
);
cache.clone().increment_active_listeners().await;
tokio::time::advance(Duration::from_millis(100)).await;
cache.insert_role_secret(
(&project_id).into(),
(&endpoint_id).into(),
(&user2).into(),
secret2.clone(),
);
// let project_id: ProjectId = "project".into();
// let endpoint_id: EndpointId = "endpoint".into();
// let user1: RoleName = "user1".into();
// let user2: RoleName = "user2".into();
// let secret1 = Some(AuthSecret::Scram(ServerSecret::mock([1; 32])));
// let secret2 = Some(AuthSecret::Scram(ServerSecret::mock([2; 32])));
// let allowed_ips = Arc::new(vec![
// "127.0.0.1".parse().unwrap(),
// "127.0.0.2".parse().unwrap(),
// ]);
// cache.insert_role_secret(
// (&project_id).into(),
// (&endpoint_id).into(),
// (&user1).into(),
// secret1.clone(),
// );
// cache.clone().increment_active_listeners().await;
// tokio::time::advance(Duration::from_millis(100)).await;
// cache.insert_role_secret(
// (&project_id).into(),
// (&endpoint_id).into(),
// (&user2).into(),
// secret2.clone(),
// );
// Added before ttl was disabled + ttl should be still cached.
let cached = cache.get_role_secret(&endpoint_id, &user1).unwrap();
assert!(cached.cached());
let cached = cache.get_role_secret(&endpoint_id, &user2).unwrap();
assert!(cached.cached());
// // Added before ttl was disabled + ttl should be still cached.
// let cached = cache.get_role_secret(&endpoint_id, &user1).unwrap();
// assert!(cached.cached());
// let cached = cache.get_role_secret(&endpoint_id, &user2).unwrap();
// assert!(cached.cached());
tokio::time::advance(Duration::from_secs(1)).await;
// Added before ttl was disabled + ttl should expire.
assert!(cache.get_role_secret(&endpoint_id, &user1).is_none());
assert!(cache.get_role_secret(&endpoint_id, &user2).is_none());
// tokio::time::advance(Duration::from_secs(1)).await;
// // Added before ttl was disabled + ttl should expire.
// assert!(cache.get_role_secret(&endpoint_id, &user1).is_none());
// assert!(cache.get_role_secret(&endpoint_id, &user2).is_none());
// Added after ttl was disabled + ttl should not be cached.
cache.insert_allowed_ips(
(&project_id).into(),
(&endpoint_id).into(),
allowed_ips.clone(),
);
let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
assert!(!cached.cached());
// // Added after ttl was disabled + ttl should not be cached.
// cache.insert_allowed_ips(
// (&project_id).into(),
// (&endpoint_id).into(),
// allowed_ips.clone(),
// );
// let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
// assert!(!cached.cached());
tokio::time::advance(Duration::from_secs(1)).await;
// Added before ttl was disabled + ttl still should expire.
assert!(cache.get_role_secret(&endpoint_id, &user1).is_none());
assert!(cache.get_role_secret(&endpoint_id, &user2).is_none());
// Shouldn't be invalidated.
// tokio::time::advance(Duration::from_secs(1)).await;
// // Added before ttl was disabled + ttl still should expire.
// assert!(cache.get_role_secret(&endpoint_id, &user1).is_none());
// assert!(cache.get_role_secret(&endpoint_id, &user2).is_none());
// // Shouldn't be invalidated.
let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
assert!(!cached.cached());
assert_eq!(cached.value, allowed_ips);
}
}
// let cached = cache.get_allowed_ips(&endpoint_id).unwrap();
// assert!(!cached.cached());
// assert_eq!(cached.value, allowed_ips);
// }
// }

4
proxy/src/control_plane/client/cplane_proxy_v1.rs
View File

@@ -323,7 +323,7 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
self.caches.project_info.insert_allowed_ips(
project_id,
normalized_ep_int,
Arc::new(auth_info.allowed_ips),
Arc::new((auth_info.allowed_ips, auth_info.auth_rules)),
);
ctx.set_project_id(project_id);
}
@@ -349,7 +349,7 @@ impl super::ControlPlaneApi for NeonControlPlaneClient {
.allowed_ips_cache_misses
.inc(CacheOutcome::Miss);
let auth_info = self.do_get_auth_info(ctx, user_info).await?;
let allowed_ips = Arc::new(auth_info.allowed_ips);
let allowed_ips = Arc::new((auth_info.allowed_ips, auth_info.auth_rules));
let user = &user_info.user;
if let Some(project_id) = auth_info.project_id {
let normalized_ep_int = normalized_ep.into();

5
proxy/src/control_plane/client/mock.rs
View File

@@ -230,10 +230,9 @@ impl super::ControlPlaneApi for MockControlPlane {
_ctx: &RequestContext,
user_info: &ComputeUserInfo,
) -> Result<(CachedAllowedIps, Option<CachedRoleSecret>), GetAuthInfoError> {
let res = self.do_get_auth_info(user_info).await?;
Ok((
Cached::new_uncached(Arc::new(
self.do_get_auth_info(user_info).await?.allowed_ips,
)),
Cached::new_uncached(Arc::new((res.allowed_ips, res.auth_rules))),
None,
))
}

1
proxy/src/control_plane/messages.rs
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>>,
#[serde(default)]
pub(crate) jwks: Vec<JwksSettings>,
}

2
proxy/src/control_plane/mod.rs
View File

@@ -100,7 +100,7 @@ pub(crate) type NodeInfoCache =
TimedLru<EndpointCacheKey, Result<NodeInfo, Box<ControlPlaneErrorMessage>>>;
pub(crate) type CachedNodeInfo = Cached<&'static NodeInfoCache, NodeInfo>;
pub(crate) type CachedRoleSecret = Cached<&'static ProjectInfoCacheImpl, Option<AuthSecret>>;
pub(crate) type CachedAllowedIps = Cached<&'static ProjectInfoCacheImpl, Arc<Vec<IpPattern>>>;
pub(crate) type CachedAllowedIps = Cached<&'static ProjectInfoCacheImpl, Arc<(Vec<IpPattern>, Vec<AuthRule>)>>;
/// This will allocate per each call, but the http requests alone
/// already require a few allocations, so it should be fine.

7
proxy/src/proxy/mod.rs
View File

@@ -23,6 +23,7 @@ use tracing::{debug, error, info, warn, Instrument};
use self::connect_compute::{connect_to_compute, TcpMechanism};
use self::passthrough::ProxyPassthrough;
use crate::auth::backend::jwt::JwkCache;
use crate::cancellation::{self, CancellationHandlerMain, CancellationHandlerMainInternal};
use crate::config::{ProxyConfig, ProxyProtocolV2, TlsConfig};
use crate::context::RequestContext;
@@ -71,6 +72,8 @@ pub async fn task_main(
let connections = tokio_util::task::task_tracker::TaskTracker::new();
let cancellations = tokio_util::task::task_tracker::TaskTracker::new();
let jwks_cache = Arc::new(JwkCache::default());
while let Some(accept_result) =
run_until_cancelled(listener.accept(), &cancellation_token).await
{
@@ -84,6 +87,7 @@ pub async fn task_main(
let session_id = uuid::Uuid::new_v4();
let cancellation_handler = Arc::clone(&cancellation_handler);
let cancellations = cancellations.clone();
let jwks_cache = jwks_cache.clone();
debug!(protocol = "tcp", %session_id, "accepted new TCP connection");
let endpoint_rate_limiter2 = endpoint_rate_limiter.clone();
@@ -136,6 +140,7 @@ pub async fn task_main(
endpoint_rate_limiter2,
conn_gauge,
cancellations,
jwks_cache,
)
.instrument(ctx.span())
.boxed()
@@ -249,6 +254,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
endpoint_rate_limiter: Arc<EndpointRateLimiter>,
conn_gauge: NumClientConnectionsGuard<'static>,
cancellations: tokio_util::task::task_tracker::TaskTracker,
jwks_cache: Arc<JwkCache>,
) -> Result<Option<ProxyPassthrough<CancellationHandlerMainInternal, S>>, ClientRequestError> {
debug!(
protocol = %ctx.protocol(),
@@ -319,6 +325,7 @@ pub(crate) async fn handle_client<S: AsyncRead + AsyncWrite + Unpin>(
mode.allow_cleartext(),
&config.authentication_config,
endpoint_rate_limiter,
jwks_cache,
)
.await
{

5
proxy/src/serverless/backend.rs
View File

@@ -57,9 +57,10 @@ impl PoolingBackend {
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 (x, maybe_secret) = backend.get_allowed_ips_and_secret(ctx).await?;
let (allowed_ips, _) = &**x;
if self.config.authentication_config.ip_allowlist_check_enabled
&& !check_peer_addr_is_in_list(&ctx.peer_addr(), &allowed_ips)
&& !check_peer_addr_is_in_list(&ctx.peer_addr(), allowed_ips)
{
return Err(AuthError::ip_address_not_allowed(ctx.peer_addr()));
}

6
proxy/src/serverless/mod.rs
View File

@@ -43,6 +43,7 @@ use tokio_util::task::TaskTracker;
use tracing::{info, warn, Instrument};
use utils::http::error::ApiError;
use crate::auth::backend::jwt::JwkCache;
use crate::cancellation::CancellationHandlerMain;
use crate::config::{ProxyConfig, ProxyProtocolV2};
use crate::context::RequestContext;
@@ -331,6 +332,8 @@ async fn connection_handler(
let http_cancellation_token = CancellationToken::new();
let _cancel_connection = http_cancellation_token.clone().drop_guard();
let jwks_cache = Arc::new(JwkCache::default());
let conn_info2 = conn_info.clone();
let server = Builder::new(TokioExecutor::new());
let conn = server.serve_connection_with_upgrades(
@@ -371,6 +374,7 @@ async fn connection_handler(
http_request_token,
endpoint_rate_limiter.clone(),
cancellations,
jwks_cache.clone(),
)
.in_current_span()
.map_ok_or_else(api_error_into_response, |r| r),
@@ -419,6 +423,7 @@ async fn request_handler(
http_cancellation_token: CancellationToken,
endpoint_rate_limiter: Arc<EndpointRateLimiter>,
cancellations: TaskTracker,
jwks_cache: Arc<JwkCache>,
) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ApiError> {
let host = request
.headers()
@@ -456,6 +461,7 @@ async fn request_handler(
endpoint_rate_limiter,
host,
cancellations,
jwks_cache,
)
.await
{

3
proxy/src/serverless/websocket.rs
View File

@@ -12,6 +12,7 @@ use pin_project_lite::pin_project;
use tokio::io::{self, AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
use tracing::warn;
use crate::auth::backend::jwt::JwkCache;
use crate::cancellation::CancellationHandlerMain;
use crate::config::ProxyConfig;
use crate::context::RequestContext;
@@ -133,6 +134,7 @@ pub(crate) async fn serve_websocket(
endpoint_rate_limiter: Arc<EndpointRateLimiter>,
hostname: Option<String>,
cancellations: tokio_util::task::task_tracker::TaskTracker,
jwks_cache: Arc<JwkCache>,
) -> anyhow::Result<()> {
let websocket = websocket.await?;
let websocket = WebSocketServer::after_handshake(TokioIo::new(websocket));
@@ -152,6 +154,7 @@ pub(crate) async fn serve_websocket(
endpoint_rate_limiter,
conn_gauge,
cancellations,
jwks_cache,
))
.await;