jwt-verify 0.1.0

use base64::Engine;
use jwt_verify::{CognitoJwtVerifier, JwtError, JwtVerifier, VerifierConfig};
use std::time::{SystemTime, UNIX_EPOCH};

// Helper function to create a test token
fn create_test_token(header: &str, payload: &str, signature: &str) -> String {
    format!("{}.{}.{}", header, payload, signature)
}

// Helper function to create a base64 encoded string
fn base64_encode(data: &str) -> String {
    base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(data)
}

// Helper function to get current timestamp
fn current_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

// Helper function to create a test ID token
fn create_test_id_token(issuer: &str, kid: &str, exp_offset: i64) -> String {
    let now = current_timestamp();
    let exp = now + exp_offset as u64;

    let header = base64_encode(&format!(r#"{{"alg":"RS256","kid":"{}","typ":"JWT"}}"#, kid));
    let payload = base64_encode(&format!(
        r#"{{
        "sub": "test-user-id",
        "iss": "{}",
        "client_id": "test-client-id",
        "token_use": "id",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-id",
        "email": "test@example.com",
        "email_verified": true,
        "cognito:username": "testuser",
        "aud": "test-client-id"
    }}"#,
        issuer,
        now - 300,
        exp,
        now
    ));

    // For testing purposes, we use a dummy signature
    let signature = "test-signature";

    create_test_token(&header, &payload, signature)
}

// Helper function to create a tampered token with modified payload
fn create_tampered_token(issuer: &str, kid: &str, tamper_type: &str) -> String {
    let now = current_timestamp();
    let exp = now + 3600;

    let header = base64_encode(&format!(r#"{{"alg":"RS256","kid":"{}","typ":"JWT"}}"#, kid));
    
    // Create different tampered payloads based on the tamper type
    let payload = match tamper_type {
        "modified_sub" => base64_encode(&format!(
            r#"{{
            "sub": "hacker-user-id",
            "iss": "{}",
            "client_id": "test-client-id",
            "token_use": "id",
            "auth_time": {},
            "exp": {},
            "iat": {},
            "jti": "test-jwt-id",
            "email": "test@example.com"
        }}"#,
            issuer,
            now - 300,
            exp,
            now
        )),
        "elevated_permissions" => base64_encode(&format!(
            r#"{{
            "sub": "test-user-id",
            "iss": "{}",
            "client_id": "test-client-id",
            "token_use": "id",
            "auth_time": {},
            "exp": {},
            "iat": {},
            "jti": "test-jwt-id",
            "email": "test@example.com",
            "cognito:groups": ["admin", "superuser"]
        }}"#,
            issuer,
            now - 300,
            exp,
            now
        )),
        "modified_iat" => base64_encode(&format!(
            r#"{{
            "sub": "test-user-id",
            "iss": "{}",
            "client_id": "test-client-id",
            "token_use": "id",
            "auth_time": {},
            "exp": {},
            "iat": {},
            "jti": "test-jwt-id",
            "email": "test@example.com"
        }}"#,
            issuer,
            now - 300,
            exp,
            now - 86400 // Set issued at to 24 hours ago
        )),
        _ => base64_encode(&format!(
            r#"{{
            "sub": "test-user-id",
            "iss": "{}",
            "client_id": "test-client-id",
            "token_use": "id",
            "auth_time": {},
            "exp": {},
            "iat": {},
            "jti": "test-jwt-id",
            "email": "test@example.com"
        }}"#,
            issuer,
            now - 300,
            exp,
            now
        )),
    };

    // Original signature from a valid token
    let signature = "test-signature";

    create_test_token(&header, &payload, signature)
}

// Helper function to create a token with none algorithm (algorithm stripping attack)
fn create_none_algorithm_token(issuer: &str) -> String {
    let now = current_timestamp();
    let exp = now + 3600;

    // Use "none" algorithm in header
    let header = base64_encode(r#"{"alg":"none","typ":"JWT"}"#);
    
    let payload = base64_encode(&format!(
        r#"{{
        "sub": "test-user-id",
        "iss": "{}",
        "client_id": "test-client-id",
        "token_use": "id",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-id",
        "email": "test@example.com",
        "cognito:groups": ["admin"]
    }}"#,
        issuer,
        now - 300,
        exp,
        now
    ));

    // Empty signature for "none" algorithm
    let signature = "";

    create_test_token(&header, &payload, signature)
}

// Helper function to create a token with invalid signature
fn create_invalid_signature_token(issuer: &str, kid: &str) -> String {
    let now = current_timestamp();
    let exp = now + 3600;

    let header = base64_encode(&format!(r#"{{"alg":"RS256","kid":"{}","typ":"JWT"}}"#, kid));
    let payload = base64_encode(&format!(
        r#"{{
        "sub": "test-user-id",
        "iss": "{}",
        "client_id": "test-client-id",
        "token_use": "id",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-id",
        "email": "test@example.com"
    }}"#,
        issuer,
        now - 300,
        exp,
        now
    ));

    // Invalid signature
    let signature = "invalid-signature";

    create_test_token(&header, &payload, signature)
}

// Helper function to create an expired token
fn create_expired_token(issuer: &str, kid: &str, expired_seconds_ago: u64) -> String {
    let now = current_timestamp();
    let exp = now - expired_seconds_ago; // Token expired in the past

    let header = base64_encode(&format!(r#"{{"alg":"RS256","kid":"{}","typ":"JWT"}}"#, kid));
    let payload = base64_encode(&format!(
        r#"{{
        "sub": "test-user-id",
        "iss": "{}",
        "client_id": "test-client-id",
        "token_use": "id",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-id",
        "email": "test@example.com"
    }}"#,
        issuer,
        now - expired_seconds_ago - 3600, // Auth time is before expiration
        exp,
        now - expired_seconds_ago - 3600  // Issued at is before expiration
    ));

    // For testing purposes, we use a dummy signature
    let signature = "test-signature";

    create_test_token(&header, &payload, signature)
}

// Helper function to create a token with wrong issuer
fn create_wrong_issuer_token(kid: &str) -> String {
    let now = current_timestamp();
    let exp = now + 3600;

    let header = base64_encode(&format!(r#"{{"alg":"RS256","kid":"{}","typ":"JWT"}}"#, kid));
    let payload = base64_encode(&format!(
        r#"{{
        "sub": "test-user-id",
        "iss": "https://wrong-issuer.com",
        "client_id": "test-client-id",
        "token_use": "id",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-id",
        "email": "test@example.com"
    }}"#,
        now - 300,
        exp,
        now
    ));

    // For testing purposes, we use a dummy signature
    let signature = "test-signature";

    create_test_token(&header, &payload, signature)
}

// Helper function to create a token with future issued at time (clock manipulation attack)
fn create_future_iat_token(issuer: &str, kid: &str, seconds_in_future: u64) -> String {
    let now = current_timestamp();
    let future_time = now + seconds_in_future;
    let exp = future_time + 3600;

    let header = base64_encode(&format!(r#"{{"alg":"RS256","kid":"{}","typ":"JWT"}}"#, kid));
    let payload = base64_encode(&format!(
        r#"{{
        "sub": "test-user-id",
        "iss": "{}",
        "client_id": "test-client-id",
        "token_use": "id",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-id",
        "email": "test@example.com"
    }}"#,
        issuer,
        future_time,
        exp,
        future_time
    ));

    // For testing purposes, we use a dummy signature
    let signature = "test-signature";

    create_test_token(&header, &payload, signature)
}

#[tokio::test]
async fn test_tampered_tokens() {
    // Create a verifier config
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
        None,
    )
    .unwrap();

    // Create a verifier
    let verifier = CognitoJwtVerifier::new(vec![config]).unwrap();

    // Create issuer
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    
    // Test with tampered subject
    let tampered_sub_token = create_tampered_token(issuer, "test-key-1", "modified_sub");
    let result = verifier.verify_id_token(&tampered_sub_token).await;
    assert!(result.is_err());
    
    // Test with elevated permissions
    let elevated_perms_token = create_tampered_token(issuer, "test-key-1", "elevated_permissions");
    let result = verifier.verify_id_token(&elevated_perms_token).await;
    assert!(result.is_err());
    
    // Test with modified issued-at time
    let modified_iat_token = create_tampered_token(issuer, "test-key-1", "modified_iat");
    let result = verifier.verify_id_token(&modified_iat_token).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_none_algorithm_attack() {
    // Create a verifier config
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
        None,
    )
    .unwrap();

    // Create a verifier
    let verifier = CognitoJwtVerifier::new(vec![config]).unwrap();

    // Create issuer
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    
    // Test with none algorithm
    let none_alg_token = create_none_algorithm_token(issuer);
    let result = verifier.verify_id_token(&none_alg_token).await;
    assert!(result.is_err());
    
    // The error should be related to invalid algorithm or parsing
    match result.unwrap_err() {
        JwtError::InvalidClaim { claim, .. } => {
            assert_eq!(claim, "alg");
        },
        JwtError::ParseError { .. } => {
            // This is also acceptable
        },
        err => panic!("Expected InvalidClaim or ParseError, got: {:?}", err),
    }
}

#[tokio::test]
async fn test_invalid_signature() {
    // Create a verifier config
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
        None,
    )
    .unwrap();

    // Create a verifier
    let verifier = CognitoJwtVerifier::new(vec![config]).unwrap();

    // Create issuer
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    
    // Test with invalid signature
    let invalid_sig_token = create_invalid_signature_token(issuer, "test-key-1");
    let result = verifier.verify_id_token(&invalid_sig_token).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_expired_tokens() {
    // Create a verifier config
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
        None,
    )
    .unwrap();

    // Create a verifier
    let verifier = CognitoJwtVerifier::new(vec![config]).unwrap();

    // Create issuer
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    
    // Test with token expired 1 hour ago
    let expired_token = create_expired_token(issuer, "test-key-1", 3600);
    let result = verifier.verify_id_token(&expired_token).await;
    assert!(result.is_err());
    
    // Test with token expired just now (1 second ago)
    let just_expired_token = create_expired_token(issuer, "test-key-1", 1);
    let result = verifier.verify_id_token(&just_expired_token).await;
    assert!(result.is_err());
    
    // Test with token expired long ago (1 year)
    let long_expired_token = create_expired_token(issuer, "test-key-1", 31536000); // 365 days
    let result = verifier.verify_id_token(&long_expired_token).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_wrong_issuer() {
    // Create a verifier config
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
        None,
    )
    .unwrap();

    // Create a verifier
    let verifier = CognitoJwtVerifier::new(vec![config]).unwrap();
    
    // Test with wrong issuer
    let wrong_issuer_token = create_wrong_issuer_token("test-key-1");
    let result = verifier.verify_id_token(&wrong_issuer_token).await;
    assert!(result.is_err());
    
    // The error should be related to configuration (issuer not found)
    match result.unwrap_err() {
        JwtError::ConfigurationError { parameter, .. } => {
            assert_eq!(parameter, Some("issuer".to_string()));
        },
        err => panic!("Expected ConfigurationError, got: {:?}", err),
    }
}

#[tokio::test]
async fn test_future_issued_at() {
    // Create a verifier config
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
        None,
    )
    .unwrap();

    // Create a verifier
    let verifier = CognitoJwtVerifier::new(vec![config]).unwrap();

    // Create issuer
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    
    // Test with token issued 1 hour in the future (clock manipulation attack)
    let future_token = create_future_iat_token(issuer, "test-key-1", 3600);
    let result = verifier.verify_id_token(&future_token).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_cross_pool_token_rejection() {
    // Create verifier configs for two different user pools
    let config1 = VerifierConfig::new(
        "us-east-1",
        "us-east-1_pool1",
        &["client-id-1".to_string()],
        None,
    )
    .unwrap();
    
    let config2 = VerifierConfig::new(
        "us-west-2",
        "us-west-2_pool2",
        &["client-id-2".to_string()],
        None,
    )
    .unwrap();

    // Create a verifier with both pools
    let verifier = CognitoJwtVerifier::new(vec![config1, config2]).unwrap();

    // Create issuers for both pools
    let issuer1 = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_pool1";
    let issuer2 = "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_pool2";
    
    // Create tokens for both pools
    let token1 = create_test_id_token(issuer1, "test-key-1", 3600);
    let token2 = create_test_id_token(issuer2, "test-key-1", 3600);
    
    // Try to verify token1 with pool2's issuer (should fail)
    let modified_token1 = token1.replace(issuer1, issuer2);
    let result = verifier.verify_id_token(&modified_token1).await;
    assert!(result.is_err());
    
    // Try to verify token2 with pool1's issuer (should fail)
    let modified_token2 = token2.replace(issuer2, issuer1);
    let result = verifier.verify_id_token(&modified_token2).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_token_replay_detection() {
    // This test is more conceptual since we don't have a real implementation of replay detection
    // In a real system, you would check if the token has been used before
    
    // Create a verifier config
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
        None,
    )
    .unwrap();

    // Create a verifier
    let verifier = CognitoJwtVerifier::new(vec![config]).unwrap();

    // Create issuer
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    
    // Create a token
    let token = create_test_id_token(issuer, "test-key-1", 3600);
    
    // First verification should fail due to signature/key issues, but not due to replay
    let result1 = verifier.verify_id_token(&token).await;
    assert!(result1.is_err());
    
    // Second verification should also fail for the same reason
    let result2 = verifier.verify_id_token(&token).await;
    assert!(result2.is_err());
    
    // In a real implementation with replay detection, the second verification would fail
    // with a specific error indicating that the token has been used before
}