jwt-verify 0.1.2

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 test access token
fn create_test_access_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": "access",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-id",
        "scope": "openid profile email",
        "version": 2
    }}"#,
        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 test refresh token
fn create_test_refresh_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": "refresh",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-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 test token with invalid signature
fn create_test_token_with_invalid_signature(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"
    }}"#,
        issuer,
        now - 300,
        exp,
        now
    ));

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

    create_test_token(&header, &payload, signature)
}

// Helper function to create a test token with invalid issuer
fn create_test_token_with_invalid_issuer(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://invalid-issuer.com",
        "client_id": "test-client-id",
        "token_use": "id",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-id"
    }}"#,
        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 test token with invalid client ID
fn create_test_token_with_invalid_client_id(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": "invalid-client-id",
        "token_use": "id",
        "auth_time": {},
        "exp": {},
        "iat": {},
        "jti": "test-jwt-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 test token with expired timestamp
fn create_test_expired_token(issuer: &str, kid: &str) -> String {
    let now = current_timestamp();
    let exp = now - 3600; // Expired 1 hour ago

    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"
    }}"#,
        issuer,
        now - 7200,
        exp,
        now - 7200
    ));

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

    create_test_token(&header, &payload, signature)
}

#[tokio::test]
async fn test_missing_token() {
    // 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 missing token
    let result = verifier.verify_id_token("").await;
    assert!(result.is_err());
    match result.unwrap_err() {
        JwtError::MissingToken => {
            // This is the expected error
        }
        err => panic!("Expected MissingToken error, got: {:?}", err),
    }
}

#[tokio::test]
async fn test_invalid_token_format() {
    // 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 invalid token format
    let result = verifier.verify_id_token("invalid.token").await;
    assert!(result.is_err());
    match result.unwrap_err() {
        JwtError::ParseError { .. } => {
            // This is the expected error
        }
        err => panic!("Expected ParseError, got: {:?}", err),
    }
}

#[tokio::test]
async fn test_invalid_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 invalid issuer
    let token = create_test_token_with_invalid_issuer("test-key-1");
    let result = verifier.verify_id_token(&token).await;
    assert!(result.is_err());
    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_different_token_types() {
    // 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 test tokens
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    let id_token = create_test_id_token(issuer, "test-key-1", 3600);
    let access_token = create_test_access_token(issuer, "test-key-1", 3600);
    let refresh_token = create_test_refresh_token(issuer, "test-key-1");

    // Verify ID token
    let id_result = verifier.verify_id_token(&id_token).await;
    assert!(id_result.is_err()); // Will fail due to signature or key not found

    // Verify access token
    let access_result = verifier.verify_access_token(&access_token).await;
    assert!(access_result.is_err()); // Will fail due to signature or key not found

    // Try to verify a refresh token (should be rejected)
    let refresh_result = verifier.verify_id_token(&refresh_token).await;
    assert!(refresh_result.is_err());
}

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

    let config2 = VerifierConfig::new(
        "us-west-2",
        "us-west-2_example2",
        &["test-client-id-2".to_string()],
        None,
    )
    .unwrap();

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

    // Create test tokens for each user pool
    let issuer1 = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example1";
    let issuer2 = "https://cognito-idp.us-west-2.amazonaws.com/us-west-2_example2";
    let token1 = create_test_id_token(issuer1, "test-key-1", 3600);
    let token2 = create_test_id_token(issuer2, "test-key-1", 3600);

    // Verify tokens
    let result1 = verifier.verify_id_token(&token1).await;
    let result2 = verifier.verify_id_token(&token2).await;

    // Both will fail due to signature or key not found, but we're testing that the correct user pool is selected
    assert!(result1.is_err());
    assert!(result2.is_err());
}
#[tokio::test]
async fn test_add_and_remove_user_pool() {
    // Create a verifier with no user pools
    let mut verifier = CognitoJwtVerifier::new(vec![]).unwrap();
    
    // Verify that there are no user pools initially
    let pool_ids = verifier.get_user_pool_ids();
    assert_eq!(pool_ids.len(), 0);
    
    // Add a user pool
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
        None,
    ).unwrap();
    
    let result = verifier.add_user_pool("test-pool", config);
    assert!(result.is_ok());
    
    // Verify that the user pool was added
    let pool_ids = verifier.get_user_pool_ids();
    assert_eq!(pool_ids.len(), 1);
    assert!(pool_ids.contains(&"test-pool".to_string()));
    
    // Add another user pool using the convenience method
    let result = verifier.add_user_pool_with_params(
        "test-pool-2",
        "us-west-2",
        "us-west-2_example",
        &["test-client-id-2".to_string()],
    );
    assert!(result.is_ok());
    
    // Verify that the second user pool was added
    let pool_ids = verifier.get_user_pool_ids();
    assert_eq!(pool_ids.len(), 2);
    assert!(pool_ids.contains(&"test-pool".to_string()));
    assert!(pool_ids.contains(&"test-pool-2".to_string()));
    
    // Remove a user pool
    let result = verifier.remove_user_pool("test-pool");
    assert!(result.is_ok());
    
    // Verify that the user pool was removed
    let pool_ids = verifier.get_user_pool_ids();
    assert_eq!(pool_ids.len(), 1);
    assert!(!pool_ids.contains(&"test-pool".to_string()));
    assert!(pool_ids.contains(&"test-pool-2".to_string()));
    
    // Try to remove a non-existent user pool
    let result = verifier.remove_user_pool("non-existent-pool");
    assert!(result.is_err());
    match result.unwrap_err() {
        JwtError::ConfigurationError { parameter, .. } => {
            assert_eq!(parameter, Some("pool_id".to_string()));
        },
        err => panic!("Expected ConfigurationError, got: {:?}", err),
    }
}

#[tokio::test]
async fn test_new_single_pool() {
    // Create a verifier with a single user pool
    let verifier = CognitoJwtVerifier::new_single_pool(
        "us-east-1",
        "us-east-1_example",
        &["test-client-id".to_string()],
    ).unwrap();
    
    // Verify that the user pool was added
    let pool_ids = verifier.get_user_pool_ids();
    assert_eq!(pool_ids.len(), 1);
    assert!(pool_ids.contains(&"us-east-1_us-east-1_example".to_string()));
}

#[tokio::test]
async fn test_prefetch_all() {
    // Create a verifier with multiple user pools
    let config1 = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example1",
        &["test-client-id-1".to_string()],
        None,
    ).unwrap();
    
    let config2 = VerifierConfig::new(
        "us-west-2",
        "us-west-2_example2",
        &["test-client-id-2".to_string()],
        None,
    ).unwrap();
    
    let verifier = CognitoJwtVerifier::new(vec![config1, config2]).unwrap();
    
    // Prefetch JWKs for all user pools
    let results = verifier.hydrate().await;
    
    // Verify that we got results for both user pools
    assert_eq!(results.len(), 2);
    
    // The prefetch will likely fail since we're not mocking the HTTP responses,
    // but we're just testing that the method returns results for all user pools
    let pool_ids: Vec<String> = results.iter().map(|(id, _)| id.clone()).collect();
    assert!(pool_ids.contains(&"us-east-1_us-east-1_example1".to_string()));
    assert!(pool_ids.contains(&"us-west-2_us-west-2_example2".to_string()));
}

#[tokio::test]
async fn test_invalid_token_use() {
    // Since we can't directly access TokenUse enum in integration tests,
    // we'll skip testing token use restrictions and just test with the default configuration
    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 an access token
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    let access_token = create_test_access_token(issuer, "test-key-1", 3600);
    
    // Try to verify the access token
    let result = verifier.verify_access_token(&access_token).await;
    
    // The verification will fail because the token use is not allowed
    assert!(result.is_err());
    
    // Since we can't directly access TokenUse enum in integration tests,
    // we'll skip testing token use restrictions and just test with the default configuration
    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 an ID token
    let id_token = create_test_id_token(issuer, "test-key-1", 3600);
    
    // Try to verify the ID token
    let result = verifier.verify_id_token(&id_token).await;
    
    // The verification will fail because the token use is not allowed
    assert!(result.is_err());
}

#[tokio::test]
async fn test_refresh_token_rejection() {
    // 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 a refresh token
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    let refresh_token = create_test_refresh_token(issuer, "test-key-1");
    
    // Try to verify the refresh token as an ID token
    let result = verifier.verify_id_token(&refresh_token).await;
    
    // The verification should fail with an UnsupportedTokenType error
    assert!(result.is_err());
    match result.unwrap_err() {
        JwtError::ConfigurationError { .. } => {
            // This is expected since the issuer lookup will fail before token type validation
        },
        err => println!("Got error: {:?}", err),
    }
}

#[tokio::test]
async fn test_invalid_signature_token() {
    // 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 a token with invalid signature
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    let token = create_test_token_with_invalid_signature(issuer, "test-key-1");
    
    // Try to verify the token
    let result = verifier.verify_id_token(&token).await;
    
    // The verification will fail due to invalid signature
    assert!(result.is_err());
}

#[tokio::test]
async fn test_expired_token() {
    // 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 an expired token
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    let token = create_test_expired_token(issuer, "test-key-1");
    
    // Try to verify the token
    let result = verifier.verify_id_token(&token).await;
    
    // The verification will fail due to expired token
    assert!(result.is_err());
}

#[tokio::test]
async fn test_invalid_client_id() {
    // Create a verifier config
    let config = VerifierConfig::new(
        "us-east-1",
        "us-east-1_example",
        &["valid-client-id".to_string()],
        None,
    ).unwrap();
    
    // Create a verifier
    let verifier = CognitoJwtVerifier::new(vec![config]).unwrap();
    
    // Create a token with invalid client ID
    let issuer = "https://cognito-idp.us-east-1.amazonaws.com/us-east-1_example";
    let token = create_test_token_with_invalid_client_id(issuer, "test-key-1");
    
    // Try to verify the token
    let result = verifier.verify_id_token(&token).await;
    
    // The verification will fail due to invalid client ID
    assert!(result.is_err());
}