klieo-auth-oauth 3.4.0

//! `OAuthAuthenticator` impl of [`klieo_auth_common::Authenticator`].
//!
//! Validates RS256/ES256 JWT bearer tokens against a JWKS cache.
//! Standard claims validated: `iss` (required), `aud` (required),
//! `exp`, `nbf`. Scope claim parsed per RFC 6749 §3.3 (space-delimited
//! string OR JSON array) from `scope` or `scp` (Azure AD).
//! `authorize_method` consults an operator-supplied
//! method-to-required-scopes map.

use crate::cache::TokenCache;
use crate::introspect::IntrospectionConfig;
use crate::jwks::Cache;
use async_trait::async_trait;
use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
use klieo_auth_common::{AuthError, Authenticator, Headers, Identity};
use serde::Deserialize;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

/// OAuth 2.0 bearer-token Authenticator. Construct via
/// [`OAuthAuthenticator::builder`].
#[derive(Clone)]
pub struct OAuthAuthenticator {
    pub(crate) issuer: String,
    pub(crate) audience: String,
    pub(crate) allowed_algs: Vec<Algorithm>,
    pub(crate) cache: Cache,
    pub(crate) scope_map: Arc<HashMap<String, Vec<String>>>,
    pub(crate) http: reqwest::Client,
    pub(crate) introspection: Option<IntrospectionConfig>,
    pub(crate) token_cache: Option<Arc<TokenCache>>,
}

impl std::fmt::Debug for OAuthAuthenticator {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // Audience may carry tenant-id-ish data — omit the value and
        // surface only the algorithm allowlist. Scope-map keys are
        // method names (safe to surface); values can be
        // operator-defined identifiers, so list the method names only.
        f.debug_struct("OAuthAuthenticator")
            .field("issuer", &self.issuer)
            .field("allowed_algs", &self.allowed_algs)
            .field(
                "scope_map_methods",
                &self.scope_map.keys().collect::<Vec<_>>(),
            )
            .field("introspection_configured", &self.introspection.is_some())
            .field("token_cache_enabled", &self.token_cache.is_some())
            .finish()
    }
}

/// JWT claims this authenticator cares about. `iss`/`exp`/`nbf`/`aud`
/// are validated by `jsonwebtoken` via [`Validation`]; we deserialise
/// them here so we can also lift `sub` + scope claims.
#[derive(Deserialize)]
struct Claims {
    sub: Option<String>,
    // RFC 6749 §3.3: space-delimited string. Some issuers emit an array;
    // accept both via the untagged enum below.
    scope: Option<ScopeClaim>,
    // Azure AD compat: same shape, different key.
    scp: Option<ScopeClaim>,
    // Expiry timestamp (seconds since epoch). `jsonwebtoken` validates
    // it inside `decode`; we lift it here so the token cache can
    // honour the same value when sizing its entry TTL.
    exp: Option<u64>,
}

#[derive(Deserialize, Debug)]
#[serde(untagged)]
enum ScopeClaim {
    Space(String),
    Array(Vec<String>),
}

impl ScopeClaim {
    fn into_set(self) -> HashSet<String> {
        match self {
            Self::Space(s) => s.split_ascii_whitespace().map(String::from).collect(),
            Self::Array(v) => v.into_iter().collect(),
        }
    }
}

/// `(sub, scopes, exp_unix_opt)` — verifier output shared across the
/// JWT and introspection branches. `exp_unix_opt` flows into the
/// token cache to size entry TTLs; it is intentionally never surfaced
/// to callers (the cache is the only consumer).
type VerifiedToken = (String, HashSet<String>, Option<u64>);

impl OAuthAuthenticator {
    /// Verify a bearer: cache check first, then JWT path, then RFC 7662
    /// introspection fallback if configured. Cache hits short-circuit
    /// before any CPU-bound JWT verify or network introspection
    /// roundtrip; cache misses fall through and successful verifies
    /// re-populate the cache for the next presentation of the same
    /// bearer.
    async fn verify_token(&self, bearer: &str) -> Result<(String, HashSet<String>), AuthError> {
        if let Some(cache) = self.token_cache.as_deref() {
            if let Some(hit) = cache.get(bearer) {
                return Ok(hit);
            }
        }
        let verified = self.verify_uncached(bearer).await?;
        let (sub, scopes, exp_unix) = verified;
        if let Some(cache) = self.token_cache.as_deref() {
            cache.insert(bearer, sub.clone(), scopes.clone(), exp_unix);
        }
        Ok((sub, scopes))
    }

    /// Run the uncached verify pipeline: JWT first, introspection
    /// fallback when configured. Both branches surface the token's
    /// `exp` so the cache can size its entry TTL.
    async fn verify_uncached(&self, bearer: &str) -> Result<VerifiedToken, AuthError> {
        match self.try_jwt_verify(bearer).await {
            Ok(verified) => return Ok(verified),
            Err(jwt_err) => {
                // Fall through to introspection only when configured.
                // Signature-failure retries already happen inside
                // `try_jwt_verify` via the one-shot JWKS refresh.
                tracing::debug!(
                    target: "klieo.auth",
                    error = ?jwt_err,
                    "JWT verify failed; checking introspection fallback"
                );
            }
        }
        let Some(cfg) = self.introspection.as_ref() else {
            return Err(AuthError::Rejected("token validation failed".into()));
        };
        self.verify_via_introspection(cfg, bearer).await
    }

    /// RFC 7662 introspection branch. Active-false / missing-sub
    /// surface as [`AuthError::Rejected`]; endpoint failures surface
    /// as [`OAuthError::IntrospectionFailed`] mapped to `Rejected`.
    #[tracing::instrument(
        skip_all,
        fields(klieo.auth.mode = "oauth_introspect"),
        level = "debug",
    )]
    async fn verify_via_introspection(
        &self,
        cfg: &IntrospectionConfig,
        bearer: &str,
    ) -> Result<VerifiedToken, AuthError> {
        let resp = crate::introspect::call(&self.http, cfg, bearer)
            .await
            .map_err(|e| AuthError::Rejected(format!("token validation failed: {e}")))?;
        if !resp.active {
            return Err(AuthError::Rejected("token inactive".into()));
        }
        // The JWT branch pins `iss` + `aud` via `Validation`; the introspection
        // branch must too. An `active` token at a shared IdP may have been
        // minted for a different issuer or resource server — accepting it here
        // is audience/issuer confusion (RFC 7662, OAuth 2.0 Security BCP). A
        // missing claim is rejected, mirroring the JWT path's required_spec.
        if resp.iss.as_deref() != Some(self.issuer.as_str()) {
            return Err(AuthError::Rejected("token validation failed".into()));
        }
        if !resp
            .aud
            .as_ref()
            .is_some_and(|aud| aud.contains(&self.audience))
        {
            return Err(AuthError::Rejected("token validation failed".into()));
        }
        let sub = resp
            .sub
            .ok_or_else(|| AuthError::Rejected("token validation failed".into()))?;
        let scopes = resp
            .scope
            .map(|s| s.split_ascii_whitespace().map(String::from).collect())
            .unwrap_or_default();
        Ok((sub, scopes, resp.exp))
    }

    /// Verify the JWT signature + standard claims and return
    /// `(sub, scopes)` on success. Retries once with a JWKS refresh if
    /// the first attempt fails — this transparently absorbs issuer key
    /// rotation without proactive polling.
    #[tracing::instrument(
        skip_all,
        fields(klieo.auth.kid = tracing::field::Empty),
        level = "debug",
    )]
    async fn try_jwt_verify(&self, bearer: &str) -> Result<VerifiedToken, AuthError> {
        let header = decode_header(bearer)
            .map_err(|_| AuthError::Rejected("malformed JWT header".into()))?;
        let kid = header
            .kid
            .ok_or_else(|| AuthError::Rejected("JWT header missing kid".into()))?;
        tracing::Span::current().record("klieo.auth.kid", kid.as_str());

        let validation = self.build_validation(header.alg)?;

        let data = match self.try_verify(bearer, &kid, &validation).await {
            Ok(data) => data,
            Err(first_err) => {
                tracing::debug!(
                    target: "klieo_auth_oauth",
                    reason = ?first_err,
                    "first JWT verify attempt failed; refreshing JWKS and retrying"
                );
                // One-shot JWKS refresh — covers issuer key rotation
                // without proactive polling. If the refresh itself
                // fails, surface the same redacted Rejected as a
                // verification failure.
                self.cache
                    .refresh()
                    .await
                    .map_err(|_| AuthError::Rejected("token validation failed".into()))?;
                self.try_verify(bearer, &kid, &validation).await?
            }
        };

        let claims = data.claims;
        let sub = claims
            .sub
            .ok_or_else(|| AuthError::Rejected("token missing sub claim".into()))?;
        let scopes = claims
            .scope
            .or(claims.scp)
            .map(ScopeClaim::into_set)
            .unwrap_or_default();
        Ok((sub, scopes, claims.exp))
    }

    fn build_validation(&self, alg: Algorithm) -> Result<Validation, AuthError> {
        if !self.allowed_algs.contains(&alg) {
            return Err(AuthError::Rejected("algorithm not permitted".into()));
        }
        let mut validation = Validation::new(alg);
        validation.set_issuer(&[self.issuer.as_str()]);
        validation.set_audience(&[self.audience.as_str()]);
        validation.validate_nbf = true;
        Ok(validation)
    }

    async fn try_verify(
        &self,
        bearer: &str,
        kid: &str,
        validation: &Validation,
    ) -> Result<jsonwebtoken::TokenData<Claims>, VerifyAttemptError> {
        let key: DecodingKey = self
            .cache
            .get(kid)
            .await
            .ok_or(VerifyAttemptError::KidMissing)?;
        decode::<Claims>(bearer, &key, validation).map_err(VerifyAttemptError::Decode)
    }
}

/// Internal: distinguishes "kid not in cache" from "decode failed" so
/// the retry path can refresh JWKS and then re-look-up the key. The
/// `Decode` payload feeds `tracing::debug!` on the retry path; it is
/// intentionally never surfaced to callers (decoder diagnostics can
/// leak token-shape detail).
#[allow(dead_code)] // Debug-formatted into trace logs on retry path.
#[derive(Debug)]
enum VerifyAttemptError {
    KidMissing,
    Decode(jsonwebtoken::errors::Error),
}

impl From<VerifyAttemptError> for AuthError {
    fn from(_err: VerifyAttemptError) -> Self {
        // Never leak decoder diagnostics to the wire; they're already
        // logged at the server boundary.
        AuthError::Rejected("token validation failed".into())
    }
}

#[async_trait]
impl Authenticator for OAuthAuthenticator {
    #[tracing::instrument(
        skip_all,
        fields(
            klieo.auth.principal_hash = tracing::field::Empty,
            klieo.auth.scopes_count = tracing::field::Empty,
            klieo.auth.mode = "oauth_jwt",
        ),
        err,
    )]
    async fn authenticate(
        &self,
        headers: &dyn Headers,
        _payload: &[u8],
    ) -> Result<Identity, AuthError> {
        let bearer = extract_bearer(headers)?;
        let (sub, scopes) = self.verify_token(&bearer).await?;
        let identity = Identity::with_scopes(sub, scopes);
        let span = tracing::Span::current();
        span.record(
            "klieo.auth.principal_hash",
            klieo_core::principal_hash(identity.as_str()).as_str(),
        );
        span.record("klieo.auth.scopes_count", identity.scopes().len());
        Ok(identity)
    }

    #[tracing::instrument(
        skip_all,
        fields(
            rpc.method = %method,
            klieo.auth.principal_hash = %klieo_core::principal_hash(identity.as_str()),
            klieo.auth.scopes_count = identity.scopes().len(),
        ),
        err,
    )]
    async fn authorize_method(&self, identity: &Identity, method: &str) -> Result<(), AuthError> {
        let Some(required) = self.scope_map.get(method) else {
            return Err(AuthError::Rejected("method not authorized".into()));
        };
        if required.is_empty() {
            return Ok(());
        }
        let granted = identity.scopes();
        if required.iter().any(|scope| granted.contains(scope)) {
            return Ok(());
        }
        Err(AuthError::Rejected("scope mismatch".into()))
    }

    fn allows_anonymous(&self) -> bool {
        false
    }
}

/// Extract a non-empty bearer token from the `Authorization` header.
/// Scheme matching is case-insensitive per RFC 7235 §2.1.
fn extract_bearer(headers: &dyn Headers) -> Result<String, AuthError> {
    let raw = headers.get("authorization").ok_or(AuthError::Missing)?;
    let token = strip_bearer_prefix(raw).ok_or(AuthError::Malformed)?;
    Ok(token.to_string())
}

/// Mirror of `klieo_auth_common::impls::strip_bearer_prefix` (private there).
/// W2.A14: case-insensitive scheme + reject empty tokens.
fn strip_bearer_prefix(header: &str) -> Option<&str> {
    let (scheme, rest) = header.split_once(|c: char| c.is_ascii_whitespace())?;
    if !scheme.eq_ignore_ascii_case("bearer") {
        return None;
    }
    let token = rest.trim_start();
    if token.is_empty() {
        None
    } else {
        Some(token)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::jwks::Cache;

    /// Build a minimal `OAuthAuthenticator` for unit tests that do not
    /// need a live JWKS endpoint. `Cache::empty()` is used; tests that
    /// exercise the verification path should use the integration suite.
    fn test_authenticator_with_algs(allowed_algs: Vec<Algorithm>) -> OAuthAuthenticator {
        OAuthAuthenticator {
            issuer: "https://issuer.example.test".into(),
            audience: "test-audience".into(),
            allowed_algs,
            cache: Cache::empty(),
            scope_map: Arc::new(std::collections::HashMap::new()),
            http: reqwest::Client::new(),
            introspection: None,
            token_cache: None,
        }
    }

    #[test]
    fn build_validation_rejects_algorithm_not_in_allowlist() {
        let authn = test_authenticator_with_algs(vec![Algorithm::RS256, Algorithm::ES256]);
        let err = authn
            .build_validation(Algorithm::HS256)
            .expect_err("HS256 must be rejected when not in allowlist");
        assert!(matches!(err, AuthError::Rejected(_)));
    }

    #[test]
    fn build_validation_accepts_algorithm_in_allowlist() {
        let authn = test_authenticator_with_algs(vec![Algorithm::RS256, Algorithm::ES256]);
        authn
            .build_validation(Algorithm::RS256)
            .expect("RS256 in allowlist must produce Validation");
        authn
            .build_validation(Algorithm::ES256)
            .expect("ES256 in allowlist must produce Validation");
    }

    #[test]
    fn validate_nbf_is_enabled() {
        let authn = test_authenticator_with_algs(vec![Algorithm::RS256]);
        let validation = authn
            .build_validation(Algorithm::RS256)
            .expect("RS256 accepted");
        assert!(
            validation.validate_nbf,
            "validate_nbf must be true after build_validation"
        );
    }

    #[test]
    fn scope_claim_space_delimited_parses() {
        let claim = ScopeClaim::Space("klieo:read klieo:write klieo:admin".into());
        let set = claim.into_set();
        assert_eq!(set.len(), 3);
        assert!(set.contains("klieo:read"));
        assert!(set.contains("klieo:write"));
        assert!(set.contains("klieo:admin"));
    }

    #[test]
    fn scope_claim_array_form_parses() {
        let claim = ScopeClaim::Array(vec!["a".into(), "b".into()]);
        let set = claim.into_set();
        assert_eq!(set.len(), 2);
        assert!(set.contains("a"));
        assert!(set.contains("b"));
    }

    #[test]
    fn scope_claim_empty_space_string_yields_empty_set() {
        let claim = ScopeClaim::Space("   ".into());
        let set = claim.into_set();
        assert!(set.is_empty());
    }

    #[test]
    fn strip_bearer_rejects_non_bearer_scheme() {
        assert!(strip_bearer_prefix("Basic abc").is_none());
    }

    #[test]
    fn strip_bearer_accepts_case_insensitive_scheme() {
        assert_eq!(strip_bearer_prefix("bearer abc"), Some("abc"));
        assert_eq!(strip_bearer_prefix("BEARER abc"), Some("abc"));
    }

    #[test]
    fn strip_bearer_rejects_empty_token() {
        assert!(strip_bearer_prefix("Bearer ").is_none());
    }
}