From 3d8fcdd5286e77e52fcf741c3fadbba608bb9c1b Mon Sep 17 00:00:00 2001 From: Dawid Nowak Date: Fri, 4 Sep 2026 19:57:43 +0100 Subject: [PATCH 1/6] AuthNZ refactoring and fixing bugs Signed-off-by: Dawid Nowak --- .../src/authorization/jwks/jwks.rs | 28 +-- .../authorization/jwks/jwks_authorization.rs | 117 ++++++----- .../src/authorization/jwks/mod.rs | 1 - .../src/authorization/jwks/principal.rs | 19 -- .../src/authorization/mod.rs | 71 ++----- .../cel_principal_extractor.rs | 185 ++++++++++++++++++ .../authorization/principal_extractor/mod.rs | 3 + .../principal_extractor/principal.rs | 31 +++ .../src/layers/mod.rs | 1 + .../src/layers/principal_extractor.rs | 97 +++++++-- crates/contextforge-data-plane-lib/src/lib.rs | 22 ++- .../contextforge-data-plane-lib/src/tools.rs | 38 +++- .../tests/gateway/harness/auth.rs | 23 ++- 13 files changed, 452 insertions(+), 184 deletions(-) delete mode 100644 crates/contextforge-data-plane-lib/src/authorization/jwks/principal.rs create mode 100644 crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs create mode 100644 crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs create mode 100644 crates/contextforge-data-plane-lib/src/authorization/principal_extractor/principal.rs diff --git a/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks.rs b/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks.rs index 33e16fc9..783b75cb 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks.rs @@ -13,10 +13,7 @@ use tokio::sync::RwLock; use tracing::debug; use typed_builder::TypedBuilder; -use crate::authorization::{ - AuthorizationClaims, AuthorizationError, - jwks::principal::{DefaultPrincipalExtractor, PrincipalExtractor}, -}; +use crate::authorization::{AuthorizationClaims, AuthorizationError}; pub const JWKS_CACHE_TTL: Duration = Duration::from_mins(5); pub const JWKS_CACHE_KEY: &str = "jwks"; @@ -24,10 +21,7 @@ pub const JWKS_CACHE_KEY: &str = "jwks"; const JWKS_MAX_RESPONSE_BYTES: usize = 1024 * 1024; #[derive(TypedBuilder)] -pub(super) struct Jwks -where - T: PrincipalExtractor, -{ +pub(super) struct Jwks { client: reqwest::Client, url: Url, #[builder(default = RwLock::new(LruCache::with_expiry_duration(JWKS_CACHE_TTL)))] @@ -38,10 +32,9 @@ where validate_expiry: bool, #[builder(default = true)] validate_not_before: bool, - principal_extractor: T, } -impl Jwks { +impl Jwks { fn validation(&self, alg: Algorithm) -> Validation { let mut validation = Validation::new(alg); validation.required_spec_claims.clear(); @@ -54,17 +47,18 @@ impl Jwks { pub async fn validate(&self, token: &str, header: &Header) -> Option { { let cache = self.cache.read().await; + if let Some(keys) = cache.peek(JWKS_CACHE_KEY) && keys.iter().any(|key| key.matches(header)) { - return self.validate_with_keys(keys, token, header, &self.validation(header.alg)); + return Self::validate_with_keys(keys, token, header, &self.validation(header.alg)); } } match fetch_jwks(&self.client, &self.url).await { Ok(keys) => { let key_count = keys.len(); - let claims = self.validate_with_keys(&keys, token, header, &self.validation(header.alg)); + let claims = Self::validate_with_keys(&keys, token, header, &self.validation(header.alg)); self.cache.write().await.insert(JWKS_CACHE_KEY.to_owned(), keys); tracing::info!("validate: SaaS JWKS cache refreshed {key_count}"); @@ -78,7 +72,6 @@ impl Jwks { } fn validate_with_keys( - &self, keys: &[VerificationKey], token: &str, header: &Header, @@ -86,26 +79,23 @@ impl Jwks { ) -> Option { keys.iter() .filter(|key| key.matches(header)) - .find_map(|key| self.validate_and_decode_claims(token, &key.decoding_key, validation)) + .find_map(|key| Self::validate_and_decode_claims(token, &key.decoding_key, validation)) } fn validate_and_decode_claims( - &self, token: &str, key: &DecodingKey, validation: &Validation, ) -> Option { + println!("Validation {validation:?}"); let claims = decode::(token, key, validation) .inspect_err(|e| { debug!("validate_and_decode_claims: problem {e:?}"); }) .ok()? .claims; - let claims = claims.as_object()?; - let user_id = self.principal_extractor.user_id(claims)?; - let tenant_id = self.principal_extractor.tenant_id(claims)?; - Some(AuthorizationClaims::new(user_id, tenant_id)) + Some(AuthorizationClaims::from(claims)) } } diff --git a/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks_authorization.rs b/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks_authorization.rs index 392f0166..403ee476 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks_authorization.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks_authorization.rs @@ -1,5 +1,4 @@ use crate::authorization::jwks::jwks::Jwks; -use crate::authorization::jwks::principal::DefaultPrincipalExtractor; use crate::authorization::{AuthorizationClaims, AuthorizationError, AuthorizationService}; use async_trait::async_trait; use jsonwebtoken::decode_header; @@ -14,7 +13,7 @@ const JWKS_CONNECT_TIMEOUT: Duration = Duration::from_secs(5); const JWKS_READ_TIMEOUT: Duration = Duration::from_secs(5); pub struct JwtAuthorizationService { - jwks: Jwks, + jwks: Jwks, } impl JwtAuthorizationService { @@ -31,9 +30,7 @@ impl JwtAuthorizationService { client = client.tls_certs_only(load_ca_certificates(ca_cert_path)?); } let client = client.build().map_err(AuthorizationError::JwksRequest)?; - Ok(Self { - jwks: Jwks::builder().client(client).url(url).principal_extractor(DefaultPrincipalExtractor {}).build(), - }) + Ok(Self { jwks: Jwks::builder().client(client).url(url).build() }) } async fn authorize_token(&self, token: &str) -> Option { @@ -58,6 +55,7 @@ impl AuthorizationService for JwtAuthorizationService { let token = str::from_utf8(token).ok()?; let claims = self.authorize_token(token).await; + println!("got claims {claims:?}"); if claims.is_none() { tracing::debug!("validate_saas_jwt SaaS JWT was rejected"); } @@ -97,7 +95,6 @@ mod test { JwtAuthorizationService, jwks::{JWKS_CACHE_KEY, Jwks, VerificationKey}, jwks_authorization::{JWKS_CONNECT_TIMEOUT, JWKS_READ_TIMEOUT, JWKS_REQUEST_TIMEOUT}, - principal::DefaultPrincipalExtractor, }, }; use crate::{ @@ -114,6 +111,7 @@ mod test { use http::{HeaderMap, Request, StatusCode}; use jsonwebtoken::{Algorithm, DecodingKey, EncodingKey, Header, encode}; use lru_time_cache::LruCache; + use serde_json::json; use std::sync::{Arc, Once}; use std::{str::FromStr, time::Duration}; @@ -132,6 +130,23 @@ mod test { } } + impl AuthorizationClaims { + fn clear(&mut self, name: &str) { + if let Some(value) = self.value.get_mut(name) { + *value = serde_json::Value::Null; + } + } + + fn set(&mut self, name: &str, new_value: serde_json::Value) { + if let Some(value) = self.value.get_mut(name) { + *value = new_value; + } + } + fn get(&mut self, name: &str) -> Option<&serde_json::Value> { + self.value.get(name) + } + } + impl JwtAuthorizationService { pub async fn from_keys(verification_keys: Vec) -> Result { let url: Url = Url::from_str("http://127.0.0.1:0/").expect("this should work"); @@ -150,14 +165,7 @@ mod test { guard.insert(JWKS_CACHE_KEY.to_owned(), verification_keys); drop(guard); - Ok(Self { - jwks: Jwks::builder() - .cache(cache) - .client(client) - .url(url) - .principal_extractor(DefaultPrincipalExtractor {}) - .build(), - }) + Ok(Self { jwks: Jwks::builder().cache(cache).client(client).url(url).build() }) } } @@ -172,40 +180,36 @@ mod test { let now = now_epoch_seconds(); let user_id = "11111111-1111-1111-1111-111111111111".to_owned(); - AuthorizationClaims { - iss: GATEWAY_ISSUER.to_owned(), - sub: user_id.clone(), - aud: GATEWAY_AUDIENCE.to_owned(), - exp: now + Duration::from_hours(1).as_secs(), - nbf: Some(now - Duration::from_mins(1).as_secs()), - iat: Some(now), - jti: Uuid::new_v4().to_string(), - token_use: Some("api".to_owned()), - teams: Some(vec!["team_awesome".to_owned()]), - user: Some( - crate::authorization::User::builder() - .tenant_id("team_awesome".to_owned()) - .user_id(user_id.clone()) - .build(), - ), - scopes: Some( - Scopes::builder() - .server_id(Some("my_id".to_owned())) - .ip_restrictions(vec!["192.169.1.0/24".to_owned()]) - .permissions(vec!["tools.read".to_owned(), "servers.use".to_owned()]) - .time_restrictions(None) - .build(), - ), - tenant_id: "tenant".to_owned(), - ..Default::default() - } + let map = json!( { + "iss": GATEWAY_ISSUER.to_owned(), + "sub": user_id.clone(), + "aud": GATEWAY_AUDIENCE.to_owned(), + "exp": now + Duration::from_hours(1).as_secs(), + "nbf": now - Duration::from_mins(1).as_secs(), + "iat": now, + "jti": Uuid::new_v4().to_string(), + "token_use": Some("api".to_owned()), + "teams": vec!["team_awesome".to_owned()], + "user": crate::authorization::User::builder() + .tenant_id("team_awesome".to_owned()) + .user_id(user_id.clone()) + .build(), + "scopes": Scopes::builder() + .server_id(Some("my_id".to_owned())) + .ip_restrictions(vec!["192.169.1.0/24".to_owned()]) + .permissions(vec!["tools.read".to_owned(), "servers.use".to_owned()]) + .time_restrictions(None) + .build(), + "tenant_id": "tenant".to_owned(), + }); + AuthorizationClaims::from(map) } fn get_hmac_token_for_claims(claims: &AuthorizationClaims) -> String { let key = EncodingKey::from_secret(HMAC_SECRET); let header = Header::new(Algorithm::HS256); - - encode::(&header, claims, &key).expect("Expecting this to work") + let claims = claims.value.clone(); + encode::(&header, &claims, &key).expect("Expecting this to work") } struct MockedUserConfigStore; @@ -220,6 +224,14 @@ mod test { } } + #[test] + fn test_active_token() { + let mut claims = active_test_claims(); + assert_ne!(claims.get("exp").and_then(serde_json::Value::as_i64), Some(0_i64)); + claims.set("exp", 0.into()); + assert_eq!(claims.get("exp").and_then(serde_json::Value::as_i64), Some(0_i64)); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[allow(clippy::items_after_statements)] #[test_log::test] @@ -269,7 +281,7 @@ mod test { } let mut claims = active_test_claims(); - claims.scopes = None; + claims.clear("scopes"); let token = get_hmac_token_for_claims(&claims); let decoding_key = DecodingKey::from_secret(HMAC_SECRET); @@ -307,11 +319,18 @@ mod test { } let mut claims = active_test_claims(); - claims.token_use = None; - - claims.user = Some( - crate::authorization::User::builder().tenant_id("team_awesome".to_owned()).user_id(user_id.clone()).build(), + claims.clear("token_use"); + claims.set( + "user", + serde_json::to_value( + crate::authorization::User::builder() + .tenant_id("team_awesome".to_owned()) + .user_id(user_id.clone()) + .build(), + ) + .expect("should work"), ); + let token = get_hmac_token_for_claims(&claims); let decoding_key = DecodingKey::from_secret(HMAC_SECRET); @@ -349,7 +368,7 @@ mod test { } let mut claims = active_test_claims(); - claims.exp = 0; + claims.set("exp", 1000.into()); let token = get_hmac_token_for_claims(&claims); let decoding_key = DecodingKey::from_secret(HMAC_SECRET); diff --git a/crates/contextforge-data-plane-lib/src/authorization/jwks/mod.rs b/crates/contextforge-data-plane-lib/src/authorization/jwks/mod.rs index 02cf598f..9c94a029 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/jwks/mod.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/jwks/mod.rs @@ -1,6 +1,5 @@ #[allow(clippy::module_inception)] mod jwks; mod jwks_authorization; -mod principal; pub use jwks_authorization::JwtAuthorizationService; diff --git a/crates/contextforge-data-plane-lib/src/authorization/jwks/principal.rs b/crates/contextforge-data-plane-lib/src/authorization/jwks/principal.rs deleted file mode 100644 index dc7575da..00000000 --- a/crates/contextforge-data-plane-lib/src/authorization/jwks/principal.rs +++ /dev/null @@ -1,19 +0,0 @@ -use serde_json::Value; - -pub struct DefaultPrincipalExtractor {} - -pub trait PrincipalExtractor { - fn user_id<'a>(&self, claims: &'a serde_json::Map) -> Option<&'a str> { - ["sub", "user_id", "UserId"].into_iter().find_map(|claim| claims.get(claim)).and_then(non_empty_string) - } - - fn tenant_id<'a>(&self, claims: &'a serde_json::Map) -> Option<&'a str> { - ["tenantId", "tenant_id"].into_iter().find_map(|claim| claims.get(claim).and_then(non_empty_string)) - } -} - -impl PrincipalExtractor for DefaultPrincipalExtractor {} - -fn non_empty_string(value: &Value) -> Option<&str> { - value.as_str().filter(|value| !value.trim().is_empty()) -} diff --git a/crates/contextforge-data-plane-lib/src/authorization/mod.rs b/crates/contextforge-data-plane-lib/src/authorization/mod.rs index 6d00c433..2a6bf608 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/mod.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/mod.rs @@ -1,7 +1,7 @@ use std::{path::PathBuf, sync::Arc}; use async_trait::async_trait; -use chrono::Duration; + use http::HeaderValue; use serde::{Deserialize, Serialize}; use typed_builder::TypedBuilder; @@ -9,9 +9,9 @@ use typed_builder::TypedBuilder; use crate::Config; mod jwks; +mod principal_extractor; -pub const AUDIENCE: &str = "audience"; -pub const ISSUER: &str = "issuer"; +pub use principal_extractor::{DefaultPrincipalExtractor, PrincipalExtractor}; pub fn get_authorization_service( config: &Config, @@ -64,12 +64,6 @@ pub struct User { pub tenant_id: String, } -impl From for User { - fn from(claims: AuthorizationClaims) -> Self { - Self { user_id: claims.idp_unique_id.clone(), tenant_id: claims.tenant_id.clone() } - } -} - #[derive(Clone, Debug, Serialize, Deserialize, TypedBuilder, PartialEq)] pub struct Scopes { server_id: Option, @@ -85,51 +79,26 @@ pub struct Idp { iss: String, } -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Default, TypedBuilder)] +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, TypedBuilder)] #[serde(rename_all = "camelCase")] pub struct AuthorizationClaims { - pub iss: String, - pub jti: String, - pub aud: String, - pub exp: u64, - pub iat: Option, - pub nbf: Option, - pub tenant_id: String, - pub subscription_id: String, - pub sub: String, - pub entity_type: String, - pub email: Option, - pub name: Option, - pub displayname: Option, - pub idp: Option, - pub groups: Option>, - pub roles: Option>, - pub idp_unique_id: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub teams: Option>, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub user: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub scopes: Option, - pub token_use: Option, + value: serde_json::Value, +} + +impl From for AuthorizationClaims { + fn from(value: serde_json::Value) -> Self { + Self { value } + } +} + +impl From<&AuthorizationClaims> for serde_json::Value { + fn from(val: &AuthorizationClaims) -> Self { + val.value.clone() + } } -impl AuthorizationClaims { - pub fn new(user_id: &str, tenant_id: &str) -> Self { - let audience = AUDIENCE.to_owned(); - let start = std::time::SystemTime::now(); - let now = start.duration_since(std::time::UNIX_EPOCH).expect("Time went backwards").as_secs(); - Self { - iss: ISSUER.to_owned(), - sub: user_id.to_owned(), - aud: audience, - exp: now + Duration::hours(1).num_seconds().cast_unsigned(), - iat: Some(now), - nbf: Some(now - Duration::minutes(5).num_seconds().cast_unsigned()), - idp_unique_id: user_id.to_owned(), - tenant_id: tenant_id.to_owned(), - groups: Some(vec!["team_awesome".to_owned()]), - ..Default::default() - } +impl From for serde_json::Value { + fn from(val: AuthorizationClaims) -> Self { + val.value } } diff --git a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs new file mode 100644 index 00000000..e2ad6c2c --- /dev/null +++ b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs @@ -0,0 +1,185 @@ +use std::fs; +use std::path::Path; +use std::sync::Arc; + +use cel::{Context, Program, objects::Key}; +use serde_json::Value as JsonValue; +use thiserror::Error; +use tracing::debug; + +use crate::{AuthorizationClaims, authorization::jwks::principal::PrincipalExtractor, layers::AuthorizedPrincipal}; + +#[derive(Error, Debug)] +pub enum CelPrincipalExtractorError { + #[error("Failed to read CEL expression file: {0}")] + FileReadError(#[from] std::io::Error), + + #[error("Failed to compile CEL expression: {0}")] + CompilationError(String), + + #[error("Failed to evaluate CEL expression: {0}")] + EvaluationError(String), + + #[error("CEL expression did not return a map: {0:?}")] + InvalidReturnType(JsonValue), + + #[error("Missing required field in CEL result: {0}")] + MissingRequiredField(String), + + #[error("Invalid field type in CEL result: field={0}, expected={1}")] + InvalidFieldType(String, String), +} + +/// A CEL-based principal extractor that evaluates a CEL expression to extract +/// principal information from authorization claims. +/// +/// The CEL expression should return a map with the following fields: +/// - `user_id` (string, required): The user identifier +/// - `tenant_id` (string, required): The tenant identifier +/// - `scopes` (list of strings, optional): The user's scopes/permissions +/// +/// The CEL expression has access to the following variables: +/// - `claims`: A map containing all the authorization claims +/// - `sub`: The subject claim (shorthand for claims.sub) +/// - `tenant_id`: The tenant_id claim (shorthand for claims.tenant_id) +/// +/// Example CEL expression: +/// ```cel +/// { +/// "user_id": claims.sub, +/// "tenant_id": claims.tenant_id, +/// "scopes": [] +/// } +/// ``` +#[derive(Clone)] +pub struct CelPrincipalExtractor { + program: Arc, +} + +impl CelPrincipalExtractor { + /// Creates a new CEL principal extractor from a file containing a CEL expression. + pub fn from_file>(path: P) -> Result { + let expression = fs::read_to_string(path)?; + Self::from_expression(&expression) + } + + /// Creates a new CEL principal extractor from a CEL expression string. + pub fn from_expression(expression: &str) -> Result { + let program = + Program::compile(expression).map_err(|e| CelPrincipalExtractorError::CompilationError(e.to_string()))?; + + Ok(Self { program: Arc::new(program) }) + } +} + +impl PrincipalExtractor for CelPrincipalExtractor { + fn extract<'a>( + &self, + claims: &'a serde_json::Map, + ) -> Result, Box> { + let mut context = Context::default(); + + context + .add_variable("claims", claims) + .map_err(|e| CelPrincipalExtractorError::EvaluationError(format!("Failed to add claims variable: {e}")))?; + + // Evaluate the CEL expression + let result = + self.program.execute(&context).map_err(|e| CelPrincipalExtractorError::EvaluationError(e.to_string()))?; + + debug!("CEL expression evaluated to: {:?}", result); + + Ok(Some(AuthorizedPrincipal::try_from(result)?)) + } +} + +impl TryFrom for AuthorizedPrincipal { + type Error = CelPrincipalExtractorError; + + fn try_from(value: cel::Value) -> Result { + match value { + cel::Value::Map(map) => { + if let Some(cel::Value::String(user_id)) = map.get(&Key::from("user_id".to_owned())) + && let Some(cel::Value::String(tenant_id)) = map.get(&Key::from("tenant_id".to_owned())) + { + let user_id = (**user_id).clone(); + let tenant_id = (**tenant_id).clone(); + Ok(AuthorizedPrincipal::builder().user_id(user_id).tenant_id(tenant_id).scopes(vec![]).build()) + } else { + Err(CelPrincipalExtractorError::InvalidReturnType(serde_json::Value::Null)) + } + }, + _ => Err(CelPrincipalExtractorError::InvalidReturnType(serde_json::Value::Null)), + } + } +} + +// #[cfg(test)] +// mod tests { +// use super::*; +// use crate::AuthorizationClaims; + +// fn create_test_claims() -> AuthorizationClaims { +// AuthorizationClaims { +// sub: "user123".to_string(), +// tenant_id: "tenant456".to_string(), +// iss: "https://auth.example.com".to_string(), +// aud: "api".to_string(), +// exp: 1234567890, +// nbf: None, +// iat: Some(1234567800), +// ..Default::default() +// } +// } + +// #[test] +// fn test_extractor_creation() { +// // Test that the extractor can be created from a valid CEL expression +// let expression = r#" +// { +// "user_id": claims.sub, +// "tenant_id": claims.tenant_id, +// "scopes": [] +// } +// "#; + +// let result = CelPrincipalExtractor::from_expression(expression); +// assert!(result.is_ok(), "Should compile valid CEL expression"); +// } + +// #[test] +// fn test_parse_valid_expression() { +// let expression = r#"{"user_id": "test", "tenant_id": "tenant"}"#; +// let result = CelPrincipalExtractor::from_expression(expression); +// assert!(result.is_ok()); +// } + +// #[test] +// fn test_missing_required_field() { +// let expression = r#" +// { +// "user_id": claims.sub, +// "scopes": [] +// } +// "#; + +// let extractor = CelPrincipalExtractor::from_expression(expression).expect("Should compile"); +// let claims = create_test_claims(); +// let result = extractor.extract(&claims); + +// assert!(result.is_err()); +// assert!(matches!(result.unwrap_err(), CelPrincipalExtractorError::MissingRequiredField(_))); +// } + +// #[test] +// fn test_invalid_return_type() { +// let expression = r#""not a map""#; + +// let extractor = CelPrincipalExtractor::from_expression(expression).expect("Should compile"); +// let claims = create_test_claims(); +// let result = extractor.extract(&claims); + +// assert!(result.is_err()); +// assert!(matches!(result.unwrap_err(), CelPrincipalExtractorError::InvalidReturnType(_))); +// } +// } diff --git a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs new file mode 100644 index 00000000..40137c9d --- /dev/null +++ b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs @@ -0,0 +1,3 @@ +mod principal; + +pub use principal::{DefaultPrincipalExtractor, PrincipalExtractor}; diff --git a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/principal.rs b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/principal.rs new file mode 100644 index 00000000..1017b834 --- /dev/null +++ b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/principal.rs @@ -0,0 +1,31 @@ +use crate::layers::AuthorizedPrincipal; + +#[derive(Debug, Clone)] +pub struct DefaultPrincipalExtractor {} + +pub trait PrincipalExtractor { + fn extract( + &self, + claims: &serde_json::Value, + ) -> Result>; +} + +impl PrincipalExtractor for DefaultPrincipalExtractor { + fn extract( + &self, + claims: &serde_json::Value, + ) -> Result> { + let user_id = + ["sub", "user_id", "UserId"].into_iter().find_map(|claim| claims.get(claim)).and_then(|v| v.as_str()); + let tenant_id = + ["tenantId", "tenant_id"].into_iter().find_map(|claim| claims.get(claim)).and_then(|v| v.as_str()); + match (user_id, tenant_id) { + (Some(user_id), Some(tenant_id)) => Ok(AuthorizedPrincipal::builder() + .user_id(user_id.to_owned()) + .tenant_id(tenant_id.to_owned()) + .scopes(vec![]) + .build()), + _ => Err("Can't create principal".into()), + } + } +} diff --git a/crates/contextforge-data-plane-lib/src/layers/mod.rs b/crates/contextforge-data-plane-lib/src/layers/mod.rs index 995a8a8c..22fde56e 100644 --- a/crates/contextforge-data-plane-lib/src/layers/mod.rs +++ b/crates/contextforge-data-plane-lib/src/layers/mod.rs @@ -7,3 +7,4 @@ pub mod virtual_host_config; pub mod virtual_host_id; pub(crate) use principal_extractor::AuthorizedPrincipal; +pub use principal_extractor::PrincipalExtractorLayer; diff --git a/crates/contextforge-data-plane-lib/src/layers/principal_extractor.rs b/crates/contextforge-data-plane-lib/src/layers/principal_extractor.rs index 4a84d627..3c6a9c76 100644 --- a/crates/contextforge-data-plane-lib/src/layers/principal_extractor.rs +++ b/crates/contextforge-data-plane-lib/src/layers/principal_extractor.rs @@ -1,10 +1,11 @@ -use axum::{extract::Request, middleware::Next, response::Response}; +use axum::{extract::Request, response::Response}; use contextforge_data_plane_apis::User; +use futures::future::BoxFuture; use tracing::debug; use typed_builder::TypedBuilder; -use crate::{AuthorizationClaims, errors::unauthorized_response}; +use crate::{AuthorizationClaims, authorization::PrincipalExtractor, errors::unauthorized_response}; #[derive(Debug, Clone, TypedBuilder)] #[allow(dead_code)] @@ -20,27 +21,81 @@ impl<'a> From<&'a AuthorizedPrincipal> for User<'a> { } } -impl TryFrom<&AuthorizationClaims> for AuthorizedPrincipal { - type Error = Box; +// pub async fn principal_extractor_layer(request: http::Request, next: Next) -> Response { +// let maybe_claims = request.extensions().get::(); +// let Some(Ok(authorized_principal)) = maybe_claims.map(|claims| { +// AuthorizedPrincipal::try_from(claims).inspect_err(|e| debug!("Can't extract the principal {e:?}")) +// }) else { +// return unauthorized_response("Invalid token. Unable to extract the principal from claims"); +// }; +// let (mut parts, body) = request.into_parts(); +// parts.extensions.insert(authorized_principal); +// let request = Request::from_parts(parts, body); +// next.run(request).await +// } - fn try_from(value: &AuthorizationClaims) -> Result { - Ok(AuthorizedPrincipal::builder() - .user_id(value.sub.clone()) - .tenant_id(value.tenant_id.clone()) - .scopes(vec![]) - .build()) +use std::task::{Context, Poll}; +use tower::{Layer, Service}; + +#[derive(Clone)] +pub struct PrincipalExtractorLayer { + principal_extractor: E, +} + +impl PrincipalExtractorLayer { + pub fn new(principal_extractor: E) -> Self { + Self { principal_extractor } + } +} + +impl Layer for PrincipalExtractorLayer +where + E: Clone, +{ + type Service = PrincipalExtractorMiddleware; + + fn layer(&self, inner: S) -> Self::Service { + PrincipalExtractorMiddleware { inner, principal_extractor: self.principal_extractor.clone() } } } -pub async fn principal_extractor_layer(request: http::Request, next: Next) -> Response { - let maybe_claims = request.extensions().get::(); - let Some(Ok(authorized_principal)) = maybe_claims.map(|claims| { - AuthorizedPrincipal::try_from(claims).inspect_err(|e| debug!("Can't extract the principal {e:?}")) - }) else { - return unauthorized_response("Invalid token. Unable to extract the principal from claims"); - }; - let (mut parts, body) = request.into_parts(); - parts.extensions.insert(authorized_principal); - let request = Request::from_parts(parts, body); - next.run(request).await +#[derive(Clone)] +pub struct PrincipalExtractorMiddleware { + principal_extractor: E, + inner: S, +} + +impl Service for PrincipalExtractorMiddleware +where + S: Service + Send + 'static + Clone, + S::Future: Send + 'static, + E: PrincipalExtractor + Send + Clone + 'static, +{ + type Response = S::Response; + type Error = S::Error; + // `BoxFuture` is a type alias for `Pin>` + type Future = BoxFuture<'static, Result>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, request: Request) -> Self::Future { + let mut inner = self.inner.clone(); + let principal_extractor = self.principal_extractor.clone(); + Box::pin(async move { + let maybe_authorization_claims = request.extensions().get::(); + let Some(Ok(authorized_principal)) = maybe_authorization_claims.map(|authorization_claims| { + principal_extractor + .extract(&authorization_claims.into()) + .inspect_err(|e| debug!("Can't extract the principal {e:?}")) + }) else { + return Ok(unauthorized_response("Invalid token. Unable to extract the principal from claims")); + }; + let (mut parts, body) = request.into_parts(); + parts.extensions.insert(authorized_principal); + let request = Request::from_parts(parts, body); + inner.call(request).await + }) + } } diff --git a/crates/contextforge-data-plane-lib/src/lib.rs b/crates/contextforge-data-plane-lib/src/lib.rs index e268ad11..d06bfe37 100644 --- a/crates/contextforge-data-plane-lib/src/lib.rs +++ b/crates/contextforge-data-plane-lib/src/lib.rs @@ -39,14 +39,16 @@ pub use crate::common::*; pub type Error = Box; pub type Result = std::result::Result; -use crate::layers::{ - claims_id::claims_layer, - mcp_header_limits::{StandardHeaderLimits, mcp_header_limits_layer}, - mcp_origin::mcp_origin_layer, - principal_extractor::principal_extractor_layer, - user_config_store::user_config_store_layer, - virtual_host_config::virtual_host_config_layer, - virtual_host_id::virtual_host_id_layer, +use crate::{ + authorization::DefaultPrincipalExtractor, + layers::{ + claims_id::claims_layer, + mcp_header_limits::{StandardHeaderLimits, mcp_header_limits_layer}, + mcp_origin::mcp_origin_layer, + user_config_store::user_config_store_layer, + virtual_host_config::virtual_host_config_layer, + virtual_host_id::virtual_host_id_layer, + }, }; pub use authorization::{AuthorizationClaims, AuthorizationService, get_authorization_service}; @@ -143,11 +145,13 @@ impl Gateway { }; let mcp_standard_header_limits = StandardHeaderLimits::from(&config); + let principal_extractor_layer = layers::PrincipalExtractorLayer::new(DefaultPrincipalExtractor {}); + let app = axum::Router::new() .nest_service("/servers/{virtual_host_name}/mcp", mcp_service) .layer(middleware::from_fn(virtual_host_config_layer)) .layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), user_config_store_layer)) - .layer(middleware::from_fn(principal_extractor_layer)) + .layer(principal_extractor_layer) .layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), claims_layer)) .layer(middleware::from_fn(virtual_host_id_layer)) // Keep this outside auth/config/RMCP work so oversized MCP headers diff --git a/crates/contextforge-data-plane-lib/src/tools.rs b/crates/contextforge-data-plane-lib/src/tools.rs index 13741b69..635445db 100644 --- a/crates/contextforge-data-plane-lib/src/tools.rs +++ b/crates/contextforge-data-plane-lib/src/tools.rs @@ -12,9 +12,11 @@ use http::{ }; use jsonwebtoken::jwk::{Jwk, JwkSet}; use serde::Deserialize; -use std::fs; +use serde_json::json; +use std::{fs, time::Duration}; +use uuid::Uuid; -use crate::{authorization::AuthorizationClaims, common::ContextForgeDataPlaneAppState}; +use crate::{Scopes, authorization::AuthorizationClaims, common::ContextForgeDataPlaneAppState}; const DEFAULT_TOKEN_EMAIL: &str = "admin@example.com"; const JWKS_CACHE_CONTROL: &str = "public, max-age=300, must-revalidate"; @@ -71,11 +73,37 @@ pub async fn get_token( .expect("Expecting this to work"); let user_email = query.email.as_deref().unwrap_or(DEFAULT_TOKEN_EMAIL); - let mut claims = AuthorizationClaims::new(&user_id, user_email); - claims.tenant_id = tenant_id; + let now = + std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).expect("Time went backwards").as_secs(); + + let map = json!( { + "iss": "contexforge-dataplane", + "sub": user_id.clone(), + "aud": "contexforge-dataplane-audience", + "exp": now + Duration::from_hours(1).as_secs(), + "nbf": now - Duration::from_mins(1).as_secs(), + "iat": now, + "jti": Uuid::new_v4().to_string(), + "token_use": Some("api".to_owned()), + "teams": vec!["team_awesome".to_owned()], + "user": crate::authorization::User::builder() + .tenant_id("team_awesome".to_owned()) + .user_id(user_id.clone()) + .build(), + "scopes": Scopes::builder() + .server_id(Some("my_id".to_owned())) + .ip_restrictions(vec!["192.169.1.0/24".to_owned()]) + .permissions(vec!["tools.read".to_owned(), "servers.use".to_owned()]) + .time_restrictions(None) + .build(), + "tenant_id": tenant_id, + "user_email": user_email + }); + + let claims = serde_json::Value::from(AuthorizationClaims::from(map)); let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256); header.kid = Some("test".to_owned()); - let token = jsonwebtoken::encode::(&header, &claims, &key).expect("Expecting this to work"); + let token = jsonwebtoken::encode::(&header, &claims, &key).expect("Expecting this to work"); token.into_response() } diff --git a/crates/contextforge-data-plane-lib/tests/gateway/harness/auth.rs b/crates/contextforge-data-plane-lib/tests/gateway/harness/auth.rs index a73fdc7d..0a1495fc 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/harness/auth.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/harness/auth.rs @@ -13,14 +13,12 @@ use super::TEST_USER_EMAIL; const TEST_TOKEN_TTL_SECS: u64 = 60 * 60; -pub(crate) fn token(user_id: &str) -> String { - let key = EncodingKey::from_rsa_pem(&fs::read("../../assets/jwt.key").expect("jwt key")).expect("encoding key"); - let mut header = Header::new(Algorithm::RS256); - header.kid = Some("test".to_owned()); +fn default_claims(user_id: &str) -> serde_json::Value { let now = SystemTime::now().duration_since(UNIX_EPOCH).expect("system clock").as_secs(); - let claims = json!({ + json!({ "iss": "mcpgateway", "sub": user_id, + "tenant_id": "test_tenant", "aud": "mcpgateway-api", "exp": now + TEST_TOKEN_TTL_SECS, "iat": now, @@ -39,8 +37,14 @@ pub(crate) fn token(user_id: &str) -> String { "ip_restrictions": ["192.169.1.0/24"], "time_restrictions": null }, - }); - encode(&header, &claims, &key).expect("jwt token") + }) +} + +pub(crate) fn token(user_id: &str) -> String { + let key = EncodingKey::from_rsa_pem(&fs::read("../../assets/jwt.key").expect("jwt key")).expect("encoding key"); + let mut header = Header::new(Algorithm::RS256); + header.kid = Some("test".to_owned()); + encode(&header, &default_claims(user_id), &key).expect("jwt token") } #[derive(Debug)] @@ -53,11 +57,10 @@ impl AlwaysAllowAuthorizatioService { Self { user } } } + #[async_trait] impl AuthorizationService for AlwaysAllowAuthorizatioService { async fn authorize(&self, _: &HeaderValue) -> Option { - let mut claims = AuthorizationClaims::default(); - claims.sub.clone_from(&self.user); - Some(claims) + Some(AuthorizationClaims::from(default_claims(&self.user))) } } From 656e5d383faaf99efb1a95622c290dea95f50864 Mon Sep 17 00:00:00 2001 From: Dawid Nowak Date: Fri, 4 Sep 2026 20:14:33 +0100 Subject: [PATCH 2/6] AuthNZ refactoring and CEL evaluator Signed-off-by: Dawid Nowak --- Cargo.lock | 139 ++++++++++++++++++ crates/contextforge-data-plane-lib/Cargo.toml | 2 +- .../src/authorization/jwks/jwks.rs | 1 - .../authorization/jwks/jwks_authorization.rs | 1 - .../src/authorization/mod.rs | 4 +- .../cel_principal_extractor.rs | 10 +- ...ipal.rs => default_principal_extractor.rs} | 9 +- .../authorization/principal_extractor/mod.rs | 29 +++- .../contextforge-data-plane-lib/src/common.rs | 5 + .../src/layers/mod.rs | 1 - .../src/layers/principal_extractor.rs | 16 -- .../src/layers/user_config_store.rs | 3 +- crates/contextforge-data-plane-lib/src/lib.rs | 9 +- .../tests/gateway/harness/mod.rs | 1 + 14 files changed, 190 insertions(+), 40 deletions(-) rename crates/contextforge-data-plane-lib/src/authorization/principal_extractor/{principal.rs => default_principal_extractor.rs} (80%) diff --git a/Cargo.lock b/Cargo.lock index 0a2824f2..f07352a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -97,6 +97,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "antlr4rust" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "093d520274bfff7278d776f7ea12981a0a0a6f96db90964658e0f38fc6e9a6a6" +dependencies = [ + "better_any", + "bit-set", + "byteorder", + "lazy_static", + "murmur3", + "once_cell", + "parking_lot", + "typed-arena", + "uuid", +] + [[package]] name = "anyhow" version = "1.0.104" @@ -343,6 +360,27 @@ version = "0.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" +[[package]] +name = "better_any" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4372b9543397a4b86050cc5e7ee36953edf4bac9518e8a774c2da694977fb6e4" + +[[package]] +name = "bit-set" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" + [[package]] name = "bitflags" version = "2.13.1" @@ -385,6 +423,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.12.1" @@ -403,6 +447,22 @@ dependencies = [ "shlex", ] +[[package]] +name = "cel" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f93082a93da8fd78394852602ced1e2e7754ed8c29dc813fa80b01a1baf9032" +dependencies = [ + "antlr4rust", + "chrono", + "lazy_static", + "nom", + "pastey", + "regex", + "serde", + "thiserror 2.0.19", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -602,6 +662,7 @@ dependencies = [ "axum-otel-metrics", "axum-server", "base64 0.22.1", + "cel", "chrono", "clap", "contextforge-data-plane-apis", @@ -1581,6 +1642,15 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.33" @@ -1636,6 +1706,12 @@ dependencies = [ "unicase", ] +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -1657,6 +1733,25 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "murmur3" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a198f9589efc03f544388dfc4a19fe8af4323662b62f598b8dcfdac62c14771c" +dependencies = [ + "byteorder", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -1845,6 +1940,29 @@ version = "2.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "pastey" version = "0.2.3" @@ -2164,6 +2282,15 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "ref-cast" version = "1.0.26" @@ -2486,6 +2613,12 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "security-framework" version = "3.7.0" @@ -3270,6 +3403,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typed-arena" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a" + [[package]] name = "typed-builder" version = "0.23.2" diff --git a/crates/contextforge-data-plane-lib/Cargo.toml b/crates/contextforge-data-plane-lib/Cargo.toml index 7c2df933..72f57f30 100644 --- a/crates/contextforge-data-plane-lib/Cargo.toml +++ b/crates/contextforge-data-plane-lib/Cargo.toml @@ -46,7 +46,7 @@ rustls-pki-types = { version = "1.14.1", features = ["std", "alloc"] } tokio-rustls = "0.26.4" typed-builder = "0.23.2" url = { workspace = true, features = ["serde"] } - +cel = "0.14.4" diff --git a/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks.rs b/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks.rs index 783b75cb..5482d25b 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks.rs @@ -87,7 +87,6 @@ impl Jwks { key: &DecodingKey, validation: &Validation, ) -> Option { - println!("Validation {validation:?}"); let claims = decode::(token, key, validation) .inspect_err(|e| { debug!("validate_and_decode_claims: problem {e:?}"); diff --git a/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks_authorization.rs b/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks_authorization.rs index 403ee476..b3de29c4 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks_authorization.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/jwks/jwks_authorization.rs @@ -55,7 +55,6 @@ impl AuthorizationService for JwtAuthorizationService { let token = str::from_utf8(token).ok()?; let claims = self.authorize_token(token).await; - println!("got claims {claims:?}"); if claims.is_none() { tracing::debug!("validate_saas_jwt SaaS JWT was rejected"); } diff --git a/crates/contextforge-data-plane-lib/src/authorization/mod.rs b/crates/contextforge-data-plane-lib/src/authorization/mod.rs index 2a6bf608..c8d1ef38 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/mod.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/mod.rs @@ -11,7 +11,9 @@ use crate::Config; mod jwks; mod principal_extractor; -pub use principal_extractor::{DefaultPrincipalExtractor, PrincipalExtractor}; +pub use principal_extractor::{ + AuthorizedPrincipal, CelPrincipalExtractor, DefaultPrincipalExtractor, PrincipalExtractor, +}; pub fn get_authorization_service( config: &Config, diff --git a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs index e2ad6c2c..a335bce9 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs @@ -7,7 +7,7 @@ use serde_json::Value as JsonValue; use thiserror::Error; use tracing::debug; -use crate::{AuthorizationClaims, authorization::jwks::principal::PrincipalExtractor, layers::AuthorizedPrincipal}; +use crate::authorization::{AuthorizedPrincipal, PrincipalExtractor}; #[derive(Error, Debug)] pub enum CelPrincipalExtractorError { @@ -73,10 +73,10 @@ impl CelPrincipalExtractor { } impl PrincipalExtractor for CelPrincipalExtractor { - fn extract<'a>( + fn extract( &self, - claims: &'a serde_json::Map, - ) -> Result, Box> { + claims: &serde_json::Value, + ) -> Result> { let mut context = Context::default(); context @@ -89,7 +89,7 @@ impl PrincipalExtractor for CelPrincipalExtractor { debug!("CEL expression evaluated to: {:?}", result); - Ok(Some(AuthorizedPrincipal::try_from(result)?)) + Ok(AuthorizedPrincipal::try_from(result)?) } } diff --git a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/principal.rs b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/default_principal_extractor.rs similarity index 80% rename from crates/contextforge-data-plane-lib/src/authorization/principal_extractor/principal.rs rename to crates/contextforge-data-plane-lib/src/authorization/principal_extractor/default_principal_extractor.rs index 1017b834..a8a9488a 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/principal.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/default_principal_extractor.rs @@ -1,15 +1,8 @@ -use crate::layers::AuthorizedPrincipal; +use crate::authorization::{AuthorizedPrincipal, PrincipalExtractor}; #[derive(Debug, Clone)] pub struct DefaultPrincipalExtractor {} -pub trait PrincipalExtractor { - fn extract( - &self, - claims: &serde_json::Value, - ) -> Result>; -} - impl PrincipalExtractor for DefaultPrincipalExtractor { fn extract( &self, diff --git a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs index 40137c9d..c38b376e 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/mod.rs @@ -1,3 +1,28 @@ -mod principal; +mod cel_principal_extractor; +mod default_principal_extractor; +use contextforge_data_plane_apis::User; +use typed_builder::TypedBuilder; -pub use principal::{DefaultPrincipalExtractor, PrincipalExtractor}; +pub use cel_principal_extractor::CelPrincipalExtractor; +pub use default_principal_extractor::DefaultPrincipalExtractor; + +#[derive(Debug, Clone, TypedBuilder)] +#[allow(dead_code)] +pub struct AuthorizedPrincipal { + user_id: String, + tenant_id: String, + scopes: Vec, +} + +impl<'a> From<&'a AuthorizedPrincipal> for User<'a> { + fn from(value: &'a AuthorizedPrincipal) -> Self { + Self::new(&value.user_id) + } +} + +pub trait PrincipalExtractor { + fn extract( + &self, + claims: &serde_json::Value, + ) -> Result>; +} diff --git a/crates/contextforge-data-plane-lib/src/common.rs b/crates/contextforge-data-plane-lib/src/common.rs index eed6b8f4..20cb026b 100644 --- a/crates/contextforge-data-plane-lib/src/common.rs +++ b/crates/contextforge-data-plane-lib/src/common.rs @@ -269,6 +269,9 @@ pub struct Config { #[cfg(feature = "with_tools")] #[arg(long)] pub token_verification_private_key: PathBuf, + + #[arg(long)] + pub cel_principal_extractor_path: PathBuf, } pub const DEFAULT_MCP_STANDARD_HEADER_MAX_COUNT: usize = 32; @@ -460,6 +463,8 @@ mod tests { log_rotation: None, mcp_allowed_origins: None, mcp_allowed_hosts: None, + cel_principal_extractor_path: PathBuf::from_str("./assets/principal_extractor.cel") + .expect("This should work"), #[cfg(feature = "with_tools")] token_verification_private_key: PathBuf::from_str("./assets/jwt.key").expect("This should work"), diff --git a/crates/contextforge-data-plane-lib/src/layers/mod.rs b/crates/contextforge-data-plane-lib/src/layers/mod.rs index 22fde56e..9b96028c 100644 --- a/crates/contextforge-data-plane-lib/src/layers/mod.rs +++ b/crates/contextforge-data-plane-lib/src/layers/mod.rs @@ -6,5 +6,4 @@ pub mod user_config_store; pub mod virtual_host_config; pub mod virtual_host_id; -pub(crate) use principal_extractor::AuthorizedPrincipal; pub use principal_extractor::PrincipalExtractorLayer; diff --git a/crates/contextforge-data-plane-lib/src/layers/principal_extractor.rs b/crates/contextforge-data-plane-lib/src/layers/principal_extractor.rs index 3c6a9c76..87e18c1c 100644 --- a/crates/contextforge-data-plane-lib/src/layers/principal_extractor.rs +++ b/crates/contextforge-data-plane-lib/src/layers/principal_extractor.rs @@ -1,26 +1,10 @@ use axum::{extract::Request, response::Response}; -use contextforge_data_plane_apis::User; use futures::future::BoxFuture; use tracing::debug; -use typed_builder::TypedBuilder; use crate::{AuthorizationClaims, authorization::PrincipalExtractor, errors::unauthorized_response}; -#[derive(Debug, Clone, TypedBuilder)] -#[allow(dead_code)] -pub struct AuthorizedPrincipal { - user_id: String, - tenant_id: String, - scopes: Vec, -} - -impl<'a> From<&'a AuthorizedPrincipal> for User<'a> { - fn from(value: &'a AuthorizedPrincipal) -> Self { - Self::new(&value.user_id) - } -} - // pub async fn principal_extractor_layer(request: http::Request, next: Next) -> Response { // let maybe_claims = request.extensions().get::(); // let Some(Ok(authorized_principal)) = maybe_claims.map(|claims| { diff --git a/crates/contextforge-data-plane-lib/src/layers/user_config_store.rs b/crates/contextforge-data-plane-lib/src/layers/user_config_store.rs index facea51e..abab56d0 100644 --- a/crates/contextforge-data-plane-lib/src/layers/user_config_store.rs +++ b/crates/contextforge-data-plane-lib/src/layers/user_config_store.rs @@ -4,6 +4,7 @@ use contextforge_data_plane_apis::User; use tracing::{debug, info, warn}; use crate::{ + authorization::AuthorizedPrincipal, common::ContextForgeDataPlaneAppState, errors::{bad_request, internal_server_error}, user_config_store::ConfigStoreError, @@ -16,7 +17,7 @@ pub async fn user_config_store_layer( ) -> Response { let method = request.method().clone(); let path = request.uri().path().to_owned(); - let maybe_principal = request.extensions().get::(); + let maybe_principal = request.extensions().get::(); if let Some(principal) = maybe_principal { debug!( "user_config_store_layer - getting user config for principal {principal:?} method = {method} path = {path}" diff --git a/crates/contextforge-data-plane-lib/src/lib.rs b/crates/contextforge-data-plane-lib/src/lib.rs index d06bfe37..71d7c6b7 100644 --- a/crates/contextforge-data-plane-lib/src/lib.rs +++ b/crates/contextforge-data-plane-lib/src/lib.rs @@ -40,7 +40,7 @@ pub type Error = Box; pub type Result = std::result::Result; use crate::{ - authorization::DefaultPrincipalExtractor, + authorization::{CelPrincipalExtractor, DefaultPrincipalExtractor}, layers::{ claims_id::claims_layer, mcp_header_limits::{StandardHeaderLimits, mcp_header_limits_layer}, @@ -145,13 +145,16 @@ impl Gateway { }; let mcp_standard_header_limits = StandardHeaderLimits::from(&config); - let principal_extractor_layer = layers::PrincipalExtractorLayer::new(DefaultPrincipalExtractor {}); + let cel_principal_extractor_layer = layers::PrincipalExtractorLayer::new(CelPrincipalExtractor::from_file( + config.cel_principal_extractor_path.clone(), + )?); + let default_principal_extractor_layer = layers::PrincipalExtractorLayer::new(DefaultPrincipalExtractor {}); let app = axum::Router::new() .nest_service("/servers/{virtual_host_name}/mcp", mcp_service) .layer(middleware::from_fn(virtual_host_config_layer)) .layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), user_config_store_layer)) - .layer(principal_extractor_layer) + .layer(cel_principal_extractor_layer) .layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), claims_layer)) .layer(middleware::from_fn(virtual_host_id_layer)) // Keep this outside auth/config/RMCP work so oversized MCP headers diff --git a/crates/contextforge-data-plane-lib/tests/gateway/harness/mod.rs b/crates/contextforge-data-plane-lib/tests/gateway/harness/mod.rs index a718e9d7..6cb2ae8c 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/harness/mod.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/harness/mod.rs @@ -79,6 +79,7 @@ pub fn create_default_config() -> Config { log_rotation: None, mcp_allowed_origins: None, mcp_allowed_hosts: None, + cel_principal_extractor_path: PathBuf::from_str("./assets/principal_extractor.cel").expect("This should work"), #[cfg(feature = "with_tools")] token_verification_private_key: PathBuf::from_str("./assets/jwt.key").expect("This should work"), From ab64c1ec4a8c6e3e6a31375bb43d33902527ac1e Mon Sep 17 00:00:00 2001 From: Dawid Nowak Date: Fri, 4 Sep 2026 20:14:51 +0100 Subject: [PATCH 3/6] AuthNZ refactoring and CEL evaluator Signed-off-by: Dawid Nowak --- assets/principal_extractor.cel | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 assets/principal_extractor.cel diff --git a/assets/principal_extractor.cel b/assets/principal_extractor.cel new file mode 100644 index 00000000..6755c1ec --- /dev/null +++ b/assets/principal_extractor.cel @@ -0,0 +1,5 @@ +{ + "user_id": claims.sub, + "tenant_id": claims.tenant_id, + "scopes": [] +} \ No newline at end of file From 7cf11763988fcf220b0c6614948e3b5292b1ba46 Mon Sep 17 00:00:00 2001 From: Dawid Nowak Date: Mon, 7 Sep 2026 10:24:55 +0100 Subject: [PATCH 4/6] AuthNZ refactoring and CEL evaluator Signed-off-by: Dawid Nowak --- .../contextforge-data-plane-lib/src/common.rs | 5 ++-- crates/contextforge-data-plane-lib/src/lib.rs | 22 ++++++++-------- .../contextforge-data-plane-lib/src/tools.rs | 25 +++++++++++++++++++ .../tests/gateway/harness/mod.rs | 2 +- 4 files changed, 39 insertions(+), 15 deletions(-) diff --git a/crates/contextforge-data-plane-lib/src/common.rs b/crates/contextforge-data-plane-lib/src/common.rs index 20cb026b..c727cad2 100644 --- a/crates/contextforge-data-plane-lib/src/common.rs +++ b/crates/contextforge-data-plane-lib/src/common.rs @@ -271,7 +271,7 @@ pub struct Config { pub token_verification_private_key: PathBuf, #[arg(long)] - pub cel_principal_extractor_path: PathBuf, + pub cel_principal_extractor_path: Option, } pub const DEFAULT_MCP_STANDARD_HEADER_MAX_COUNT: usize = 32; @@ -463,8 +463,7 @@ mod tests { log_rotation: None, mcp_allowed_origins: None, mcp_allowed_hosts: None, - cel_principal_extractor_path: PathBuf::from_str("./assets/principal_extractor.cel") - .expect("This should work"), + cel_principal_extractor_path: None, #[cfg(feature = "with_tools")] token_verification_private_key: PathBuf::from_str("./assets/jwt.key").expect("This should work"), diff --git a/crates/contextforge-data-plane-lib/src/lib.rs b/crates/contextforge-data-plane-lib/src/lib.rs index 71d7c6b7..b18e2496 100644 --- a/crates/contextforge-data-plane-lib/src/lib.rs +++ b/crates/contextforge-data-plane-lib/src/lib.rs @@ -145,24 +145,24 @@ impl Gateway { }; let mcp_standard_header_limits = StandardHeaderLimits::from(&config); - let cel_principal_extractor_layer = layers::PrincipalExtractorLayer::new(CelPrincipalExtractor::from_file( - config.cel_principal_extractor_path.clone(), - )?); - let default_principal_extractor_layer = layers::PrincipalExtractorLayer::new(DefaultPrincipalExtractor {}); - let app = axum::Router::new() .nest_service("/servers/{virtual_host_name}/mcp", mcp_service) .layer(middleware::from_fn(virtual_host_config_layer)) - .layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), user_config_store_layer)) - .layer(cel_principal_extractor_layer) + .layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), user_config_store_layer)); + + let app = if let Some(cel_principal_extractor_path) = config.cel_principal_extractor_path.as_ref() { + app.layer(layers::PrincipalExtractorLayer::new(CelPrincipalExtractor::from_file( + cel_principal_extractor_path, + )?)) + } else { + app.layer(layers::PrincipalExtractorLayer::new(DefaultPrincipalExtractor {})) + }; + + let app = app .layer(middleware::from_fn_with_state(mcp_gateway_state.clone(), claims_layer)) .layer(middleware::from_fn(virtual_host_id_layer)) - // Keep this outside auth/config/RMCP work so oversized MCP headers - // are rejected before JWT validation or body parsing. .layer(middleware::from_fn_with_state(mcp_standard_header_limits, mcp_header_limits_layer)) .layer(cors_layer) - // mcp_origin_layer is the outermost wrapper: fires before JWT auth, - // session creation, and backend fan-out. .layer(middleware::from_fn_with_state(config.clone(), mcp_origin_layer)); #[cfg(feature = "with_tools")] diff --git a/crates/contextforge-data-plane-lib/src/tools.rs b/crates/contextforge-data-plane-lib/src/tools.rs index 635445db..b95f6342 100644 --- a/crates/contextforge-data-plane-lib/src/tools.rs +++ b/crates/contextforge-data-plane-lib/src/tools.rs @@ -49,6 +49,7 @@ async fn get_jwks(State(state): State) -> Respons pub fn add_tools(router: Router) -> Router { router .route(TOKEN_PATH, get(get_token)) + .route(TOKEN_PATH, post(get_custom_token)) .route(JWKS_PATH, get(get_jwks)) .route(CONFIGURE_USER_PATH, post(configure_user)) .route("/health", get(health)) @@ -62,6 +63,30 @@ pub async fn health() -> Response { .expect("Expecting this to work") } +pub async fn get_custom_token( + State(state): State, + Json(mut claims): Json, +) -> Response { + let key = jsonwebtoken::EncodingKey::from_rsa_pem( + &fs::read(&state.config.token_verification_private_key).expect("Expecting this to work"), + ) + .expect("Expecting this to work"); + + let now = + std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).expect("Time went backwards").as_secs(); + + claims["exp"] = (now + Duration::from_hours(1).as_secs()).into(); + claims["nbf"] = (now - Duration::from_mins(1).as_secs()).into(); + claims["iat"] = (now).into(); + claims["jti"] = Uuid::new_v4().to_string().into(); + + let mut header = jsonwebtoken::Header::new(jsonwebtoken::Algorithm::RS256); + header.kid = Some("test".to_owned()); + let token = jsonwebtoken::encode::(&header, &claims, &key).expect("Expecting this to work"); + + token.into_response() +} + pub async fn get_token( State(state): State, Path((tenant_id, user_id)): Path<(String, String)>, diff --git a/crates/contextforge-data-plane-lib/tests/gateway/harness/mod.rs b/crates/contextforge-data-plane-lib/tests/gateway/harness/mod.rs index 6cb2ae8c..31ce2063 100644 --- a/crates/contextforge-data-plane-lib/tests/gateway/harness/mod.rs +++ b/crates/contextforge-data-plane-lib/tests/gateway/harness/mod.rs @@ -79,7 +79,7 @@ pub fn create_default_config() -> Config { log_rotation: None, mcp_allowed_origins: None, mcp_allowed_hosts: None, - cel_principal_extractor_path: PathBuf::from_str("./assets/principal_extractor.cel").expect("This should work"), + cel_principal_extractor_path: None, #[cfg(feature = "with_tools")] token_verification_private_key: PathBuf::from_str("./assets/jwt.key").expect("This should work"), From feec17c7614eec08c9b5be73467c3c1b5b059152 Mon Sep 17 00:00:00 2001 From: Dawid Nowak Date: Mon, 7 Sep 2026 10:36:55 +0100 Subject: [PATCH 5/6] AuthNZ refactoring and CEL evaluator.Unit tests Signed-off-by: Dawid Nowak --- .../cel_principal_extractor.rs | 138 +++++++++--------- 1 file changed, 66 insertions(+), 72 deletions(-) diff --git a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs index a335bce9..aaf9d8f0 100644 --- a/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs +++ b/crates/contextforge-data-plane-lib/src/authorization/principal_extractor/cel_principal_extractor.rs @@ -57,13 +57,11 @@ pub struct CelPrincipalExtractor { } impl CelPrincipalExtractor { - /// Creates a new CEL principal extractor from a file containing a CEL expression. pub fn from_file>(path: P) -> Result { let expression = fs::read_to_string(path)?; Self::from_expression(&expression) } - /// Creates a new CEL principal extractor from a CEL expression string. pub fn from_expression(expression: &str) -> Result { let program = Program::compile(expression).map_err(|e| CelPrincipalExtractorError::CompilationError(e.to_string()))?; @@ -83,7 +81,6 @@ impl PrincipalExtractor for CelPrincipalExtractor { .add_variable("claims", claims) .map_err(|e| CelPrincipalExtractorError::EvaluationError(format!("Failed to add claims variable: {e}")))?; - // Evaluate the CEL expression let result = self.program.execute(&context).map_err(|e| CelPrincipalExtractorError::EvaluationError(e.to_string()))?; @@ -114,72 +111,69 @@ impl TryFrom for AuthorizedPrincipal { } } -// #[cfg(test)] -// mod tests { -// use super::*; -// use crate::AuthorizationClaims; - -// fn create_test_claims() -> AuthorizationClaims { -// AuthorizationClaims { -// sub: "user123".to_string(), -// tenant_id: "tenant456".to_string(), -// iss: "https://auth.example.com".to_string(), -// aud: "api".to_string(), -// exp: 1234567890, -// nbf: None, -// iat: Some(1234567800), -// ..Default::default() -// } -// } - -// #[test] -// fn test_extractor_creation() { -// // Test that the extractor can be created from a valid CEL expression -// let expression = r#" -// { -// "user_id": claims.sub, -// "tenant_id": claims.tenant_id, -// "scopes": [] -// } -// "#; - -// let result = CelPrincipalExtractor::from_expression(expression); -// assert!(result.is_ok(), "Should compile valid CEL expression"); -// } - -// #[test] -// fn test_parse_valid_expression() { -// let expression = r#"{"user_id": "test", "tenant_id": "tenant"}"#; -// let result = CelPrincipalExtractor::from_expression(expression); -// assert!(result.is_ok()); -// } - -// #[test] -// fn test_missing_required_field() { -// let expression = r#" -// { -// "user_id": claims.sub, -// "scopes": [] -// } -// "#; - -// let extractor = CelPrincipalExtractor::from_expression(expression).expect("Should compile"); -// let claims = create_test_claims(); -// let result = extractor.extract(&claims); - -// assert!(result.is_err()); -// assert!(matches!(result.unwrap_err(), CelPrincipalExtractorError::MissingRequiredField(_))); -// } - -// #[test] -// fn test_invalid_return_type() { -// let expression = r#""not a map""#; - -// let extractor = CelPrincipalExtractor::from_expression(expression).expect("Should compile"); -// let claims = create_test_claims(); -// let result = extractor.extract(&claims); - -// assert!(result.is_err()); -// assert!(matches!(result.unwrap_err(), CelPrincipalExtractorError::InvalidReturnType(_))); -// } -// } +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + fn create_test_claims() -> serde_json::Value { + json!( { + "sub": "user123", + "tenant_id": "tenant456", + "iss": "https://auth.example.com", + "aud": "api", + "exp": "1234567890", + "iat": "1234567800" + }) + } + + #[test] + fn test_extractor_creation() { + // Test that the extractor can be created from a valid CEL expression + let expression = r#" + { + "user_id": claims.sub, + "tenant_id": claims.tenant_id, + "scopes": [] + } + "#; + + let result = CelPrincipalExtractor::from_expression(expression); + assert!(result.is_ok(), "Should compile valid CEL expression"); + } + + #[test] + fn test_parse_valid_expression() { + let expression = r#"{"user_id": "test", "tenant_id": "tenant"}"#; + let result = CelPrincipalExtractor::from_expression(expression); + assert!(result.is_ok()); + } + + #[test] + fn test_missing_required_field() { + let expression = r#" + { + "user_id": claims.sub, + "scopes": [] + } + "#; + + let extractor = CelPrincipalExtractor::from_expression(expression).expect("Should compile"); + let claims = create_test_claims(); + let result = extractor.extract(&claims); + + assert!(result.is_err()); + } + + #[test] + fn test_invalid_return_type() { + let expression = r#""not a map""#; + + let extractor = CelPrincipalExtractor::from_expression(expression).expect("Should compile"); + let claims = create_test_claims(); + let result = extractor.extract(&claims); + + assert!(result.is_err()); + } +} From 6c05e0d959491e79897132be49944070efc29e28 Mon Sep 17 00:00:00 2001 From: Dawid Nowak Date: Mon, 7 Sep 2026 12:57:05 +0100 Subject: [PATCH 6/6] AuthNZ refactoring and CEL evaluator.Unit tests Signed-off-by: Dawid Nowak --- Cargo.lock | 1 - crates/contextforge-data-plane-lib/Cargo.toml | 1 - 2 files changed, 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f07352a9..51bfc1e4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -663,7 +663,6 @@ dependencies = [ "axum-server", "base64 0.22.1", "cel", - "chrono", "clap", "contextforge-data-plane-apis", "contextforge-data-plane-cpex", diff --git a/crates/contextforge-data-plane-lib/Cargo.toml b/crates/contextforge-data-plane-lib/Cargo.toml index 72f57f30..40f2e574 100644 --- a/crates/contextforge-data-plane-lib/Cargo.toml +++ b/crates/contextforge-data-plane-lib/Cargo.toml @@ -29,7 +29,6 @@ tower = "0.5.3" http.workspace = true futures = { version = "0.3", features = ["std", "alloc"] } jsonwebtoken.workspace = true -chrono = "0.4.44" redis.workspace = true clap.workspace = true thiserror.workspace = true