diff options
Diffstat (limited to 'src/main.rs')
| -rw-r--r-- | src/main.rs | 176 |
1 files changed, 175 insertions, 1 deletions
diff --git a/src/main.rs b/src/main.rs index 25d56ea..9b6b0ea 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,8 +12,8 @@ fn main() { #[cfg(test)] mod tests { - use super::*; use std::collections::HashMap; + use base64::Engine; #[test] fn test_oauth_server_creation() { @@ -74,4 +74,178 @@ mod tests { assert!(result.is_err()); assert!(result.unwrap_err().contains("unsupported_response_type")); } + + #[test] + fn test_jwks_endpoint() { + let config = sts::Config::from_env(); + let oauth_server = sts::OAuthServer::new(&config).expect("Failed to create OAuth server"); + + let jwks_json = oauth_server.get_jwks(); + assert!(!jwks_json.is_empty()); + + // Parse the JSON to verify structure + let jwks: serde_json::Value = serde_json::from_str(&jwks_json).expect("Invalid JWKS JSON"); + assert!(jwks["keys"].is_array()); + assert!(jwks["keys"].as_array().unwrap().len() > 0); + + let key = &jwks["keys"][0]; + assert_eq!(key["kty"], "RSA"); + assert_eq!(key["use"], "sig"); + assert_eq!(key["alg"], "RS256"); + assert!(key["kid"].is_string()); + assert!(key["n"].is_string()); + assert!(key["e"].is_string()); + } + + #[test] + fn test_full_oauth_flow_with_rsa_tokens() { + let config = sts::Config::from_env(); + let oauth_server = sts::OAuthServer::new(&config).expect("Failed to create OAuth server"); + + // Step 1: Authorization request + let mut auth_params = HashMap::new(); + auth_params.insert("client_id".to_string(), "test_client".to_string()); + auth_params.insert("redirect_uri".to_string(), "http://localhost:3000/callback".to_string()); + auth_params.insert("response_type".to_string(), "code".to_string()); + auth_params.insert("state".to_string(), "test_state".to_string()); + + let auth_result = oauth_server.handle_authorize(&auth_params).expect("Authorization failed"); + + // Extract the authorization code from redirect URL + let redirect_url = url::Url::parse(&auth_result).expect("Invalid redirect URL"); + let auth_code = redirect_url + .query_pairs() + .find(|(key, _)| key == "code") + .map(|(_, value)| value.to_string()) + .expect("No authorization code in redirect"); + + // Step 2: Token request + let mut token_params = HashMap::new(); + token_params.insert("grant_type".to_string(), "authorization_code".to_string()); + token_params.insert("code".to_string(), auth_code); + token_params.insert("client_id".to_string(), "test_client".to_string()); + + let token_result = oauth_server.handle_token(&token_params).expect("Token request failed"); + + // Parse token response + let token_response: serde_json::Value = serde_json::from_str(&token_result) + .expect("Invalid token response JSON"); + + assert_eq!(token_response["token_type"], "Bearer"); + assert_eq!(token_response["expires_in"], 3600); + assert!(token_response["access_token"].is_string()); + + let access_token = token_response["access_token"].as_str().unwrap(); + + // Step 3: Verify the JWT token has RSA signature and key ID + let header = jsonwebtoken::decode_header(access_token).expect("Failed to decode JWT header"); + assert_eq!(header.alg, jsonwebtoken::Algorithm::RS256); + assert!(header.kid.is_some()); + assert!(!header.kid.as_ref().unwrap().is_empty()); + } + + #[test] + fn test_token_validation_with_jwks() { + let config = sts::Config::from_env(); + let oauth_server = sts::OAuthServer::new(&config).expect("Failed to create OAuth server"); + + // Generate a token + let mut auth_params = HashMap::new(); + auth_params.insert("client_id".to_string(), "test_client".to_string()); + auth_params.insert("redirect_uri".to_string(), "http://localhost:3000/callback".to_string()); + auth_params.insert("response_type".to_string(), "code".to_string()); + + let auth_result = oauth_server.handle_authorize(&auth_params).expect("Authorization failed"); + let redirect_url = url::Url::parse(&auth_result).expect("Invalid redirect URL"); + let auth_code = redirect_url + .query_pairs() + .find(|(key, _)| key == "code") + .map(|(_, value)| value.to_string()) + .expect("No authorization code in redirect"); + + let mut token_params = HashMap::new(); + token_params.insert("grant_type".to_string(), "authorization_code".to_string()); + token_params.insert("code".to_string(), auth_code); + token_params.insert("client_id".to_string(), "test_client".to_string()); + + let token_result = oauth_server.handle_token(&token_params).expect("Token request failed"); + let token_response: serde_json::Value = serde_json::from_str(&token_result) + .expect("Invalid token response JSON"); + let access_token = token_response["access_token"].as_str().unwrap(); + + // Get the JWKS + let jwks_json = oauth_server.get_jwks(); + let jwks: serde_json::Value = serde_json::from_str(&jwks_json).expect("Invalid JWKS JSON"); + + // Decode the token header to get the key ID + let header = jsonwebtoken::decode_header(access_token).expect("Failed to decode JWT header"); + let kid = header.kid.as_ref().expect("No key ID in token"); + + // Find the matching key in JWKS + let matching_key = jwks["keys"] + .as_array() + .unwrap() + .iter() + .find(|key| key["kid"] == *kid) + .expect("Key ID not found in JWKS"); + + assert_eq!(matching_key["kty"], "RSA"); + assert_eq!(matching_key["alg"], "RS256"); + } + + #[test] + fn test_token_contains_proper_claims() { + let config = sts::Config::from_env(); + let oauth_server = sts::OAuthServer::new(&config).expect("Failed to create OAuth server"); + + // Generate a token through the full flow + let mut auth_params = HashMap::new(); + auth_params.insert("client_id".to_string(), "test_client".to_string()); + auth_params.insert("redirect_uri".to_string(), "http://localhost:3000/callback".to_string()); + auth_params.insert("response_type".to_string(), "code".to_string()); + auth_params.insert("scope".to_string(), "openid profile".to_string()); + + let auth_result = oauth_server.handle_authorize(&auth_params).expect("Authorization failed"); + let redirect_url = url::Url::parse(&auth_result).expect("Invalid redirect URL"); + let auth_code = redirect_url + .query_pairs() + .find(|(key, _)| key == "code") + .map(|(_, value)| value.to_string()) + .expect("No authorization code in redirect"); + + let mut token_params = HashMap::new(); + token_params.insert("grant_type".to_string(), "authorization_code".to_string()); + token_params.insert("code".to_string(), auth_code); + token_params.insert("client_id".to_string(), "test_client".to_string()); + + let token_result = oauth_server.handle_token(&token_params).expect("Token request failed"); + let token_response: serde_json::Value = serde_json::from_str(&token_result) + .expect("Invalid token response JSON"); + let access_token = token_response["access_token"].as_str().unwrap(); + + // Decode the token without verification to check claims + let _token_data = jsonwebtoken::decode::<serde_json::Value>( + access_token, + &jsonwebtoken::DecodingKey::from_secret(b"dummy"), // We're not validating, just parsing + &jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::HS256) + ); + + // Since we can't validate with a dummy key, we'll just verify the structure + // by decoding the payload manually + let parts: Vec<&str> = access_token.split('.').collect(); + assert_eq!(parts.len(), 3); // header.payload.signature + + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(parts[1]) + .expect("Failed to decode payload"); + let claims: serde_json::Value = serde_json::from_slice(&payload).expect("Invalid claims JSON"); + + assert!(claims["sub"].is_string()); + assert!(claims["iss"].is_string()); + assert!(claims["aud"].is_string()); + assert!(claims["exp"].is_number()); + assert!(claims["iat"].is_number()); + assert_eq!(claims["aud"], "test_client"); + assert_eq!(claims["scope"], "openid profile"); + } } |