PoC: downstream role mapping and MCP permissions

crates/contextforge-data-plane-lib/src/authorization/mod.rs

@@ -15,6 +15,12 @@ pub use principal_extractor::{
     AuthorizedPrincipal, CelPrincipalExtractor, DefaultPrincipalExtractor, PrincipalExtractor,
 };
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub enum Permission {
+    Admin,
+    MCPUser,
+}
+
 pub fn get_authorization_service(
     config: &JwksConfig,
 ) -> Result<Arc<dyn AuthorizationService + Send + Sync>, AuthorizationError> {
@@ -27,6 +33,12 @@ pub trait AuthorizationService: std::fmt::Debug {
     async fn authorize(&self, authorization_token: &HeaderValue) -> Option<AuthorizationClaims>;
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
+pub enum AuthenticationError {
+    #[error("invalid bearer token")]
+    InvalidToken,
+}
+
 #[derive(Debug, thiserror::Error)]
 #[allow(dead_code)]
 pub enum AuthorizationError {
@@ -86,6 +98,12 @@ pub struct AuthorizationClaims {
     value: serde_json::Value,
 }
 
+impl AuthorizationClaims {
+    pub fn as_value(&self) -> &serde_json::Value {
+        &self.value
+    }
+}
+
 impl From<serde_json::Value> for AuthorizationClaims {
     fn from(value: serde_json::Value) -> Self {
         Self { value }

crates/contextforge-data-plane-lib/src/authorization/principal_extractor/default_principal_extractor.rs

@@ -13,8 +13,10 @@ impl PrincipalExtractor for DefaultPrincipalExtractor {
     ) -> Result<AuthorizedPrincipal, Box<dyn std::error::Error + Send + Sync>> {
         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());
+        let tenant_id = ["tenantId", "tenant_id", "woTenantId"]
+            .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())
@@ -25,3 +27,44 @@ impl PrincipalExtractor for DefaultPrincipalExtractor {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn supported_identity_aliases_are_preserved() {
+        for user_claim in ["sub", "user_id", "UserId"] {
+            for tenant_claim in ["tenantId", "tenant_id", "woTenantId"] {
+                let claims = json!({user_claim: "user", tenant_claim: "tenant"});
+                let principal = DefaultPrincipalExtractor {}.extract(&claims).unwrap();
+                assert_eq!(principal.user_id, "user");
+                assert_eq!(principal.tenant_id, "tenant");
+                assert!(principal.scopes.is_empty());
+            }
+        }
+    }
+
+    #[test]
+    fn existing_alias_precedence_is_preserved() {
+        let claims = json!({
+            "sub": "subject", "user_id": "alternate", "UserId": "other",
+            "tenantId": "tenant", "tenant_id": "alternate", "woTenantId": "other"
+        });
+        let principal = DefaultPrincipalExtractor {}.extract(&claims).unwrap();
+        assert_eq!(principal.user_id, "subject");
+        assert_eq!(principal.tenant_id, "tenant");
+    }
+
+    #[test]
+    fn missing_identity_is_rejected() {
+        for claims in [
+            json!({"sub": "user"}),
+            json!({"tenant_id": "tenant"}),
+            json!({"woUserId": "user", "woTenantId": "tenant"}),
+        ] {
+            assert!(DefaultPrincipalExtractor {}.extract(&claims).is_err());
+        }
+    }
+}

crates/contextforge-data-plane-lib/src/layers/mod.rs

@@ -6,3 +6,5 @@ pub mod virtual_host_config;
 pub mod virtual_host_id;
 
 pub use principal_extractor::PrincipalExtractorLayer;
+
+pub mod permission;

crates/contextforge-data-plane-lib/src/layers/permission.rs

@@ -0,0 +1,128 @@
+use crate::{
+    authorization::{AuthenticationError, AuthorizationClaims, AuthorizedPrincipal, Permission},
+    errors::{custom_error, unauthorized_response},
+};
+use axum::{
+    extract::{Request, State},
+    middleware::Next,
+    response::Response,
+};
+use http::StatusCode;
+
+pub async fn require_permission(State(permission): State<Permission>, request: Request, next: Next) -> Response {
+    if request.extensions().get::<AuthorizedPrincipal>().is_none() {
+        return unauthorized_response("Missing verified identity");
+    }
+    let Some(claims) = request.extensions().get::<AuthorizationClaims>() else {
+        return unauthorized_response("Missing verified claims");
+    };
+    match has_permission(claims, permission) {
+        Ok(true) => next.run(request).await,
+        Ok(false) => custom_error(StatusCode::FORBIDDEN, "Insufficient permission"),
+        Err(_) => unauthorized_response("Invalid role claims"),
+    }
+}
+
+/// Test role mapping; awaiting confirmation from WxO.
+fn has_permission(claims: &AuthorizationClaims, permission: Permission) -> Result<bool, AuthenticationError> {
+    let allows = |role: &str| match role {
+        "admin" => true,
+        "builder" | "user" => permission == Permission::MCPUser,
+        _ => false,
+    };
+    let claims = claims.as_value();
+    let mut granted = false;
+    if let Some(roles) = claims.get("roles") {
+        for role in roles.as_array().ok_or(AuthenticationError::InvalidToken)? {
+            granted |= allows(role.as_str().ok_or(AuthenticationError::InvalidToken)?);
+        }
+    }
+    if let Some(role) = claims.get("role") {
+        granted |= allows(role.as_str().ok_or(AuthenticationError::InvalidToken)?);
+    }
+    Ok(granted)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::authorization::{CelPrincipalExtractor, DefaultPrincipalExtractor, PrincipalExtractor};
+    use axum::{Router, body::Body, middleware, routing::get};
+    use serde_json::json;
+    use tower::ServiceExt;
+
+    fn app(permission: Permission) -> Router {
+        Router::new()
+            .route("/", get(|| async { StatusCode::NO_CONTENT }))
+            .layer(middleware::from_fn_with_state(permission, require_permission))
+    }
+
+    #[tokio::test]
+    async fn guards_require_both_verified_claims_and_identity() {
+        let claims = json!({"sub":"user", "tenant_id":"tenant", "roles":["admin"]});
+        let principal = DefaultPrincipalExtractor {}.extract(&claims).unwrap();
+        for (include_claims, include_principal) in [(false, false), (true, false), (false, true)] {
+            let mut request = Request::new(Body::empty());
+            if include_claims {
+                request.extensions_mut().insert(AuthorizationClaims::from(claims.clone()));
+            }
+            if include_principal {
+                request.extensions_mut().insert(principal.clone());
+            }
+            let response = app(Permission::MCPUser).oneshot(request).await.unwrap();
+            assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
+        }
+    }
+
+    #[tokio::test]
+    async fn guards_map_verified_roles_to_the_requested_permission() {
+        use StatusCode as S;
+        for (role_claims, admin_status, mcp_status) in [
+            (json!({"roles":["admin"]}), S::NO_CONTENT, S::NO_CONTENT),
+            (json!({"roles":["builder"]}), S::FORBIDDEN, S::NO_CONTENT),
+            (json!({"roles":["user"]}), S::FORBIDDEN, S::NO_CONTENT),
+            (json!({"role":"admin"}), S::NO_CONTENT, S::NO_CONTENT),
+            (json!({"roles":["user"], "role":"admin"}), S::NO_CONTENT, S::NO_CONTENT),
+            (json!({"roles":["unknown", "user"]}), S::FORBIDDEN, S::NO_CONTENT),
+            (json!({"roles":["ServiceAdmin"]}), S::FORBIDDEN, S::FORBIDDEN),
+            (json!({"roles":[]}), S::FORBIDDEN, S::FORBIDDEN),
+            (json!({}), S::FORBIDDEN, S::FORBIDDEN),
+            (json!({"roles":"admin"}), S::UNAUTHORIZED, S::UNAUTHORIZED),
+            (json!({"role":["admin"]}), S::UNAUTHORIZED, S::UNAUTHORIZED),
+            (json!({"roles":["admin", 42]}), S::UNAUTHORIZED, S::UNAUTHORIZED),
+            (json!({"roles":["admin"], "role":42}), S::UNAUTHORIZED, S::UNAUTHORIZED),
+        ] {
+            let mut claims = json!({"sub":"user", "tenant_id":"tenant"});
+            claims.as_object_mut().unwrap().extend(role_claims.as_object().unwrap().clone());
+            let principal = DefaultPrincipalExtractor {}.extract(&claims).unwrap();
+            for (permission, expected) in [(Permission::Admin, admin_status), (Permission::MCPUser, mcp_status)] {
+                let mut request = Request::new(Body::empty());
+                request.extensions_mut().insert(principal.clone());
+                request.extensions_mut().insert(AuthorizationClaims::from(claims.clone()));
+                let response = app(permission).oneshot(request).await.unwrap();
+                assert_eq!(response.status(), expected, "{permission:?}: {role_claims}");
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn cel_identity_mapping_cannot_grant_roles_absent_from_the_token() {
+        let extractor = CelPrincipalExtractor::from_expression(
+            r#"{"user_id": claims.sub, "tenant_id": claims.woTenantId, "role": "admin", "scopes": ["Admin"]}"#,
+        )
+        .unwrap();
+        for (roles, mcp_status) in [(json!(["user"]), StatusCode::NO_CONTENT), (json!([]), StatusCode::FORBIDDEN)] {
+            let claims = json!({"sub":"user", "woTenantId":"tenant", "roles":roles});
+            let principal = extractor.extract(&claims).unwrap();
+            for (permission, expected) in
+                [(Permission::Admin, StatusCode::FORBIDDEN), (Permission::MCPUser, mcp_status)]
+            {
+                let mut request = Request::new(Body::empty());
+                request.extensions_mut().insert(principal.clone());
+                request.extensions_mut().insert(AuthorizationClaims::from(claims.clone()));
+                let response = app(permission).oneshot(request).await.unwrap();
+                assert_eq!(response.status(), expected);
+            }
+        }
+    }
+}

crates/contextforge-data-plane-lib/src/lib.rs

@@ -42,7 +42,7 @@ pub type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
 pub type Result<T> = std::result::Result<T, Error>;
 
 use crate::{
-    authorization::{CelPrincipalExtractor, DefaultPrincipalExtractor},
+    authorization::{CelPrincipalExtractor, DefaultPrincipalExtractor, Permission},
     config_stores::RedisStore,
     layers::{
         claims_id::claims_layer,
@@ -53,6 +53,7 @@ use crate::{
     },
 };
 pub use authorization::{AuthorizationClaims, AuthorizationService, get_authorization_service};
+pub use layers::permission::require_permission;
 
 #[derive(Clone)]
 pub enum UserConfigStoreType {
@@ -143,7 +144,8 @@ impl Gateway {
         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_with_state(mcp_gateway_state.clone(), user_config_store_layer))
+            .layer(middleware::from_fn_with_state(Permission::MCPUser, require_permission));
 
         let app = if let Some(cel_principal_extractor_path) = config.cel_principal_extractor_path.as_ref() {
             app.layer(layers::PrincipalExtractorLayer::new(CelPrincipalExtractor::from_file(

crates/contextforge-data-plane-lib/src/tools.rs

@@ -98,6 +98,8 @@ pub async fn get_token(
         "iss": "contexforge-dataplane",
         "sub": user_id.clone(),
         "aud": "contexforge-dataplane-audience",
+        "role": "user",
+        "woUserId": user_id.clone(),
         "exp": now + Duration::from_hours(1).as_secs(),
         "nbf": now - Duration::from_mins(1).as_secs(),
         "iat": now,

crates/contextforge-data-plane-lib/tests/gateway.rs

@@ -17,3 +17,6 @@ mod resources;
 mod subscriptions;
 #[path = "gateway/tools.rs"]
 mod tools;
+
+#[path = "gateway/downstream_auth.rs"]
+mod downstream_auth;