core/server/src/http/jwks.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! JWKS key resolver for trusted-issuer (A2A) JWT verification.
//!
//! Ported from `server::http::jwt::jwks`. Keys are fetched lazily per
//! `{issuer, kid}` and cached in a `DashMap`; a refresh re-reads the issuer's
//! key set and evicts cached kids no longer present (rotation cleanup).

use std::collections::HashSet;
use std::hash::Hash;
use std::sync::Arc;
use std::time::{Duration, Instant};

use dashmap::DashMap;
use iggy_common::IggyError;
use jsonwebtoken::DecodingKey;
use serde::Deserialize;
use strum::EnumString;
use tokio::sync::Mutex;

/// Minimum wall-clock gap between outbound JWKS fetches for one trusted issuer.
/// Inside this window a cache miss is authoritative: the issuer's key set was
/// just read, so an absent `kid` is genuinely absent and no fetch is issued.
/// This bounds a pre-auth caller replaying unknown `kid`s to at most one
/// outbound request per issuer per window; the trade is that a freshly rotated
/// key is only discoverable once the window elapses.
const JWKS_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(10);

thread_local! {
    // cyper's `Client` is `!Send`/`!Sync` (`Rc`-backed) and `new()` returns a
    // `Result`, so it cannot be a global `OnceLock`. compio is thread-per-core;
    // keep one client per thread. The `Rc` inner makes cloning cheap, so callers
    // take an owned handle.
    static HTTP_CLIENT: cyper::Client =
        cyper::Client::new().expect("failed to build cyper HTTP client for JWKS");
}

fn get_http_client() -> cyper::Client {
    HTTP_CLIENT.with(cyper::Client::clone)
}

/// JWK key type enumeration
#[derive(Debug, Clone, Copy, Deserialize)]
#[serde(rename_all = "UPPERCASE")]
enum JwkKeyType {
    /// RSA key type
    Rsa,
    /// EC (Elliptic Curve) key type
    Ec,
}

/// EC curve type enumeration
#[derive(Debug, Clone, Copy, EnumString)]
#[strum(serialize_all = "UPPERCASE")]
enum EcCurve {
    /// P-256 curve
    #[strum(serialize = "P-256")]
    P256,
    /// P-384 curve
    #[strum(serialize = "P-384")]
    P384,
    /// P-521 curve
    #[strum(serialize = "P-521")]
    P521,
}

#[derive(Debug, Deserialize)]
struct Jwk {
    kty: JwkKeyType,
    kid: Option<String>,
    n: Option<String>,
    e: Option<String>,
    x: Option<String>,
    y: Option<String>,
    crv: Option<String>,
}

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

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
struct CacheKey {
    issuer: String,
    kid: String,
}

#[derive(Debug, Clone, Default)]
pub struct JwksClient {
    cache: DashMap<CacheKey, DecodingKey>,
    /// Per-issuer single-flight and rate-limit guard. The async mutex serialises
    /// refresh attempts for one issuer so a burst of concurrent misses collapses
    /// onto a single fetch; the inner instant is when that issuer was last
    /// fetched and gates [`JWKS_REFRESH_MIN_INTERVAL`]. Keyed only by issuer
    /// (operator-configured, bounded), never by the attacker-controlled `kid`.
    refresh_guards: DashMap<String, Arc<Mutex<Option<Instant>>>>,
}

impl JwksClient {
    /// Resolve the decoding key for `{issuer, kid}`, fetching and caching the
    /// issuer's JWKS on a cache miss. Returns the rich [`IggyError`] built during
    /// the fetch so the caller can log why a trusted-issuer token could not be
    /// verified instead of collapsing every failure into an opaque miss.
    ///
    /// A cache miss is served under a per-issuer guard: concurrent misses fetch
    /// once, and within [`JWKS_REFRESH_MIN_INTERVAL`] of the last fetch a miss is
    /// treated as a known-absent `kid` and rejected without an outbound request.
    /// This keeps an unauthenticated caller replaying unknown `kid`s from
    /// amplifying into unbounded fetches against the issuer's JWKS endpoint.
    // The per-issuer guard is deliberately held across the JWKS fetch: that hold
    // is what serialises concurrent misses onto a single outbound request. Drop-
    // tightening would release it before the await and defeat the single-flight.
    #[allow(clippy::significant_drop_tightening)]
    pub async fn get_key(
        &self,
        issuer: &str,
        jwks_url: &str,
        kid: &str,
    ) -> Result<DecodingKey, IggyError> {
        let cache_key = CacheKey {
            issuer: issuer.to_string(),
            kid: kid.to_string(),
        };

        // Positive-cache fast path: no lock, no fetch.
        if let Some(key) = self.cache.get(&cache_key) {
            return Ok(key.clone());
        }

        // Take the per-issuer guard so concurrent misses serialise onto one
        // fetch. Clone the Arc out and drop the DashMap entry lock before the
        // await, so no shard lock is held across the network I/O.
        let entry = self.refresh_guards.entry(issuer.to_string()).or_default();
        let guard = entry.value().clone();
        drop(entry);
        let mut last_fetch = guard.lock().await;

        // A prior holder of the guard may have populated our kid while we waited.
        if let Some(key) = self.cache.get(&cache_key) {
            return Ok(key.clone());
        }

        // Inside the refresh window the last fetch's key set still stands, so a
        // miss here means the kid is genuinely absent: reject without touching
        // the network. One per-issuer timestamp both negative-caches unknown kids
        // and rate-limits outbound fetches, with no attacker-keyed state.
        if let Some(fetched_at) = *last_fetch
            && fetched_at.elapsed() < JWKS_REFRESH_MIN_INTERVAL
        {
            return Err(IggyError::InvalidAccessToken);
        }

        // Stale or first contact: fetch. Record the attempt up front so a failing
        // issuer is rate-limited too, not re-hit on every miss.
        *last_fetch = Some(Instant::now());
        self.refresh_keys(issuer, jwks_url).await?;

        self.cache
            .get(&cache_key)
            .map(|entry| entry.clone())
            .ok_or(IggyError::InvalidAccessToken)
    }

    async fn refresh_keys(&self, issuer: &str, jwks_url: &str) -> Result<(), IggyError> {
        // The cyper client is `!Send`; callers reached from the axum extractor
        // wrap this future in `SendWrapper` (see `http::extractor`), so it is
        // free to await cyper directly here.
        let client = get_http_client();
        let request = client
            .get(jwks_url)
            .map_err(|e| IggyError::CannotFetchJwks(format!("Failed to build request: {e}")))?
            .build();
        let response = client
            .execute(request)
            .await
            .map_err(|e| IggyError::CannotFetchJwks(format!("HTTP request failed: {e}")))?;

        let body = response.text().await.map_err(|e| {
            IggyError::CannotFetchJwks(format!("Failed to read response body: {e}"))
        })?;

        let jwks: JwkSet = serde_json::from_str(&body)
            .map_err(|e| IggyError::CannotFetchJwks(format!("Failed to parse JWKS: {e}")))?;

        // Collect all current kids from the JWKS response
        let current_kids: HashSet<String> =
            jwks.keys.iter().filter_map(|key| key.kid.clone()).collect();

        // Remove cached keys for this issuer that are no longer in the JWKS response
        // Security fix: Clean up revoked/rotated keys to prevent accepting tokens signed with old keys
        let keys_to_remove: Vec<CacheKey> = self
            .cache
            .iter()
            .filter(|entry| {
                entry.key().issuer == issuer && !current_kids.contains(&entry.key().kid)
            })
            .map(|entry| entry.key().clone())
            .collect();

        for key in keys_to_remove {
            self.cache.remove(&key);
        }

        for key in jwks.keys {
            if let Some(kid) = key.kid {
                let decoding_key: DecodingKey = match key.kty {
                    JwkKeyType::Rsa => {
                        if let (Some(n), Some(e)) = (key.n.as_deref(), key.e.as_deref()) {
                            DecodingKey::from_rsa_components(n, e).map_err(|e| {
                                IggyError::CannotFetchJwks(format!("Invalid RSA key: {e}"))
                            })?
                        } else {
                            continue;
                        }
                    }
                    JwkKeyType::Ec => {
                        if let (Some(x), Some(y), Some(crv_str)) =
                            (key.x.as_deref(), key.y.as_deref(), key.crv.as_deref())
                        {
                            if let Ok(_curve) = crv_str.parse::<EcCurve>() {
                                DecodingKey::from_ec_components(x, y).map_err(|e| {
                                    IggyError::CannotFetchJwks(format!("Invalid EC key: {e}"))
                                })?
                            } else {
                                continue;
                            }
                        } else {
                            continue;
                        }
                    }
                };

                let cache_key = CacheKey {
                    issuer: issuer.to_string(),
                    kid,
                };
                self.cache.insert(cache_key, decoding_key);
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use jsonwebtoken::DecodingKey;

    const TEST_ISSUER: &str = "https://test-issuer.com";
    const TEST_KID: &str = "test-key";

    fn create_test_decoding_key() -> DecodingKey {
        // Use HMAC secret to create a simple test DecodingKey
        // Note: This is only for testing cache logic, not a real RSA/EC key
        DecodingKey::from_secret(b"test-secret-key-for-cache-testing-only")
    }

    #[test]
    fn test_cache_key_equality() {
        let key1 = CacheKey {
            issuer: TEST_ISSUER.to_string(),
            kid: TEST_KID.to_string(),
        };
        let key2 = CacheKey {
            issuer: TEST_ISSUER.to_string(),
            kid: TEST_KID.to_string(),
        };
        assert_eq!(key1, key2);
    }

    #[test]
    fn test_cache_key_different_issuer() {
        let key1 = CacheKey {
            issuer: "issuer1".to_string(),
            kid: TEST_KID.to_string(),
        };
        let key2 = CacheKey {
            issuer: "issuer2".to_string(),
            kid: TEST_KID.to_string(),
        };
        assert_ne!(key1, key2);
    }

    #[test]
    fn test_cache_key_different_kid() {
        let key1 = CacheKey {
            issuer: TEST_ISSUER.to_string(),
            kid: "kid1".to_string(),
        };
        let key2 = CacheKey {
            issuer: TEST_ISSUER.to_string(),
            kid: "kid2".to_string(),
        };
        assert_ne!(key1, key2);
    }

    #[test]
    fn test_jwks_client_default() {
        let client = JwksClient::default();
        assert!(client.cache.is_empty());
    }

    #[test]
    fn test_cache_insert_and_get() {
        let client = JwksClient::default();
        let cache_key = CacheKey {
            issuer: TEST_ISSUER.to_string(),
            kid: TEST_KID.to_string(),
        };
        let decoding_key = create_test_decoding_key();

        client.cache.insert(cache_key.clone(), decoding_key);

        assert!(client.cache.get(&cache_key).is_some());
    }

    #[test]
    fn test_cache_multiple_keys() {
        let client = JwksClient::default();

        let key1 = CacheKey {
            issuer: "issuer1".to_string(),
            kid: "kid1".to_string(),
        };
        let key2 = CacheKey {
            issuer: "issuer2".to_string(),
            kid: "kid2".to_string(),
        };

        let decoding_key1 = create_test_decoding_key();
        let decoding_key2 = create_test_decoding_key();

        client.cache.insert(key1.clone(), decoding_key1);
        client.cache.insert(key2.clone(), decoding_key2);

        assert_eq!(client.cache.len(), 2);
        assert!(client.cache.get(&key1).is_some());
        assert!(client.cache.get(&key2).is_some());
    }
}