ppoppo_sdk_core/token_cache/
mod.rs1use 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#[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 #[error("{0}")]
43 Source(#[source] Box<dyn std::error::Error + Send + Sync>),
44}
45
46#[async_trait]
51pub trait TokenSource: Send + Sync {
52 async fn fetch_token(&self) -> Result<(String, Duration), TokenCacheError>;
57}
58
59#[derive(Debug, Clone)]
61pub struct TokenCacheConfig {
62 pub refresh_skew: Duration,
66}
67
68impl Default for TokenCacheConfig {
69 fn default() -> Self {
70 #[allow(clippy::cast_sign_loss)] 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
81pub struct TokenCache {
85 data: Mutex<CacheInner>,
86 refresh: Mutex<()>,
90 config: TokenCacheConfig,
91 source: Box<dyn TokenSource>,
92 clock: ArcClock,
93}
94
95impl TokenCache {
96 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 #[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 pub async fn get(&self) -> Result<String, TokenCacheError> {
141 {
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 let _refresh_guard = self.refresh.lock().await;
152
153 {
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 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 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 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), token: "tok-1",
242 };
243 let cache = TokenCache::new(Box::new(source), TokenCacheConfig::default());
244
245 let _ = cache.get().await.unwrap(); let _ = cache.get().await.unwrap(); assert_eq!(count.load(Ordering::SeqCst), 1, "second get should hit cache");
248
249 cache.invalidate().await;
251
252 let _ = cache.get().await.unwrap(); 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 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 assert_eq!(count.load(Ordering::SeqCst), 1, "source called more than once (single-flight broken)");
326 }
327}