poimen-memory/tests/it_authorized_pipeline.rs

/// Integration Tests: Authorized Pipeline (RBAC + JWT)
///
/// Tests two-level access control:
/// 1. Capability level: memory:read / memory:write (checked at HTTP layer)
/// 2. Resource level: project/skill policies (checked in AuthorizedPipeline)
///
/// Access flow:
/// JWT → validate → check project access → execute pipeline → filter by document access

use std::collections::BTreeMap;
use std::sync::Arc;

use mem_core::{GlobalTfIdfScorer, SemanticScorer};
use mem_ingest::wiki_link::WikiLinkGraph;
use mem_cli::{
    AuthorizedPipeline, AuthorizedPipelineBuilder, AccessStats,
    FullPipeline, PipelineConfig,
    rbac::{AccessPolicy, MockPolicyProvider, OidcClaims, LegacyNoOpAuditLogger as NoOpAuditLogger},
};

// ============================================================================
// Test Fixtures
// ============================================================================

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);
    vocab.insert("secret".to_string(), 0.9);
    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_test_wiki_graph() -> WikiLinkGraph {
    let mut graph = WikiLinkGraph::new("test");
    graph.add_link("index.md", "docs/public.md");
    graph.add_link("index.md", "docs/internal.md");
    graph.add_link("docs/internal.md", "shared/skills/SKILL-secret/SKILL.md");
    graph
}

fn create_test_candidates() -> Vec<(String, String)> {
    vec![
        ("index.md".to_string(), "# Index\nKubernetes documentation.".to_string()),
        ("docs/public.md".to_string(), "# Public Docs\nPublic kubernetes guide.".to_string()),
        ("docs/internal.md".to_string(), "# Internal\nInternal pod documentation.".to_string()),
        ("shared/skills/SKILL-secret/SKILL.md".to_string(), "# Secret Skill\nConfidential deployment skill.".to_string()),
    ]
}

// Policy helpers
fn public_policy() -> AccessPolicy {
    AccessPolicy {
        access_level: "public".to_string(),
        owner_group: "".to_string(),
        allowed_groups: vec![],
        required_role: None,
        required_permission: None,
    }
}

fn group_policy(groups: Vec<&str>) -> AccessPolicy {
    AccessPolicy {
        access_level: "group".to_string(),
        owner_group: "".to_string(),
        allowed_groups: groups.into_iter().map(String::from).collect(),
        required_role: None,
        required_permission: None,
    }
}

fn 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 role_policy(role: &str) -> AccessPolicy {
    AccessPolicy {
        access_level: "public".to_string(),
        owner_group: "".to_string(),
        allowed_groups: vec![],
        required_role: Some(role.to_string()),
        required_permission: None,
    }
}

fn permission_policy(perm: &str) -> AccessPolicy {
    AccessPolicy {
        access_level: "public".to_string(),
        owner_group: "".to_string(),
        allowed_groups: vec![],
        required_role: None,
        required_permission: Some(perm.to_string()),
    }
}

// ============================================================================
// Project Access Tests
// ============================================================================

#[tokio::test]
async fn test_public_project_anyone_can_access() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "docs", public_policy())
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    // Anyone can access public project
    let claims = OidcClaims {
        sub: "anonymous".to_string(),
        groups: vec![],
        roles: vec![],
        permissions: vec![],
    };

    let allowed = auth.check_project_access(&claims, "docs").await.unwrap();
    assert!(allowed, "Public project should allow anyone");
}

#[tokio::test]
async fn test_group_project_member_can_access() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "internal", group_policy(vec!["engineering", "devops"]))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    // Engineering member can access
    let eng_claims = OidcClaims {
        sub: "alice".to_string(),
        groups: vec!["engineering".to_string()],
        roles: vec![],
        permissions: vec![],
    };
    assert!(auth.check_project_access(&eng_claims, "internal").await.unwrap());

    // DevOps member can access
    let devops_claims = OidcClaims {
        sub: "bob".to_string(),
        groups: vec!["devops".to_string()],
        roles: vec![],
        permissions: vec![],
    };
    assert!(auth.check_project_access(&devops_claims, "internal").await.unwrap());

    // Sales cannot access
    let sales_claims = OidcClaims {
        sub: "charlie".to_string(),
        groups: vec!["sales".to_string()],
        roles: vec![],
        permissions: vec![],
    };
    assert!(!auth.check_project_access(&sales_claims, "internal").await.unwrap());
}

#[tokio::test]
async fn test_private_project_owner_only() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "secret", private_policy("ml-team"))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    // Owner group can access
    let owner_claims = OidcClaims {
        sub: "ml-researcher".to_string(),
        groups: vec!["ml-team".to_string()],
        roles: vec![],
        permissions: vec![],
    };
    assert!(auth.check_project_access(&owner_claims, "secret").await.unwrap());

    // Non-owner denied
    let other_claims = OidcClaims {
        sub: "engineer".to_string(),
        groups: vec!["engineering".to_string()],
        roles: vec![],
        permissions: vec![],
    };
    assert!(!auth.check_project_access(&other_claims, "secret").await.unwrap());
}

#[tokio::test]
async fn test_role_required_project() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "admin-only", role_policy("admin"))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    // Admin can access
    let admin_claims = OidcClaims {
        sub: "admin-user".to_string(),
        groups: vec![],
        roles: vec!["admin".to_string()],
        permissions: vec![],
    };
    assert!(auth.check_project_access(&admin_claims, "admin-only").await.unwrap());

    // Viewer cannot access
    let viewer_claims = OidcClaims {
        sub: "viewer-user".to_string(),
        groups: vec![],
        roles: vec!["viewer".to_string()],
        permissions: vec![],
    };
    assert!(!auth.check_project_access(&viewer_claims, "admin-only").await.unwrap());
}

#[tokio::test]
async fn test_permission_required_project() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "special", permission_policy("project:special:read"))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    // User with permission can access
    let permitted_claims = OidcClaims {
        sub: "special-user".to_string(),
        groups: vec![],
        roles: vec![],
        permissions: vec!["project:special:read".to_string()],
    };
    assert!(auth.check_project_access(&permitted_claims, "special").await.unwrap());

    // User without permission denied
    let regular_claims = OidcClaims {
        sub: "regular-user".to_string(),
        groups: vec![],
        roles: vec![],
        permissions: vec!["memory:read".to_string()],
    };
    assert!(!auth.check_project_access(&regular_claims, "special").await.unwrap());
}

// ============================================================================
// Document/Skill Filtering Tests
// ============================================================================

#[tokio::test]
async fn test_skill_filtered_by_access() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "docs", public_policy())
            .with_policy("skill", "SKILL-secret", private_policy("ml-team"))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    let claims = OidcClaims {
        sub: "engineer".to_string(),
        groups: vec!["engineering".to_string()], // Not ml-team
        roles: vec![],
        permissions: vec![],
    };

    let graph = create_test_wiki_graph();
    let candidates = create_test_candidates();

    let result = auth
        .execute_with_claims(&claims, "kubernetes", "docs", &graph, candidates)
        .await
        .unwrap();

    // Project access allowed
    assert!(result.access_stats.project_access_allowed);

    // SKILL-secret should be filtered out
    let has_secret_skill = result.result.chunks
        .iter()
        .any(|c| c.id.contains("SKILL-secret"));
    assert!(!has_secret_skill, "Secret skill should be filtered out");

    // Should have some denied chunks
    assert!(result.access_stats.chunks_denied >= 0);
}

#[tokio::test]
async fn test_skill_visible_to_owner() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "docs", public_policy())
            .with_policy("skill", "SKILL-secret", private_policy("ml-team"))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    // ML team member should see the skill
    let claims = OidcClaims {
        sub: "ml-researcher".to_string(),
        groups: vec!["ml-team".to_string()],
        roles: vec![],
        permissions: vec![],
    };

    let graph = create_test_wiki_graph();
    let candidates = create_test_candidates();

    let result = auth
        .execute_with_claims(&claims, "secret skill", "docs", &graph, candidates)
        .await
        .unwrap();

    // No chunks should be denied for owner
    assert_eq!(result.access_stats.chunks_denied, 0);
}

// ============================================================================
// End-to-End Pipeline Tests
// ============================================================================

#[tokio::test]
async fn test_full_pipeline_with_access_control() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "poimen", group_policy(vec!["platform-team"]))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    let claims = OidcClaims {
        sub: "platform-engineer".to_string(),
        groups: vec!["platform-team".to_string()],
        roles: vec!["engineer".to_string()],
        permissions: vec!["memory:read".to_string(), "memory:write".to_string()],
    };

    let graph = create_test_wiki_graph();
    let candidates = create_test_candidates();

    let result = auth
        .execute_with_claims(&claims, "kubernetes pod", "poimen", &graph, candidates)
        .await
        .unwrap();

    // Access stats should be populated
    assert!(result.access_stats.project_access_checked);
    assert!(result.access_stats.project_access_allowed);
    assert!(result.access_stats.chunks_before_filter >= 0);

    // User ID should be captured
    assert_eq!(result.user_id, "platform-engineer");
}

#[tokio::test]
async fn test_denied_project_returns_error() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "secret", private_policy("ml-team"))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    let claims = OidcClaims {
        sub: "outsider".to_string(),
        groups: vec!["other-team".to_string()],
        roles: vec![],
        permissions: vec![],
    };

    let graph = create_test_wiki_graph();
    let candidates = create_test_candidates();

    let result = auth
        .execute_with_claims(&claims, "query", "secret", &graph, candidates)
        .await;

    assert!(result.is_err());
    let err = result.unwrap_err().to_string();
    assert!(err.contains("Access denied"));
    assert!(err.contains("secret"));
}

#[tokio::test]
async fn test_direct_execution_with_access() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "public-docs", public_policy())
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    let claims = OidcClaims {
        sub: "anyone".to_string(),
        groups: vec![],
        roles: vec![],
        permissions: vec![],
    };

    let candidates = create_test_candidates();

    let result = auth
        .execute_direct_with_claims(&claims, "kubernetes", "public-docs", candidates)
        .await
        .unwrap();

    assert!(result.access_stats.project_access_allowed);
}

// ============================================================================
// Builder Tests
// ============================================================================

#[tokio::test]
async fn test_builder_full_config() {
    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 audit = Arc::new(NoOpAuditLogger);

    let auth = AuthorizedPipelineBuilder::new()
        .with_scorers(tfidf, semantic)
        .with_policy_provider(provider)
        .with_audit_logger(audit)
        .with_pipeline_config(PipelineConfig {
            budget_bytes: 4096,
            ..PipelineConfig::default()
        })
        .build()
        .unwrap();

    assert_eq!(auth.pipeline().config().budget_bytes, 4096);
}

#[tokio::test]
async fn test_builder_minimal() {
    let vocab = create_test_vocab();
    let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab));
    let semantic = Arc::new(SemanticScorer::new());
    let provider = Arc::new(MockPolicyProvider::new());

    // Minimal config: scorers + policy provider (audit defaults to NoOp)
    let auth = AuthorizedPipelineBuilder::new()
        .with_scorers(tfidf, semantic)
        .with_policy_provider(provider)
        .build()
        .unwrap();

    // Should work
    assert!(auth.pipeline().config().budget_bytes > 0);
}

// ============================================================================
// Access Stats Tests
// ============================================================================

#[test]
fn test_access_stats_initialization() {
    let stats = AccessStats::new();
    
    assert!(!stats.project_access_checked);
    assert!(!stats.project_access_allowed);
    assert_eq!(stats.chunks_before_filter, 0);
    assert_eq!(stats.chunks_after_filter, 0);
    assert_eq!(stats.chunks_denied, 0);
    assert!(stats.denied_reasons.is_empty());
}

#[tokio::test]
async fn test_access_stats_populated() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "test", public_policy())
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    let claims = OidcClaims {
        sub: "user".to_string(),
        groups: vec![],
        roles: vec![],
        permissions: vec![],
    };

    let graph = create_test_wiki_graph();
    let candidates = create_test_candidates();

    let result = auth
        .execute_with_claims(&claims, "kubernetes", "test", &graph, candidates)
        .await
        .unwrap();

    let stats = &result.access_stats;
    
    assert!(stats.project_access_checked);
    assert!(stats.project_access_allowed);
    // chunks_before_filter should be >= chunks_after_filter
    assert!(stats.chunks_before_filter >= stats.chunks_after_filter);
}

// ============================================================================
// Multi-Group Membership Tests
// ============================================================================

#[tokio::test]
async fn test_user_with_multiple_groups() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "devops", group_policy(vec!["devops"]))
            .with_policy("project", "engineering", group_policy(vec!["engineering"]))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    // User in both groups
    let claims = OidcClaims {
        sub: "sre".to_string(),
        groups: vec!["devops".to_string(), "engineering".to_string()],
        roles: vec![],
        permissions: vec![],
    };

    // Can access both projects
    assert!(auth.check_project_access(&claims, "devops").await.unwrap());
    assert!(auth.check_project_access(&claims, "engineering").await.unwrap());
}

#[tokio::test]
async fn test_multiple_permissions_combined() {
    let pipeline = create_test_pipeline();
    let provider = Arc::new(
        MockPolicyProvider::new()
            .with_policy("project", "special", permission_policy("special:access"))
    );
    let audit = Arc::new(NoOpAuditLogger);
    let auth = AuthorizedPipeline::new(pipeline, provider, audit, None);

    // User with multiple permissions including required one
    let claims = OidcClaims {
        sub: "power-user".to_string(),
        groups: vec![],
        roles: vec![],
        permissions: vec![
            "memory:read".to_string(),
            "memory:write".to_string(),
            "special:access".to_string(),
        ],
    };

    assert!(auth.check_project_access(&claims, "special").await.unwrap());
}