b6f7f73
ZhenYi
refactor: reorganize configuration files and update module paths for improved structure
b6f7f73
use std::sync::atomic::{AtomicI64, Ordering};
use std::sync::{Arc, OnceLock, RwLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use config::jwt::JwtConfig;
use crate::claims::AppClaims;
use crate::codec;
use crate::error::JwtError;
use crate::jwks::Jwks;
pub const PLATFORM_AUDIENCE: &str = "gitdata.ai";
const REFRESH_SECS: u64 = 3600;
const REFRESH_COOLDOWN_SECS: i64 = 5;
const FETCH_TIMEOUT_SECS: u64 = 5;
/// `reqwest::Url` 与真正发请求的解析器是同一实现，避免手写切分与 WHATWG 解析分叉。
fn parsed(url: &str) -> Result<reqwest::Url, JwtError> {
reqwest::Url::parse(url).map_err(|error| JwtError::Config(format!("invalid url: {error}")))
}
fn is_loopback(url: &str) -> bool {
let Ok(url) = reqwest::Url::parse(url) else {
return false;
};
url.host_str().is_some_and(config::is_loopback_host)
}
static INSTALLED: OnceLock<JwtState> = OnceLock::new();
#[derive(Clone)]
pub struct JwtState {
inner: Arc<Inner>,
}
struct Inner {
profile_url: Option<String>,
keys: RwLock<Arc<Jwks>>,
issuer: String,
audience: String,
leeway_secs: i64,
source: Option<Source>,
last_refresh: AtomicI64,
}
struct Source {
url: String,
pin: Option<String>,
client: reqwest::Client,
}
impl JwtState {
pub async fn load(config: &JwtConfig) -> Result<Self, JwtError> {
let url = trimmed(config.jwks_url.as_deref());
let pin = trimmed(config.jwks_sha256.as_deref());
if let Some(url) = url.as_deref() {
let parsed = parsed(url)?;
require_secure(&parsed, "jwt.jwks_url")?;
if pin.is_none() && !is_loopback(url) {
return Err(JwtError::Config(
"jwt.jwks_sha256 is required when the document is fetched remotely".to_owned(),
));
}
}
if let Some(authority) = config.authority_grpc_url.as_deref() {
let parsed = parsed(authority)?;
require_secure(&parsed, "jwt.authority_grpc_url")?;
}
if !config.retired.is_empty()
&& config.retired.iter().any(|key| key.key_id == config.key_id)
{
return Err(JwtError::Config(
"jwt.retired must not reuse the active jwt.key_id".to_owned(),
));
}
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(FETCH_TIMEOUT_SECS))
.build()
.map_err(|error| JwtError::Internal(format!("building the jwks client: {error}")))?;
let source = url.map(|url| Source {
url,
pin: pin.clone(),
client,
});
let keys = match derived(config)? {
Some(keys) => {
if let Some(pin) = pin.as_deref() {
check_pin(&keys, pin)?;
}
keys
}
None => match &source {
Some(source) => source.fetch(None).await?,
None => {
return Err(JwtError::Config(
"jwt.jwks_url is required when jwt.private_key is absent".to_owned(),
));
}
},
};
let state = Self {
inner: Arc::new(Inner {
profile_url: config.profile_url.clone(),
keys: RwLock::new(Arc::new(keys)),
issuer: config.issuer.clone(),
audience: trimmed(config.audience.as_deref())
.unwrap_or_else(|| PLATFORM_AUDIENCE.to_owned()),
leeway_secs: config.leeway_secs.max(0),
source,
last_refresh: AtomicI64::new(now_secs()),
}),
};
match INSTALLED.set(state.clone()) {
// 已经装过一个实例：返回已装的那个，避免调用方拿到未安装副本并重复启动刷新任务。
Ok(()) => {
state.spawn_refresh_loop();
Ok(state)
}
Err(existing) => {
if existing.snapshot().digest()? != state.snapshot().digest()? {
return Err(JwtError::Config(
"another jwt verifier with different keys is already installed".to_owned(),
));
}
Ok(existing)
}
}
}
pub fn snapshot(&self) -> Arc<Jwks> {
// 毒锁回退：keys 快照丢失比进程崩溃代价小，取上次快照（或空表）继续服务。
self.inner
.keys
.read()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.clone()
}
pub fn jwks(&self) -> Arc<Jwks> {
self.snapshot()
}
pub fn issuer(&self) -> &str {
&self.inner.issuer
}
pub fn profile_url(&self) -> Option<&str> {
self.inner.profile_url.as_deref()
}
pub fn audience(&self) -> &str {
&self.inner.audience
}
pub fn leeway_secs(&self) -> i64 {
self.inner.leeway_secs
}
pub fn parse(&self, token: &str) -> Result<AppClaims, JwtError> {
self.parse_for_audience(token, &self.inner.audience)
}
pub fn parse_for_audience(&self, token: &str, audience: &str) -> Result<AppClaims, JwtError> {
let keys = self.snapshot();
let result = codec::verify_app(
&keys,
token,
&self.inner.issuer,
audience,
self.inner.leeway_secs,
);
if matches!(result, Err(JwtError::UnknownKey)) {
self.refresh_soon();
}
result
}
#[tracing::instrument(skip_all, fields(request_id = %request_id.unwrap_or_default()))]
async fn refresh(&self, request_id: Option<&str>) -> Result<(), JwtError> {
let Some(source) = self.inner.source.as_ref() else {
return Ok(());
};
let keys = source.fetch(request_id).await?;
// 写入毒锁：另一个线程在刷新中 panic 时，新钥匙依然替换成功。
*self
.inner
.keys
.write()
.unwrap_or_else(|poisoned| poisoned.into_inner()) = Arc::new(keys);
self.inner.last_refresh.store(now_secs(), Ordering::Relaxed);
Ok(())
}
fn refresh_soon(&self) {
if self.inner.source.is_none() {
return;
}
let now = now_secs();
let last = self.inner.last_refresh.load(Ordering::Relaxed);
if now - last < REFRESH_COOLDOWN_SECS {
return;
}
if self
.inner
.last_refresh
.compare_exchange(last, now, Ordering::AcqRel, Ordering::Relaxed)
.is_err()
{
return;
}
if tokio::runtime::Handle::try_current().is_err() {
return;
}
let state = self.clone();
tokio::spawn(async move {
if let Err(error) = state.refresh(None).await {
tracing::warn!(%error, "refreshing the jwks document failed");
}
});
}
fn spawn_refresh_loop(&self) {
if self.inner.source.is_none() {
return;
}
if tokio::runtime::Handle::try_current().is_err() {
return;
}
let state = self.clone();
tokio::spawn(async move {
let mut ticker = tokio::time::interval(Duration::from_secs(REFRESH_SECS));
ticker.tick().await;
loop {
ticker.tick().await;
if let Err(error) = state.refresh(None).await {
tracing::warn!(%error, "refreshing the jwks document failed");
}
}
});
}
}
impl Source {
#[tracing::instrument(skip_all, fields(url = %self.url, request_id = %_request_id.unwrap_or_default()))]
async fn fetch(&self, _request_id: Option<&str>) -> Result<Jwks, JwtError> {
let response = self.client.get(&self.url).send().await.map_err(|error| {
JwtError::Internal(format!("fetching jwks from {}: {error}", self.url))
})?;
if !response.status().is_success() {
return Err(JwtError::Internal(format!(
"fetching jwks from {}: status {}",
self.url,
response.status()
)));
}
let document: Jwks = response
.json()
.await
.map_err(|error| JwtError::Internal(format!("parsing jwks: {error}")))?;
if document.keys.is_empty() {
return Err(JwtError::Config("jwks document carries no keys".to_owned()));
}
for key in &document.keys {
key.verifying_key()?;
}
if let Some(pin) = self.pin.as_deref() {
check_pin(&document, pin)?;
}
Ok(document)
}
}
pub fn installed() -> Option<&'static JwtState> {
INSTALLED.get()
}
pub fn require_installed() -> Result<&'static JwtState, JwtError> {
installed().ok_or_else(|| JwtError::Config("jwt verifier is not installed".to_owned()))
}
pub fn verify(token: &str) -> Result<AppClaims, JwtError> {
require_installed()?.parse(token)
}
pub fn verify_for_audience(token: &str, audience: &str) -> Result<AppClaims, JwtError> {
require_installed()?.parse_for_audience(token, audience)
}
fn derived(config: &JwtConfig) -> Result<Option<Jwks>, JwtError> {
let Some(material) = trimmed(config.private_key.as_deref()) else {
return Ok(None);
};
let signing = codec::signing_key(&material)?;
Ok(Some(Jwks::single(
config.key_id.clone(),
&signing.verifying_key(),
)))
}
/// 非 loopback 地址一律要求 https；`http` 大小写、`http:/host` 这类写法都由同一解析器归一。
fn require_secure(url: &reqwest::Url, field: &str) -> Result<(), JwtError> {
match url.scheme() {
"https" => Ok(()),
"http" if url.host_str().is_some_and(config::is_loopback_host) => Ok(()),
"http" => Err(JwtError::Config(format!(
"{field} must use https when it is not a loopback address"
))),
_ => Err(JwtError::Config(format!("{field} must use http(s)"))),
}
}
fn check_pin(keys: &Jwks, expected: &str) -> Result<(), JwtError> {
let actual = keys.digest()?;
if expected != actual {
return Err(JwtError::Config(
"jwt.jwks_sha256 does not match the published public keys".to_owned(),
));
}
Ok(())
}
fn trimmed(value: Option<&str>) -> Option<String> {
value
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_owned)
}
fn now_secs() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|elapsed| elapsed.as_secs() as i64)
.unwrap_or_default()
}