soma_auth/upstream/
refresh.rs

1//! Single-flight refresh coordination for upstream OAuth clients.
2//!
3//! `RefreshLocks` prevents concurrent callers for the same `(upstream, subject)` pair
4//! from issuing simultaneous token refresh requests against the authorization server.
5//! One caller wins the lock and executes `get_access_token()` (which internally handles
6//! proactive refresh); all others wait and then return the already-refreshed token.
7//!
8//! **Scope:** This module handles *proactive* refresh triggered before making an MCP call.
9//! Reactive 401-retry logic is wired in Task 4 (`dispatch/gateway/`).
10//!
11//! ## rmcp refresh semantics
12//!
13//! `AuthorizationManager::get_access_token()` refreshes the token when fewer than 30 s
14//! remain before expiry.  It does **not** react to 401 responses from the resource server.
15//! A 401 with a locally-still-valid token requires an explicit `refresh_token()` call
16//! followed by a retry — that is the Task 4 responsibility.
17
18use std::sync::Arc;
19use std::time::{Duration, Instant};
20
21use dashmap::DashMap;
22use tokio::sync::Mutex;
23
24/// Per-`(upstream_name, subject)` mutex pool.
25///
26/// Entries are created lazily on first access and are never removed (the number of
27/// distinct `(upstream, subject)` pairs is bounded by the number of configured upstreams
28/// times the number of users, which is small in a homelab context).
29#[derive(Default)]
30pub struct RefreshLocks(DashMap<(String, String), Arc<Mutex<()>>>);
31
32impl RefreshLocks {
33    pub fn new() -> Self {
34        Self(DashMap::new())
35    }
36
37    /// Return the mutex for `(upstream_name, subject)`, creating it if absent.
38    pub fn acquire(&self, upstream_name: &str, subject: &str) -> Arc<Mutex<()>> {
39        let key = (upstream_name.to_string(), subject.to_string());
40        self.0
41            .entry(key)
42            .or_insert_with(|| Arc::new(Mutex::new(())))
43            .clone()
44    }
45}
46
47/// How long a confirmed refresh failure suppresses further live retries for the
48/// same `(upstream, subject)` pair. Chosen to be well short of any human patience
49/// window (so a fix shows up promptly) while still cutting a dead credential's
50/// call volume against the authorization server by roughly two orders of
51/// magnitude versus retrying on every single request.
52pub const REFRESH_FAILURE_COOLDOWN: Duration = Duration::from_secs(300);
53
54/// Per-`(upstream_name, subject)` "is this credential known-broken right now"
55/// cache.
56///
57/// Without this, a dead refresh token (revoked, expired, `invalid_grant`, ...)
58/// gets retried against the authorization server on every single request
59/// forever — `TokenRefreshState::refresh_due()` is purely time-based and has
60/// no memory of prior outcomes. That wastes latency on every real request
61/// touching the upstream, and can itself contribute to the authorization
62/// server rate-limiting or flagging the client_id, which is especially bad
63/// when multiple upstreams share one client_id (see `labby-auth::upstream`
64/// module docs).
65#[derive(Default)]
66pub struct RefreshFailureCache(DashMap<(String, String), Instant>);
67
68impl RefreshFailureCache {
69    pub fn new() -> Self {
70        Self(DashMap::new())
71    }
72
73    /// Record that a refresh just failed for `(upstream_name, subject)`.
74    pub fn record_failure(&self, upstream_name: &str, subject: &str) {
75        self.0.insert(
76            (upstream_name.to_string(), subject.to_string()),
77            Instant::now(),
78        );
79    }
80
81    /// Clear any recorded failure for `(upstream_name, subject)` — call this on
82    /// any successful refresh, a fresh authorization completing, or explicit
83    /// credential clearing, so a fix is picked up immediately instead of
84    /// waiting out the cooldown.
85    pub fn clear(&self, upstream_name: &str, subject: &str) {
86        self.0
87            .remove(&(upstream_name.to_string(), subject.to_string()));
88    }
89
90    /// Whether `(upstream_name, subject)` failed recently enough that a live
91    /// retry should be skipped.
92    pub fn recently_failed(&self, upstream_name: &str, subject: &str) -> bool {
93        self.0
94            .get(&(upstream_name.to_string(), subject.to_string()))
95            .is_some_and(|entry| entry.elapsed() < REFRESH_FAILURE_COOLDOWN)
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::RefreshFailureCache;
102
103    #[test]
104    fn fresh_cache_has_no_recent_failures() {
105        let cache = RefreshFailureCache::new();
106        assert!(!cache.recently_failed("google-drive", "gateway"));
107    }
108
109    #[test]
110    fn recorded_failure_is_recently_failed_until_cleared() {
111        let cache = RefreshFailureCache::new();
112        cache.record_failure("google-drive", "gateway");
113        assert!(cache.recently_failed("google-drive", "gateway"));
114
115        cache.clear("google-drive", "gateway");
116        assert!(!cache.recently_failed("google-drive", "gateway"));
117    }
118
119    #[test]
120    fn failures_are_scoped_per_upstream_and_subject() {
121        let cache = RefreshFailureCache::new();
122        cache.record_failure("google-drive", "gateway");
123
124        assert!(!cache.recently_failed("google-gmail", "gateway"));
125        assert!(!cache.recently_failed("google-drive", "alice"));
126    }
127}
soma_auth/upstream/refresh.rs

soma_auth/upstream/
refresh.rs
