Skip to content
Clone
dd4650a
252 lines 8.4 KB Rust
use std::sync::Arc;
use std::time::{Duration, Instant};
use config::oauth::OAuthConfig;
use jsonwebtoken::jwk::JwkSet;
use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};
use tokio::sync::{Mutex, RwLock};
use crate::claims::{AccessClaims, IdClaims, ServiceClaims};
use crate::error::SessionError;
const ACCESS_TYP: &str = "at+jwt";
const SERVICE_TYP: &str = "at+jwt";
struct CachedJwks {
set: JwkSet,
fetched_at: Instant,
}
pub struct JwtVerifier {
identity: Option<kenoauth::auth::AuthSvc>,
issuer: String,
audience: String,
chat_audience: String,
client_id: String,
session_cookie: String,
leeway: u64,
jwks_url: String,
cache_ttl: Duration,
min_refresh: Duration,
http: reqwest::Client,
cache: RwLock<Option<CachedJwks>>,
refresh: Mutex<()>,
}
impl JwtVerifier {
pub fn new(config: &OAuthConfig) -> Result<Arc<Self>, SessionError> {
Self::with_identity(config, None)
}
pub fn with_identity(
config: &OAuthConfig,
identity: Option<kenoauth::auth::AuthSvc>,
) -> Result<Arc<Self>, SessionError> {
if config.issuer.is_empty() {
return Err(SessionError::Verifier("oauth.issuer is empty".into()));
}
let http = reqwest::Client::builder()
.timeout(Duration::from_millis(config.http_timeout_ms.max(200)))
.build()
.map_err(|err| SessionError::Verifier(err.to_string()))?;
Ok(Arc::new(Self {
identity,
issuer: config.issuer.trim_end_matches('/').to_string(),
audience: config.audience.clone(),
chat_audience: config.chat_audience.clone(),
client_id: config.client_id.clone(),
session_cookie: config.session_cookie.clone(),
leeway: config.leeway_secs,
jwks_url: config.jwks_url(),
cache_ttl: Duration::from_secs(config.jwks_cache_secs.max(1)),
min_refresh: Duration::from_secs(config.jwks_min_refresh_secs.max(1)),
http,
cache: RwLock::new(None),
refresh: Mutex::new(()),
}))
}
pub fn jwks_url(&self) -> &str {
&self.jwks_url
}
pub fn session_cookie(&self) -> &str {
&self.session_cookie
}
pub async fn warm_up(&self) {
if let Err(error) = self.fetch_jwks().await {
tracing::error!(target: "gitdata::jwks", error = %error, url = %self.jwks_url, "failed to load JWKS");
}
}
pub async fn verify_access(&self, token: &str) -> Result<AccessClaims, SessionError> {
let (claims, typ) = self.verify::<AccessClaims>(token, &self.audience).await?;
if typ != ACCESS_TYP {
return Err(SessionError::Verifier(format!("unexpected typ {typ:?}")));
}
if claims.token_type != "access" {
return Err(SessionError::Verifier(format!(
"unexpected token_type {:?}",
claims.token_type
)));
}
if claims.client_id != self.client_id {
return Err(SessionError::Verifier(format!(
"unexpected client_id {:?}",
claims.client_id
)));
}
Ok(claims)
}
pub async fn verify_id(&self, token: &str) -> Result<IdClaims, SessionError> {
let (claims, _) = self.verify::<IdClaims>(token, &self.client_id).await?;
Ok(claims)
}
pub async fn verify_service(
&self,
token: &str,
audience: &str,
) -> Result<ServiceClaims, SessionError> {
let (claims, typ) = self.verify::<ServiceClaims>(token, audience).await?;
if typ != SERVICE_TYP && typ != "JWT" {
return Err(SessionError::Verifier(format!("unexpected typ {typ:?}")));
}
if claims.token_type != "service" {
return Err(SessionError::Verifier(format!(
"unexpected token_type {:?}",
claims.token_type
)));
}
Ok(claims)
}
pub fn chat_audience(&self) -> &str {
&self.chat_audience
}
async fn verify<T: serde::de::DeserializeOwned>(
&self,
token: &str,
audience: &str,
) -> Result<(T, String), SessionError> {
let header = decode_header(token).map_err(|err| SessionError::Verifier(err.to_string()))?;
if header.alg != Algorithm::EdDSA {
return Err(SessionError::Verifier(format!(
"unexpected algorithm {:?}",
header.alg
)));
}
let kid = header
.kid
.clone()
.filter(|value| !value.is_empty())
.ok_or_else(|| SessionError::Verifier("token has no kid".into()))?;
let key = self.decoding_key(&kid).await?;
let mut validation = Validation::new(Algorithm::EdDSA);
validation.leeway = self.leeway;
validation.set_issuer(&[self.issuer.as_str()]);
validation.set_audience(&[audience]);
validation.required_spec_claims = ["exp", "iss", "aud", "sub"]
.into_iter()
.map(str::to_string)
.collect();
let data = decode::<T>(token, &key, &validation).map_err(|err| {
if matches!(
err.kind(),
jsonwebtoken::errors::ErrorKind::ExpiredSignature
) {
SessionError::Expired
} else {
SessionError::Verifier(err.to_string())
}
})?;
let typ = header.typ.unwrap_or_default();
Ok((data.claims, typ))
}
async fn decoding_key(&self, kid: &str) -> Result<DecodingKey, SessionError> {
{
let cache = self.cache.read().await;
if let Some(cached) = cache.as_ref()
&& cached.fetched_at.elapsed() < self.cache_ttl
&& let Some(key) = decoding_key_for(&cached.set, kid)?
{
return Ok(key);
}
}
let _guard = self.refresh.lock().await;
{
let cache = self.cache.read().await;
if let Some(cached) = cache.as_ref()
&& cached.fetched_at.elapsed() < self.min_refresh
{
return decoding_key_for(&cached.set, kid)?.ok_or_else(|| {
SessionError::Verifier(format!("unknown kid {kid:?} (refresh rate limited)"))
});
}
}
match self.fetch_jwks().await {
Ok(set) => {
let key = decoding_key_for(&set, kid)?;
*self.cache.write().await = Some(CachedJwks {
set,
fetched_at: Instant::now(),
});
key.ok_or_else(|| SessionError::Verifier(format!("unknown kid {kid:?}")))
}
Err(error) => {
let cache = self.cache.read().await;
if let Some(cached) = cache.as_ref()
&& let Some(key) = decoding_key_for(&cached.set, kid)?
{
tracing::error!(target: "gitdata::jwks", error = %error, "JWKS refresh failed; using cached keys");
return Ok(key);
}
Err(error)
}
}
}
async fn fetch_jwks(&self) -> Result<JwkSet, SessionError> {
if let Some(service) = &self.identity {
let value = serde_json::to_value(service.signing().jwks())
.map_err(|error| SessionError::Verifier(error.to_string()))?;
return serde_json::from_value(value)
.map_err(|error| SessionError::Verifier(error.to_string()));
}
let response = self
.http
.get(&self.jwks_url)
.send()
.await
.map_err(|err| SessionError::Verifier(format!("jwks request failed: {err}")))?;
if !response.status().is_success() {
return Err(SessionError::Verifier(format!(
"jwks endpoint returned {}",
response.status()
)));
}
response
.json::<JwkSet>()
.await
.map_err(|err| SessionError::Verifier(format!("invalid jwks document: {err}")))
}
}
fn decoding_key_for(set: &JwkSet, kid: &str) -> Result<Option<DecodingKey>, SessionError> {
match set.find(kid) {
Some(jwk) => DecodingKey::from_jwk(jwk)
.map(Some)
.map_err(|err| SessionError::Verifier(err.to_string())),
None => Ok(None),
}
}