i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-12 04:09:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(provider): harden Agent Identity OAuth lifecycle

Browse Source
This commit is contained in:
elky
2026-07-23 09:33:00 +08:00
parent e49024d33b
commit 3606290ac8
84 changed files with 8063 additions and 1104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-provider/transport/src/oauth_refresh/mod.rs
+474 -16
View File
@@ -1,13 +1,14 @@
use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_oauth::core::OAuthError;
use aether_oauth::network::{
OAuthHttpExecutor, OAuthHttpRequest, OAuthHttpResponse, OAuthNetworkContext,
};
use aether_oauth::provider::ProviderOAuthTransportContext;
use aether_runtime_state::RuntimeState;
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
use async_trait::async_trait;
use serde_json::Value;
use thiserror::Error;
@@ -40,6 +41,12 @@ pub struct LocalOAuthResolution {
pub auth: Option<LocalResolvedOAuthRequestAuth>,
pub refreshed_entry: Option<CachedOAuthEntry>,
pub refresh_in_flight: bool,
/// Indicates that a forced caller reused a newer completed refresh rather
/// than producing a new entry that needs persistence.
pub reused_refresh: bool,
/// Held until the caller persists `refreshed_entry`. The lease TTL remains
/// the cancellation fallback if the caller is dropped.
pub distributed_lease: Option<RuntimeLockLease>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -49,6 +56,8 @@ pub struct CachedOAuthEntry {
pub auth_header_value: String,
pub expires_at_unix_secs: Option<u64>,
pub metadata: Option<Value>,
/// Non-secret fingerprint of the credential/configuration that produced it.
pub source_fingerprint: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
@@ -297,6 +306,28 @@ pub trait LocalOAuthRefreshAdapter: Send + Sync {
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth>;
/// Resolves a cache entry that is known to have advanced the caller's
/// refresh fence. Agent task rotation can safely use the winner even while
/// the caller still holds the pre-refresh transport snapshot.
fn resolve_fenced_cached(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
self.resolve_cached(transport, entry)
}
/// Resolves the entry returned by this adapter's immediately preceding
/// refresh. Unlike a reusable cache entry, this entry is expected to have
/// advanced the transport generation.
fn resolve_refreshed(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
self.resolve_cached(transport, entry)
}
fn resolve_without_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
@@ -308,6 +339,36 @@ pub trait LocalOAuthRefreshAdapter: Send + Sync {
entry: Option<&CachedOAuthEntry>,
) -> bool;
/// Identifies the credential/configuration generation used by a refresh.
/// Adapters that support fencing override this method.
fn refresh_fingerprint(
&self,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> Option<String> {
None
}
/// Reconstructs a cache entry from an already-persisted transport after a
/// distributed refresh waiter reloads the winner.
fn cached_entry_from_transport(
&self,
_transport: &GatewayProviderTransportSnapshot,
) -> Option<CachedOAuthEntry> {
None
}
/// Enables bounded negative backoff for transient refresh failures.
fn should_backoff_after_error(&self, _error: &LocalOAuthRefreshError) -> bool {
false
}
/// Agent task registration is a non-idempotent external mutation and must
/// not continue unlocked when a configured distributed lock is unavailable.
fn requires_distributed_refresh_lock(&self) -> bool {
false
}
async fn refresh(
&self,
executor: &dyn LocalOAuthHttpExecutor,
@@ -320,6 +381,13 @@ pub struct LocalOAuthRefreshCoordinator {
adapters: Vec<Arc<dyn LocalOAuthRefreshAdapter>>,
cache: Mutex<BTreeMap<String, CachedOAuthEntry>>,
key_locks: Mutex<BTreeMap<String, Arc<Mutex<()>>>>,
refresh_backoff: Mutex<BTreeMap<String, RefreshBackoffState>>,
}
#[derive(Debug, Clone)]
struct RefreshBackoffState {
failures: u32,
retry_after: Instant,
}
impl fmt::Debug for LocalOAuthRefreshCoordinator {
@@ -337,7 +405,10 @@ impl Default for LocalOAuthRefreshCoordinator {
}
impl LocalOAuthRefreshCoordinator {
const DISTRIBUTED_REFRESH_LOCK_TTL_MS: u64 = 30_000;
// Keep the lease alive through the 30s upstream HTTP timeout and the
// subsequent encrypted DB CAS/persistence step. Cancellation still relies
// on expiry as the last-resort release path.
const DISTRIBUTED_REFRESH_LOCK_TTL_MS: u64 = 120_000;
pub fn new() -> Self {
Self {
@@ -349,6 +420,7 @@ impl LocalOAuthRefreshCoordinator {
],
cache: Mutex::new(BTreeMap::new()),
key_locks: Mutex::new(BTreeMap::new()),
refresh_backoff: Mutex::new(BTreeMap::new()),
}
}
@@ -373,7 +445,9 @@ impl LocalOAuthRefreshCoordinator {
}
pub async fn invalidate_cached_entry(&self, key_id: &str) -> bool {
self.cache.lock().await.remove(key_id).is_some()
let removed = self.cache.lock().await.remove(key_id).is_some();
self.clear_refresh_backoff(key_id).await;
removed
}
pub async fn resolve_with_result(
@@ -389,6 +463,7 @@ impl LocalOAuthRefreshCoordinator {
distributed_lock,
distributed_owner,
false,
None,
)
.await
}
@@ -399,6 +474,27 @@ impl LocalOAuthRefreshCoordinator {
transport: &GatewayProviderTransportSnapshot,
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
self.force_refresh_with_result_fenced(
executor,
transport,
distributed_lock,
distributed_owner,
None,
)
.await
}
/// Force a refresh unless another request has already advanced the supplied
/// refresh fence. This prevents a distributed waiter from re-registering a
/// task after the winner has persisted it.
pub async fn force_refresh_with_result_fenced(
&self,
executor: &dyn LocalOAuthHttpExecutor,
transport: &GatewayProviderTransportSnapshot,
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
expected_refresh_fingerprint: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
self.resolve_with_result_mode(
executor,
@@ -406,6 +502,7 @@ impl LocalOAuthRefreshCoordinator {
distributed_lock,
distributed_owner,
true,
expected_refresh_fingerprint,
)
.await
}
@@ -417,6 +514,7 @@ impl LocalOAuthRefreshCoordinator {
distributed_lock: Option<&RuntimeState>,
distributed_owner: Option<&str>,
force_refresh: bool,
expected_refresh_fingerprint: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
let Some(adapter) = self
.adapters
@@ -450,10 +548,37 @@ impl LocalOAuthRefreshCoordinator {
return Ok(None);
}
if force_refresh {
if let Some(resolution) = Self::resolve_if_refresh_fence_advanced(
adapter.as_ref(),
transport,
cached_entry.as_ref(),
expected_refresh_fingerprint,
) {
return Ok(Some(resolution));
}
}
if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await {
return Err(error);
}
let key_lock = self.lock_for_key(key_id).await;
let _key_guard = key_lock.lock().await;
let cached_entry = self.cached_entry(key_id).await;
if force_refresh {
if let Some(resolution) = Self::resolve_if_refresh_fence_advanced(
adapter.as_ref(),
transport,
cached_entry.as_ref(),
expected_refresh_fingerprint,
) {
return Ok(Some(resolution));
}
}
if let Some(error) = self.backoff_error(key_id, adapter.provider_type()).await {
return Err(error);
}
if !force_refresh {
if let Some(auth) = cached_entry
.as_ref()
@@ -488,6 +613,16 @@ impl LocalOAuthRefreshCoordinator {
error = ?err,
"gateway local oauth refresh distributed lock unavailable"
);
if adapter.requires_distributed_refresh_lock() {
let error = LocalOAuthRefreshError::TransportMessage {
provider_type: adapter.provider_type(),
message: "distributed refresh lock is unavailable".to_string(),
};
if adapter.should_backoff_after_error(&error) {
self.record_refresh_failure(key_id).await;
}
return Err(error);
}
None
}
}
@@ -501,22 +636,147 @@ impl LocalOAuthRefreshCoordinator {
// came from the original transport snapshot.
let refresh_entry = cached_entry.as_ref();
let refresh_result = adapter.refresh(executor, transport, refresh_entry).await;
if let (Some(lock), Some(lease)) = (distributed_lock, distributed_lease.as_ref()) {
if let Err(err) = lock.lock_release(lease).await {
tracing::warn!(
key_id = %key_id,
provider_type = adapter.provider_type(),
error = ?err,
"gateway local oauth refresh distributed lock release failed"
);
let refreshed_entry = match refresh_result {
Ok(Some(entry)) => {
self.clear_refresh_backoff(key_id).await;
entry
}
Ok(None) => {
Self::release_distributed_lease(
distributed_lock,
distributed_lease.as_ref(),
key_id,
adapter.provider_type(),
)
.await;
return Ok(None);
}
Err(error) => {
if adapter.should_backoff_after_error(&error) {
self.record_refresh_failure(key_id).await;
}
Self::release_distributed_lease(
distributed_lock,
distributed_lease.as_ref(),
key_id,
adapter.provider_type(),
)
.await;
return Err(error);
}
};
// In production the distributed lease is held through the gateway's
// DB CAS. Do not publish a provisional task before that CAS succeeds;
// otherwise a waiter could consume an assertion that loses the CAS.
// Lock-free/test callers retain the historical in-memory behavior.
if distributed_lease.is_none() {
self.insert_cached_entry(key_id, refreshed_entry.clone())
.await;
}
let Some(refreshed_entry) = refresh_result? else {
let Some(auth) = adapter.resolve_refreshed(transport, &refreshed_entry) else {
Self::release_distributed_lease(
distributed_lock,
distributed_lease.as_ref(),
key_id,
adapter.provider_type(),
)
.await;
return Ok(None);
};
Ok(adapter
.resolve_cached(transport, &refreshed_entry)
.map(|auth| LocalOAuthResolution::resolved(auth, Some(refreshed_entry))))
Ok(Some(LocalOAuthResolution::refreshed(
auth,
refreshed_entry,
distributed_lease,
)))
}
fn resolve_if_refresh_fence_advanced(
adapter: &dyn LocalOAuthRefreshAdapter,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
expected_refresh_fingerprint: Option<&str>,
) -> Option<LocalOAuthResolution> {
let expected = expected_refresh_fingerprint?;
if adapter.refresh_fingerprint(transport, entry).as_deref() == Some(expected) {
return None;
}
entry
.and_then(|entry| adapter.resolve_fenced_cached(transport, entry))
.map(|auth| {
LocalOAuthResolution::reused(
auth,
entry.expect("cached auth is required when a refresh fence advanced"),
)
})
.or_else(|| {
adapter
.cached_entry_from_transport(transport)
.and_then(|entry| {
adapter
.resolve_cached(transport, &entry)
.map(|auth| LocalOAuthResolution::reused(auth, &entry))
})
})
.or_else(|| {
adapter
.resolve_without_refresh(transport)
.map(|auth| LocalOAuthResolution::resolved(auth, None))
})
}
async fn backoff_error(
&self,
key_id: &str,
provider_type: &'static str,
) -> Option<LocalOAuthRefreshError> {
let backoff = self.refresh_backoff.lock().await;
let state = backoff.get(key_id)?;
let remaining = state.retry_after.checked_duration_since(Instant::now())?;
Some(LocalOAuthRefreshError::InvalidResponse {
provider_type,
message: format!(
"refresh temporarily backed off after {} failed attempts (retry in {}ms)",
state.failures,
remaining.as_millis()
),
})
}
async fn record_refresh_failure(&self, key_id: &str) {
let mut backoff = self.refresh_backoff.lock().await;
let state = backoff
.entry(key_id.to_string())
.or_insert(RefreshBackoffState {
failures: 0,
retry_after: Instant::now(),
});
state.failures = state.failures.saturating_add(1);
let exponent = state.failures.saturating_sub(1).min(4);
let delay = Duration::from_millis(500u64.saturating_mul(1u64 << exponent));
state.retry_after = Instant::now() + delay.min(Duration::from_secs(8));
}
async fn clear_refresh_backoff(&self, key_id: &str) {
self.refresh_backoff.lock().await.remove(key_id);
}
async fn release_distributed_lease(
distributed_lock: Option<&RuntimeState>,
lease: Option<&RuntimeLockLease>,
key_id: &str,
provider_type: &'static str,
) {
let (Some(lock), Some(lease)) = (distributed_lock, lease) else {
return;
};
if let Err(err) = lock.lock_release(lease).await {
tracing::warn!(
key_id = %key_id,
provider_type,
error = ?err,
"gateway local oauth refresh distributed lock release failed"
);
}
}
pub fn with_adapters_for_tests(adapters: Vec<Arc<dyn LocalOAuthRefreshAdapter>>) -> Self {
@@ -524,6 +784,7 @@ impl LocalOAuthRefreshCoordinator {
adapters,
cache: Mutex::new(BTreeMap::new()),
key_locks: Mutex::new(BTreeMap::new()),
refresh_backoff: Mutex::new(BTreeMap::new()),
}
}
}
@@ -537,6 +798,37 @@ impl LocalOAuthResolution {
auth: Some(auth),
refreshed_entry,
refresh_in_flight: false,
reused_refresh: false,
distributed_lease: None,
}
}
fn refreshed(
auth: LocalResolvedOAuthRequestAuth,
refreshed_entry: CachedOAuthEntry,
distributed_lease: Option<RuntimeLockLease>,
) -> Self {
Self {
auth: Some(auth),
refreshed_entry: Some(refreshed_entry),
refresh_in_flight: false,
reused_refresh: false,
distributed_lease,
}
}
fn reused(auth: LocalResolvedOAuthRequestAuth, entry: &CachedOAuthEntry) -> Self {
let mut refreshed_entry = entry.clone();
if let LocalResolvedOAuthRequestAuth::Header { name, value } = &auth {
refreshed_entry.auth_header_name = name.clone();
refreshed_entry.auth_header_value = value.clone();
}
Self {
auth: Some(auth),
refreshed_entry: Some(refreshed_entry),
refresh_in_flight: false,
reused_refresh: true,
distributed_lease: None,
}
}
@@ -545,6 +837,8 @@ impl LocalOAuthResolution {
auth: None,
refreshed_entry: None,
refresh_in_flight: true,
reused_refresh: false,
distributed_lease: None,
}
}
}
@@ -560,7 +854,7 @@ pub fn supports_local_oauth_request_auth_resolution(
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
@@ -580,6 +874,12 @@ mod tests {
refresh_with_entry_hits: Arc<AtomicUsize>,
}
#[derive(Debug)]
struct FencedTestAdapter {
refresh_hits: Arc<AtomicUsize>,
fail_refresh: Arc<AtomicBool>,
}
#[async_trait]
impl LocalOAuthRefreshAdapter for TestAdapter {
fn provider_type(&self) -> &'static str {
@@ -634,6 +934,77 @@ mod tests {
auth_header_value: "Bearer refreshed-token".to_string(),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
source_fingerprint: None,
}))
}
}
#[async_trait]
impl LocalOAuthRefreshAdapter for FencedTestAdapter {
fn provider_type(&self) -> &'static str {
"test-oauth"
}
fn resolve_cached(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
Some(LocalResolvedOAuthRequestAuth::Header {
name: entry.auth_header_name.clone(),
value: "fresh-winner-assertion".to_string(),
})
}
fn resolve_without_refresh(
&self,
_transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
None
}
fn should_refresh(
&self,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> bool {
true
}
fn refresh_fingerprint(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Option<String> {
entry
.and_then(|entry| entry.source_fingerprint.clone())
.or_else(|| Some("generation-1".to_string()))
}
fn should_backoff_after_error(&self, _error: &LocalOAuthRefreshError) -> bool {
true
}
async fn refresh(
&self,
_executor: &dyn LocalOAuthHttpExecutor,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
self.refresh_hits.fetch_add(1, Ordering::SeqCst);
if self.fail_refresh.load(Ordering::SeqCst) {
return Err(LocalOAuthRefreshError::TransportMessage {
provider_type: "test-oauth",
message: "temporary failure".to_string(),
});
}
Ok(Some(CachedOAuthEntry {
provider_type: "test-oauth".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: "stale-winner-cache-value".to_string(),
expires_at_unix_secs: None,
metadata: None,
source_fingerprint: Some("generation-2".to_string()),
}))
}
}
@@ -739,8 +1110,11 @@ mod tests {
auth_header_value: "Bearer refreshed-token".to_string(),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
source_fingerprint: None,
}),
refresh_in_flight: false,
reused_refresh: false,
distributed_lease: None,
})
);
assert_eq!(
@@ -752,6 +1126,8 @@ mod tests {
}),
refreshed_entry: None,
refresh_in_flight: false,
reused_refresh: false,
distributed_lease: None,
})
);
}
@@ -791,4 +1167,86 @@ mod tests {
assert_eq!(refresh_hits.load(Ordering::SeqCst), 2);
assert_eq!(refresh_with_entry_hits.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn fenced_force_refresh_reuses_the_winner() {
let refresh_hits = Arc::new(AtomicUsize::new(0));
let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(
FencedTestAdapter {
refresh_hits: Arc::clone(&refresh_hits),
fail_refresh: Arc::new(AtomicBool::new(false)),
},
)]);
let transport = sample_transport();
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
let first = coordinator
.force_refresh_with_result_fenced(
&executor,
&transport,
None,
None,
Some("generation-1"),
)
.await
.expect("first refresh should succeed")
.expect("first refresh should resolve");
let waiter = coordinator
.force_refresh_with_result_fenced(
&executor,
&transport,
None,
None,
Some("generation-1"),
)
.await
.expect("waiter should reuse winner")
.expect("waiter should resolve");
assert!(first.refreshed_entry.is_some());
assert!(waiter.refreshed_entry.is_some());
assert!(waiter.reused_refresh);
assert_eq!(
waiter
.refreshed_entry
.as_ref()
.expect("reused entry")
.auth_header_value,
"fresh-winner-assertion"
);
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn refresh_failure_enters_bounded_negative_backoff() {
let refresh_hits = Arc::new(AtomicUsize::new(0));
let fail_refresh = Arc::new(AtomicBool::new(true));
let coordinator = LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(
FencedTestAdapter {
refresh_hits: Arc::clone(&refresh_hits),
fail_refresh: Arc::clone(&fail_refresh),
},
)]);
let transport = sample_transport();
let executor = ReqwestLocalOAuthHttpExecutor::new(reqwest::Client::new());
assert!(coordinator
.force_refresh_with_result(&executor, &transport, None, None)
.await
.is_err());
let second = coordinator
.force_refresh_with_result(&executor, &transport, None, None)
.await
.expect_err("second refresh should be backed off");
assert!(second.to_string().contains("temporarily backed off"));
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
fail_refresh.store(false, Ordering::SeqCst);
coordinator.invalidate_cached_entry("key-1").await;
assert!(coordinator
.force_refresh_with_result(&executor, &transport, None, None)
.await
.expect("replacement should refresh immediately")
.is_some());
assert_eq!(refresh_hits.load(Ordering::SeqCst), 2);
}
}