hive-router 0.2.7

use std::collections::HashMap;
use std::time::Duration;

use http::HeaderMap;
use jsonwebtoken::Algorithm;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

use crate::config::primitives::{file_path::FilePath, http_header::HttpHeaderName};

#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
#[serde(deny_unknown_fields)]
#[non_exhaustive]
pub struct JwtAuthConfig {
    #[serde(default = "default_enabled")]
    pub enabled: bool,
    /// A list of JWKS providers to use for verifying the JWT signature.
    /// Can be either a path to a local JSON of the file-system, or a URL to a remote JWKS provider.
    pub jwks_providers: Vec<JwksProviderSourceConfig>,
    /// Specify the [principal](https://tools.ietf.org/html/rfc7519#section-4.1.1) that issued the JWT, usually a URL or an email address.
    /// If specified, it has to match the `iss` field in JWT, otherwise the token's `iss` field is not checked.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub issuers: Option<Vec<String>>,
    /// The list of [JWT audiences](https://tools.ietf.org/html/rfc7519#section-4.1.3) are allowed to access.
    /// If this field is set, the token's `aud` field must be one of the values in this list, otherwise the token's `aud` field is not checked.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audiences: Option<Vec<String>>,
    /// A list of locations to look up for the JWT token in the incoming HTTP request.
    /// The first one that is found will be used.
    #[serde(
        default = "default_lookup_location",
        skip_serializing_if = "Vec::is_empty"
    )]
    pub lookup_locations: Vec<JwtAuthPluginLookupLocation>,
    /// If set to `true`, the entire request will be rejected if the JWT token is not present in the request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub require_authentication: Option<bool>,
    /// List of allowed algorithms for verifying the JWT signature.
    /// If not specified, the default list of all supported algorithms in [`jsonwebtoken` crate](https://crates.io/crates/jsonwebtoken) are used.
    #[serde(
        skip_serializing_if = "Option::is_none",
        default = "default_allowed_algorithms"
    )]
    #[schemars(with = "Option<Vec<String>>")]
    pub allowed_algorithms: Option<Vec<Algorithm>>,
    #[serde(default = "default_forward_claims_to_upstream_extensions")]
    /// Forward the JWT claims to the upstream service using GraphQL's `.extensions`.
    pub forward_claims_to_upstream_extensions: JwtClaimsForwardingConfig,
    /// The name of the JWT claim to read authorization scopes from, used to evaluate the `@requiresScopes` directive.
    /// If not specified, the `scope` claim is used (falling back to `scopes` if `scope` is not present), each holding
    /// either a space-delimited string or an array of strings.
    ///
    /// This is useful for identity providers that grant authorization data under a different claim name, such as
    /// Microsoft Entra ID, which issues app roles under a `roles` claim.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub scopes_claim: Option<String>,
}

impl JwtAuthConfig {
    pub fn is_jwt_extensions_forwarding_enabled(&self) -> bool {
        self.is_jwt_auth_enabled() && self.forward_claims_to_upstream_extensions.enabled
    }

    pub fn is_jwt_auth_enabled(&self) -> bool {
        self.enabled
    }

    pub fn is_jwt_auth_disabled(&self) -> bool {
        !self.is_jwt_auth_enabled()
    }
}

fn default_enabled() -> bool {
    false
}

impl Default for JwtAuthConfig {
    fn default() -> Self {
        JwtAuthConfig {
            enabled: default_enabled(),
            require_authentication: None,
            lookup_locations: vec![],
            jwks_providers: vec![],
            forward_claims_to_upstream_extensions: default_forward_claims_to_upstream_extensions(),
            audiences: None,
            issuers: None,
            allowed_algorithms: None,
            scopes_claim: None,
        }
    }
}

#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
#[non_exhaustive]
pub struct JwtClaimsForwardingConfig {
    pub enabled: bool,
    pub field_name: String,
    /// Only forward these root-level claim keys, instead of the entire JWT payload.
    /// If not specified, all claims are forwarded.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub include_claims: Option<Vec<String>>,
}

fn default_forward_claims_to_upstream_extensions() -> JwtClaimsForwardingConfig {
    JwtClaimsForwardingConfig {
        enabled: false,
        field_name: "jwt".to_string(),
        include_claims: None,
    }
}

#[derive(Deserialize, Serialize, Debug, Clone, JsonSchema)]
#[serde(tag = "source")]
#[non_exhaustive]
pub enum JwksProviderSourceConfig {
    /// A local file on the file-system. This file will be read once on startup and cached.
    #[serde(rename = "file")]
    #[schemars(title = "file")]
    File {
        #[serde(rename = "path")]
        /// A path to a local file on the file-system. Relative to the location of the root configuration file.
        file: FilePath,
    },
    /// A remote JWKS provider. The JWKS will be fetched via HTTP/HTTPS and cached.
    #[serde(rename = "remote")]
    #[schemars(title = "remote")]
    Remote {
        /// The URL to fetch the JWKS key set from, via HTTP/HTTPS.
        url: String,
        #[serde(
            deserialize_with = "humantime_serde::deserialize",
            serialize_with = "humantime_serde::serialize",
            default = "default_polling_interval"
        )]
        #[schemars(with = "String")]
        /// How often the JWKS should be polled for updates.
        polling_interval: Option<Duration>,
        /// If set to `true`, the JWKS will be fetched on startup and cached. In case of invalid JWKS, the error will be ignored and the plugin will try to fetch again when server receives the first request.
        /// If set to `false`, the JWKS will be fetched on-demand, when the first request comes in.
        prefetch: Option<bool>,
        /// Additional HTTP headers to send with the JWKS request.
        ///
        /// Useful when the JWKS endpoint is not reachable as a plain public URL:
        /// an identity provider that selects its tenant from the `Host` header
        /// while being addressed internally, or a JWKS endpoint behind a gateway
        /// that expects an API key.
        ///
        /// Keys must be valid HTTP header names (RFC 7230) and values must be
        /// valid HTTP header values.
        ///
        /// `Host` is honored: it is applied to the request and overrides the
        /// authority derived from `url`, without changing where the request is
        /// sent.
        ///
        /// Example:
        /// ```yaml
        /// headers:
        ///   Host: "auth.example.com"
        /// ```
        #[serde(
            default,
            with = "http_serde::header_map",
            skip_serializing_if = "HeaderMap::is_empty"
        )]
        #[schemars(with = "HashMap<String, String>")]
        headers: HeaderMap,
    },
}

fn default_polling_interval() -> Option<Duration> {
    // Some providers like MS Azure have rate limit configured. So let's use 10 minutes, like Envoy does.
    // and allow users to adjust it if needed.
    // See https://community.auth0.com/t/caching-jwks-signing-key/17654/2
    Some(Duration::from_secs(10 * 60))
}

pub fn default_lookup_location() -> Vec<JwtAuthPluginLookupLocation> {
    vec![JwtAuthPluginLookupLocation::Header {
        name: "Authorization".into(),
        prefix: Some("Bearer".to_string()),
    }]
}

pub fn default_allowed_algorithms() -> Option<Vec<Algorithm>> {
    Some(vec![
        Algorithm::HS256,
        Algorithm::HS384,
        Algorithm::HS512,
        Algorithm::RS256,
        Algorithm::RS384,
        Algorithm::RS512,
        Algorithm::ES256,
        Algorithm::ES384,
        Algorithm::PS256,
        Algorithm::PS384,
        Algorithm::PS512,
        Algorithm::EdDSA,
    ])
}

#[derive(Deserialize, Serialize, Debug, Clone, JsonSchema)]
#[serde(tag = "source")]
#[non_exhaustive]
pub enum JwtAuthPluginLookupLocation {
    #[serde(rename = "header")]
    #[schemars(title = "header")]
    Header {
        name: HttpHeaderName,
        prefix: Option<String>,
    },
    #[serde(rename = "cookies")]
    #[schemars(title = "cookies")]
    Cookie { name: String },
}

#[cfg(test)]
mod tests {
    use super::{JwksProviderSourceConfig, JwtAuthConfig};

    #[test]
    fn scopes_claim_defaults_to_none() {
        let config = JwtAuthConfig::default();
        assert_eq!(config.scopes_claim, None);
    }

    #[test]
    fn scopes_claim_is_none_when_not_specified_in_config() {
        let config = serde_json::from_str::<JwtAuthConfig>(r#"{"jwks_providers":[]}"#)
            .expect("config should parse");
        assert_eq!(config.scopes_claim, None);
    }

    #[test]
    fn scopes_claim_parses_configured_claim_name() {
        let config = serde_json::from_str::<JwtAuthConfig>(
            r#"{"jwks_providers":[],"scopes_claim":"roles"}"#,
        )
        .expect("config should parse");
        assert_eq!(config.scopes_claim, Some("roles".to_string()));
    }

    #[test]
    fn remote_jwks_headers_default_to_empty() {
        let config = serde_json::from_str::<JwtAuthConfig>(
            r#"{"jwks_providers":[{"source":"remote","url":"https://example.com/jwks.json"}]}"#,
        )
        .expect("config should parse");

        match &config.jwks_providers[0] {
            JwksProviderSourceConfig::Remote { headers, .. } => assert!(headers.is_empty()),
            other => panic!("expected a remote provider, got {other:?}"),
        }
    }

    #[test]
    fn remote_jwks_headers_are_parsed() {
        let config = serde_json::from_str::<JwtAuthConfig>(
            r#"{"jwks_providers":[{"source":"remote","url":"http://idp.internal/jwks.json","headers":{"Host":"auth.example.com","X-Api-Key":"secret"}}]}"#,
        )
        .expect("config should parse");

        match &config.jwks_providers[0] {
            JwksProviderSourceConfig::Remote { headers, .. } => {
                assert_eq!(headers.len(), 2);
                assert_eq!(headers.get("host").unwrap(), "auth.example.com");
                assert_eq!(headers.get("x-api-key").unwrap(), "secret");
            }
            other => panic!("expected a remote provider, got {other:?}"),
        }
    }

    #[test]
    fn remote_jwks_rejects_invalid_header_name() {
        let result = serde_json::from_str::<JwtAuthConfig>(
            r#"{"jwks_providers":[{"source":"remote","url":"https://example.com/jwks.json","headers":{"Invalid Header":"value"}}]}"#,
        );
        assert!(result.is_err(), "a header name with a space must not parse");
    }
}