easy-auth-sdk 0.2.0

mod claims;
mod error;
mod jwks;

pub use claims::Claims;
pub use error::AuthError;
pub use jsonwebtoken::Algorithm;

use std::sync::RwLock;

use jsonwebtoken::{decode, decode_header, DecodingKey, Validation};

pub struct EasyAuth {
    decoding_keys: RwLock<Vec<(Option<String>, DecodingKey)>>,
    validation: Validation,
}

impl EasyAuth {
    /// Create an EasyAuth instance from a JWKS JSON string.
    ///
    /// # Example
    /// ```ignore
    /// let jwks_json = r#"{"keys":[...]}"#;
    /// let auth = EasyAuth::from_jwks_json(jwks_json)?;
    /// ```
    pub fn from_jwks_json(jwks_json: &str) -> Result<Self, AuthError> {
        let keys = jwks::parse_jwks(jwks_json)?;

        let mut validation = Validation::new(Algorithm::RS256);
        validation.validate_exp = true;

        Ok(Self {
            decoding_keys: RwLock::new(keys),
            validation,
        })
    }

    /// Create an EasyAuth instance from a PEM-encoded public key.
    ///
    /// # Example
    /// ```ignore
    /// let pem = "-----BEGIN PUBLIC KEY-----\n...";
    /// let auth = EasyAuth::from_pem(pem)?;
    /// ```
    pub fn from_pem(pem: &str) -> Result<Self, AuthError> {
        let key = DecodingKey::from_rsa_pem(pem.as_bytes())
            .map_err(|e| AuthError::InvalidKey(format!("Failed to parse PEM: {}", e)))?;

        let mut validation = Validation::new(Algorithm::RS256);
        validation.validate_exp = true;

        Ok(Self {
            decoding_keys: RwLock::new(vec![(None, key)]),
            validation,
        })
    }

    /// Hot-swap the JWKS keys without reconstructing the `EasyAuth` instance.
    ///
    /// Call this when a `KeyNotFound` error indicates the signing keys have
    /// been rotated and the current key set is stale.
    pub fn update_jwks(&self, jwks_json: &str) -> Result<(), AuthError> {
        let keys = jwks::parse_jwks(jwks_json)?;
        let mut guard = self.decoding_keys.write().expect("decoding_keys poisoned");
        *guard = keys;
        Ok(())
    }

    /// Validate the JWT token and return the claims.
    ///
    /// Verifies the token signature and expiration, then returns the decoded claims.
    ///
    /// # Example
    /// ```ignore
    /// let claims = auth.validate(&token)?;
    /// println!("User: {}", claims.sub);
    /// println!("Roles: {:?}", claims.domain_roles);
    /// ```
    pub fn validate(&self, token: &str) -> Result<Claims, AuthError> {
        self.decode_token(token)
    }

    fn decode_token(&self, token: &str) -> Result<Claims, AuthError> {
        let header = decode_header(token)?;
        let kid = header.kid.as_deref();

        let guard = self.decoding_keys.read().expect("decoding_keys poisoned");

        let decoding_key = Self::find_key(&guard, kid)?;
        let token_data = decode::<Claims>(token, decoding_key, &self.validation)?;

        Ok(token_data.claims)
    }

    fn find_key<'a>(
        keys: &'a [(Option<String>, DecodingKey)],
        kid: Option<&str>,
    ) -> Result<&'a DecodingKey, AuthError> {
        if keys.is_empty() {
            return Err(AuthError::InvalidKey("No keys available".to_string()));
        }

        match kid {
            Some(kid) => {
                // Look for exact kid match
                for (key_kid, key) in keys {
                    if key_kid.as_deref() == Some(kid) {
                        return Ok(key);
                    }
                }
                // No match — if all stored keys have kids, this is a key-not-found
                // (the signing key rotated and we don't have the new one).
                // If stored keys have no kids (e.g. PEM), fall back to first key.
                let all_keys_have_kids = keys.iter().all(|(k, _)| k.is_some());
                if all_keys_have_kids {
                    Err(AuthError::KeyNotFound(kid.to_string()))
                } else {
                    Ok(&keys[0].1)
                }
            }
            None => Ok(&keys[0].1),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
    use jsonwebtoken::{encode, EncodingKey, Header};
    use rand::rngs::OsRng;
    use rsa::pkcs1::EncodeRsaPrivateKey;
    use rsa::pkcs8::EncodePublicKey;
    use rsa::traits::PublicKeyParts;
    use rsa::RsaPrivateKey;
    use serde::Serialize;
    use std::time::{SystemTime, UNIX_EPOCH};

    #[derive(Debug, Serialize)]
    struct TestClaims {
        sub: String,
        domain_roles: Vec<String>,
        exp: u64,
        iat: u64,
    }

    struct TestKeys {
        encoding_key: EncodingKey,
        pem_public: String,
        jwks_json: String,
    }

    fn generate_test_keys() -> TestKeys {
        let mut rng = OsRng;
        let private_key = RsaPrivateKey::new(&mut rng, 2048).unwrap();
        let public_key = private_key.to_public_key();

        let private_pem = private_key.to_pkcs1_pem(Default::default()).unwrap();
        let public_pem = public_key.to_public_key_pem(Default::default()).unwrap();

        let encoding_key = EncodingKey::from_rsa_pem(private_pem.as_bytes()).unwrap();

        let n = URL_SAFE_NO_PAD.encode(private_key.n().to_bytes_be());
        let e = URL_SAFE_NO_PAD.encode(private_key.e().to_bytes_be());

        let jwks_json = format!(
            r#"{{"keys":[{{"kty":"RSA","kid":"test-key","use":"sig","alg":"RS256","n":"{}","e":"{}"}}]}}"#,
            n, e
        );

        TestKeys {
            encoding_key,
            pem_public: public_pem,
            jwks_json,
        }
    }

    fn create_token(keys: &TestKeys, claims: &TestClaims) -> String {
        let mut header = Header::new(Algorithm::RS256);
        header.kid = Some("test-key".to_string());
        encode(&header, claims, &keys.encoding_key).unwrap()
    }

    fn now_secs() -> u64 {
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs()
    }

    #[test]
    fn test_allowed_domain_roles_with_matching_role() {
        let keys = generate_test_keys();
        let auth = EasyAuth::from_jwks_json(&keys.jwks_json).unwrap();

        let test_claims = TestClaims {
            sub: "user-123".to_string(),
            domain_roles: vec!["moon:user".to_string(), "example:admin".to_string()],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        let token = create_token(&keys, &test_claims);
        let claims = auth.validate(&token).unwrap();
        assert!(claims.allowed_domain_roles(&["moon:user"]));
        assert_eq!(claims.sub, "user-123");
    }

    #[test]
    fn test_allowed_domain_roles_without_matching_role() {
        let keys = generate_test_keys();
        let auth = EasyAuth::from_jwks_json(&keys.jwks_json).unwrap();

        let test_claims = TestClaims {
            sub: "user-123".to_string(),
            domain_roles: vec!["example:viewer".to_string()],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        let token = create_token(&keys, &test_claims);
        let claims = auth.validate(&token).unwrap();
        assert!(!claims.allowed_domain_roles(&["moon:admin"]));
    }

    #[test]
    fn test_is_subject_matching() {
        let keys = generate_test_keys();
        let auth = EasyAuth::from_pem(&keys.pem_public).unwrap();

        let test_claims = TestClaims {
            sub: "295fafbb-7da3-4881-858f-e6ea5d2b65ae".to_string(),
            domain_roles: vec![],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        let mut header = Header::new(Algorithm::RS256);
        header.kid = None;
        let token = encode(&header, &test_claims, &keys.encoding_key).unwrap();

        let claims = auth.validate(&token).unwrap();
        assert!(claims.is_subject("295fafbb-7da3-4881-858f-e6ea5d2b65ae"));
    }

    #[test]
    fn test_is_subject_not_matching() {
        let keys = generate_test_keys();
        let auth = EasyAuth::from_pem(&keys.pem_public).unwrap();

        let test_claims = TestClaims {
            sub: "user-123".to_string(),
            domain_roles: vec![],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        let mut header = Header::new(Algorithm::RS256);
        header.kid = None;
        let token = encode(&header, &test_claims, &keys.encoding_key).unwrap();

        let claims = auth.validate(&token).unwrap();
        assert!(!claims.is_subject("different-user"));
    }

    #[test]
    fn test_validate() {
        let keys = generate_test_keys();
        let auth = EasyAuth::from_jwks_json(&keys.jwks_json).unwrap();

        let test_claims = TestClaims {
            sub: "user-456".to_string(),
            domain_roles: vec!["test:role".to_string()],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        let token = create_token(&keys, &test_claims);
        let claims = auth.validate(&token).unwrap();
        assert_eq!(claims.sub, "user-456");
        assert_eq!(claims.domain_roles, vec!["test:role".to_string()]);
    }

    #[test]
    fn test_combined_checks() {
        let keys = generate_test_keys();
        let auth = EasyAuth::from_jwks_json(&keys.jwks_json).unwrap();

        let test_claims = TestClaims {
            sub: "user-789".to_string(),
            domain_roles: vec!["api:read".to_string(), "api:write".to_string()],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        let token = create_token(&keys, &test_claims);

        // Validate once, check multiple times
        let claims = auth.validate(&token).unwrap();
        assert!(claims.allowed_domain_roles(&["api:read"]));
        assert!(claims.is_subject("user-789"));

        // OR logic: allow if subject matches OR has admin role
        assert!(claims.is_subject("user-789") || claims.allowed_domain_roles(&["admin"]));
    }

    #[test]
    fn test_expired_token() {
        let keys = generate_test_keys();
        let auth = EasyAuth::from_jwks_json(&keys.jwks_json).unwrap();

        let test_claims = TestClaims {
            sub: "user-123".to_string(),
            domain_roles: vec!["moon:user".to_string()],
            exp: now_secs() - 3600, // Expired 1 hour ago
            iat: now_secs() - 7200,
        };

        let token = create_token(&keys, &test_claims);
        let result = auth.validate(&token);

        assert!(matches!(result, Err(AuthError::TokenExpired)));
    }

    #[test]
    fn test_invalid_signature() {
        let keys1 = generate_test_keys();
        let keys2 = generate_test_keys();

        // Create auth with keys1
        let auth = EasyAuth::from_jwks_json(&keys1.jwks_json).unwrap();

        // Create token with keys2 (different key)
        let test_claims = TestClaims {
            sub: "user-123".to_string(),
            domain_roles: vec!["moon:user".to_string()],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };
        let token = create_token(&keys2, &test_claims);

        let result = auth.validate(&token);
        assert!(matches!(result, Err(AuthError::InvalidSignature)));
    }

    #[test]
    fn test_malformed_token() {
        let keys = generate_test_keys();
        let auth = EasyAuth::from_jwks_json(&keys.jwks_json).unwrap();

        let result = auth.validate("not.a.valid.token");
        assert!(matches!(result, Err(AuthError::InvalidToken(_))));
    }

    #[test]
    fn test_invalid_jwks() {
        let result = EasyAuth::from_jwks_json("not valid json");
        assert!(matches!(result, Err(AuthError::JsonError(_))));
    }

    #[test]
    fn test_empty_jwks() {
        let result = EasyAuth::from_jwks_json(r#"{"keys":[]}"#);
        assert!(matches!(result, Err(AuthError::InvalidKey(_))));
    }

    fn generate_test_keys_with_kid(kid: &str) -> TestKeys {
        let mut rng = OsRng;
        let private_key = RsaPrivateKey::new(&mut rng, 2048).unwrap();
        let public_key = private_key.to_public_key();

        let private_pem = private_key.to_pkcs1_pem(Default::default()).unwrap();
        let public_pem = public_key.to_public_key_pem(Default::default()).unwrap();

        let encoding_key = EncodingKey::from_rsa_pem(private_pem.as_bytes()).unwrap();

        let n = URL_SAFE_NO_PAD.encode(private_key.n().to_bytes_be());
        let e = URL_SAFE_NO_PAD.encode(private_key.e().to_bytes_be());

        let jwks_json = format!(
            r#"{{"keys":[{{"kty":"RSA","kid":"{}","use":"sig","alg":"RS256","n":"{}","e":"{}"}}]}}"#,
            kid, n, e
        );

        TestKeys {
            encoding_key,
            pem_public: public_pem,
            jwks_json,
        }
    }

    fn create_token_with_kid(keys: &TestKeys, claims: &TestClaims, kid: &str) -> String {
        let mut header = Header::new(Algorithm::RS256);
        header.kid = Some(kid.to_string());
        encode(&header, claims, &keys.encoding_key).unwrap()
    }

    #[test]
    fn test_key_not_found() {
        let keys = generate_test_keys_with_kid("old-key");
        let auth = EasyAuth::from_jwks_json(&keys.jwks_json).unwrap();

        let test_claims = TestClaims {
            sub: "user-123".to_string(),
            domain_roles: vec![],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        // Sign with a kid that doesn't exist in the JWKS
        let token = create_token_with_kid(&keys, &test_claims, "rotated-new-key");
        let result = auth.validate(&token);
        assert!(
            matches!(result, Err(AuthError::KeyNotFound(ref kid)) if kid == "rotated-new-key"),
            "Expected KeyNotFound for unknown kid, got: {:?}",
            result
        );
    }

    #[test]
    fn test_update_jwks() {
        let old_keys = generate_test_keys_with_kid("old-key");
        let new_keys = generate_test_keys_with_kid("new-key");
        let auth = EasyAuth::from_jwks_json(&old_keys.jwks_json).unwrap();

        let test_claims = TestClaims {
            sub: "user-123".to_string(),
            domain_roles: vec![],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        // Token signed with new key initially fails
        let token = create_token_with_kid(&new_keys, &test_claims, "new-key");
        assert!(matches!(
            auth.validate(&token),
            Err(AuthError::KeyNotFound(_))
        ));

        // After updating JWKS with new keys, validation succeeds
        auth.update_jwks(&new_keys.jwks_json).unwrap();
        let claims = auth.validate(&token).unwrap();
        assert_eq!(claims.sub, "user-123");
    }

    #[test]
    fn test_pem_fallback_no_key_not_found() {
        // PEM keys have no kid — should fall back to first key, not KeyNotFound
        let keys = generate_test_keys();
        let auth = EasyAuth::from_pem(&keys.pem_public).unwrap();

        let test_claims = TestClaims {
            sub: "user-123".to_string(),
            domain_roles: vec![],
            exp: now_secs() + 3600,
            iat: now_secs(),
        };

        // Token with an unknown kid still works because PEM keys have no kids
        let token = create_token_with_kid(&keys, &test_claims, "any-kid");
        let claims = auth.validate(&token).unwrap();
        assert_eq!(claims.sub, "user-123");
    }
}