remora/crates/remora-identity/src/oidc.rs

//! OIDC bearer-token validation against discovery + JWKS.
//!
//! Validates signature, `iss`, `exp`, and a mandatory expected `aud`. Role
//! claims are taken only from explicitly configured sources: Keycloak
//! `resource_access[<client>]` and exact ZITADEL project-role claim names.
//! Remora authority is encoded only in namespaced `remora:v1:…` roles.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
use remora_domain::ports::IdentityProvider;
use serde::Deserialize;
use serde_json::Value;
use tokio::sync::RwLock;

use crate::error::IdentityError;
use crate::principal::ExternalIdentity;
use crate::roles::{DecodedRoles, decode_roles};

/// Default JWKS cache lifetime before re-fetch (key rotation / removal).
const DEFAULT_JWKS_CACHE_TTL: Duration = Duration::from_secs(300);

/// Well-known ZITADEL claim that carries project role keys when role assertion
/// is enabled for the application.
pub const ZITADEL_PROJECT_ROLES_CLAIM: &str = "urn:zitadel:iam:org:project:roles";

#[derive(Clone, Debug)]
pub struct OidcConfig {
    pub issuer: String,
    /// Required expected `aud` claim (typically the OIDC client id).
    pub audience: String,
    /// Keycloak client id whose `resource_access` roles are trusted.
    /// Defaults to `audience` when not overridden.
    pub resource_access_client: String,
    /// Exact JWT claim names that carry ZITADEL (or similar) project role maps.
    pub project_role_claims: Vec<String>,
    pub reserved_segments: Vec<String>,
    pub http_timeout: Duration,
    /// How long a successful JWKS fetch remains valid before revalidation.
    pub jwks_cache_ttl: Duration,
}

impl OidcConfig {
    /// Production constructor: non-empty issuer and audience are mandatory.
    pub fn new(
        issuer: impl Into<String>,
        audience: impl Into<String>,
    ) -> Result<Self, IdentityError> {
        let issuer = issuer.into();
        let audience = audience.into();
        if issuer.is_empty() {
            return Err(IdentityError::Config("issuer must be non-empty".into()));
        }
        if audience.is_empty() {
            return Err(IdentityError::Config("audience must be non-empty".into()));
        }
        Ok(Self {
            resource_access_client: audience.clone(),
            issuer,
            audience,
            project_role_claims: Vec::new(),
            reserved_segments: Vec::new(),
            http_timeout: Duration::from_secs(10),
            jwks_cache_ttl: DEFAULT_JWKS_CACHE_TTL,
        })
    }

    /// Trust Keycloak client roles under a specific client id (defaults to audience).
    #[must_use]
    pub fn with_resource_access_client(mut self, client_id: impl Into<String>) -> Self {
        self.resource_access_client = client_id.into();
        self
    }

    /// Trust an exact ZITADEL project-role claim name (repeatable).
    #[must_use]
    pub fn with_project_role_claim(mut self, claim: impl Into<String>) -> Self {
        let claim = claim.into();
        if !self
            .project_role_claims
            .iter()
            .any(|existing| existing == &claim)
        {
            self.project_role_claims.push(claim);
        }
        self
    }

    /// Convenience: trust the standard ZITADEL project roles claim.
    #[must_use]
    pub fn with_zitadel_project_roles(self) -> Self {
        self.with_project_role_claim(ZITADEL_PROJECT_ROLES_CLAIM)
    }

    #[must_use]
    pub fn with_reserved_segments(mut self, segments: impl IntoIterator<Item = String>) -> Self {
        self.reserved_segments = segments.into_iter().collect();
        self
    }

    #[must_use]
    pub fn with_jwks_cache_ttl(mut self, ttl: Duration) -> Self {
        self.jwks_cache_ttl = ttl;
        self
    }
}

#[derive(Clone)]
pub struct OidcIdentityProvider {
    config: OidcConfig,
    client: reqwest::Client,
    state: Arc<RwLock<ProviderState>>,
}

#[derive(Default)]
struct ProviderState {
    jwks_uri: Option<String>,
    keys: HashMap<String, DecodingKey>,
    keys_fetched_at: Option<Instant>,
    static_keys: bool,
}

impl ProviderState {
    fn is_stale(&self, ttl: Duration) -> bool {
        if self.static_keys {
            return false;
        }
        match self.keys_fetched_at {
            Some(at) => at.elapsed() >= ttl,
            None => true,
        }
    }
}

#[derive(Debug, Deserialize)]
struct DiscoveryDocument {
    issuer: String,
    jwks_uri: String,
}

#[derive(Debug, Deserialize)]
struct JwksDocument {
    keys: Vec<Jwk>,
}

#[derive(Debug, Deserialize)]
struct Jwk {
    kid: Option<String>,
    kty: String,
    #[serde(rename = "use")]
    use_: Option<String>,
    alg: Option<String>,
    n: Option<String>,
    e: Option<String>,
}

#[derive(Debug, Deserialize)]
struct TokenClaims {
    iss: String,
    sub: String,
    #[serde(default)]
    aud: Audience,
    #[allow(dead_code)]
    exp: i64,
    #[serde(flatten)]
    extra: HashMap<String, Value>,
}

#[derive(Debug, Default, Deserialize)]
#[serde(untagged)]
enum Audience {
    String(String),
    List(Vec<String>),
    #[default]
    Missing,
}

impl OidcIdentityProvider {
    pub fn new(config: OidcConfig) -> Result<Self, IdentityError> {
        if config.audience.is_empty() {
            return Err(IdentityError::Config("audience must be non-empty".into()));
        }
        let client = reqwest::Client::builder()
            .timeout(config.http_timeout)
            .build()
            .map_err(|error| IdentityError::Config(error.to_string()))?;
        Ok(Self {
            config,
            client,
            state: Arc::new(RwLock::new(ProviderState::default())),
        })
    }

    /// Test and offline harness helper: install fixed RS256 keys without discovery.
    pub fn with_static_keys(
        config: OidcConfig,
        keys: HashMap<String, DecodingKey>,
    ) -> Result<Self, IdentityError> {
        let provider = Self::new(config)?;
        let mut state = provider
            .state
            .try_write()
            .map_err(|_| IdentityError::Config("provider state lock busy".into()))?;
        state.keys = keys;
        state.jwks_uri = Some("static://test-jwks".into());
        state.keys_fetched_at = Some(Instant::now());
        state.static_keys = true;
        drop(state);
        Ok(provider)
    }

    pub async fn authenticate_token(
        &self,
        bearer_token: &str,
    ) -> Result<ExternalIdentity, IdentityError> {
        let token = bearer_token
            .strip_prefix("Bearer ")
            .or_else(|| bearer_token.strip_prefix("bearer "))
            .unwrap_or(bearer_token)
            .trim();
        if token.is_empty() {
            return Err(IdentityError::InvalidToken);
        }

        let header = decode_header(token).map_err(|_| IdentityError::InvalidToken)?;
        let kid = header.kid.ok_or(IdentityError::InvalidToken)?;
        let key = self.decoding_key(&kid).await?;

        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_issuer(std::slice::from_ref(&self.config.issuer));
        validation.set_audience(std::slice::from_ref(&self.config.audience));
        validation.leeway = 30;

        let data = decode::<TokenClaims>(token, &key, &validation).map_err(|error| {
            if is_provider_side(error.kind()) {
                IdentityError::ProviderUnavailable(error.to_string())
            } else {
                IdentityError::InvalidToken
            }
        })?;
        let claims = data.claims;
        if claims.iss != self.config.issuer {
            return Err(IdentityError::InvalidToken);
        }
        if !audience_matches(&claims.aud, &self.config.audience) {
            return Err(IdentityError::InvalidToken);
        }
        if claims.sub.is_empty() {
            return Err(IdentityError::InvalidToken);
        }

        let role_names = extract_role_names(
            &claims.extra,
            &self.config.resource_access_client,
            &self.config.project_role_claims,
        );
        let reserved: Vec<&str> = self
            .config
            .reserved_segments
            .iter()
            .map(String::as_str)
            .collect();
        let roles = decode_roles(role_names.iter().map(String::as_str), &reserved)?;
        Ok(ExternalIdentity {
            issuer: claims.iss,
            subject: claims.sub,
            roles,
        })
    }

    async fn decoding_key(&self, kid: &str) -> Result<DecodingKey, IdentityError> {
        {
            let state = self.state.read().await;
            if state.static_keys {
                return state
                    .keys
                    .get(kid)
                    .cloned()
                    .ok_or(IdentityError::InvalidToken);
            }
            if !state.is_stale(self.config.jwks_cache_ttl)
                && let Some(key) = state.keys.get(kid)
            {
                return Ok(key.clone());
            }
        }
        self.refresh_jwks().await?;
        let state = self.state.read().await;
        state
            .keys
            .get(kid)
            .cloned()
            .ok_or(IdentityError::InvalidToken)
    }

    async fn refresh_jwks(&self) -> Result<(), IdentityError> {
        let jwks_uri = {
            let state = self.state.read().await;
            if state.static_keys {
                return Ok(());
            }
            state.jwks_uri.clone()
        };
        let jwks_uri = match jwks_uri {
            Some(uri) => uri,
            None => self.discover().await?,
        };
        let document = self
            .client
            .get(&jwks_uri)
            .send()
            .await
            .map_err(|error| IdentityError::ProviderUnavailable(error.to_string()))?
            .error_for_status()
            .map_err(|error| IdentityError::ProviderUnavailable(error.to_string()))?
            .json::<JwksDocument>()
            .await
            .map_err(|error| IdentityError::ProviderUnavailable(error.to_string()))?;

        let mut keys = HashMap::new();
        for jwk in document.keys {
            if jwk.kty != "RSA" {
                continue;
            }
            if let Some(use_) = &jwk.use_
                && use_ != "sig"
            {
                continue;
            }
            if let Some(alg) = &jwk.alg
                && alg != "RS256"
            {
                continue;
            }
            let (Some(n), Some(e), Some(kid)) = (jwk.n, jwk.e, jwk.kid) else {
                continue;
            };
            let key = DecodingKey::from_rsa_components(&n, &e)
                .map_err(|error| IdentityError::ProviderUnavailable(error.to_string()))?;
            keys.insert(kid, key);
        }
        if keys.is_empty() {
            return Err(IdentityError::ProviderUnavailable(
                "JWKS contained no usable RS256 signing keys".into(),
            ));
        }
        let mut state = self.state.write().await;
        state.jwks_uri = Some(jwks_uri);
        state.keys = keys;
        state.keys_fetched_at = Some(Instant::now());
        state.static_keys = false;
        Ok(())
    }

    async fn discover(&self) -> Result<String, IdentityError> {
        let url = format!(
            "{}/.well-known/openid-configuration",
            self.config.issuer.trim_end_matches('/')
        );
        let document = self
            .client
            .get(url)
            .send()
            .await
            .map_err(|error| IdentityError::ProviderUnavailable(error.to_string()))?
            .error_for_status()
            .map_err(|error| IdentityError::ProviderUnavailable(error.to_string()))?
            .json::<DiscoveryDocument>()
            .await
            .map_err(|error| IdentityError::ProviderUnavailable(error.to_string()))?;
        if document.issuer != self.config.issuer {
            return Err(IdentityError::ProviderUnavailable(format!(
                "discovery issuer {} does not match configured {}",
                document.issuer, self.config.issuer
            )));
        }
        let mut state = self.state.write().await;
        state.jwks_uri = Some(document.jwks_uri.clone());
        Ok(document.jwks_uri)
    }
}

impl IdentityProvider for OidcIdentityProvider {
    type Error = IdentityError;
    type Principal = ExternalIdentity;

    async fn authenticate(&self, bearer_token: &str) -> Result<Self::Principal, Self::Error> {
        self.authenticate_token(bearer_token).await
    }
}

fn audience_matches(aud: &Audience, expected: &str) -> bool {
    match aud {
        Audience::String(value) => value == expected,
        Audience::List(values) => values.iter().any(|value| value == expected),
        Audience::Missing => false,
    }
}

/// Extract Remora role *names* only from configured Keycloak client roles and
/// exact project-role claim maps. Plain `roles`, realm roles, and unlisted
/// project claims are ignored.
fn extract_role_names(
    extra: &HashMap<String, Value>,
    resource_access_client: &str,
    project_role_claims: &[String],
) -> Vec<String> {
    let mut roles = Vec::new();
    if let Some(Value::Object(resource_access)) = extra.get("resource_access")
        && let Some(Value::Object(client_roles)) = resource_access.get(resource_access_client)
        && let Some(Value::Array(values)) = client_roles.get("roles")
    {
        push_string_array(&mut roles, values);
    }
    for claim in project_role_claims {
        if let Some(Value::Object(map)) = extra.get(claim) {
            for role_key in map.keys() {
                roles.push(role_key.clone());
            }
        }
    }
    roles.sort();
    roles.dedup();
    roles
}

fn push_string_array(roles: &mut Vec<String>, values: &[Value]) {
    for value in values {
        if let Value::String(role) = value {
            roles.push(role.clone());
        }
    }
}

fn is_provider_side(kind: &jsonwebtoken::errors::ErrorKind) -> bool {
    matches!(
        kind,
        jsonwebtoken::errors::ErrorKind::InvalidKeyFormat
            | jsonwebtoken::errors::ErrorKind::InvalidRsaKey(_)
            | jsonwebtoken::errors::ErrorKind::InvalidEcdsaKey
    )
}

pub fn roles_from_names(
    roles: &[&str],
    reserved_segments: &[&str],
) -> Result<DecodedRoles, IdentityError> {
    decode_roles(roles.iter().copied(), reserved_segments)
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn claim_map(value: Value) -> HashMap<String, Value> {
        match value {
            Value::Object(map) => map.into_iter().collect(),
            _ => panic!("expected object"),
        }
    }

    #[test]
    fn only_configured_client_resource_access_is_trusted() {
        let map = claim_map(json!({
            "resource_access": {
                "remora-client": { "roles": ["remora:v1:admin"] },
                "evil-client": { "roles": ["remora:v1:admin"] }
            },
            "realm_access": { "roles": ["admin", "remora:v1:admin"] },
            "roles": ["remora:v1:admin"]
        }));
        let roles = extract_role_names(&map, "remora-client", &[]);
        assert_eq!(roles, vec!["remora:v1:admin".to_owned()]);
    }

    #[test]
    fn only_configured_project_role_claims_are_trusted() {
        let map = claim_map(json!({
            "urn:zitadel:iam:org:project:roles": {
                "remora:v1:admin": { "org": "x" }
            },
            "urn:zitadel:iam:org:project:999:roles": {
                "remora:v1:admin": { "org": "x" }
            }
        }));
        let allowed = vec![ZITADEL_PROJECT_ROLES_CLAIM.to_owned()];
        let roles = extract_role_names(&map, "unused", &allowed);
        assert_eq!(roles, vec!["remora:v1:admin".to_owned()]);
    }

    #[test]
    fn config_rejects_empty_audience() {
        assert!(OidcConfig::new("https://issuer.example", "").is_err());
    }
}