Skip to main content

ppoppo_sdk_core/token_cache/
mod.rs

1//! JWT/token retention layer for Ppoppo client SDKs.
2//!
3//! [`TokenCache`] wraps a [`TokenSource`] and adds near-expiry refresh with
4//! single-flight semantics: when multiple tasks call [`TokenCache::get`]
5//! concurrently during an expiry window, only one performs the fetch; the
6//! rest queue on the refresh lock and then reuse the result.
7//!
8//! ## Typical wiring
9//!
10//! ```ignore
11//! use ppoppo_sdk_core::token_cache::{
12//!     ClientCredentialsSource, TokenCache, TokenCacheConfig,
13//! };
14//!
15//! let source = ClientCredentialsSource::new(token_url, client_id, client_secret);
16//! let cache = Arc::new(TokenCache::new(Box::new(source), TokenCacheConfig::default()));
17//! let token: String = cache.get().await?;
18//! ```
19
20use std::sync::Arc;
21use std::time::Duration;
22
23use async_trait::async_trait;
24use ppoppo_clock::ArcClock;
25use ppoppo_clock::native::WallClock;
26use tokio::sync::Mutex;
27
28/// Errors returned by [`TokenCache`] and [`TokenSource`] implementations.
29#[derive(Debug, thiserror::Error)]
30#[non_exhaustive]
31pub enum TokenCacheError {
32    #[error("token fetch failed: {0}")]
33    Fetch(String),
34    #[error("token response is malformed: {0}")]
35    Malformed(String),
36    /// A source-specific failure, preserved intact for the SDK layer that
37    /// installed the source. A custom [`TokenSource`] wraps its own typed
38    /// error here instead of flattening it to a string, so the SDK facade
39    /// can `downcast` back to its native error type (e.g. distinguish
40    /// "not authenticated — no durable credential" from a transport
41    /// failure) after the error has passed through the cache.
42    #[error("{0}")]
43    Source(#[source] Box<dyn std::error::Error + Send + Sync>),
44}
45
46/// Object-safe async port for acquiring a fresh `(token, ttl)` pair.
47///
48/// Implement this for custom token acquisition strategies. The built-in
49/// implementation is [`ClientCredentialsSource`].
50#[async_trait]
51pub trait TokenSource: Send + Sync {
52    /// Fetch a fresh token. Returns the raw JWT string and its lifetime.
53    /// The cache subtracts [`TokenCacheConfig::refresh_skew`] from the TTL
54    /// before storing the expiry, so the token is treated as stale slightly
55    /// before the server considers it expired.
56    async fn fetch_token(&self) -> Result<(String, Duration), TokenCacheError>;
57}
58
59/// Configuration for [`TokenCache`].
60#[derive(Debug, Clone)]
61pub struct TokenCacheConfig {
62    /// How far before a token's actual expiry to treat it as stale.
63    /// Defaults to 60 s — guards against clock skew and network latency on
64    /// the refresh call.
65    pub refresh_skew: Duration,
66}
67
68impl Default for TokenCacheConfig {
69    fn default() -> Self {
70        // Single skew SSOT: the same 60 s constant drives the pure
71        // `is_expired` helper and the cache's stored-expiry discount.
72        #[allow(clippy::cast_sign_loss)] // compile-time positive constant
73        Self { refresh_skew: Duration::from_secs(jwt_expiry::EXPIRY_SKEW_SECS as u64) }
74    }
75}
76
77struct CacheInner {
78    cached: Option<(String, i64)>,
79}
80
81/// Thread-safe JWT cache with near-expiry refresh and single-flight semantics.
82///
83/// Wrap in `Arc<TokenCache>` and share across tasks.
84pub struct TokenCache {
85    data: Mutex<CacheInner>,
86    /// Serialises concurrent refreshes (single-flight pattern).
87    /// Only the task that acquires this lock performs the HTTP round-trip;
88    /// the rest wait, then double-check the cache before deciding to fetch.
89    refresh: Mutex<()>,
90    config: TokenCacheConfig,
91    source: Box<dyn TokenSource>,
92    clock: ArcClock,
93}
94
95impl TokenCache {
96    /// Create a new cache backed by `source`.
97    pub fn new(source: Box<dyn TokenSource>, config: TokenCacheConfig) -> Self {
98        Self {
99            data: Mutex::new(CacheInner { cached: None }),
100            refresh: Mutex::new(()),
101            config,
102            source,
103            clock: Arc::new(WallClock),
104        }
105    }
106
107    #[must_use]
108    pub fn with_clock(mut self, clock: ArcClock) -> Self {
109        self.clock = clock;
110        self
111    }
112
113    /// Seed the cache with an already-held token (e.g. loaded from disk),
114    /// builder-style, before the cache is shared.
115    ///
116    /// The stored expiry is derived from the token's JWT `exp` claim (via
117    /// [`parse_expiry`]) minus [`TokenCacheConfig::refresh_skew`] — the same
118    /// discount the fetch path applies. Tokens without a parseable `exp`
119    /// are ignored, and an already-stale seed is filtered by `get()`'s
120    /// freshness check — either way the first [`get`](Self::get) falls
121    /// through to the source, which is the safe direction.
122    ///
123    /// This exists so a seeded cache never confuses "still-valid stored
124    /// token" with "freshly minted token": seeding happens exactly once at
125    /// construction, so [`invalidate`](Self::invalidate) always forces a
126    /// genuine re-mint (a source that re-served stored state would hand the
127    /// same dead token back after a 401).
128    #[must_use]
129    pub fn with_initial_token(mut self, token: impl Into<String>) -> Self {
130        let token = token.into();
131        if let Some(exp) = jwt_expiry::parse_expiry(&token) {
132            let skew_ms = i64::try_from(self.config.refresh_skew.as_millis()).unwrap_or(i64::MAX);
133            let exp_ms = exp.as_millisecond().saturating_sub(skew_ms);
134            self.data.get_mut().cached = Some((token, exp_ms));
135        }
136        self
137    }
138
139    /// Return the current token, refreshing via the source if near-expired.
140    pub async fn get(&self) -> Result<String, TokenCacheError> {
141        // Fast path: cache is fresh.
142        {
143            let inner = self.data.lock().await;
144            if let Some((token, exp)) = &inner.cached
145                && self.clock.now_unix_millis() < *exp {
146                    return Ok(token.clone());
147                }
148        }
149
150        // Slow path: acquire the refresh lock (serialises concurrent refreshes).
151        let _refresh_guard = self.refresh.lock().await;
152
153        // Double-check: a previous holder may have already refreshed.
154        {
155            let inner = self.data.lock().await;
156            if let Some((token, exp)) = &inner.cached
157                && self.clock.now_unix_millis() < *exp {
158                    return Ok(token.clone());
159                }
160        }
161
162        // We are the refresh leader. Fetch a fresh token.
163        let (token, ttl) = self.source.fetch_token().await?;
164        let skewed_ttl = ttl.saturating_sub(self.config.refresh_skew);
165        let exp = self.clock.now_unix_millis() + skewed_ttl.as_millis() as i64;
166        self.data.lock().await.cached = Some((token.clone(), exp));
167
168        Ok(token)
169    }
170
171    /// Discard the cached token so the next [`get`](Self::get) re-mints.
172    ///
173    /// Call this when the server rejects the current token (a 401 /
174    /// `Unauthenticated`): the token is *dead* — revoked, session-epoch-bumped,
175    /// or rotated out from under a still-valid clock TTL — so clock-based expiry
176    /// is the wrong staleness signal. Forcing a fetch lets a caller recover on
177    /// its next request instead of waiting out the remaining TTL.
178    pub async fn invalidate(&self) {
179        self.data.lock().await.cached = None;
180    }
181}
182
183#[cfg(feature = "client-credentials")]
184mod credentials;
185mod jwt_expiry;
186
187#[cfg(feature = "client-credentials")]
188pub use credentials::ClientCredentialsSource;
189pub use jwt_expiry::{EXPIRY_SKEW_SECS, is_expired, parse_expiry};
190
191#[cfg(test)]
192#[allow(clippy::unwrap_used, clippy::expect_used)]
193mod tests {
194    use super::*;
195    use std::sync::{
196        Arc,
197        atomic::{AtomicUsize, Ordering},
198    };
199
200    struct CountingSource {
201        count: Arc<AtomicUsize>,
202        ttl: Duration,
203        token: &'static str,
204    }
205
206    #[async_trait]
207    impl TokenSource for CountingSource {
208        async fn fetch_token(&self) -> Result<(String, Duration), TokenCacheError> {
209            self.count.fetch_add(1, Ordering::SeqCst);
210            // Simulate a brief network hop so concurrent callers actually
211            // overlap on the refresh lock.
212            tokio::time::sleep(Duration::from_millis(10)).await;
213            Ok((self.token.to_string(), self.ttl))
214        }
215    }
216
217    #[tokio::test]
218    async fn token_cache_returns_cached_until_near_expiry() {
219        let count = Arc::new(AtomicUsize::new(0));
220        let source = CountingSource {
221            count: Arc::clone(&count),
222            ttl: Duration::from_secs(3600),
223            token: "tok-abc",
224        };
225        let cache = TokenCache::new(Box::new(source), TokenCacheConfig::default());
226
227        let t1 = cache.get().await.unwrap();
228        let t2 = cache.get().await.unwrap();
229
230        assert_eq!(t1, "tok-abc");
231        assert_eq!(t2, "tok-abc");
232        assert_eq!(count.load(Ordering::SeqCst), 1, "source called more than once for valid cache");
233    }
234
235    #[tokio::test]
236    async fn invalidate_forces_refetch_despite_valid_ttl() {
237        let count = Arc::new(AtomicUsize::new(0));
238        let source = CountingSource {
239            count: Arc::clone(&count),
240            ttl: Duration::from_secs(3600), // long TTL: clock says "still fresh"
241            token: "tok-1",
242        };
243        let cache = TokenCache::new(Box::new(source), TokenCacheConfig::default());
244
245        let _ = cache.get().await.unwrap(); // fetch #1, cached
246        let _ = cache.get().await.unwrap(); // served from cache
247        assert_eq!(count.load(Ordering::SeqCst), 1, "second get should hit cache");
248
249        // The server rejected this token — it is dead despite a valid clock TTL.
250        cache.invalidate().await;
251
252        let _ = cache.get().await.unwrap(); // must re-fetch, not serve the dead token
253        assert_eq!(
254            count.load(Ordering::SeqCst),
255            2,
256            "invalidate() must force the next get() to re-mint",
257        );
258    }
259
260    fn fake_jwt(exp: jiff::Timestamp) -> String {
261        jwt_expiry::tests::make_fake_token(exp)
262    }
263
264    #[tokio::test]
265    async fn with_initial_token_serves_seed_without_fetch() {
266        let count = Arc::new(AtomicUsize::new(0));
267        let source = CountingSource {
268            count: Arc::clone(&count),
269            ttl: Duration::from_secs(3600),
270            token: "tok-fresh",
271        };
272        let seed = fake_jwt(jiff::Timestamp::now() + jiff::SignedDuration::from_hours(1));
273        let cache = TokenCache::new(Box::new(source), TokenCacheConfig::default())
274            .with_initial_token(seed.clone());
275
276        assert_eq!(cache.get().await.unwrap(), seed);
277        assert_eq!(count.load(Ordering::SeqCst), 0, "a valid seed must not trigger a fetch");
278
279        // A 401-invalidate must force a genuine re-mint, never re-serve the seed.
280        cache.invalidate().await;
281        assert_eq!(cache.get().await.unwrap(), "tok-fresh");
282        assert_eq!(count.load(Ordering::SeqCst), 1);
283    }
284
285    #[tokio::test]
286    async fn with_initial_token_ignores_unparseable_and_stale_seeds() {
287        for seed in [
288            "not-a-jwt".to_string(),
289            fake_jwt(jiff::Timestamp::now() - jiff::SignedDuration::from_hours(1)),
290        ] {
291            let count = Arc::new(AtomicUsize::new(0));
292            let source = CountingSource {
293                count: Arc::clone(&count),
294                ttl: Duration::from_secs(3600),
295                token: "tok-fresh",
296            };
297            let cache = TokenCache::new(Box::new(source), TokenCacheConfig::default())
298                .with_initial_token(seed);
299
300            assert_eq!(cache.get().await.unwrap(), "tok-fresh");
301            assert_eq!(count.load(Ordering::SeqCst), 1, "unusable seed must fall through to fetch");
302        }
303    }
304
305    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
306    async fn token_cache_single_flights_concurrent_refresh() {
307        let count = Arc::new(AtomicUsize::new(0));
308        let source = CountingSource {
309            count: Arc::clone(&count),
310            ttl: Duration::from_secs(3600),
311            token: "tok-xyz",
312        };
313        let cache = Arc::new(TokenCache::new(Box::new(source), TokenCacheConfig::default()));
314
315        let mut handles = Vec::new();
316        for _ in 0..8 {
317            let cache = Arc::clone(&cache);
318            handles.push(tokio::spawn(async move { cache.get().await }));
319        }
320        for h in handles {
321            assert_eq!(h.await.unwrap().unwrap(), "tok-xyz");
322        }
323
324        // All 8 concurrent callers must share one fetch (single-flight).
325        assert_eq!(count.load(Ordering::SeqCst), 1, "source called more than once (single-flight broken)");
326    }
327}