use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant};
#[derive(Clone)]
struct JwksEntry {
jwks_json: String,
kids: Vec<String>,
fetched_at: Instant,
}
#[derive(Clone, Default)]
pub struct OidcJwksCache {
inner: Arc<RwLock<HashMap<String, JwksEntry>>>,
}
#[derive(Debug, Clone)]
pub struct JwksResolution {
pub jwks_json: String,
pub kids: Vec<String>,
pub from_cache: bool,
}
fn cache_key(tenant_id: &str, provider_id: &str) -> String {
format!("{tenant_id}\u{1f}{provider_id}")
}
fn jwks_ttl_secs() -> u64 {
std::env::var("UDB_OIDC_JWKS_TTL_SECS")
.ok()
.and_then(|v| v.parse::<u64>().ok())
.filter(|v| *v > 0)
.unwrap_or(3600)
}
impl OidcJwksCache {
pub fn new() -> Self {
Self::default()
}
fn cached(&self, tenant_id: &str, provider_id: &str) -> Option<JwksEntry> {
let ttl = Duration::from_secs(jwks_ttl_secs());
let guard = self.inner.read().ok()?;
let entry = guard.get(&cache_key(tenant_id, provider_id))?;
if entry.fetched_at.elapsed() <= ttl {
Some(entry.clone())
} else {
None
}
}
fn store(&self, tenant_id: &str, provider_id: &str, jwks_json: String, kids: Vec<String>) {
if let Ok(mut guard) = self.inner.write() {
guard.insert(
cache_key(tenant_id, provider_id),
JwksEntry {
jwks_json,
kids,
fetched_at: Instant::now(),
},
);
}
}
pub fn invalidate(&self, tenant_id: &str, provider_id: &str) {
if let Ok(mut guard) = self.inner.write() {
guard.remove(&cache_key(tenant_id, provider_id));
}
}
pub async fn resolve(
&self,
tenant_id: &str,
provider_id: &str,
issuer: &str,
jwks_url: &str,
force_refresh: bool,
) -> Result<JwksResolution, String> {
if !force_refresh {
if let Some(entry) = self.cached(tenant_id, provider_id) {
return Ok(JwksResolution {
jwks_json: entry.jwks_json,
kids: entry.kids,
from_cache: true,
});
}
}
let resolved_url = if !jwks_url.trim().is_empty() {
jwks_url.trim().to_string()
} else {
discover_jwks_url(issuer).await?
};
let jwks_json = fetch_url(&resolved_url).await?;
let kids = parse_jwks_kids(&jwks_json);
self.store(tenant_id, provider_id, jwks_json.clone(), kids.clone());
Ok(JwksResolution {
jwks_json,
kids,
from_cache: false,
})
}
}
async fn discover_jwks_url(issuer: &str) -> Result<String, String> {
if issuer.trim().is_empty() {
return Err("OIDC issuer is required to discover the JWKS endpoint".to_string());
}
let well_known = format!(
"{}/.well-known/openid-configuration",
issuer.trim_end_matches('/')
);
let body = fetch_url(&well_known).await?;
let doc: serde_json::Value =
serde_json::from_str(&body).map_err(|e| format!("invalid discovery document: {e}"))?;
doc.get("jwks_uri")
.and_then(|v| v.as_str())
.filter(|s| !s.is_empty())
.map(ToString::to_string)
.ok_or_else(|| "discovery document has no jwks_uri".to_string())
}
pub fn parse_jwks_kids(jwks_json: &str) -> Vec<String> {
let doc: serde_json::Value = match serde_json::from_str(jwks_json) {
Ok(v) => v,
Err(_) => return Vec::new(),
};
doc.get("keys")
.and_then(|k| k.as_array())
.map(|keys| {
keys.iter()
.filter_map(|k| {
k.get("kid")
.and_then(|v| v.as_str())
.map(ToString::to_string)
})
.collect()
})
.unwrap_or_default()
}
#[cfg(feature = "http-client")]
async fn fetch_url(url: &str) -> Result<String, String> {
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(10))
.build()
.map_err(|e| format!("http client build failed: {e}"))?;
let resp = client
.get(url)
.send()
.await
.map_err(|e| format!("GET {url} failed: {e}"))?;
if !resp.status().is_success() {
return Err(format!("GET {url} returned HTTP {}", resp.status()));
}
resp.text()
.await
.map_err(|e| format!("reading {url} body failed: {e}"))
}
#[cfg(not(feature = "http-client"))]
async fn fetch_url(_url: &str) -> Result<String, String> {
Err("OIDC discovery/JWKS requires the `http-client` feature".to_string())
}
pub async fn fetch_text(url: &str) -> Result<String, String> {
fetch_url(url).await
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_jwks_kids_and_handles_garbage() {
let jwks = r#"{"keys":[{"kid":"k1","kty":"RSA"},{"kid":"k2"},{"kty":"oct"}]}"#;
assert_eq!(parse_jwks_kids(jwks), vec!["k1", "k2"]);
assert!(parse_jwks_kids("not json").is_empty());
assert!(parse_jwks_kids("{}").is_empty());
}
#[test]
fn cache_round_trips_within_ttl() {
let cache = OidcJwksCache::new();
cache.store("t1", "p1", "{\"keys\":[]}".into(), vec!["a".into()]);
let hit = cache.cached("t1", "p1").expect("cache hit");
assert_eq!(hit.kids, vec!["a".to_string()]);
cache.invalidate("t1", "p1");
assert!(cache.cached("t1", "p1").is_none());
}
}