From 0e62453b927b04c68a8e23f43e46068174e75dec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Danilo=20Trnini=C4=87?= Date: Fri, 7 Aug 2026 17:05:01 +0000 Subject: [PATCH 1/2] [Rust] Harden OAuth token caching MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Parse the OAuth `expires_in` field from a quoted integer ("3600") in addition to a plain JSON integer, so a token endpoint returning it as a JSON string no longer drops the TTL and re-mints on every stream creation. Add deterministic TokenCache tests covering concurrency, invalidation, cancellation, and expiry. Fixes #607. Signed-off-by: Danilo Trninić --- rust/NEXT_CHANGELOG.md | 3 + rust/sdk/src/default_token_factory.rs | 55 ++++-- rust/sdk/src/token_cache.rs | 275 +++++++++++++++++++++++++- 3 files changed, 314 insertions(+), 19 deletions(-) diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index bff28992..22b806c5 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -22,6 +22,9 @@ only after pending replay succeeds, while initial supervisor handoff and failed or cancelled replay promptly drop redundant senders instead of retaining incomplete `DoPut` request channels until later teardown. +- The OAuth `expires_in` field is now parsed from a quoted integer (`"3600"`) in + addition to a plain JSON integer. A value that is missing or does not represent + a positive integer still yields no token lifetime, as before. ### Documentation diff --git a/rust/sdk/src/default_token_factory.rs b/rust/sdk/src/default_token_factory.rs index 0b178cd7..60ecca2d 100644 --- a/rust/sdk/src/default_token_factory.rs +++ b/rust/sdk/src/default_token_factory.rs @@ -239,13 +239,23 @@ impl DefaultTokenFactory { } /// Parses the OAuth `expires_in` field (token lifetime in seconds) into a - /// `Duration`. It is optional in the OAuth spec; if it is missing or not a - /// positive integer the token has no known TTL and must not be cached. + /// `Duration`. A plain JSON integer (`3600`) and a quoted one (`"3600"`) are + /// both accepted. + /// + /// `expires_in` is optional in the OAuth spec; a missing value, or one that + /// is not a positive integer, yields `None`. fn parse_expires_in(body: &serde_json::Value) -> Option { - body["expires_in"] - .as_u64() - .filter(|secs| *secs > 0) - .map(Duration::from_secs) + let secs = match &body["expires_in"] { + serde_json::Value::Number(n) => n.as_u64(), + serde_json::Value::String(s) => s.trim().parse::().ok(), + _ => None, + }?; + + if secs == 0 { + return None; + } + + Some(Duration::from_secs(secs)) } /// Classifies HTTP status codes as retryable or non-retryable errors. @@ -355,20 +365,43 @@ mod tests { #[test] fn test_parse_expires_in() { - let with_ttl = serde_json::json!({ "expires_in": 3600 }); + // A JSON integer parses to that many seconds. + let integer = serde_json::json!({ "expires_in": 3600 }); + assert_eq!( + DefaultTokenFactory::parse_expires_in(&integer), + Some(Duration::from_secs(3600)) + ); + + // A quoted integer parses to the same value. + let quoted = serde_json::json!({ "expires_in": "3600" }); assert_eq!( - DefaultTokenFactory::parse_expires_in(&with_ttl), + DefaultTokenFactory::parse_expires_in("ed), Some(Duration::from_secs(3600)) ); - let missing = serde_json::json!({ "access_token": "abc" }); + // Surrounding whitespace in the string is trimmed. + let padded = serde_json::json!({ "expires_in": " 3600 " }); + assert_eq!( + DefaultTokenFactory::parse_expires_in(&padded), + Some(Duration::from_secs(3600)) + ); + + // Absent, zero, and negative all yield no TTL. + let missing = serde_json::json!({}); assert_eq!(DefaultTokenFactory::parse_expires_in(&missing), None); let zero = serde_json::json!({ "expires_in": 0 }); assert_eq!(DefaultTokenFactory::parse_expires_in(&zero), None); - // A string value (non-integer) is not usable and yields no TTL. - let non_numeric = serde_json::json!({ "expires_in": "3600" }); + let negative = serde_json::json!({ "expires_in": -1 }); + assert_eq!(DefaultTokenFactory::parse_expires_in(&negative), None); + + // A fractional value is not a whole number of seconds. + let fractional = serde_json::json!({ "expires_in": 3600.9 }); + assert_eq!(DefaultTokenFactory::parse_expires_in(&fractional), None); + + // A non-numeric string yields no TTL. + let non_numeric = serde_json::json!({ "expires_in": "abc" }); assert_eq!(DefaultTokenFactory::parse_expires_in(&non_numeric), None); } diff --git a/rust/sdk/src/token_cache.rs b/rust/sdk/src/token_cache.rs index 8b99ce67..30859e07 100644 --- a/rust/sdk/src/token_cache.rs +++ b/rust/sdk/src/token_cache.rs @@ -285,6 +285,47 @@ mod tests { assert_eq!(calls.load(Ordering::SeqCst), 2); } + #[tokio::test] + async fn refresh_installs_new_expiry() { + // A proactive refresh must re-stabilize the cache: once it returns a + // token with a healthy TTL, the following call should hit rather than + // refresh again. The first mint uses a within-buffer TTL (30s < 60s + // buffer) to force one refresh; later mints return a healthy TTL. + let cache = TokenCache::new(true, Duration::from_secs(60)); + let calls = AtomicUsize::new(0); + + let make = |_reason| async { + let n = calls.fetch_add(1, Ordering::SeqCst); + let ttl = if n == 0 { 30 } else { 3600 }; + Ok(fetched(&format!("tok{n}"), Some(ttl))) + }; + + // Call 1 mints tok0 (within-buffer, immediately due for refresh). + let a = cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + // Call 2 refreshes to tok1 with a healthy TTL. + let b = cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + // Call 3 must be a cache hit on tok1: no further mint. + let c = cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + + assert_eq!(a, "tok0"); + assert_eq!(b, "tok1"); + assert_eq!(c, "tok1", "the refreshed token should be served from cache"); + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "refresh should install a new expiry so the third call hits cache" + ); + } + #[tokio::test] async fn separate_tables_get_separate_entries() { let cache = TokenCache::new(true, Duration::from_secs(60)); @@ -377,6 +418,96 @@ mod tests { assert_eq!(calls.load(Ordering::SeqCst), 2); } + #[tokio::test] + async fn invalidate_affects_only_its_own_key() { + let cache = TokenCache::new(true, Duration::from_secs(60)); + let calls = AtomicUsize::new(0); + + let make = |_reason| async { + let n = calls.fetch_add(1, Ordering::SeqCst); + Ok(fetched(&format!("tok{n}"), Some(3600))) + }; + + // Seed two different tables (tok0 and tok1). + cache + .get_or_fetch("id", "secret", "c.s.t1", make) + .await + .unwrap(); + cache + .get_or_fetch("id", "secret", "c.s.t2", make) + .await + .unwrap(); + + // Invalidating t1 must not disturb t2. + cache.invalidate("id", "secret", "c.s.t1").await; + + // t1 re-mints (tok2); t2 still hits its original cached token (tok1). + let t1 = cache + .get_or_fetch("id", "secret", "c.s.t1", make) + .await + .unwrap(); + let t2 = cache + .get_or_fetch("id", "secret", "c.s.t2", make) + .await + .unwrap(); + + assert_eq!(t1, "tok2", "invalidated table should re-mint"); + assert_eq!(t2, "tok1", "untouched table should still hit cache"); + assert_eq!(calls.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn invalidate_unknown_key_is_a_noop() { + let cache = TokenCache::new(true, Duration::from_secs(60)); + let calls = AtomicUsize::new(0); + + let make = |_reason| async { + calls.fetch_add(1, Ordering::SeqCst); + Ok(fetched("tok", Some(3600))) + }; + + cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + + // Invalidating a key that was never cached must leave the existing + // entry intact, so the next call still hits. + cache.invalidate("id", "secret", "other.table.here").await; + cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "invalidating an unknown key must not evict the cached token" + ); + } + + #[tokio::test] + async fn invalidate_on_disabled_cache_is_a_noop() { + // A disabled cache never stores anything, so invalidate has nothing to + // do; it must simply not panic, and fetching must keep working. + let cache = TokenCache::new(false, Duration::from_secs(60)); + let calls = AtomicUsize::new(0); + + let make = |_reason| async { + calls.fetch_add(1, Ordering::SeqCst); + Ok(fetched("tok", Some(3600))) + }; + + cache.invalidate("id", "secret", "c.s.t").await; + let token = cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + + assert_eq!(token, "tok"); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + #[tokio::test] async fn disabled_cache_always_fetches() { let cache = TokenCache::new(false, Duration::from_secs(60)); @@ -445,6 +576,35 @@ mod tests { assert_eq!(served, "valid"); } + #[tokio::test] + async fn refresh_failure_does_not_serve_expired_token() { + let cache = TokenCache::new(true, Duration::from_secs(60)); + + // Seed a token with a zero TTL: `expires_at` becomes the mint instant. By + // the second await below the monotonic clock has reached or passed it, and + // `is_expired` (`Instant::now() >= expires_at`) treats equality as expired, + // so the token reads as expired. + cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Ok(fetched("stale", Some(0))) + }) + .await + .unwrap(); + + // A retryable refresh failure would serve a still-valid cached token, but + // this one has expired, so the error must surface rather than handing the + // caller a dead token. + let result = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Err(crate::ZerobusError::TokenFetchError("blip".to_string())) + }) + .await; + assert!(matches!( + result, + Err(crate::ZerobusError::TokenFetchError(_)) + )); + } + #[tokio::test] async fn refresh_failure_propagates_non_retryable_error() { let cache = TokenCache::new(true, Duration::from_secs(60)); @@ -505,22 +665,52 @@ mod tests { assert_eq!(served, "valid"); } - #[tokio::test] + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn single_flight_mints_once_for_concurrent_callers() { let cache = Arc::new(TokenCache::new(true, Duration::from_secs(60))); let calls = Arc::new(AtomicUsize::new(0)); + const FOLLOWERS: usize = 15; - let mut handles = Vec::new(); - for _ in 0..16 { + // Keep the leader's mint in flight (blocked on `gate`) while the + // followers pile in. + let gate = Arc::new(tokio::sync::Notify::new()); + let (queued_tx, mut queued_rx) = tokio::sync::mpsc::unbounded_channel(); + + // Leader: occupies the slot and blocks inside the mint on `gate`. + let leader = { let cache = Arc::clone(&cache); let calls = Arc::clone(&calls); - handles.push(tokio::spawn(async move { + let gate = Arc::clone(&gate); + tokio::spawn(async move { + cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async move { + calls.fetch_add(1, Ordering::SeqCst); + gate.notified().await; + Ok(fetched("tok", Some(3600))) + }) + .await + .unwrap() + }) + }; + + // Wait until the leader is inside the mint (one call recorded) before + // launching followers, so they cannot win the slot first. + while calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + + // Followers: each signals that it has started, then calls get_or_fetch + // and contends for the same per-entry lock the leader holds. + let mut followers = Vec::new(); + for _ in 0..FOLLOWERS { + let cache = Arc::clone(&cache); + let calls = Arc::clone(&calls); + let queued_tx = queued_tx.clone(); + followers.push(tokio::spawn(async move { + queued_tx.send(()).unwrap(); cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { calls.fetch_add(1, Ordering::SeqCst); - // Hold the slot briefly so the other callers pile up - // behind the single-flight lock rather than racing. - tokio::time::sleep(Duration::from_millis(20)).await; Ok(fetched("tok", Some(3600))) }) .await @@ -528,7 +718,15 @@ mod tests { })); } - for handle in handles { + // Once all followers report they have started, release the leader's mint + // so it caches the single token. + for _ in 0..FOLLOWERS { + queued_rx.recv().await.unwrap(); + } + gate.notify_one(); + + assert_eq!(leader.await.unwrap(), "tok"); + for handle in followers { assert_eq!(handle.await.unwrap(), "tok"); } assert_eq!( @@ -537,4 +735,65 @@ mod tests { "single-flight must mint exactly once for concurrent same-key callers" ); } + + #[tokio::test] + async fn cancelled_mint_leaves_cache_usable() { + let cache = Arc::new(TokenCache::new(true, Duration::from_secs(60))); + let calls = Arc::new(AtomicUsize::new(0)); + + // Signals that the leader has entered the mint (and so is holding the + // per-entry lock) so we can cancel it at a known point, without relying + // on wall-clock timing. + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + + let task = { + let cache = Arc::clone(&cache); + let calls = Arc::clone(&calls); + tokio::spawn(async move { + cache + .get_or_fetch("id", "secret", "c.s.t", move |_reason| async move { + calls.fetch_add(1, Ordering::SeqCst); + let _ = started_tx.send(()); + // Never completes: the task is aborted while awaiting + // here, dropping the get_or_fetch future mid-mint. + std::future::pending::>().await + }) + .await + }) + }; + + // Wait until the mint is in flight, then cancel it. Awaiting the aborted + // task guarantees its future (and the slot guard) has been dropped. + started_rx.await.unwrap(); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + + // The cancelled leader must have released the lock and left no + // half-written entry, so the next caller mints cleanly... + let minted = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + calls.fetch_add(1, Ordering::SeqCst); + Ok(fetched("tok", Some(3600))) + }) + .await + .unwrap(); + assert_eq!(minted, "tok"); + + // ...and that freshly minted token is cached, not a phantom entry: a + // follow-up call hits without minting again. + let cached = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + calls.fetch_add(1, Ordering::SeqCst); + Ok(fetched("other", Some(3600))) + }) + .await + .unwrap(); + assert_eq!(cached, "tok"); + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "one aborted mint plus one real mint; the final call must hit cache" + ); + } } From 77565c08e7615b63a66b5d55f984816e9979b8e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Danilo=20Trnini=C4=87?= Date: Mon, 10 Aug 2026 08:42:39 +0000 Subject: [PATCH 2/2] [Rust] Format default token factory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Danilo Trninić --- rust/sdk/src/default_token_factory.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/sdk/src/default_token_factory.rs b/rust/sdk/src/default_token_factory.rs index 60ecca2d..f4f599d1 100644 --- a/rust/sdk/src/default_token_factory.rs +++ b/rust/sdk/src/default_token_factory.rs @@ -254,7 +254,7 @@ impl DefaultTokenFactory { if secs == 0 { return None; } - + Some(Duration::from_secs(secs)) }