prvng_platform/crates/control-center/tests/jwt_integration_tests.rs

//! JWT Integration Tests
//!
//! Comprehensive integration tests for JWT token system including:
//! - Token generation and validation
//! - Token rotation and refresh
//! - Token revocation and blacklist
//! - Concurrent access patterns
//! - Security edge cases

#![allow(unused_imports, clippy::needless_borrows_for_generic_args)]

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use control_center::auth::{
    jwt::{BlacklistStats, JwtService, TokenType},
    password::PasswordService,
    user::{User, UserRole, UserService},
    AuthService,
};
use control_center::services::jwt::generate_rsa_key_pair;

/// Generate RSA key pair for testing
fn generate_test_keys() -> (Vec<u8>, Vec<u8>) {
    let keys = generate_rsa_key_pair().expect("Failed to generate test RSA keys");
    (
        keys.private_key_pem.into_bytes(),
        keys.public_key_pem.into_bytes(),
    )
}

/// Create JWT service for testing
fn create_jwt_service() -> JwtService {
    let (private_key, public_key) = generate_test_keys();
    JwtService::new(
        &private_key,
        &public_key,
        "test-control-center",
        vec!["orchestrator".to_string(), "cli".to_string()],
    )
    .expect("Failed to create JWT service for tests")
}

#[test]
fn test_token_pair_generation() {
    let jwt_service = create_jwt_service();

    let token_pair = jwt_service
        .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
        .unwrap();

    assert_eq!(token_pair.token_type, "Bearer");
    assert!(token_pair.expires_in > 0);
    assert!(!token_pair.access_token.is_empty());
    assert!(!token_pair.refresh_token.is_empty());
    assert_ne!(token_pair.access_token, token_pair.refresh_token);
}

#[test]
fn test_token_validation_success() {
    let jwt_service = create_jwt_service();

    let token_pair = jwt_service
        .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
        .unwrap();

    // Validate access token
    let access_claims = jwt_service
        .validate_token(&token_pair.access_token)
        .unwrap();
    assert_eq!(access_claims.sub, "user123");
    assert_eq!(access_claims.workspace, "workspace1");
    assert_eq!(access_claims.token_type, TokenType::Access);

    // Validate refresh token
    let refresh_claims = jwt_service
        .validate_token(&token_pair.refresh_token)
        .unwrap();
    assert_eq!(refresh_claims.sub, "user123");
    assert_eq!(refresh_claims.token_type, TokenType::Refresh);
}

#[test]
fn test_token_validation_with_metadata() {
    let jwt_service = create_jwt_service();

    let mut metadata = HashMap::new();
    metadata.insert("ip_address".to_string(), "192.168.1.100".to_string());
    metadata.insert("user_agent".to_string(), "test-client/1.0".to_string());
    metadata.insert("session_id".to_string(), "sess_xyz789".to_string());

    let token_pair = jwt_service
        .generate_token_pair(
            "user456",
            "workspace2",
            "perm_hash_def",
            Some(metadata.clone()),
        )
        .unwrap();

    let claims = jwt_service
        .validate_token(&token_pair.access_token)
        .unwrap();

    assert!(claims.metadata.is_some());
    let token_metadata = claims.metadata.unwrap();
    assert_eq!(
        token_metadata.get("ip_address"),
        Some(&"192.168.1.100".to_string())
    );
    assert_eq!(
        token_metadata.get("user_agent"),
        Some(&"test-client/1.0".to_string())
    );
    assert_eq!(
        token_metadata.get("session_id"),
        Some(&"sess_xyz789".to_string())
    );
}

#[test]
fn test_token_rotation() {
    let jwt_service = create_jwt_service();

    // Generate initial token pair
    let initial_pair = jwt_service
        .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
        .unwrap();

    // Rotate using refresh token
    let rotated_pair = jwt_service
        .rotate_token(&initial_pair.refresh_token)
        .unwrap();

    // New tokens should be different
    assert_ne!(rotated_pair.access_token, initial_pair.access_token);
    assert_ne!(rotated_pair.refresh_token, initial_pair.refresh_token);

    // New tokens should be valid
    let new_claims = jwt_service
        .validate_token(&rotated_pair.access_token)
        .unwrap();
    assert_eq!(new_claims.sub, "user123");
    assert_eq!(new_claims.workspace, "workspace1");

    // Old refresh token should be revoked
    let old_refresh_validation = jwt_service.validate_token(&initial_pair.refresh_token);
    assert!(old_refresh_validation.is_err());
}

#[test]
fn test_token_revocation() {
    let jwt_service = create_jwt_service();

    let token_pair = jwt_service
        .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
        .unwrap();

    // Initially valid
    let claims = jwt_service
        .validate_token(&token_pair.access_token)
        .unwrap();
    assert!(!jwt_service.is_revoked(&claims.jti).unwrap());

    // Revoke token
    jwt_service.revoke_token(&claims.jti, claims.exp).unwrap();

    // Check revoked
    assert!(jwt_service.is_revoked(&claims.jti).unwrap());

    // Validation should fail
    let validation_result = jwt_service.validate_token(&token_pair.access_token);
    assert!(validation_result.is_err());
}

#[test]
fn test_blacklist_cleanup() {
    let jwt_service = create_jwt_service();

    // Create expired token (simulate with past timestamp)
    let expired_jti = "expired_token_123";
    let expired_exp = chrono::Utc::now().timestamp() - 3600; // 1 hour ago
    jwt_service.revoke_token(expired_jti, expired_exp).unwrap();

    // Create valid token
    let token_pair = jwt_service
        .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
        .unwrap();
    let claims = jwt_service
        .validate_token(&token_pair.access_token)
        .unwrap();
    jwt_service.revoke_token(&claims.jti, claims.exp).unwrap();

    // Check stats before cleanup
    let stats_before = jwt_service.blacklist_stats().unwrap();
    assert_eq!(stats_before.total_revoked, 2);

    // Cleanup expired tokens
    let removed_count = jwt_service.cleanup_expired_tokens().unwrap();
    assert_eq!(removed_count, 1);

    // Check stats after cleanup
    let stats_after = jwt_service.blacklist_stats().unwrap();
    assert_eq!(stats_after.total_revoked, 1);
    assert_eq!(stats_after.active_revocations, 1);
}

#[test]
fn test_multiple_rotations() {
    let jwt_service = create_jwt_service();

    let mut current_pair = jwt_service
        .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
        .unwrap();

    // Perform 5 rotations
    for i in 0..5 {
        let new_pair = jwt_service
            .rotate_token(&current_pair.refresh_token)
            .unwrap();

        // Verify new tokens work
        let claims = jwt_service.validate_token(&new_pair.access_token).unwrap();
        assert_eq!(claims.sub, "user123");

        // Old refresh token should fail
        let old_validation = jwt_service.validate_token(&current_pair.refresh_token);
        assert!(old_validation.is_err(), "Rotation {} failed", i);

        current_pair = new_pair;
    }

    // Final tokens should still work
    let final_claims = jwt_service
        .validate_token(&current_pair.access_token)
        .unwrap();
    assert_eq!(final_claims.sub, "user123");
}

#[test]
fn test_extract_token_from_header() {
    let test_cases = vec![
        ("Bearer abc123xyz", Ok("abc123xyz")),
        ("bearer abc123xyz", Err(())), // Wrong case
        ("Token abc123xyz", Err(())),  // Wrong prefix
        ("Bearer", Err(())),           // Missing token
        ("abc123xyz", Err(())),        // No prefix
    ];

    for (header, expected) in test_cases {
        let result = JwtService::extract_token_from_header(header);
        match expected {
            Ok(token) => assert_eq!(result.unwrap(), token),
            Err(_) => assert!(result.is_err()),
        }
    }
}

#[test]
fn test_concurrent_token_operations() {
    use std::sync::Arc;
    use std::thread;

    let jwt_service = Arc::new(create_jwt_service());
    let mut handles = vec![];

    // Spawn 10 threads generating and validating tokens concurrently
    for i in 0..10 {
        let service = jwt_service.clone();
        let handle = thread::spawn(move || {
            let user_id = format!("user{}", i);
            let workspace = format!("workspace{}", i);

            // Generate token pair
            let token_pair = service
                .generate_token_pair(&user_id, &workspace, "perm_hash", None)
                .unwrap();

            // Validate access token
            let claims = service.validate_token(&token_pair.access_token).unwrap();
            assert_eq!(claims.sub, user_id);
            assert_eq!(claims.workspace, workspace);

            // Rotate token
            let new_pair = service.rotate_token(&token_pair.refresh_token).unwrap();
            let new_claims = service.validate_token(&new_pair.access_token).unwrap();
            assert_eq!(new_claims.sub, user_id);
        });
        handles.push(handle);
    }

    // Wait for all threads to complete
    for handle in handles {
        handle.join().unwrap();
    }
}

#[test]
fn test_token_expiry_detection() {
    let jwt_service = create_jwt_service();

    let token_pair = jwt_service
        .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
        .unwrap();

    let claims = jwt_service
        .validate_token(&token_pair.access_token)
        .unwrap();

    // Token should not be expired immediately
    assert!(!claims.is_expired());

    // Token should have remaining validity
    let remaining = claims.remaining_validity();
    assert!(remaining.as_secs() > 0);
}

#[tokio::test]
async fn test_full_auth_flow() {
    let jwt_service = create_jwt_service();
    let password_service = PasswordService::new();
    let user_service = UserService::new();
    let auth_service = AuthService::new(jwt_service, password_service, user_service);

    // Create user
    let password = "SecurePassword123!";
    let password_hash = auth_service.password.hash_password(password).unwrap();

    let user = User::new(
        "testuser",
        "test@example.com",
        "Test User",
        password_hash,
        vec![UserRole::Developer],
    );

    auth_service.user.create_user(user).await.unwrap();

    // Login
    let token_pair = auth_service
        .login("testuser", password, "workspace1")
        .await
        .unwrap();

    // Validate access token
    let claims = auth_service.validate(&token_pair.access_token).unwrap();
    assert_eq!(claims.workspace, "workspace1");

    // Rotate tokens
    let new_pair = auth_service.refresh(&token_pair.refresh_token).unwrap();
    assert_ne!(new_pair.access_token, token_pair.access_token);

    // Logout
    auth_service.logout(&new_pair.access_token).await.unwrap();

    // Token should be revoked
    let validation_result = auth_service.validate(&new_pair.access_token);
    assert!(validation_result.is_err());
}

#[test]
fn test_invalid_signature_detection() {
    let (private_key1, public_key1) = generate_test_keys();
    let (_private_key2, public_key2) = generate_test_keys(); // Different key pair

    // Service 1 generates token
    let jwt_service1 = JwtService::new(
        &private_key1,
        &public_key1,
        "test-issuer",
        vec!["test-audience".to_string()],
    )
    .expect("Failed to create jwt_service1");

    let token_pair = jwt_service1
        .generate_token_pair("user123", "workspace1", "perm_hash", None)
        .expect("Failed to generate token pair");

    // Service 2 with different public key tries to validate
    // This should fail because the token was signed with key1 but we're validating with key2
    let jwt_service2 = JwtService::new(
        &private_key1,
        &public_key2, // Different public key!
        "test-issuer",
        vec!["test-audience".to_string()],
    )
    .expect("Failed to create jwt_service2");

    // Should fail signature verification because public keys don't match
    let validation_result = jwt_service2.validate_token(&token_pair.access_token);
    assert!(validation_result.is_err());
}

#[test]
fn test_blacklist_statistics() {
    let jwt_service = create_jwt_service();

    // Generate and revoke multiple tokens
    for i in 0..5 {
        let token_pair = jwt_service
            .generate_token_pair(&format!("user{}", i), "workspace1", "perm_hash", None)
            .unwrap();

        let claims = jwt_service
            .validate_token(&token_pair.access_token)
            .unwrap();
        jwt_service.revoke_token(&claims.jti, claims.exp).unwrap();
    }

    let stats = jwt_service.blacklist_stats().unwrap();
    assert_eq!(stats.total_revoked, 5);
    assert_eq!(stats.active_revocations, 5);
    assert_eq!(stats.expired_tokens, 0);
}

#[test]
fn test_token_cannot_rotate_with_access_token() {
    let jwt_service = create_jwt_service();

    let token_pair = jwt_service
        .generate_token_pair("user123", "workspace1", "perm_hash_abc", None)
        .unwrap();

    // Attempting to rotate with access token should fail
    let rotation_result = jwt_service.rotate_token(&token_pair.access_token);
    assert!(rotation_result.is_err());
}