use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};
use super::account::{Account, CredentialRef};
#[derive(Clone, PartialEq, Eq, Default)]
pub enum CredentialSource {
#[default]
Local,
Remote {
endpoint: String,
machine_token: String,
},
}
impl CredentialSource {
pub fn from_parts(endpoint: Option<String>, machine_token: Option<String>) -> Self {
match endpoint {
Some(e) if !e.trim().is_empty() => CredentialSource::Remote {
endpoint: e.trim().trim_end_matches('/').to_string(),
machine_token: machine_token.unwrap_or_default().trim().to_string(),
},
_ => CredentialSource::Local,
}
}
pub fn is_remote(&self) -> bool {
matches!(self, CredentialSource::Remote { .. })
}
}
impl std::fmt::Debug for CredentialSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
CredentialSource::Local => write!(f, "Local"),
CredentialSource::Remote { endpoint, .. } => f
.debug_struct("Remote")
.field("endpoint", endpoint)
.field("machine_token", &"***")
.finish(),
}
}
}
#[derive(Clone, serde::Deserialize)]
pub struct BrokerToken {
pub access_token: String,
pub expires: u64,
#[serde(default)]
pub ttl_ms: Option<u64>,
}
impl std::fmt::Debug for BrokerToken {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BrokerToken")
.field("access_token", &"[REDACTED]")
.field("expires", &self.expires)
.field("ttl_ms", &self.ttl_ms)
.finish()
}
}
fn now_millis() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
pub fn is_expired_with_margin(expires_ms: u64, margin_ms: u64) -> bool {
now_millis().saturating_add(margin_ms) >= expires_ms
}
pub const DEFAULT_MARGIN_MS: u64 = 5 * 60 * 1000;
type LocalBrokerMap = std::collections::BTreeMap<
(std::path::PathBuf, std::path::PathBuf),
Arc<super::broker::LocalBroker>,
>;
#[derive(Clone, Default)]
pub struct TokenCache {
inner: Arc<RwLock<HashMap<String, BrokerToken>>>,
pub(crate) local_brokers: Arc<std::sync::Mutex<LocalBrokerMap>>,
}
pub fn source_scope(endpoint: &str, machine_token: &str) -> String {
use sha2::{Digest, Sha256};
let principal = if machine_token.is_empty() {
"anon".to_string()
} else {
let digest = Sha256::digest(machine_token.as_bytes());
hex_prefix(&digest, 16)
};
format!("{}#{principal}", endpoint.trim_end_matches('/'))
}
fn hex_prefix(bytes: &[u8], chars: usize) -> String {
let mut out = String::with_capacity(chars);
for b in bytes {
if out.len() >= chars {
break;
}
out.push_str(&format!("{b:02x}"));
}
out.truncate(chars);
out
}
fn key_provider(key: &str) -> Option<String> {
let storage_key = key.split_once('|').map(|(k, _)| k).unwrap_or(key);
CredentialRef::parse_storage_key(storage_key)
.map(|c| c.provider.as_str().to_string())
.or_else(|| Some(storage_key.to_string()))
}
impl std::fmt::Debug for TokenCache {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let providers: Vec<String> = self
.inner
.read()
.map(|m| m.keys().cloned().collect())
.unwrap_or_default();
f.debug_struct("TokenCache")
.field("providers", &providers)
.finish()
}
}
impl TokenCache {
pub fn new() -> Self {
Self::default()
}
pub fn get_fresh(&self, provider: &str, margin_ms: u64) -> Option<BrokerToken> {
let map = self.inner.read().ok()?;
let tok = map.get(provider)?;
if is_expired_with_margin(tok.expires, margin_ms) {
None
} else {
Some(tok.clone())
}
}
pub fn put(&self, provider: &str, token: BrokerToken) {
if let Ok(mut map) = self.inner.write() {
map.insert(provider.to_string(), token);
}
}
pub fn invalidate(&self, provider: &str) {
if let Ok(mut map) = self.inner.write() {
map.retain(|key, _| key_provider(key).as_deref() != Some(provider));
}
}
pub fn scoped_key(scope: &str, cred: &CredentialRef) -> String {
format!("{}|{scope}", cred.storage_key())
}
pub fn invalidate_credential(&self, scope: &str, cred: &CredentialRef) {
if let Ok(mut map) = self.inner.write() {
map.remove(&Self::scoped_key(scope, cred));
}
}
pub fn get_unexpired(&self, provider: &str) -> Option<BrokerToken> {
let map = self.inner.read().ok()?;
let tok = map.get(provider)?;
if is_expired_with_margin(tok.expires, 0) {
None
} else {
Some(tok.clone())
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TokenFetchError {
Unauthorized,
UnknownAccount,
InvalidAccount,
UnsupportedAccount,
AccountMismatch,
Http(u16),
Other(String),
}
impl std::fmt::Display for TokenFetchError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Unauthorized => f.write_str(&super::broker_token::unauthorized_hint()),
Self::UnknownAccount => write!(f, "broker reports unknown account (404)"),
Self::InvalidAccount => write!(f, "broker rejected account label (400)"),
Self::UnsupportedAccount => write!(f, "named accounts are not supported by this fetcher"),
Self::AccountMismatch => write!(
f,
"broker served a token for a different account than requested (broker too old for named accounts?)"
),
Self::Http(status) => write!(f, "broker returned HTTP {status}"),
Self::Other(msg) => f.write_str(msg),
}
}
}
impl TokenFetchError {
pub fn into_broker_error(self, cred: &CredentialRef) -> super::broker::BrokerError {
use super::broker::BrokerError;
match self {
Self::Unauthorized => BrokerError::Unauthorized,
Self::UnknownAccount => BrokerError::UnknownAccount {
provider: cred.provider.as_str().to_string(),
label: cred.account.label_str().to_string(),
},
Self::InvalidAccount => {
BrokerError::InvalidAccount(format!("broker rejected '{}'", cred.account))
}
Self::UnsupportedAccount => BrokerError::UnsupportedAccount {
provider: cred.provider.as_str().to_string(),
label: cred.account.label_str().to_string(),
},
Self::AccountMismatch => BrokerError::Transport(format!(
"broker served a token for a different account than '{}' (no fallback)",
cred
)),
Self::Http(status) => BrokerError::Transport(format!("broker returned HTTP {status}")),
Self::Other(msg) => BrokerError::Transport(msg),
}
}
}
#[allow(async_fn_in_trait)]
pub trait TokenFetcher {
async fn fetch_token(&self, provider: &str) -> Result<BrokerToken, String>;
async fn fetch_credential(&self, cred: &CredentialRef) -> Result<BrokerToken, TokenFetchError> {
match &cred.account {
Account::Default => self
.fetch_token(cred.provider.as_str())
.await
.map_err(TokenFetchError::Other),
Account::Named(_) => Err(TokenFetchError::UnsupportedAccount),
}
}
}
pub async fn resolve_remote<F: TokenFetcher>(
fetcher: &F,
cache: &TokenCache,
provider: &str,
margin_ms: u64,
) -> Result<BrokerToken, String> {
if let Some(tok) = cache.get_fresh(provider, margin_ms) {
return Ok(tok);
}
match fetcher.fetch_token(provider).await {
Ok(tok) => {
cache.put(provider, tok.clone());
Ok(tok)
}
Err(e) => {
if let Some(tok) = cache.get_unexpired(provider) {
return Ok(tok);
}
Err(e)
}
}
}
pub async fn resolve_remote_credential<F: TokenFetcher>(
fetcher: &F,
cache: &TokenCache,
scope: &str,
cred: &CredentialRef,
margin_ms: u64,
) -> Result<BrokerToken, TokenFetchError> {
let key = TokenCache::scoped_key(scope, cred);
if let Some(tok) = cache.get_fresh(&key, margin_ms) {
return Ok(tok);
}
match fetcher.fetch_credential(cred).await {
Ok(tok) => {
cache.put(&key, tok.clone());
Ok(tok)
}
Err(e @ TokenFetchError::Other(_)) | Err(e @ TokenFetchError::Http(_)) => {
if let Some(tok) = cache.get_unexpired(&key) {
return Ok(tok);
}
Err(e)
}
Err(rejected @ TokenFetchError::AccountMismatch) => {
cache.invalidate_credential(scope, cred);
Err(rejected)
}
Err(rejected) => {
cache.invalidate_credential(scope, cred);
Err(rejected)
}
}
}
pub async fn resolve_remote_token<F: TokenFetcher>(
fetcher: &F,
cache: &TokenCache,
provider: &str,
margin_ms: u64,
) -> Result<String, String> {
Ok(resolve_remote(fetcher, cache, provider, margin_ms)
.await?
.access_token)
}
pub struct BrokerClient {
http: reqwest::Client,
endpoint: String,
machine_token: super::MachineToken,
}
impl BrokerClient {
pub fn new(endpoint: impl Into<String>, machine_token: impl Into<String>) -> Self {
let http = reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(5))
.timeout(std::time::Duration::from_secs(20))
.build()
.unwrap_or_default();
let endpoint = endpoint.into();
let machine_token = super::MachineToken::for_endpoint(&endpoint, machine_token);
Self {
http,
endpoint,
machine_token,
}
}
pub fn with_client(
endpoint: impl Into<String>,
machine_token: impl Into<String>,
http: reqwest::Client,
) -> Self {
let endpoint = endpoint.into();
let machine_token = super::MachineToken::for_endpoint(&endpoint, machine_token);
Self {
http,
endpoint,
machine_token,
}
}
pub fn with_machine_token(
endpoint: impl Into<String>,
machine_token: super::MachineToken,
http: reqwest::Client,
) -> Self {
Self {
http,
endpoint: endpoint.into(),
machine_token,
}
}
pub fn from_source(source: &CredentialSource) -> Option<Self> {
match source {
CredentialSource::Remote {
endpoint,
machine_token,
} => Some(Self::new(endpoint.clone(), machine_token.clone())),
CredentialSource::Local => None,
}
}
}
pub async fn resolve_access_token(
provider: &str,
source: &CredentialSource,
cache: &TokenCache,
http: &reqwest::Client,
) -> Result<String, String> {
let provider: super::provider::OAuthProviderId = provider.try_into()?;
super::broker::broker_from_source(source, cache, http.clone())
.access_token(provider)
.await
.map(|t| t.token)
.map_err(|e| e.to_string())
}
impl BrokerClient {
pub fn scope(&self) -> String {
source_scope(&self.endpoint, &self.machine_token.get())
}
async fn fetch_token_query(
&self,
provider: &str,
account: Option<&str>,
) -> Result<BrokerToken, TokenFetchError> {
let url = format!("{}/token", self.endpoint);
let mut query: Vec<(&str, &str)> = vec![("provider", provider)];
if let Some(account) = account {
query.push(("account", account));
}
let resp = self
.machine_token
.send(|t| self.http.get(&url).query(&query).bearer_auth(t))
.await
.map_err(|e| TokenFetchError::Other(format!("broker request failed: {e}")))?;
let status = resp.status();
match status.as_u16() {
401 => return Err(TokenFetchError::Unauthorized),
404 if account.is_some() => return Err(TokenFetchError::UnknownAccount),
400 if account.is_some() => return Err(TokenFetchError::InvalidAccount),
s if !status.is_success() => return Err(TokenFetchError::Http(s)),
_ => {}
}
#[derive(serde::Deserialize)]
struct Wire {
access_token: String,
expires: u64,
#[serde(default)]
ttl_ms: Option<u64>,
#[serde(default)]
account: Option<String>,
}
let wire = resp
.json::<Wire>()
.await
.map_err(|e| TokenFetchError::Other(format!("invalid broker token response: {e}")))?;
if let Some(requested) = account {
match wire.account.as_deref() {
Some(served) if served == requested => {}
None if requested == super::account::DEFAULT_ACCOUNT_NAME => {}
_ => return Err(TokenFetchError::AccountMismatch),
}
}
let mut tok = BrokerToken {
access_token: wire.access_token,
expires: wire.expires,
ttl_ms: wire.ttl_ms,
};
if let Some(ttl) = tok.ttl_ms {
tok.expires = now_millis().saturating_add(ttl);
}
if tok.access_token.is_empty() {
return Err(TokenFetchError::Other(
"broker returned an empty access_token".to_string(),
));
}
if tok.expires <= now_millis() {
return Err(TokenFetchError::Other(
"broker returned an already-expired token".to_string(),
));
}
Ok(tok)
}
}
impl TokenFetcher for BrokerClient {
async fn fetch_token(&self, provider: &str) -> Result<BrokerToken, String> {
self.fetch_token_query(provider, None)
.await
.map_err(|e| e.to_string())
}
async fn fetch_credential(&self, cred: &CredentialRef) -> Result<BrokerToken, TokenFetchError> {
self.fetch_token_query(cred.provider.as_str(), Some(cred.account.label_str()))
.await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn from_parts_local_when_no_endpoint() {
assert_eq!(
CredentialSource::from_parts(None, None),
CredentialSource::Local
);
assert_eq!(
CredentialSource::from_parts(Some(" ".into()), Some("m".into())),
CredentialSource::Local
);
}
#[test]
fn from_parts_remote_when_endpoint_set() {
let s =
CredentialSource::from_parts(Some("https://jade.jade:8181".into()), Some("tok".into()));
assert_eq!(
s,
CredentialSource::Remote {
endpoint: "https://jade.jade:8181".into(),
machine_token: "tok".into()
}
);
assert!(s.is_remote());
}
#[test]
fn from_parts_trims_trailing_slash_and_whitespace() {
let s = CredentialSource::from_parts(Some(" https://b/ ".into()), Some(" tok ".into()));
assert_eq!(
s,
CredentialSource::Remote {
endpoint: "https://b".into(),
machine_token: "tok".into()
}
);
}
#[test]
fn remote_with_missing_machine_token_defaults_empty() {
let s = CredentialSource::from_parts(Some("https://b".into()), None);
assert_eq!(
s,
CredentialSource::Remote {
endpoint: "https://b".into(),
machine_token: String::new()
}
);
}
#[test]
fn local_is_not_remote() {
assert!(!CredentialSource::Local.is_remote());
}
#[test]
fn expiry_far_future_not_expired() {
let far = now_millis() + 60 * 60 * 1000; assert!(!is_expired_with_margin(far, DEFAULT_MARGIN_MS));
}
#[test]
fn expiry_past_is_expired() {
let past = now_millis().saturating_sub(1000);
assert!(is_expired_with_margin(past, 0));
}
#[test]
fn expiry_within_margin_is_expired() {
let soon = now_millis() + 2 * 60 * 1000;
assert!(is_expired_with_margin(soon, DEFAULT_MARGIN_MS));
assert!(!is_expired_with_margin(soon, 60 * 1000));
}
#[test]
fn broker_token_deserializes_without_refresh_field() {
let json = r#"{"access_token":"sk-abc","expires":1750000000000}"#;
let t: BrokerToken = serde_json::from_str(json).unwrap();
assert_eq!(t.access_token, "sk-abc");
assert_eq!(t.expires, 1_750_000_000_000);
}
fn tok(expires: u64) -> BrokerToken {
BrokerToken {
access_token: "sk-live".into(),
expires,
ttl_ms: None,
}
}
#[test]
fn cache_get_fresh_returns_unexpired() {
let c = TokenCache::new();
c.put("anthropic", tok(now_millis() + 60 * 60 * 1000));
assert!(c.get_fresh("anthropic", DEFAULT_MARGIN_MS).is_some());
}
#[test]
fn cache_get_fresh_none_when_expired() {
let c = TokenCache::new();
c.put("anthropic", tok(now_millis().saturating_sub(1000)));
assert!(c.get_fresh("anthropic", 0).is_none());
}
#[test]
fn cache_invalidate_removes() {
let c = TokenCache::new();
c.put("anthropic", tok(now_millis() + 60 * 60 * 1000));
c.invalidate("anthropic");
assert!(c.get_fresh("anthropic", 0).is_none());
}
#[test]
fn cache_providers_isolated() {
let c = TokenCache::new();
c.put("anthropic", tok(now_millis() + 60 * 60 * 1000));
assert!(c.get_fresh("openai", 0).is_none());
}
use std::sync::atomic::{AtomicUsize, Ordering};
struct FakeFetcher {
token: BrokerToken,
calls: AtomicUsize,
}
impl TokenFetcher for FakeFetcher {
async fn fetch_token(&self, _provider: &str) -> Result<BrokerToken, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Ok(self.token.clone())
}
}
#[tokio::test]
async fn resolve_cache_hit_does_not_fetch() {
let cache = TokenCache::new();
cache.put("anthropic", tok(now_millis() + 60 * 60 * 1000));
let f = FakeFetcher {
token: tok(0),
calls: AtomicUsize::new(0),
};
let t = resolve_remote_token(&f, &cache, "anthropic", DEFAULT_MARGIN_MS)
.await
.unwrap();
assert_eq!(t, "sk-live");
assert_eq!(
f.calls.load(Ordering::SeqCst),
0,
"must not fetch on a cache hit"
);
}
#[tokio::test]
async fn resolve_miss_fetches_then_caches() {
let cache = TokenCache::new();
let f = FakeFetcher {
token: tok(now_millis() + 60 * 60 * 1000),
calls: AtomicUsize::new(0),
};
resolve_remote_token(&f, &cache, "anthropic", DEFAULT_MARGIN_MS)
.await
.unwrap();
resolve_remote_token(&f, &cache, "anthropic", DEFAULT_MARGIN_MS)
.await
.unwrap();
assert_eq!(
f.calls.load(Ordering::SeqCst),
1,
"second resolve should hit the cache"
);
}
struct FailFetcher;
impl TokenFetcher for FailFetcher {
async fn fetch_token(&self, _provider: &str) -> Result<BrokerToken, String> {
Err("broker unreachable".into())
}
}
#[tokio::test]
async fn resolve_serves_stale_cache_when_broker_down() {
let cache = TokenCache::new();
cache.put("anthropic", tok(now_millis() + 2 * 60 * 1000));
let t = resolve_remote(&FailFetcher, &cache, "anthropic", DEFAULT_MARGIN_MS)
.await
.unwrap();
assert_eq!(t.access_token, "sk-live");
}
#[tokio::test]
async fn resolve_errors_when_broker_down_and_no_cache() {
let cache = TokenCache::new();
let err = resolve_remote(&FailFetcher, &cache, "anthropic", DEFAULT_MARGIN_MS)
.await
.unwrap_err();
assert!(err.contains("unreachable"), "got: {err}");
}
#[tokio::test]
async fn resolve_errors_when_broker_down_and_cache_hard_expired() {
let cache = TokenCache::new();
cache.put("anthropic", tok(now_millis().saturating_sub(1000))); let err = resolve_remote(&FailFetcher, &cache, "anthropic", 0)
.await
.unwrap_err();
assert!(err.contains("unreachable"), "got: {err}");
}
async fn spawn_broker(token_json: &'static str, require_token: &'static str) -> String {
use axum::{http::HeaderMap, http::StatusCode, routing::get, Router};
let app = Router::new().route(
"/token",
get(move |headers: HeaderMap| async move {
let auth = headers
.get("authorization")
.and_then(|v| v.to_str().ok())
.unwrap_or("");
if auth == format!("Bearer {require_token}") {
(StatusCode::OK, token_json.to_string())
} else {
(StatusCode::UNAUTHORIZED, String::new())
}
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
axum::serve(listener, app).await.unwrap();
});
format!("http://{addr}")
}
#[tokio::test]
async fn broker_client_fetches_and_parses() {
let url = spawn_broker(
r#"{"access_token":"sk-from-broker","expires":9999999999999}"#,
"machine-xyz",
)
.await;
let c = BrokerClient::new(url, "machine-xyz");
let t = c.fetch_token("anthropic").await.unwrap();
assert_eq!(t.access_token, "sk-from-broker");
assert_eq!(t.expires, 9_999_999_999_999);
}
#[tokio::test]
async fn broker_client_401_on_bad_machine_auth() {
let url = spawn_broker(
r#"{"access_token":"x","expires":9999999999999}"#,
"right-token",
)
.await;
let c = BrokerClient::new(url, "WRONG-token");
let err = c.fetch_token("anthropic").await.unwrap_err();
assert!(err.contains("401"), "expected 401 error, got: {err}");
}
#[tokio::test]
async fn fetch_token_recomputes_expiry_from_ttl_ms() {
let url = spawn_broker(r#"{"access_token":"sk","expires":1,"ttl_ms":3600000}"#, "m").await;
let c = BrokerClient::new(url, "m");
let t = c.fetch_token("anthropic").await.unwrap();
assert!(
t.expires > now_millis(),
"expires must come from ttl_ms, got {}",
t.expires
);
}
#[tokio::test]
async fn fetch_token_rejects_empty_access_token() {
let url = spawn_broker(r#"{"access_token":"","expires":9999999999999}"#, "m").await;
let c = BrokerClient::new(url, "m");
assert!(c
.fetch_token("anthropic")
.await
.unwrap_err()
.contains("empty"));
}
#[tokio::test]
async fn fetch_token_rejects_already_expired() {
let url = spawn_broker(r#"{"access_token":"sk","expires":1}"#, "m").await;
let c = BrokerClient::new(url, "m");
assert!(c
.fetch_token("anthropic")
.await
.unwrap_err()
.contains("expired"));
}
#[tokio::test]
async fn resolve_access_token_remote_fetches_from_broker() {
let url = spawn_broker(
r#"{"access_token":"sk-broker","expires":9999999999999,"account":"default"}"#,
"m",
)
.await;
let source = CredentialSource::Remote {
endpoint: url,
machine_token: "m".into(),
};
let cache = TokenCache::new();
let http = reqwest::Client::new();
let t = resolve_access_token("anthropic", &source, &cache, &http)
.await
.unwrap();
assert_eq!(t, "sk-broker");
let t2 = resolve_access_token("anthropic", &source, &cache, &http)
.await
.unwrap();
assert_eq!(t2, "sk-broker");
}
#[test]
fn broker_token_structurally_drops_any_refresh_field() {
let json = r#"{"access_token":"a","expires":1,"refresh_token":"LEAK","refresh":"LEAK"}"#;
let t: BrokerToken = serde_json::from_str(json).unwrap();
assert_eq!(t.access_token, "a");
assert_eq!(t.expires, 1);
}
}