/// Authorized Pipeline: RBAC-protected retrieval /// /// Wraps FullPipeline with JWT authentication and access control: /// 1. Validate JWT token and extract claims /// 2. Check project-level access before retrieval /// 3. Filter results by document/skill access policies /// 4. Audit log all access decisions /// /// Access flow: /// ```text /// JWT Token → validate → OidcClaims /// ↓ /// check_project_access(claims, project) /// ↓ (if allowed) /// FullPipeline.execute() /// ↓ /// filter_results_by_access(claims, chunks) /// ↓ /// AuthorizedResult { chunks, access_stats } /// ``` use anyhow::{anyhow, Result}; use std::sync::Arc; use mem_core::{GlobalTfIdfScorer, SemanticScorer}; use mem_ingest::wiki_link::WikiLinkGraph; use crate::full_pipeline::{FullPipeline, PipelineConfig, PipelineResult, EnrichedChunk, PipelineMetrics}; use crate::rbac::{ AccessPolicy, PolicyProvider, AccessDecisionEngine, OidcClaims, LegacyAccessDecision as AccessDecision, LegacyAuditLogger as AuditLogger, LegacyNoOpAuditLogger as NoOpAuditLogger, }; use crate::jwt_validator::{JwtValidator, JwtClaims}; /// Access statistics for audit/metrics #[derive(Debug, Clone)] pub struct AccessStats { pub project_access_checked: bool, pub project_access_allowed: bool, pub chunks_before_filter: usize, pub chunks_after_filter: usize, pub chunks_denied: usize, pub denied_reasons: Vec<(String, String)>, // (chunk_id, reason) } impl AccessStats { pub fn new() -> Self { Self { project_access_checked: false, project_access_allowed: false, chunks_before_filter: 0, chunks_after_filter: 0, chunks_denied: 0, denied_reasons: Vec::new(), } } } /// Result with access control metadata #[derive(Debug, Clone)] pub struct AuthorizedResult { pub result: PipelineResult, pub access_stats: AccessStats, pub user_id: String, } /// Authorized Pipeline: RBAC-protected retrieval pub struct AuthorizedPipeline { pipeline: FullPipeline, access_engine: AccessDecisionEngine, policy_provider: Arc, jwt_validator: Option>, } impl AuthorizedPipeline { pub fn new( pipeline: FullPipeline, policy_provider: Arc, audit_logger: Arc, jwt_validator: Option>, ) -> Self { let access_engine = AccessDecisionEngine::new( policy_provider.clone(), audit_logger, ); Self { pipeline, access_engine, policy_provider, jwt_validator, } } /// Validate JWT and extract OIDC claims pub async fn validate_token(&self, token: &str) -> Result { let validator = self.jwt_validator .as_ref() .ok_or_else(|| anyhow!("JWT validator not configured"))?; let jwt_claims = validator.validate_token(token).await?; Ok(OidcClaims { sub: jwt_claims.sub, groups: jwt_claims.groups.unwrap_or_default(), roles: vec![], // Authentik may include roles differently permissions: jwt_claims.permissions.unwrap_or_default(), }) } /// Check if user can access a project pub async fn check_project_access( &self, claims: &OidcClaims, project: &str, ) -> Result { self.access_engine .check_access(claims, "project", project) .await } /// Check if user can access a specific document/skill pub async fn check_document_access( &self, claims: &OidcClaims, doc_id: &str, ) -> Result { // Extract resource type from doc_id let (resource_type, resource_name) = self.parse_doc_id(doc_id); // Try to get policy, default to project policy if not found match self.policy_provider.get_policy(&resource_type, &resource_name).await { Ok(_) => { self.access_engine .check_access(claims, &resource_type, &resource_name) .await } Err(_) => { // No specific policy, allow if project access was granted Ok(true) } } } /// Parse doc_id to determine resource type fn parse_doc_id(&self, doc_id: &str) -> (String, String) { if doc_id.contains("SKILL-") || doc_id.contains("/skills/") { // Extract skill name let skill_name = doc_id .split('/') .find(|s| s.starts_with("SKILL-")) .unwrap_or(doc_id) .to_string(); ("skill".to_string(), skill_name) } else { // Regular document, use project-level access ("document".to_string(), doc_id.to_string()) } } /// Filter chunks by access policies pub async fn filter_by_access( &self, claims: &OidcClaims, chunks: Vec, ) -> (Vec, Vec<(String, String)>) { let mut allowed = Vec::new(); let mut denied = Vec::new(); for chunk in chunks { match self.check_document_access(claims, &chunk.id).await { Ok(true) => allowed.push(chunk), Ok(false) => denied.push((chunk.id.clone(), "access_denied".to_string())), Err(e) => denied.push((chunk.id.clone(), format!("error: {}", e))), } } (allowed, denied) } /// Execute with JWT token authentication pub async fn execute_with_token( &self, token: &str, query: &str, project: &str, wiki_graph: &WikiLinkGraph, candidates: Vec<(String, String)>, ) -> Result { // Step 1: Validate JWT let claims = self.validate_token(token).await?; // Step 2: Execute with claims self.execute_with_claims(&claims, query, project, wiki_graph, candidates).await } /// Execute with pre-validated claims (for internal use or testing) pub async fn execute_with_claims( &self, claims: &OidcClaims, query: &str, project: &str, wiki_graph: &WikiLinkGraph, candidates: Vec<(String, String)>, ) -> Result { let mut access_stats = AccessStats::new(); access_stats.project_access_checked = true; // Step 1: Check project access let project_allowed = self.check_project_access(claims, project).await?; access_stats.project_access_allowed = project_allowed; if !project_allowed { return Err(anyhow!( "Access denied to project '{}' for user '{}'", project, claims.sub )); } // Step 2: Execute pipeline let result = self.pipeline .execute_with_wiki(query, wiki_graph, candidates) .await?; access_stats.chunks_before_filter = result.chunks.len(); // Step 3: Filter by document access let (allowed_chunks, denied) = self.filter_by_access(claims, result.chunks).await; access_stats.chunks_after_filter = allowed_chunks.len(); access_stats.chunks_denied = denied.len(); access_stats.denied_reasons = denied; // Step 4: Build authorized result let authorized_result = PipelineResult { query: result.query, query_intent: result.query_intent, chunks: allowed_chunks, metrics: result.metrics, }; Ok(AuthorizedResult { result: authorized_result, access_stats, user_id: claims.sub.clone(), }) } /// Execute direct (no wiki) with token pub async fn execute_direct_with_token( &self, token: &str, query: &str, project: &str, candidates: Vec<(String, String)>, ) -> Result { let claims = self.validate_token(token).await?; self.execute_direct_with_claims(&claims, query, project, candidates).await } /// Execute direct with claims pub async fn execute_direct_with_claims( &self, claims: &OidcClaims, query: &str, project: &str, candidates: Vec<(String, String)>, ) -> Result { let mut access_stats = AccessStats::new(); access_stats.project_access_checked = true; // Check project access let project_allowed = self.check_project_access(claims, project).await?; access_stats.project_access_allowed = project_allowed; if !project_allowed { return Err(anyhow!( "Access denied to project '{}' for user '{}'", project, claims.sub )); } // Execute pipeline let result = self.pipeline.execute_direct(query, candidates).await?; access_stats.chunks_before_filter = result.chunks.len(); // Filter by access let (allowed_chunks, denied) = self.filter_by_access(claims, result.chunks).await; access_stats.chunks_after_filter = allowed_chunks.len(); access_stats.chunks_denied = denied.len(); access_stats.denied_reasons = denied; let authorized_result = PipelineResult { query: result.query, query_intent: result.query_intent, chunks: allowed_chunks, metrics: result.metrics, }; Ok(AuthorizedResult { result: authorized_result, access_stats, user_id: claims.sub.clone(), }) } pub fn pipeline(&self) -> &FullPipeline { &self.pipeline } } /// Builder for AuthorizedPipeline pub struct AuthorizedPipelineBuilder { tfidf_scorer: Option>, semantic_scorer: Option>, policy_provider: Option>, audit_logger: Option>, jwt_validator: Option>, pipeline_config: PipelineConfig, } impl AuthorizedPipelineBuilder { pub fn new() -> Self { Self { tfidf_scorer: None, semantic_scorer: None, policy_provider: None, audit_logger: None, jwt_validator: None, pipeline_config: PipelineConfig::default(), } } pub fn with_scorers( mut self, tfidf: Arc, semantic: Arc, ) -> Self { self.tfidf_scorer = Some(tfidf); self.semantic_scorer = Some(semantic); self } pub fn with_policy_provider(mut self, provider: Arc) -> Self { self.policy_provider = Some(provider); self } pub fn with_audit_logger(mut self, logger: Arc) -> Self { self.audit_logger = Some(logger); self } pub fn with_jwt_validator(mut self, validator: Arc) -> Self { self.jwt_validator = Some(validator); self } pub fn with_pipeline_config(mut self, config: PipelineConfig) -> Self { self.pipeline_config = config; self } pub fn build(self) -> Result { let tfidf = self.tfidf_scorer .ok_or_else(|| anyhow!("TF-IDF scorer required"))?; let semantic = self.semantic_scorer .ok_or_else(|| anyhow!("Semantic scorer required"))?; let policy_provider = self.policy_provider .ok_or_else(|| anyhow!("PolicyProvider required"))?; let audit_logger = self.audit_logger .unwrap_or_else(|| Arc::new(NoOpAuditLogger)); let pipeline = FullPipeline::new(tfidf, semantic, self.pipeline_config); Ok(AuthorizedPipeline::new( pipeline, policy_provider, audit_logger, self.jwt_validator, )) } } #[cfg(test)] mod tests { use super::*; use std::collections::BTreeMap; use crate::rbac::MockPolicyProvider; fn create_test_vocab() -> Arc> { let mut vocab = BTreeMap::new(); vocab.insert("kubernetes".to_string(), 0.8); vocab.insert("pod".to_string(), 0.7); Arc::new(vocab) } fn create_test_pipeline() -> FullPipeline { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); FullPipeline::new(tfidf, semantic, PipelineConfig::default()) } fn create_public_policy() -> AccessPolicy { AccessPolicy { access_level: "public".to_string(), owner_group: "".to_string(), allowed_groups: vec![], required_role: None, required_permission: None, } } fn create_group_policy(groups: Vec<&str>) -> AccessPolicy { AccessPolicy { access_level: "group".to_string(), owner_group: "".to_string(), allowed_groups: groups.into_iter().map(|s| s.to_string()).collect(), required_role: None, required_permission: None, } } fn create_private_policy(owner: &str) -> AccessPolicy { AccessPolicy { access_level: "private".to_string(), owner_group: owner.to_string(), allowed_groups: vec![], required_role: None, required_permission: None, } } fn create_test_wiki_graph() -> WikiLinkGraph { let mut graph = WikiLinkGraph::new("test"); graph.add_link("index.md", "docs/guide.md"); graph } #[test] fn test_access_stats_new() { let stats = AccessStats::new(); assert!(!stats.project_access_checked); assert!(!stats.project_access_allowed); assert_eq!(stats.chunks_before_filter, 0); } #[test] fn test_parse_doc_id_skill() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let pipeline = FullPipeline::new(tfidf, semantic, PipelineConfig::default()); let provider = Arc::new(MockPolicyProvider::new()); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let (rtype, rname) = auth_pipeline.parse_doc_id("shared/skills/SKILL-kubernetes-debug/SKILL.md"); assert_eq!(rtype, "skill"); assert!(rname.contains("SKILL-")); } #[test] fn test_parse_doc_id_document() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let pipeline = FullPipeline::new(tfidf, semantic, PipelineConfig::default()); let provider = Arc::new(MockPolicyProvider::new()); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let (rtype, _) = auth_pipeline.parse_doc_id("docs/guide.md"); assert_eq!(rtype, "document"); } #[tokio::test] async fn test_check_project_access_public() { let pipeline = create_test_pipeline(); let provider = Arc::new( MockPolicyProvider::new() .with_policy("project", "poimen", create_public_policy()) ); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let claims = OidcClaims { sub: "anyone".to_string(), groups: vec![], roles: vec![], permissions: vec![], }; let allowed = auth_pipeline.check_project_access(&claims, "poimen").await.unwrap(); assert!(allowed); } #[tokio::test] async fn test_check_project_access_group_allowed() { let pipeline = create_test_pipeline(); let provider = Arc::new( MockPolicyProvider::new() .with_policy("project", "poimen", create_group_policy(vec!["platform-team"])) ); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let claims = OidcClaims { sub: "charlie".to_string(), groups: vec!["platform-team".to_string()], roles: vec![], permissions: vec![], }; let allowed = auth_pipeline.check_project_access(&claims, "poimen").await.unwrap(); assert!(allowed); } #[tokio::test] async fn test_check_project_access_group_denied() { let pipeline = create_test_pipeline(); let provider = Arc::new( MockPolicyProvider::new() .with_policy("project", "poimen", create_group_policy(vec!["platform-team"])) ); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let claims = OidcClaims { sub: "alice".to_string(), groups: vec!["data-team".to_string()], roles: vec![], permissions: vec![], }; let allowed = auth_pipeline.check_project_access(&claims, "poimen").await.unwrap(); assert!(!allowed); } #[tokio::test] async fn test_check_project_access_private() { let pipeline = create_test_pipeline(); let provider = Arc::new( MockPolicyProvider::new() .with_policy("project", "secret", create_private_policy("ml-team")) ); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); // Owner has access let owner_claims = OidcClaims { sub: "bob".to_string(), groups: vec!["ml-team".to_string()], roles: vec![], permissions: vec![], }; assert!(auth_pipeline.check_project_access(&owner_claims, "secret").await.unwrap()); // Non-owner denied let other_claims = OidcClaims { sub: "charlie".to_string(), groups: vec!["platform-team".to_string()], roles: vec![], permissions: vec![], }; assert!(!auth_pipeline.check_project_access(&other_claims, "secret").await.unwrap()); } #[tokio::test] async fn test_execute_with_claims_allowed() { let pipeline = create_test_pipeline(); let provider = Arc::new( MockPolicyProvider::new() .with_policy("project", "poimen", create_public_policy()) ); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let claims = OidcClaims { sub: "charlie".to_string(), groups: vec!["platform-team".to_string()], roles: vec![], permissions: vec![], }; let graph = create_test_wiki_graph(); let candidates = vec![ ("index.md".to_string(), "kubernetes guide".to_string()), ]; let result = auth_pipeline .execute_with_claims(&claims, "kubernetes", "poimen", &graph, candidates) .await .unwrap(); assert!(result.access_stats.project_access_allowed); assert_eq!(result.user_id, "charlie"); } #[tokio::test] async fn test_execute_with_claims_denied() { let pipeline = create_test_pipeline(); let provider = Arc::new( MockPolicyProvider::new() .with_policy("project", "secret", create_private_policy("ml-team")) ); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let claims = OidcClaims { sub: "charlie".to_string(), groups: vec!["platform-team".to_string()], // Not ml-team roles: vec![], permissions: vec![], }; let graph = create_test_wiki_graph(); let candidates = vec![ ("index.md".to_string(), "secret content".to_string()), ]; let result = auth_pipeline .execute_with_claims(&claims, "query", "secret", &graph, candidates) .await; assert!(result.is_err()); assert!(result.unwrap_err().to_string().contains("Access denied")); } #[tokio::test] async fn test_filter_by_access_skill() { let pipeline = create_test_pipeline(); let provider = Arc::new( MockPolicyProvider::new() .with_policy("project", "poimen", create_public_policy()) .with_policy("skill", "SKILL-private", create_private_policy("ml-team")) ); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let claims = OidcClaims { sub: "charlie".to_string(), groups: vec!["platform-team".to_string()], roles: vec![], permissions: vec![], }; let chunks = vec![ EnrichedChunk { id: "docs/public.md".to_string(), text: "public content".to_string(), tfidf_score: 0.5, semantic_score: 0.5, rrf_score: 0.5, pre_boost_score: 0.5, final_score: 0.5, category: crate::chunk_metadata::ChunkCategory::Concept, heading: None, key_terms: vec![], metadata_boost: 0.0, query_intent_match: false, wiki_distance: None, cache_slot: 0, cache_priority: 0.5, }, EnrichedChunk { id: "shared/skills/SKILL-private/SKILL.md".to_string(), text: "private skill".to_string(), tfidf_score: 0.8, semantic_score: 0.8, rrf_score: 0.8, pre_boost_score: 0.8, final_score: 0.8, category: crate::chunk_metadata::ChunkCategory::Tool, heading: None, key_terms: vec![], metadata_boost: 0.0, query_intent_match: false, wiki_distance: None, cache_slot: 0, cache_priority: 0.8, }, ]; let (allowed, denied) = auth_pipeline.filter_by_access(&claims, chunks).await; // Public doc allowed, private skill denied assert_eq!(allowed.len(), 1); assert_eq!(denied.len(), 1); assert_eq!(allowed[0].id, "docs/public.md"); assert!(denied[0].0.contains("SKILL-private")); } #[tokio::test] async fn test_builder() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let provider = Arc::new(MockPolicyProvider::new()); let auth_pipeline = AuthorizedPipelineBuilder::new() .with_scorers(tfidf, semantic) .with_policy_provider(provider) .build() .unwrap(); // Should build successfully assert!(auth_pipeline.jwt_validator.is_none()); } #[tokio::test] async fn test_builder_missing_provider() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let result = AuthorizedPipelineBuilder::new() .with_scorers(tfidf, semantic) // Missing policy_provider .build(); assert!(result.is_err()); } #[tokio::test] async fn test_execute_direct_with_claims() { let pipeline = create_test_pipeline(); let provider = Arc::new( MockPolicyProvider::new() .with_policy("project", "test", create_public_policy()) ); let audit = Arc::new(NoOpAuditLogger); let auth_pipeline = AuthorizedPipeline::new(pipeline, provider, audit, None); let claims = OidcClaims { sub: "user".to_string(), groups: vec![], roles: vec![], permissions: vec![], }; let candidates = vec![ ("doc.md".to_string(), "content".to_string()), ]; let result = auth_pipeline .execute_direct_with_claims(&claims, "query", "test", candidates) .await .unwrap(); assert!(result.access_stats.project_access_allowed); } }