poimen-memory/crates/mem-cli/src/authorized_pipeline.rs

/// 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<dyn PolicyProvider>,
    jwt_validator: Option<Arc<JwtValidator>>,
}

impl AuthorizedPipeline {
    pub fn new(
        pipeline: FullPipeline,
        policy_provider: Arc<dyn PolicyProvider>,
        audit_logger: Arc<dyn AuditLogger>,
        jwt_validator: Option<Arc<JwtValidator>>,
    ) -> 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<OidcClaims> {
        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<bool> {
        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<bool> {
        // 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<EnrichedChunk>,
    ) -> (Vec<EnrichedChunk>, 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<AuthorizedResult> {
        // 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<AuthorizedResult> {
        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<AuthorizedResult> {
        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<AuthorizedResult> {
        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<Arc<GlobalTfIdfScorer>>,
    semantic_scorer: Option<Arc<SemanticScorer>>,
    policy_provider: Option<Arc<dyn PolicyProvider>>,
    audit_logger: Option<Arc<dyn AuditLogger>>,
    jwt_validator: Option<Arc<JwtValidator>>,
    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<GlobalTfIdfScorer>,
        semantic: Arc<SemanticScorer>,
    ) -> Self {
        self.tfidf_scorer = Some(tfidf);
        self.semantic_scorer = Some(semantic);
        self
    }

    pub fn with_policy_provider(mut self, provider: Arc<dyn PolicyProvider>) -> Self {
        self.policy_provider = Some(provider);
        self
    }

    pub fn with_audit_logger(mut self, logger: Arc<dyn AuditLogger>) -> Self {
        self.audit_logger = Some(logger);
        self
    }

    pub fn with_jwt_validator(mut self, validator: Arc<JwtValidator>) -> 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<AuthorizedPipeline> {
        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<BTreeMap<String, f32>> {
        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);
    }
}