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/4] [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/4] [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)) } From 4b16b8338e90315b72f63fc9bb8370f8e7100970 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Danilo=20Trnini=C4=87?= Date: Mon, 17 Aug 2026 10:41:35 +0000 Subject: [PATCH 3/4] [Rust] Harden OAuth token refresh handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the OAuth token cache's proactive refresh robust so it never fails or hangs stream creation while a usable token is still cached: - On any refresh error (including non-retryable ones such as revoked credentials), serve the still-valid cached token; surface the error only when there is nothing valid to fall back to. - Anchor token lifetime to request start, not response arrival, so a slow response can't extend perceived validity. Treat a token returned already past its start-anchored expiry as dead on arrival. - Bound a proactive refresh at half the stream's recovery_timeout_ms, but only while the cached token outlives that cap; otherwise, and on a cold miss, run the mint unbounded. - Pace post-fallback re-attempts with a refresh-buffer-scaled backoff that shrinks near expiry, preventing a token-endpoint mint stampede. Add deterministic TokenCache tests (Tokio virtual time) for these paths. Signed-off-by: Danilo Trninić --- rust/NEXT_CHANGELOG.md | 23 ++ rust/README.md | 2 + rust/sdk/Cargo.toml | 2 + rust/sdk/src/builder/stream_builder.rs | 35 +- rust/sdk/src/headers_provider.rs | 46 ++- rust/sdk/src/stream/arrow/options.rs | 6 + rust/sdk/src/stream_configuration.rs | 6 + rust/sdk/src/token_cache.rs | 546 +++++++++++++++++++++---- 8 files changed, 564 insertions(+), 102 deletions(-) diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index 22b806c5..c7fe8025 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -25,6 +25,29 @@ - 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. +- OAuth token lifetime is now anchored to when the token request starts rather + than when the response arrives, so a slow response no longer makes the SDK treat + a token as valid longer than the issuer does. A token already past its + (start-anchored) expiry is dead on arrival and, on a cold miss, surfaces a + retryable `TokenFetchError`. +- A stalled proactive token refresh no longer hangs stream creation. For streams + built with `.oauth(...)`, the refresh is capped at half the stream's configured + setup budget (`recovery_timeout_ms`) — but only when the cached token has more + life left than that cap, so a hung endpoint becomes a prompt failure and the + still-valid cached token is served in time. When too little of the token's life + remains for that fallback to help, the refresh runs unbounded (like a cold miss) + so a slow-but-working mint isn't cut off early. +- When a proactive token refresh fails — for any error, including a non-retryable + one such as revoked credentials — the still-valid cached token is served rather + than failing the caller. The token was validly issued and hasn't expired, and the + server re-validates it on every connection, so the error surfaces only when there + is no still-valid token to fall back to. A refresh that returns a dead-on-arrival + token falls back the same way. +- After a proactive refresh falls back to the cached token, a short backoff paces + further attempts so a burst of stream creations cannot turn one token-endpoint + failure into a mint stampede. The interval scales with the configured token + refresh buffer and shrinks as the token nears expiry, down to a floor, and never + extends past expiry. ### Documentation diff --git a/rust/README.md b/rust/README.md index f5f2fc0e..561bf948 100644 --- a/rust/README.md +++ b/rust/README.md @@ -325,6 +325,8 @@ let sdk = ZerobusSdk::builder() # Ok::<(), databricks_zerobus_ingest_sdk::ZerobusError>(()) ``` +When a proactive refresh fails, the SDK keeps serving the still-valid cached token rather than failing stream creation, including on non-retryable failures such as revoked credentials. The cached token is served until it actually expires, after which the next call mints a fresh one and any failure then surfaces. A proactive refresh is also capped at half the stream's [`recovery_timeout_ms`](#configuration-options), but only when the cached token has more life left than that cap, so a stalled endpoint falls back to the cached token in time. When too little of the token's life remains, the refresh runs unbounded instead, like a cold miss. + ### Custom Authentication For advanced use cases, you can implement the `HeadersProvider` trait to supply your own authentication headers. This is useful for integrating with a different OAuth provider, using a centralized token caching service, or implementing alternative authentication mechanisms. diff --git a/rust/sdk/Cargo.toml b/rust/sdk/Cargo.toml index b5897888..0affbe9d 100644 --- a/rust/sdk/Cargo.toml +++ b/rust/sdk/Cargo.toml @@ -51,6 +51,8 @@ futures = { workspace = true, optional = true } [dev-dependencies] tracing-subscriber.workspace = true +# Deterministic virtual time (start_paused / advance) for timing-sensitive tests. +tokio = { workspace = true, features = ["test-util"] } # zeroparser e2e + bench dev-deps criterion = { version = "0.5", default-features = false, features = ["html_reports"] } plotters = { version = "0.3", default-features = false, features = ["svg_backend"] } diff --git a/rust/sdk/src/builder/stream_builder.rs b/rust/sdk/src/builder/stream_builder.rs index 5e406870..36f59180 100644 --- a/rust/sdk/src/builder/stream_builder.rs +++ b/rust/sdk/src/builder/stream_builder.rs @@ -231,6 +231,11 @@ impl<'a> StreamBuilder<'a> { } /// Set the timeout in milliseconds for each recovery attempt. + /// + /// For streams authenticated with [`oauth`](Self::oauth), this also caps a + /// proactive OAuth token refresh at half this value, but only when the cached + /// token has more life left than that cap, so a stalled endpoint falls back + /// before the attempt deadline. Very near expiry the refresh runs unbounded. pub fn recovery_timeout_ms(mut self, ms: u64) -> Self { self.grpc_config.recovery_timeout_ms = ms; #[cfg(feature = "arrow-flight")] @@ -383,20 +388,32 @@ impl<'a> StreamBuilder<'a> { Ok(()) } - /// Resolve the headers provider from the stored auth config. + /// Resolve the headers provider from the stored auth config. For OAuth, a + /// proactive token refresh is bounded by half the stream's recovery timeout. fn resolve_headers_provider(&self) -> ZerobusResult> { match self.auth.as_ref() { Some(AuthConfig::OAuth { client_id, client_secret, - }) => Ok(Arc::new(OAuthHeadersProvider::with_cache( - client_id.clone(), - client_secret.clone(), - self.table_name.clone(), - self.sdk.workspace_id.clone(), - self.sdk.unity_catalog_url.clone(), - Arc::clone(&self.sdk.token_cache), - ))), + }) => { + // Give a proactive refresh half the recovery timeout so a stalled + // refresh falls back to the cached token before the setup deadline + // cancels the request. Both recovery-timeout setters write the gRPC + // and Arrow configs in lockstep, so grpc_config is a valid single + // source even for an Arrow stream; revisit if a per-transport + // recovery timeout is ever added. + let refresh_timeout = + std::time::Duration::from_millis(self.grpc_config.recovery_timeout_ms) / 2; + Ok(Arc::new(OAuthHeadersProvider::with_cache( + client_id.clone(), + client_secret.clone(), + self.table_name.clone(), + self.sdk.workspace_id.clone(), + self.sdk.unity_catalog_url.clone(), + Arc::clone(&self.sdk.token_cache), + Some(refresh_timeout), + ))) + } Some(AuthConfig::HeadersProvider(p)) => Ok(Arc::clone(p)), #[cfg(feature = "testing")] Some(AuthConfig::NoAuth) => Ok(Arc::new(NoAuthHeadersProvider)), diff --git a/rust/sdk/src/headers_provider.rs b/rust/sdk/src/headers_provider.rs index 9c361f04..88b7c325 100644 --- a/rust/sdk/src/headers_provider.rs +++ b/rust/sdk/src/headers_provider.rs @@ -4,6 +4,7 @@ use crate::ZerobusResult; use async_trait::async_trait; use std::collections::HashMap; use std::sync::Arc; +use std::time::Duration; /// A trait for providing custom headers for gRPC requests. /// @@ -68,6 +69,9 @@ pub struct OAuthHeadersProvider { workspace_id: String, unity_catalog_url: String, token_cache: Arc, + /// How long a proactive refresh may run before it's treated as failed and the + /// cached token is served; `None` leaves it unbounded. + refresh_timeout: Option, } impl OAuthHeadersProvider { @@ -92,6 +96,7 @@ impl OAuthHeadersProvider { workspace_id, unity_catalog_url, Arc::new(TokenCache::new(true, DEFAULT_REFRESH_BUFFER)), + None, ) } @@ -106,6 +111,7 @@ impl OAuthHeadersProvider { workspace_id: String, unity_catalog_url: String, token_cache: Arc, + refresh_timeout: Option, ) -> Self { Self { client_id, @@ -114,6 +120,7 @@ impl OAuthHeadersProvider { workspace_id, unity_catalog_url, token_cache, + refresh_timeout, } } } @@ -121,24 +128,39 @@ impl OAuthHeadersProvider { #[async_trait] impl HeadersProvider for OAuthHeadersProvider { async fn get_headers(&self) -> ZerobusResult> { - let token = self - .token_cache - .get_or_fetch( + let fetch = |reason| { + DefaultTokenFactory::fetch_token( + &self.unity_catalog_url, + &self.table_name, &self.client_id, &self.client_secret, - &self.table_name, - |reason| { - DefaultTokenFactory::fetch_token( - &self.unity_catalog_url, + &self.workspace_id, + reason, + ) + }; + let token = match self.refresh_timeout { + Some(refresh_timeout) => { + self.token_cache + .get_or_fetch_within( + &self.client_id, + &self.client_secret, &self.table_name, + refresh_timeout, + fetch, + ) + .await? + } + None => { + self.token_cache + .get_or_fetch( &self.client_id, &self.client_secret, - &self.workspace_id, - reason, + &self.table_name, + fetch, ) - }, - ) - .await?; + .await? + } + }; let mut headers = HashMap::new(); headers.insert("authorization", format!("Bearer {}", token)); headers.insert("x-databricks-zerobus-table-name", self.table_name.clone()); diff --git a/rust/sdk/src/stream/arrow/options.rs b/rust/sdk/src/stream/arrow/options.rs index 49189e43..ed9603b4 100644 --- a/rust/sdk/src/stream/arrow/options.rs +++ b/rust/sdk/src/stream/arrow/options.rs @@ -50,6 +50,12 @@ pub struct ArrowStreamConfigurationOptions { /// Values whose absolute deadline cannot be represented by the platform's /// monotonic clock are rejected when the stream is built. /// + /// For OAuth-authenticated streams this also caps a proactive token refresh at + /// half its value, but only when the cached token has more life left than that + /// cap, so a stalled endpoint falls back to the cached token before the attempt + /// deadline. When too little of the token's life remains, the refresh runs + /// unbounded, like a cold miss. + /// /// Default: 15,000 (15 seconds) pub recovery_timeout_ms: u64, diff --git a/rust/sdk/src/stream_configuration.rs b/rust/sdk/src/stream_configuration.rs index 108d0674..62369f4c 100644 --- a/rust/sdk/src/stream_configuration.rs +++ b/rust/sdk/src/stream_configuration.rs @@ -47,6 +47,12 @@ pub struct StreamConfigurationOptions { /// /// If a recovery attempt takes longer than this, it will be retried. /// + /// For OAuth-authenticated streams this also caps a proactive token refresh at + /// half its value, but only when the cached token has more life left than that + /// cap, so a stalled endpoint falls back to the cached token before the attempt + /// deadline. When too little of the token's life remains, the refresh runs + /// unbounded, like a cold miss. + /// /// Default: 15,000 (15 seconds) pub recovery_timeout_ms: u64, diff --git a/rust/sdk/src/token_cache.rs b/rust/sdk/src/token_cache.rs index 30859e07..253a4647 100644 --- a/rust/sdk/src/token_cache.rs +++ b/rust/sdk/src/token_cache.rs @@ -23,21 +23,57 @@ use tokio::time::Instant; use tracing::{debug, warn}; use crate::default_token_factory::{FetchedToken, MintReason}; -use crate::ZerobusResult; +use crate::{ZerobusError, ZerobusResult}; /// Default lead time before expiry at which a cached token is refreshed. pub(crate) const DEFAULT_REFRESH_BUFFER: Duration = Duration::from_secs(300); +/// Post-failure backoff bounds, as fractions of the refresh window +/// (`refresh_buffer`) so they scale with it rather than being a fixed value. +/// `MAX_BACKOFF_WINDOW_FRACTION` gives the normal backoff, +/// `MIN_BACKOFF_WINDOW_FRACTION` a near-expiry floor that keeps the `remaining/2` +/// shrink from converging to zero. +const MAX_BACKOFF_WINDOW_FRACTION: u32 = 60; +const MIN_BACKOFF_WINDOW_FRACTION: u32 = 300; + /// A cached token and the instant at which it expires. struct CachedToken { value: String, expires_at: Instant, + /// Defer the next proactive refresh until this instant, set after a refresh + /// fell back to this token (a failed mint or a dead-on-arrival token). Always + /// `<= expires_at`, so it never serves an expired token. + refresh_retry_at: Option, } impl CachedToken { fn is_expired(&self) -> bool { Instant::now() >= self.expires_at } + + /// Whether the token is inside its post-fallback backoff window, during which + /// the next proactive refresh is deferred. + fn in_backoff_window(&self) -> bool { + self.refresh_retry_at + .is_some_and(|retry_at| Instant::now() < retry_at) + } + + /// Arms the post-fallback backoff so a burst of callers reuses this cached + /// token instead of each re-attempting a refresh that just fell back. The delay + /// is half the remaining validity, clamped to the window-scaled + /// `[min_backoff, max_backoff]` and capped at expiry. + fn arm_refresh_backoff(&mut self, refresh_buffer: Duration) { + let now = Instant::now(); + let remaining = self.expires_at.saturating_duration_since(now); + let max_backoff = refresh_buffer / MAX_BACKOFF_WINDOW_FRACTION; + let min_backoff = refresh_buffer / MIN_BACKOFF_WINDOW_FRACTION; + let backoff = (remaining / 2).clamp(min_backoff, max_backoff); + let retry_at = now + .checked_add(backoff) + .unwrap_or(self.expires_at) + .min(self.expires_at); + self.refresh_retry_at = Some(retry_at); + } } /// Identifies a cache entry. The client secret is keyed by its SHA-256 digest, @@ -85,12 +121,10 @@ impl TokenCache { } } - /// Returns a valid token for the given credentials and table, fetching a new - /// one only if the cache is empty, the token has entered the refresh window, - /// or caching is disabled. - /// - /// `fetch` is invoked to mint a fresh token. It is only ever called once per - /// key at a time thanks to the per-entry lock. + /// Returns a valid token, leaving a proactive refresh unbounded. Used where + /// there is no stream-setup budget to bound by, such as the standalone + /// `OAuthHeadersProvider::new`. Callers with a setup deadline use + /// [`get_or_fetch_within`](Self::get_or_fetch_within) instead. pub(crate) async fn get_or_fetch( &self, client_id: &str, @@ -98,6 +132,50 @@ impl TokenCache { table_name: &str, fetch: F, ) -> ZerobusResult + where + F: FnOnce(MintReason) -> Fut, + Fut: std::future::Future>, + { + self.get_or_fetch_bounded(client_id, client_secret, table_name, None, fetch) + .await + } + + /// Returns a valid token, bounding a proactive refresh at `refresh_timeout` so a + /// stall surfaces as a retryable error and the still-valid cached token is served + /// before the outer setup deadline. The bound is skipped when too little of the + /// token's life remains to fall back on, and on a cold miss. + pub(crate) async fn get_or_fetch_within( + &self, + client_id: &str, + client_secret: &str, + table_name: &str, + refresh_timeout: Duration, + fetch: F, + ) -> ZerobusResult + where + F: FnOnce(MintReason) -> Fut, + Fut: std::future::Future>, + { + self.get_or_fetch_bounded( + client_id, + client_secret, + table_name, + Some(refresh_timeout), + fetch, + ) + .await + } + + /// Shared implementation. `refresh_timeout` bounds a proactive-refresh mint + /// when `Some`; `None` leaves it unbounded. + async fn get_or_fetch_bounded( + &self, + client_id: &str, + client_secret: &str, + table_name: &str, + refresh_timeout: Option, + fetch: F, + ) -> ZerobusResult where F: FnOnce(MintReason) -> Fut, Fut: std::future::Future>, @@ -125,48 +203,98 @@ impl TokenCache { if let Some(cached) = guard.as_ref() { if !self.needs_refresh(cached) { - debug!(table = %table_name, "token cache hit, reusing cached token"); + // Distinguish a healthy hit from serving a token whose refresh is in + // post-fallback backoff, since they are operationally different states. + if cached.in_backoff_window() { + debug!(table = %table_name, "serving cached token; proactive refresh in backoff after a recent fallback"); + } else { + debug!(table = %table_name, "token cache hit, reusing cached token"); + } return Ok(cached.value.clone()); } } - // A present-but-stale token means we are refreshing; an empty slot is a - // cold miss. The reason is surfaced on the mint log. - let reason = if guard.is_some() { + // Anchor the lifetime to request start, not response arrival, so a slow + // response can't make the token look valid longer than the issuer allows. + let fetch_started_at = Instant::now(); + + // The cached token's remaining validity (`None` if there is no usable + // token), measured once from the request start so the mint reason and the + // refresh bound below share a single reading. + let cached_remaining = guard.as_ref().map(|cached| { + cached + .expires_at + .saturating_duration_since(fetch_started_at) + }); + + // A still-valid token in the refresh window is a proactive refresh; an empty + // or already-expired slot is a cold miss (`MintReason::ColdMiss`). Surfaced + // on the mint log. + let reason = if cached_remaining.is_some_and(|remaining| !remaining.is_zero()) { MintReason::Refresh } else { MintReason::ColdMiss }; - let fetched = match fetch(reason).await { + let fetch_result = match (reason, refresh_timeout) { + // Bound a proactive refresh only when the cached token outlasts the + // budget: a stall then surfaces as a retryable error and the fallback + // below serves the still-valid token before the setup deadline. Otherwise + // there's nothing to fall back on, so run unbounded like a cold miss. + (MintReason::Refresh, Some(budget)) + if cached_remaining.is_some_and(|remaining| remaining > budget) => + { + match tokio::time::timeout(budget, fetch(reason)).await { + Ok(result) => result, + Err(_) => Err(ZerobusError::TokenFetchError( + "proactive token refresh timed out".to_string(), + )), + } + } + _ => fetch(reason).await, + }; + + let fetched = match fetch_result { Ok(fetched) => fetched, Err(err) => { - // On a retryable failure, serve the still-valid cached token; - // let non-retryable errors (bad/revoked creds) surface. - if err.is_retryable() { - if let Some(cached) = guard.as_ref() { - if !cached.is_expired() { - warn!(table = %table_name, "token refresh failed (retryable); serving still-valid cached token"); - return Ok(cached.value.clone()); - } - } + // On any refresh error, serve the still-valid cached token if we have + // one, arming the backoff; otherwise surface the error. + if let Some(value) = + Self::serve_valid_cached_fallback(&mut guard, self.refresh_buffer) + { + warn!(table = %table_name, error = %err, "token refresh failed; serving still-valid cached token"); + return Ok(value); } return Err(err); } }; - let token = fetched.token.clone(); - // Cache only tokens with a usable TTL. `checked_add` also drops an absurd // `expires_in` that would overflow the clock instead of panicking. let expires_at = fetched .expires_in - .and_then(|ttl| Instant::now().checked_add(ttl)); + .and_then(|ttl| fetch_started_at.checked_add(ttl)); + + // A token already past its (start-anchored) expiry is dead on arrival: serve + // an older still-valid cached token if there is one and arm the backoff, + // otherwise surface a retryable error. + if expires_at.is_some_and(|deadline| deadline <= Instant::now()) { + if let Some(value) = Self::serve_valid_cached_fallback(&mut guard, self.refresh_buffer) + { + warn!(table = %table_name, "fetched OAuth token expired on arrival; serving still-valid cached token"); + return Ok(value); + } + return Err(ZerobusError::TokenFetchError( + "fetched OAuth token expired before arrival".to_string(), + )); + } + match expires_at { Some(expires_at) => { *guard = Some(CachedToken { - value: fetched.token, + value: fetched.token.clone(), expires_at, + refresh_retry_at: None, }); } None => { @@ -179,13 +307,13 @@ impl TokenCache { } } - Ok(token) + Ok(fetched.token) } /// Drops any cached token for the given credentials and table so the next - /// `get_or_fetch` re-mints. Called when the server rejects the token (e.g. - /// it was revoked at the IdP), so the re-mint re-checks grants at UC. No-op - /// when caching is disabled or no entry exists. + /// fetch re-mints. Called when the server rejects the token (e.g. it was + /// revoked at the IdP), so the re-mint re-checks grants at UC. No-op when + /// caching is disabled or no entry exists. pub(crate) async fn invalidate(&self, client_id: &str, client_secret: &str, table_name: &str) { if !self.enabled { return; @@ -197,6 +325,11 @@ impl TokenCache { } fn needs_refresh(&self, cached: &CachedToken) -> bool { + // Within a post-fallback backoff window, don't refresh yet (see + // `arm_refresh_backoff`). + if cached.in_backoff_window() { + return false; + } // `checked_add` avoids a panic on an absurd refresh buffer (e.g. // `Duration::MAX`); an overflowing deadline means "always refresh". match Instant::now().checked_add(self.refresh_buffer) { @@ -205,6 +338,21 @@ impl TokenCache { } } + /// If a still-valid token is cached, arm its backoff and return it. This is the + /// fallback when a proactive refresh can't produce a usable token; `None` when + /// there is none to fall back to. + fn serve_valid_cached_fallback( + slot: &mut Option, + refresh_buffer: Duration, + ) -> Option { + let cached = slot.as_mut()?; + if cached.is_expired() { + return None; + } + cached.arm_refresh_backoff(refresh_buffer); + Some(cached.value.clone()) + } + /// Drops entries whose token has fully expired. Locked (in-flight) entries, /// still-valid tokens, and empty slots are kept — keeping empty slots is /// what preserves single-flight for a key being minted concurrently. @@ -551,12 +699,12 @@ mod tests { assert_eq!(ok, "tok"); } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn refresh_failure_serves_still_valid_token() { + // A refresh failure serves the still-valid cached token regardless of error + // kind: this covers both a retryable and a non-retryable (revoked-creds) one. let cache = TokenCache::new(true, Duration::from_secs(60)); - // Seed a token that is within the refresh buffer (ttl < buffer) but not - // yet expired, so the next call is due for a refresh. let seeded = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Ok(fetched("valid", Some(30))) @@ -565,71 +713,289 @@ mod tests { .unwrap(); assert_eq!(seeded, "valid"); - // The refresh mint fails; the still-valid cached token is served instead - // of surfacing the error. - let served = cache + // Count attempts to prove each failing refresh actually ran (not suppressed). + let refresh_attempts = AtomicUsize::new(0); + + let served_retryable = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { - Err(crate::ZerobusError::TokenFetchError("blip".to_string())) + refresh_attempts.fetch_add(1, Ordering::SeqCst); + Err(ZerobusError::TokenFetchError("blip".to_string())) }) .await .unwrap(); - assert_eq!(served, "valid"); + assert_eq!(served_retryable, "valid"); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 1); + + // Clear the armed backoff so the next call refreshes again. + tokio::time::advance(Duration::from_secs(2)).await; + + let served_non_retryable = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + refresh_attempts.fetch_add(1, Ordering::SeqCst); + Err(ZerobusError::InvalidUCTokenError("revoked".to_string())) + }) + .await + .unwrap(); + assert_eq!(served_non_retryable, "valid"); + assert_eq!(refresh_attempts.load(Ordering::SeqCst), 2); } #[tokio::test] - async fn refresh_failure_does_not_serve_expired_token() { + async fn fetch_error_surfaces_unchanged_when_no_valid_token() { + // With nothing to fall back to, the fetch error surfaces unchanged, so its + // retryability is preserved for callers (both a retryable and a non-retryable). 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 + let retryable = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { - Ok(fetched("stale", Some(0))) + Err(ZerobusError::TokenFetchError("blip".to_string())) }) .await - .unwrap(); + .unwrap_err(); + assert!(matches!(retryable, ZerobusError::TokenFetchError(_))); + assert!(retryable.is_retryable()); - // 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 + let non_retryable = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { - Err(crate::ZerobusError::TokenFetchError("blip".to_string())) + Err(ZerobusError::InvalidUCTokenError("revoked".to_string())) }) - .await; + .await + .unwrap_err(); assert!(matches!( - result, - Err(crate::ZerobusError::TokenFetchError(_)) + non_retryable, + ZerobusError::InvalidUCTokenError(_) )); + assert!(!non_retryable.is_retryable()); } - #[tokio::test] - async fn refresh_failure_propagates_non_retryable_error() { + #[tokio::test(start_paused = true)] + async fn cold_miss_dead_on_arrival_token_surfaces_error() { let cache = TokenCache::new(true, Duration::from_secs(60)); - // Seed a token that is within the refresh buffer but not yet expired. - cache + // The fetch outlasts the token's TTL, so it returns already expired (dead on + // arrival). A cold miss has nothing to fall back to, so a retryable error surfaces. + let result = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + tokio::time::sleep(Duration::from_secs(2)).await; + Ok(fetched("doa", Some(1))) + }) + .await; + assert!(matches!(result, Err(ZerobusError::TokenFetchError(_)))); + } + + #[tokio::test(start_paused = true)] + async fn refresh_dead_on_arrival_keeps_cached_token() { + // A refresh returning a token already past its start-anchored expiry is dead + // on arrival, so the cached token is served and the backoff armed. + let cache = TokenCache::new(true, Duration::from_secs(60)); + let mints = AtomicUsize::new(0); + + let seeded = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); Ok(fetched("valid", Some(30))) }) .await .unwrap(); + assert_eq!(seeded, "valid"); - // A non-retryable refresh error (e.g. revoked or invalid credentials) - // must surface rather than being masked by the still-valid cached token. - let result = cache + let served = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { - Err(crate::ZerobusError::InvalidUCTokenError( - "revoked".to_string(), - )) + mints.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_secs(2)).await; + Ok(fetched("doa", Some(1))) }) - .await; - assert!(matches!( - result, - Err(crate::ZerobusError::InvalidUCTokenError(_)) - )); + .await + .unwrap(); + assert_eq!(served, "valid"); + + // The armed backoff suppresses the next refresh, so the mint count stays at 2. + let reused = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); + Ok(fetched("unexpected", Some(3600))) + }) + .await + .unwrap(); + assert_eq!(reused, "valid"); + assert_eq!(mints.load(Ordering::SeqCst), 2); + } + + #[tokio::test(start_paused = true)] + async fn expired_cached_entry_re_mints_as_cold_miss() { + // An expired cached entry is a cold miss, not a proactive refresh, so the + // re-mint runs unbounded: with a 1s budget, a 2s fetch still succeeds (a + // budget-capped refresh would time out and fail). + let cache = TokenCache::new(true, Duration::from_secs(60)); + + let seeded = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Ok(fetched("stale", Some(1))) + }) + .await + .unwrap(); + assert_eq!(seeded, "stale"); + tokio::time::advance(Duration::from_secs(2)).await; + + let minted = cache + .get_or_fetch_within( + "id", + "secret", + "c.s.t", + Duration::from_secs(1), + |_reason| async { + tokio::time::sleep(Duration::from_secs(2)).await; + Ok(fetched("fresh", Some(3600))) + }, + ) + .await + .unwrap(); + assert_eq!(minted, "fresh"); + } + + #[tokio::test(start_paused = true)] + async fn failed_refresh_backoff_suppresses_repeat_mint() { + // A failed refresh arms a backoff that suppresses the next refresh, so the + // cached token is reused without re-minting (the mint count stays at 2). + let cache = TokenCache::new(true, Duration::from_secs(60)); + let mints = AtomicUsize::new(0); + + let seeded = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); + Ok(fetched("valid", Some(30))) + }) + .await + .unwrap(); + assert_eq!(seeded, "valid"); + + let served = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); + Err(ZerobusError::TokenFetchError("blip".to_string())) + }) + .await + .unwrap(); + assert_eq!(served, "valid"); + + let reused = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); + Ok(fetched("unexpected", Some(3600))) + }) + .await + .unwrap(); + assert_eq!(reused, "valid"); + assert_eq!(mints.load(Ordering::SeqCst), 2); + } + + #[tokio::test(start_paused = true)] + async fn near_expiry_backoff_shrinks_to_retry_before_cold_miss() { + // With only 3s of token life left, a flat 5s backoff (the 300s buffer's cap) + // would suppress every retry until the token died. The remaining/2 backoff + // (1.5s) instead retries while the token is still valid. + let cache = TokenCache::new(true, Duration::from_secs(300)); + let mints = AtomicUsize::new(0); + + let seeded = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); + Ok(fetched("valid", Some(3))) + }) + .await + .unwrap(); + assert_eq!(seeded, "valid"); + + let served = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); + Err(ZerobusError::TokenFetchError("blip".to_string())) + }) + .await + .unwrap(); + assert_eq!(served, "valid"); + + // Still inside the ~1.5s backoff: reused, so the mint count stays at 2. + tokio::time::advance(Duration::from_secs(1)).await; + let reused = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); + Ok(fetched("unexpected", Some(3600))) + }) + .await + .unwrap(); + assert_eq!(reused, "valid"); + assert_eq!(mints.load(Ordering::SeqCst), 2); + + // Past the backoff and still valid: the proactive retry runs and succeeds. + tokio::time::advance(Duration::from_millis(600)).await; + let refreshed = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + mints.fetch_add(1, Ordering::SeqCst); + Ok(fetched("fresh", Some(3600))) + }) + .await + .unwrap(); + assert_eq!(refreshed, "fresh"); + assert_eq!(mints.load(Ordering::SeqCst), 3); + } + + #[tokio::test(start_paused = true)] + async fn stalled_refresh_serves_cached_token() { + // A hung refresh is cut off by the 1s budget (elapsed in virtual time), so it + // becomes a retryable error and the still-valid cached token is served. + let cache = TokenCache::new(true, Duration::from_secs(60)); + + let seeded = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Ok(fetched("valid", Some(30))) + }) + .await + .unwrap(); + assert_eq!(seeded, "valid"); + + let served = cache + .get_or_fetch_within( + "id", + "secret", + "c.s.t", + Duration::from_secs(1), + |_reason| async { std::future::pending::>().await }, + ) + .await + .unwrap(); + assert_eq!(served, "valid"); + } + + #[tokio::test(start_paused = true)] + async fn refresh_runs_unbounded_when_token_would_expire_before_budget() { + // The token has less life left (3s) than the refresh budget (5s), so bounding + // is pointless: it would expire before the budget fires, with nothing to fall + // back to. The refresh runs unbounded instead, so an 8s mint still succeeds. + let cache = TokenCache::new(true, Duration::from_secs(60)); + + let seeded = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Ok(fetched("valid", Some(3))) + }) + .await + .unwrap(); + assert_eq!(seeded, "valid"); + + let minted = cache + .get_or_fetch_within( + "id", + "secret", + "c.s.t", + Duration::from_secs(5), + |_reason| async { + tokio::time::sleep(Duration::from_secs(8)).await; + Ok(fetched("fresh", Some(3600))) + }, + ) + .await + .unwrap(); + assert_eq!(minted, "fresh"); } #[tokio::test] @@ -672,9 +1038,10 @@ mod tests { const FOLLOWERS: usize = 15; // Keep the leader's mint in flight (blocked on `gate`) while the - // followers pile in. + // followers pile in. A broken single-flight — one that didn't hold the + // per-entry lock across the mint — would let a follower mint too and push + // `calls` above 1. 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 = { @@ -693,21 +1060,29 @@ mod tests { }) }; - // Wait until the leader is inside the mint (one call recorded) before - // launching followers, so they cannot win the slot first. + // Wait until the leader is inside the mint (one call recorded), so the + // slot exists in the map and the leader holds its lock. 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. + // Take the test's own clone of the occupied slot; its Arc strong count + // then reports how many callers have reached it. + let slot = { + let entries = cache.entries.lock().await; + Arc::clone( + entries + .get(&TokenKey::new("id", "secret", "c.s.t")) + .unwrap(), + ) + }; + + // Followers: each calls get_or_fetch and contends for the held lock. 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); @@ -718,11 +1093,20 @@ mod tests { })); } - // 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(); - } + // Release the leader only once every follower has cloned the occupied + // slot — i.e. entered get_or_fetch and reached the per-entry lock. The + // strong count is the map, the leader, this test handle, and each + // follower. The timeout only guards against a stuck follower (e.g. a + // single-flight regression) hanging the test; it is generous so a loaded + // CI box scheduling 15 tasks can't trip it spuriously. + tokio::time::timeout(Duration::from_secs(10), async { + while Arc::strong_count(&slot) < FOLLOWERS + 3 { + tokio::task::yield_now().await; + } + }) + .await + .expect("followers did not reach the occupied slot"); + gate.notify_one(); assert_eq!(leader.await.unwrap(), "tok"); From 8c181f23910869a406a425e1b50e00c1ab926a2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Danilo=20Trnini=C4=87?= Date: Tue, 18 Aug 2026 15:46:12 +0000 Subject: [PATCH 4/4] [Rust] Harden OAuth token invalidation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up on the #607 PR review. Make credential invalidation after an auth rejection precise and non-blocking. Each cached token now carries a monotonic generation, and invalidation records the rejected generation on the cache entry with a lock-free atomic watermark instead of taking the per-token lock. A token at or below the watermark is dropped before its next use so the next fetch re-mints, while a newer token installed by a concurrent refresh has a higher generation and is kept and reused. This removes the previous blocking invalidate and the risk of discarding a still-good token. Add an OAuth token-caching integration test suite that drives the real mint path over a loopback mock Unity Catalog endpoint (mock_oauth). It covers the provider directly (caching, quoted/absent expires_in, error classification, proactive refresh, dead-on-arrival, invalidation) and the whole stream- creation chain (shared-cache reuse, auth-rejection re-mint, one-shot auth retry, single-flight, refresh backoff, hung-refresh fallback, and per-table invalidation scoping). Signed-off-by: Danilo Trninić --- rust/NEXT_CHANGELOG.md | 7 + rust/sdk/src/headers_provider.rs | 33 +- rust/sdk/src/token_cache.rs | 466 +++++++++++--- rust/tests/Cargo.toml | 4 + rust/tests/src/mock_oauth.rs | 285 ++++++++ rust/tests/src/oauth_token_tests.rs | 964 ++++++++++++++++++++++++++++ 6 files changed, 1656 insertions(+), 103 deletions(-) create mode 100644 rust/tests/src/mock_oauth.rs create mode 100644 rust/tests/src/oauth_token_tests.rs diff --git a/rust/NEXT_CHANGELOG.md b/rust/NEXT_CHANGELOG.md index c7fe8025..bae99684 100644 --- a/rust/NEXT_CHANGELOG.md +++ b/rust/NEXT_CHANGELOG.md @@ -48,6 +48,13 @@ failure into a mint stampede. The interval scales with the configured token refresh buffer and shrinks as the token nears expiry, down to a floor, and never extends past expiry. +- OAuth credential invalidation after an auth rejection is more precise and never + waits on an in-flight mint. Each cached token carries a monotonic generation, and + invalidation records the rejected generation on the cache entry with a lock-free + atomic, so it never takes the per-token lock or detaches a mint. A token at or + below the recorded generation is dropped before its next use and is never refreshed + from, so the next fetch re-mints; a newer token installed by a concurrent refresh + has a higher generation and is kept and reused. ### Documentation diff --git a/rust/sdk/src/headers_provider.rs b/rust/sdk/src/headers_provider.rs index 88b7c325..85738fcd 100644 --- a/rust/sdk/src/headers_provider.rs +++ b/rust/sdk/src/headers_provider.rs @@ -3,6 +3,7 @@ use crate::token_cache::{TokenCache, DEFAULT_REFRESH_BUFFER}; use crate::ZerobusResult; use async_trait::async_trait; use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -48,13 +49,14 @@ pub trait HeadersProvider: Send + Sync { /// Returns a `ZerobusError` if header generation fails (e.g., token request fails). async fn get_headers(&self) -> ZerobusResult>; - /// Invalidates any cached authentication state so the next `get_headers` - /// call re-derives it from scratch. + /// Invalidates cached authentication state that the server just rejected. /// - /// The SDK calls this when the server rejects the supplied credentials with - /// an authentication error during stream creation. The default is a no-op, - /// which is correct for providers that hold no cache; the built-in OAuth - /// provider overrides it to drop its cached token so the next call re-mints. + /// The SDK calls this when the server rejects the supplied credentials with an + /// authentication error during stream creation. The default is a no-op, which is + /// correct for providers that hold no cache. The built-in OAuth provider clears + /// the rejected token from its cache, so the next `get_headers` re-mints — unless + /// a concurrent refresh already replaced it with a newer token, which is kept and + /// served without re-minting. async fn invalidate(&self) {} } @@ -72,6 +74,11 @@ pub struct OAuthHeadersProvider { /// How long a proactive refresh may run before it's treated as failed and the /// cached token is served; `None` leaves it unbounded. refresh_timeout: Option, + /// Generation of the token last returned by `get_headers`, so `invalidate` + /// rejects only that token, not a newer one from a concurrent refresh. 0 means + /// there is no cached token to reject: nothing has been served yet, or the most + /// recently served token was not cacheable (a no-TTL response or a disabled cache). + last_served_generation: AtomicU64, } impl OAuthHeadersProvider { @@ -121,6 +128,7 @@ impl OAuthHeadersProvider { unity_catalog_url, token_cache, refresh_timeout, + last_served_generation: AtomicU64::new(0), } } } @@ -138,7 +146,7 @@ impl HeadersProvider for OAuthHeadersProvider { reason, ) }; - let token = match self.refresh_timeout { + let (token, generation) = match self.refresh_timeout { Some(refresh_timeout) => { self.token_cache .get_or_fetch_within( @@ -161,6 +169,9 @@ impl HeadersProvider for OAuthHeadersProvider { .await? } }; + // Remember the served token's generation so invalidate() rejects only it. + self.last_served_generation + .store(generation, Ordering::SeqCst); let mut headers = HashMap::new(); headers.insert("authorization", format!("Bearer {}", token)); headers.insert("x-databricks-zerobus-table-name", self.table_name.clone()); @@ -168,8 +179,14 @@ impl HeadersProvider for OAuthHeadersProvider { } async fn invalidate(&self) { + let rejected_generation = self.last_served_generation.load(Ordering::SeqCst); self.token_cache - .invalidate(&self.client_id, &self.client_secret, &self.table_name) + .invalidate( + &self.client_id, + &self.client_secret, + &self.table_name, + rejected_generation, + ) .await; } } diff --git a/rust/sdk/src/token_cache.rs b/rust/sdk/src/token_cache.rs index 253a4647..2a0403e4 100644 --- a/rust/sdk/src/token_cache.rs +++ b/rust/sdk/src/token_cache.rs @@ -14,6 +14,7 @@ //! name is part of the cache key. use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -36,14 +37,23 @@ pub(crate) const DEFAULT_REFRESH_BUFFER: Duration = Duration::from_secs(300); const MAX_BACKOFF_WINDOW_FRACTION: u32 = 60; const MIN_BACKOFF_WINDOW_FRACTION: u32 = 300; +// Guard `arm_refresh_backoff` against two panics from a bad retune of these divisors: +// a zero divisor (`Duration / 0`) and an inverted pair (`clamp(min, max)` with +// min > max). A nonzero MAX plus MIN >= MAX also guarantees MIN is nonzero. +const _: () = assert!( + MAX_BACKOFF_WINDOW_FRACTION > 0 && MIN_BACKOFF_WINDOW_FRACTION >= MAX_BACKOFF_WINDOW_FRACTION +); + /// A cached token and the instant at which it expires. struct CachedToken { value: String, expires_at: Instant, - /// Defer the next proactive refresh until this instant, set after a refresh - /// fell back to this token (a failed mint or a dead-on-arrival token). Always - /// `<= expires_at`, so it never serves an expired token. + /// Instant until which the next proactive refresh is deferred, set by + /// `arm_refresh_backoff`. Always `<= expires_at`. refresh_retry_at: Option, + /// Monotonic id assigned when this token was installed, matched against the + /// slot's rejection watermark by `invalidate`. + generation: u64, } impl CachedToken { @@ -98,16 +108,25 @@ impl TokenKey { } } -/// Per-entry slot. Each key has its own mutex so that a cold-cache burst of -/// concurrent stream creations for the same table mints a single token -/// (single-flight) while creations for different tables never block each other. -type Slot = Arc>>; +/// Per-entry cache slot. Each key has its own token mutex, so a cold-cache burst for +/// one table mints once (single-flight) while other tables never block each other. +#[derive(Default)] +struct SlotInner { + token: Mutex>, + /// Highest token generation the server has rejected for this key; a cached token + /// at or below it is stale. Raised by `invalidate`. + rejected_generation: AtomicU64, +} + +type Slot = Arc; /// Caches OAuth tokens per table for the lifetime of a [`ZerobusSdk`]. /// /// Safe for concurrent use across streams created from the same SDK instance. pub(crate) struct TokenCache { entries: Mutex>, + /// Source of monotonic token generations; each installed token gets the next one. + next_generation: AtomicU64, refresh_buffer: Duration, enabled: bool, } @@ -116,6 +135,7 @@ impl TokenCache { pub(crate) fn new(enabled: bool, refresh_buffer: Duration) -> Self { Self { entries: Mutex::new(HashMap::new()), + next_generation: AtomicU64::new(0), refresh_buffer, enabled, } @@ -131,7 +151,7 @@ impl TokenCache { client_secret: &str, table_name: &str, fetch: F, - ) -> ZerobusResult + ) -> ZerobusResult<(String, u64)> where F: FnOnce(MintReason) -> Fut, Fut: std::future::Future>, @@ -151,7 +171,7 @@ impl TokenCache { table_name: &str, refresh_timeout: Duration, fetch: F, - ) -> ZerobusResult + ) -> ZerobusResult<(String, u64)> where F: FnOnce(MintReason) -> Fut, Fut: std::future::Future>, @@ -175,7 +195,7 @@ impl TokenCache { table_name: &str, refresh_timeout: Option, fetch: F, - ) -> ZerobusResult + ) -> ZerobusResult<(String, u64)> where F: FnOnce(MintReason) -> Fut, Fut: std::future::Future>, @@ -183,7 +203,7 @@ impl TokenCache { if !self.enabled { return fetch(MintReason::CacheDisabled) .await - .map(|fetched| fetched.token); + .map(|fetched| (fetched.token, 0)); } let key = TokenKey::new(client_id, client_secret, table_name); @@ -199,7 +219,15 @@ impl TokenCache { // Hold the per-entry lock across the fetch so concurrent callers for the // same key reuse a single mint instead of stampeding the token endpoint. - let mut guard = slot.lock().await; + let mut guard = slot.token.lock().await; + + // Drop a token the server has rejected (generation at or below the rejection + // watermark) so we re-mint rather than serve or refresh it. + if guard.as_ref().is_some_and(|cached| { + cached.generation <= slot.rejected_generation.load(Ordering::SeqCst) + }) { + *guard = None; + } if let Some(cached) = guard.as_ref() { if !self.needs_refresh(cached) { @@ -210,7 +238,7 @@ impl TokenCache { } else { debug!(table = %table_name, "token cache hit, reusing cached token"); } - return Ok(cached.value.clone()); + return Ok((cached.value.clone(), cached.generation)); } } @@ -259,11 +287,11 @@ impl TokenCache { Err(err) => { // On any refresh error, serve the still-valid cached token if we have // one, arming the backoff; otherwise surface the error. - if let Some(value) = - Self::serve_valid_cached_fallback(&mut guard, self.refresh_buffer) + if let Some((value, generation)) = + Self::serve_valid_cached_fallback(&slot, &mut guard, self.refresh_buffer) { warn!(table = %table_name, error = %err, "token refresh failed; serving still-valid cached token"); - return Ok(value); + return Ok((value, generation)); } return Err(err); } @@ -279,48 +307,72 @@ impl TokenCache { // an older still-valid cached token if there is one and arm the backoff, // otherwise surface a retryable error. if expires_at.is_some_and(|deadline| deadline <= Instant::now()) { - if let Some(value) = Self::serve_valid_cached_fallback(&mut guard, self.refresh_buffer) + if let Some((value, generation)) = + Self::serve_valid_cached_fallback(&slot, &mut guard, self.refresh_buffer) { warn!(table = %table_name, "fetched OAuth token expired on arrival; serving still-valid cached token"); - return Ok(value); + return Ok((value, generation)); } return Err(ZerobusError::TokenFetchError( "fetched OAuth token expired before arrival".to_string(), )); } - match expires_at { + let generation = match expires_at { Some(expires_at) => { + // A fresh token gets the next generation, above any rejection + // watermark, so it is never mistaken for a rejected token. + let generation = self.next_generation.fetch_add(1, Ordering::Relaxed) + 1; *guard = Some(CachedToken { value: fetched.token.clone(), expires_at, refresh_retry_at: None, + generation, }); + generation } None => { - // No usable TTL: keep an existing still-valid token rather than - // discarding it. - let keep_existing = guard.as_ref().is_some_and(|cached| !cached.is_expired()); + // No usable TTL: keep an existing still-valid, non-rejected token + // rather than discarding it. The returned token is uncached (gen 0). + let keep_existing = guard.as_ref().is_some_and(|cached| { + !cached.is_expired() + && cached.generation > slot.rejected_generation.load(Ordering::SeqCst) + }); if !keep_existing { *guard = None; } + 0 } - } + }; - Ok(fetched.token) + Ok((fetched.token, generation)) } - /// Drops any cached token for the given credentials and table so the next - /// fetch re-mints. Called when the server rejects the token (e.g. it was - /// revoked at the IdP), so the re-mint re-checks grants at UC. No-op when - /// caching is disabled or no entry exists. - pub(crate) async fn invalidate(&self, client_id: &str, client_secret: &str, table_name: &str) { + /// Raises the slot's rejection watermark to `rejected_generation` with a lock-free + /// `fetch_max`, so it never takes the token lock or waits on an in-flight mint. A + /// token at or below the watermark is dropped before its next use (a newer one + /// from a concurrent refresh is kept), so the next fetch re-mints. A no-op when + /// caching is off, the entry is absent, or `rejected_generation` is 0. + pub(crate) async fn invalidate( + &self, + client_id: &str, + client_secret: &str, + table_name: &str, + rejected_generation: u64, + ) { if !self.enabled { return; } let key = TokenKey::new(client_id, client_secret, table_name); - if self.entries.lock().await.remove(&key).is_some() { - debug!(table = %table_name, "token cache entry invalidated after auth rejection"); + if let Some(slot) = self.entries.lock().await.get(&key) { + let previous = slot + .rejected_generation + .fetch_max(rejected_generation, Ordering::SeqCst); + // Only log when the watermark actually advanced (not for gen 0 or a + // generation already at or below the current watermark). + if rejected_generation > previous { + debug!(table = %table_name, generation = rejected_generation, "recorded token rejection after auth failure"); + } } } @@ -338,31 +390,45 @@ impl TokenCache { } } - /// If a still-valid token is cached, arm its backoff and return it. This is the - /// fallback when a proactive refresh can't produce a usable token; `None` when - /// there is none to fall back to. + /// If a still-valid, non-rejected token is cached, arm its backoff and return it + /// with its generation. This is the fallback when a proactive refresh can't + /// produce a usable token; `None` when there is none to fall back to. A token at + /// or below the rejection watermark is refused rather than re-served. fn serve_valid_cached_fallback( - slot: &mut Option, + slot: &SlotInner, + guard: &mut Option, refresh_buffer: Duration, - ) -> Option { - let cached = slot.as_mut()?; + ) -> Option<(String, u64)> { + let cached = guard.as_ref()?; if cached.is_expired() { return None; } + if cached.generation <= slot.rejected_generation.load(Ordering::SeqCst) { + // Drop a rejected token here rather than leave it for the next lookup. + *guard = None; + return None; + } + let cached = guard.as_mut().expect("guard is Some"); cached.arm_refresh_backoff(refresh_buffer); - Some(cached.value.clone()) + Some((cached.value.clone(), cached.generation)) } - /// Drops entries whose token has fully expired. Locked (in-flight) entries, - /// still-valid tokens, and empty slots are kept — keeping empty slots is - /// what preserves single-flight for a key being minted concurrently. + /// Removes slots whose token is absent, expired, or rejected (generation at or + /// below the watermark). A slot a caller still holds (strong count > 1) is kept + /// whatever its contents, since it may be about to be minted into and single-flight + /// relies on that. fn prune_expired(entries: &mut HashMap) { - entries.retain(|_, slot| match slot.try_lock() { - Ok(guard) => match guard.as_ref() { - Some(cached) => !cached.is_expired(), - None => true, - }, - Err(_) => true, + entries.retain(|_, slot| { + if Arc::strong_count(slot) > 1 { + return true; + } + match slot.token.try_lock() { + Ok(guard) => guard.as_ref().is_some_and(|cached| { + !cached.is_expired() + && cached.generation > slot.rejected_generation.load(Ordering::SeqCst) + }), + Err(_) => true, + } }); } } @@ -389,11 +455,11 @@ mod tests { Ok(fetched("tok", Some(3600))) }; - let a = cache + let (a, _) = cache .get_or_fetch("id", "secret", "c.s.t", make) .await .unwrap(); - let b = cache + let (b, _) = cache .get_or_fetch("id", "secret", "c.s.t", make) .await .unwrap(); @@ -419,11 +485,11 @@ mod tests { Ok(fetched(&format!("tok{n}"), Some(1))) }; - let a = cache + let (a, _) = cache .get_or_fetch("id", "secret", "c.s.t", make) .await .unwrap(); - let b = cache + let (b, _) = cache .get_or_fetch("id", "secret", "c.s.t", make) .await .unwrap(); @@ -449,17 +515,17 @@ mod tests { }; // Call 1 mints tok0 (within-buffer, immediately due for refresh). - let a = cache + 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 + 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 + let (c, _) = cache .get_or_fetch("id", "secret", "c.s.t", make) .await .unwrap(); @@ -484,11 +550,11 @@ mod tests { Ok(fetched(&format!("tok{n}"), Some(3600))) }; - let a = cache + let (a, _) = cache .get_or_fetch("id", "secret", "c.s.t1", make) .await .unwrap(); - let b = cache + let (b, _) = cache .get_or_fetch("id", "secret", "c.s.t2", make) .await .unwrap(); @@ -551,13 +617,13 @@ mod tests { Ok(fetched("tok", Some(3600))) }; - cache + let (_, generation) = cache .get_or_fetch("id", "secret", "c.s.t", make) .await .unwrap(); - // Without invalidation a second call would hit the cache; invalidating - // the entry forces the next call to re-mint. - cache.invalidate("id", "secret", "c.s.t").await; + // Without invalidation a second call would hit the cache; rejecting the + // cached token's generation forces the next call to re-mint. + cache.invalidate("id", "secret", "c.s.t", generation).await; cache .get_or_fetch("id", "secret", "c.s.t", make) .await @@ -577,7 +643,7 @@ mod tests { }; // Seed two different tables (tok0 and tok1). - cache + let (_, t1_generation) = cache .get_or_fetch("id", "secret", "c.s.t1", make) .await .unwrap(); @@ -586,15 +652,17 @@ mod tests { .await .unwrap(); - // Invalidating t1 must not disturb t2. - cache.invalidate("id", "secret", "c.s.t1").await; + // Invalidating c.s.t1 must not disturb c.s.t2. + cache + .invalidate("id", "secret", "c.s.t1", t1_generation) + .await; // t1 re-mints (tok2); t2 still hits its original cached token (tok1). - let t1 = cache + let (t1, _) = cache .get_or_fetch("id", "secret", "c.s.t1", make) .await .unwrap(); - let t2 = cache + let (t2, _) = cache .get_or_fetch("id", "secret", "c.s.t2", make) .await .unwrap(); @@ -614,14 +682,16 @@ mod tests { Ok(fetched("tok", Some(3600))) }; - cache + let (_, generation) = 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 + .invalidate("id", "secret", "other.table.here", generation) + .await; cache .get_or_fetch("id", "secret", "c.s.t", make) .await @@ -646,16 +716,222 @@ mod tests { Ok(fetched("tok", Some(3600))) }; - cache.invalidate("id", "secret", "c.s.t").await; - let token = cache + let (token, generation) = cache .get_or_fetch("id", "secret", "c.s.t", make) .await .unwrap(); + // On a disabled cache invalidate is a no-op regardless of the generation. + cache.invalidate("id", "secret", "c.s.t", generation).await; assert_eq!(token, "tok"); assert_eq!(calls.load(Ordering::SeqCst), 1); } + #[tokio::test] + async fn invalidate_only_clears_the_rejected_token() { + // invalidate() clears only the token it is given; a newer token that a + // refresh installed in the meantime is left in place. + let cache = TokenCache::new(true, Duration::from_secs(60)); + let calls = AtomicUsize::new(0); + + // tok0 has a within-buffer TTL so the next call refreshes it to tok1. + 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))) + }; + + let (old, old_generation) = cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + assert_eq!(old, "tok0"); + let (refreshed, _) = cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + assert_eq!(refreshed, "tok1"); + + // Rejecting the stale tok0's generation must not evict the newer tok1. + cache + .invalidate("id", "secret", "c.s.t", old_generation) + .await; + let (served, _) = cache + .get_or_fetch("id", "secret", "c.s.t", make) + .await + .unwrap(); + assert_eq!( + served, "tok1", + "stale invalidate must not evict a newer token" + ); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn empty_slot_is_removed_on_next_miss() { + // A failed cold mint leaves an empty slot behind; a later miss on a different + // key prunes it so the map doesn't accumulate dead entries. + let cache = TokenCache::new(true, Duration::from_secs(60)); + + let err = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Err(ZerobusError::TokenFetchError("boom".to_string())) + }) + .await; + assert!(err.is_err()); + + // A miss on a different key triggers prune_expired, dropping the empty slot. + cache + .get_or_fetch("id", "secret", "other.table", |_reason| async { + Ok(fetched("tok2", Some(3600))) + }) + .await + .unwrap(); + + let entries = cache.entries.lock().await; + assert!( + !entries.contains_key(&TokenKey::new("id", "secret", "c.s.t")), + "empty slot should be removed on the next miss" + ); + assert!(entries.contains_key(&TokenKey::new("id", "secret", "other.table"))); + } + + #[tokio::test] + async fn rejected_token_entry_is_removed_on_next_miss() { + // An invalidated token that is still within its TTL is removed by + // prune_expired, not left cached until it expires. + let cache = TokenCache::new(true, Duration::from_secs(60)); + + let (_, generation) = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Ok(fetched("tok", Some(3600))) + }) + .await + .unwrap(); + cache.invalidate("id", "secret", "c.s.t", generation).await; + + // A miss on a different key triggers prune_expired, which removes the + // rejected (but not-yet-expired) entry. + cache + .get_or_fetch("id", "secret", "other.table", |_reason| async { + Ok(fetched("tok2", Some(3600))) + }) + .await + .unwrap(); + + let entries = cache.entries.lock().await; + assert!( + !entries.contains_key(&TokenKey::new("id", "secret", "c.s.t")), + "rejected token's entry should be removed on the next miss" + ); + assert!(entries.contains_key(&TokenKey::new("id", "secret", "other.table"))); + } + + #[tokio::test] + async fn invalidate_does_not_block_on_an_in_flight_mint() { + // invalidate never takes the token lock (it only raises an atomic watermark), + // so it returns promptly even while a refresh holds the lock. The refresh then + // installs a newer token, unaffected by the rejection of the older generation. + let cache = Arc::new(TokenCache::new(true, Duration::from_secs(60))); + + // Seed a within-buffer token so the leader's access refreshes it. + let (_, seeded_generation) = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Ok(fetched("valid", Some(30))) + }) + .await + .unwrap(); + + let gate = Arc::new(tokio::sync::Notify::new()); + let minting = Arc::new(tokio::sync::Notify::new()); + + let leader = { + let cache = Arc::clone(&cache); + let gate = Arc::clone(&gate); + let minting = Arc::clone(&minting); + tokio::spawn(async move { + cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async move { + minting.notify_one(); + gate.notified().await; + Ok(fetched("fresh", Some(3600))) + }) + .await + .unwrap() + }) + }; + + // Wait until the refresh holds the slot lock, then reject the seeded token. + minting.notified().await; + tokio::time::timeout( + Duration::from_secs(5), + cache.invalidate("id", "secret", "c.s.t", seeded_generation), + ) + .await + .expect("invalidate must not block on an in-flight mint"); + + gate.notify_one(); + let (token, _) = leader.await.unwrap(); + assert_eq!(token, "fresh"); + } + + #[tokio::test] + async fn invalidate_during_failing_refresh_is_honored() { + // invalidate races a refresh that then fails. The rejection is recorded + // out-of-band, so the failing refresh's fallback refuses the rejected token + // and the next get re-mints instead of serving it. + let cache = Arc::new(TokenCache::new(true, Duration::from_secs(60))); + + // Seed a within-buffer token so the next access is due for a refresh. + let (_, seeded_generation) = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Ok(fetched("valid", Some(30))) + }) + .await + .unwrap(); + + let gate = Arc::new(tokio::sync::Notify::new()); + let minting = Arc::new(tokio::sync::Notify::new()); + + // Leader: a proactive refresh that holds the slot lock, then fails. + let leader = { + let cache = Arc::clone(&cache); + let gate = Arc::clone(&gate); + let minting = Arc::clone(&minting); + tokio::spawn(async move { + cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async move { + minting.notify_one(); + gate.notified().await; + Err(ZerobusError::TokenFetchError("boom".to_string())) + }) + .await + }) + }; + + // Reject the served token while the refresh holds the lock. + minting.notified().await; + cache + .invalidate("id", "secret", "c.s.t", seeded_generation) + .await; + + // The failing refresh must not fall back to the rejected token. + gate.notify_one(); + assert!( + leader.await.unwrap().is_err(), + "failing refresh must not fall back to the rejected token" + ); + + // The next get re-mints rather than serving the rejected "valid". + let (refreshed, _) = cache + .get_or_fetch("id", "secret", "c.s.t", |_reason| async { + Ok(fetched("fresh", Some(3600))) + }) + .await + .unwrap(); + assert_eq!(refreshed, "fresh"); + } + #[tokio::test] async fn disabled_cache_always_fetches() { let cache = TokenCache::new(false, Duration::from_secs(60)); @@ -690,7 +966,7 @@ mod tests { assert!(err.is_err()); // A subsequent successful fetch should still succeed and cache. - let ok = cache + let (ok, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Ok(fetched("tok", Some(3600))) }) @@ -705,7 +981,7 @@ mod tests { // kind: this covers both a retryable and a non-retryable (revoked-creds) one. let cache = TokenCache::new(true, Duration::from_secs(60)); - let seeded = cache + let (seeded, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Ok(fetched("valid", Some(30))) }) @@ -716,7 +992,7 @@ mod tests { // Count attempts to prove each failing refresh actually ran (not suppressed). let refresh_attempts = AtomicUsize::new(0); - let served_retryable = cache + let (served_retryable, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { refresh_attempts.fetch_add(1, Ordering::SeqCst); Err(ZerobusError::TokenFetchError("blip".to_string())) @@ -729,7 +1005,7 @@ mod tests { // Clear the armed backoff so the next call refreshes again. tokio::time::advance(Duration::from_secs(2)).await; - let served_non_retryable = cache + let (served_non_retryable, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { refresh_attempts.fetch_add(1, Ordering::SeqCst); Err(ZerobusError::InvalidUCTokenError("revoked".to_string())) @@ -790,7 +1066,7 @@ mod tests { let cache = TokenCache::new(true, Duration::from_secs(60)); let mints = AtomicUsize::new(0); - let seeded = cache + let (seeded, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Ok(fetched("valid", Some(30))) @@ -799,7 +1075,7 @@ mod tests { .unwrap(); assert_eq!(seeded, "valid"); - let served = cache + let (served, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); tokio::time::sleep(Duration::from_secs(2)).await; @@ -810,7 +1086,7 @@ mod tests { assert_eq!(served, "valid"); // The armed backoff suppresses the next refresh, so the mint count stays at 2. - let reused = cache + let (reused, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Ok(fetched("unexpected", Some(3600))) @@ -828,7 +1104,7 @@ mod tests { // budget-capped refresh would time out and fail). let cache = TokenCache::new(true, Duration::from_secs(60)); - let seeded = cache + let (seeded, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Ok(fetched("stale", Some(1))) }) @@ -837,7 +1113,7 @@ mod tests { assert_eq!(seeded, "stale"); tokio::time::advance(Duration::from_secs(2)).await; - let minted = cache + let (minted, _) = cache .get_or_fetch_within( "id", "secret", @@ -860,7 +1136,7 @@ mod tests { let cache = TokenCache::new(true, Duration::from_secs(60)); let mints = AtomicUsize::new(0); - let seeded = cache + let (seeded, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Ok(fetched("valid", Some(30))) @@ -869,7 +1145,7 @@ mod tests { .unwrap(); assert_eq!(seeded, "valid"); - let served = cache + let (served, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Err(ZerobusError::TokenFetchError("blip".to_string())) @@ -878,7 +1154,7 @@ mod tests { .unwrap(); assert_eq!(served, "valid"); - let reused = cache + let (reused, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Ok(fetched("unexpected", Some(3600))) @@ -897,7 +1173,7 @@ mod tests { let cache = TokenCache::new(true, Duration::from_secs(300)); let mints = AtomicUsize::new(0); - let seeded = cache + let (seeded, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Ok(fetched("valid", Some(3))) @@ -906,7 +1182,7 @@ mod tests { .unwrap(); assert_eq!(seeded, "valid"); - let served = cache + let (served, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Err(ZerobusError::TokenFetchError("blip".to_string())) @@ -917,7 +1193,7 @@ mod tests { // Still inside the ~1.5s backoff: reused, so the mint count stays at 2. tokio::time::advance(Duration::from_secs(1)).await; - let reused = cache + let (reused, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Ok(fetched("unexpected", Some(3600))) @@ -929,7 +1205,7 @@ mod tests { // Past the backoff and still valid: the proactive retry runs and succeeds. tokio::time::advance(Duration::from_millis(600)).await; - let refreshed = cache + let (refreshed, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { mints.fetch_add(1, Ordering::SeqCst); Ok(fetched("fresh", Some(3600))) @@ -946,7 +1222,7 @@ mod tests { // becomes a retryable error and the still-valid cached token is served. let cache = TokenCache::new(true, Duration::from_secs(60)); - let seeded = cache + let (seeded, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Ok(fetched("valid", Some(30))) }) @@ -954,7 +1230,7 @@ mod tests { .unwrap(); assert_eq!(seeded, "valid"); - let served = cache + let (served, _) = cache .get_or_fetch_within( "id", "secret", @@ -974,7 +1250,7 @@ mod tests { // back to. The refresh runs unbounded instead, so an 8s mint still succeeds. let cache = TokenCache::new(true, Duration::from_secs(60)); - let seeded = cache + let (seeded, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Ok(fetched("valid", Some(3))) }) @@ -982,7 +1258,7 @@ mod tests { .unwrap(); assert_eq!(seeded, "valid"); - let minted = cache + let (minted, _) = cache .get_or_fetch_within( "id", "secret", @@ -1012,7 +1288,7 @@ mod tests { // A refresh returns a token with no TTL: the caller gets the fresh token, // but the cached valid token must not be discarded. - let fresh = cache + let (fresh, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Ok(fetched("nottl", None)) }) @@ -1022,7 +1298,7 @@ mod tests { // A later refresh failure still finds the original valid token, proving // it was retained. - let served = cache + let (served, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { Err(crate::ZerobusError::TokenFetchError("blip".to_string())) }) @@ -1109,9 +1385,9 @@ mod tests { gate.notify_one(); - assert_eq!(leader.await.unwrap(), "tok"); + assert_eq!(leader.await.unwrap().0, "tok"); for handle in followers { - assert_eq!(handle.await.unwrap(), "tok"); + assert_eq!(handle.await.unwrap().0, "tok"); } assert_eq!( calls.load(Ordering::SeqCst), @@ -1154,7 +1430,7 @@ mod tests { // The cancelled leader must have released the lock and left no // half-written entry, so the next caller mints cleanly... - let minted = cache + let (minted, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { calls.fetch_add(1, Ordering::SeqCst); Ok(fetched("tok", Some(3600))) @@ -1165,7 +1441,7 @@ mod tests { // ...and that freshly minted token is cached, not a phantom entry: a // follow-up call hits without minting again. - let cached = cache + let (cached, _) = cache .get_or_fetch("id", "secret", "c.s.t", |_reason| async { calls.fetch_add(1, Ordering::SeqCst); Ok(fetched("other", Some(3600))) diff --git a/rust/tests/Cargo.toml b/rust/tests/Cargo.toml index f76a2854..51916138 100644 --- a/rust/tests/Cargo.toml +++ b/rust/tests/Cargo.toml @@ -8,6 +8,10 @@ publish = false name = "rust_tests" path = "src/rust_tests.rs" +[[test]] +name = "oauth_token_tests" +path = "src/oauth_token_tests.rs" + [[test]] name = "proxy_tests" path = "src/proxy_tests.rs" diff --git a/rust/tests/src/mock_oauth.rs b/rust/tests/src/mock_oauth.rs new file mode 100644 index 00000000..fc989bc9 --- /dev/null +++ b/rust/tests/src/mock_oauth.rs @@ -0,0 +1,285 @@ +//! Mock Unity Catalog OAuth token endpoint for integration tests. +//! +//! Stands up a loopback HTTP/1.1 server answering `POST /oidc/v1/token` — the +//! endpoint the SDK's `DefaultTokenFactory` mints against. Tests script the +//! replies (token value, `expires_in`, HTTP status, or an indefinite hang) to +//! drive the OAuth token-cache behavior end to end through the SDK's real +//! `reqwest` client, rather than a stubbed `HeadersProvider`. +//! +//! It mirrors the loopback pattern of `mock_grpc.rs`: bind `127.0.0.1:0`, spawn +//! the server on a background task, and hand the caller back the base URL to +//! point the SDK at. The gRPC mock speaks HTTP/2 via tonic; this endpoint is a +//! plain HTTP/1.1 JSON POST, so it is hand-rolled on a raw `TcpListener` and +//! needs no extra dependency. + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::Mutex; +use tracing::{debug, error}; + +/// A scripted reply the mock returns for one token request. +#[derive(Clone)] +pub enum MockTokenResponse { + /// `200 OK` with a JSON body carrying `access_token` and, optionally, + /// `expires_in`. + Ok { + access_token: String, + /// The `expires_in` value as a raw JSON fragment inserted verbatim, so a + /// test can send an integer (`3600`) or a quoted integer (`"3600"`); + /// `None` omits the field entirely (the SDK then treats the token as + /// uncacheable). + expires_in: Option, + /// Wall-clock delay before replying, to model a slow issuer. Used to make + /// a token dead-on-arrival: a delay longer than `expires_in` means the + /// (start-anchored) token has already expired by the time it arrives. + delay: Duration, + }, + /// Reply with the given HTTP status code and body. 5xx classifies as a + /// retryable `TokenFetchError`, 4xx as a non-retryable `InvalidUCTokenError`. + Error { status: u16, body: String }, + /// Accept the connection but never reply, so the caller's own timeout fires. + Hang, +} + +impl MockTokenResponse { + /// `200 OK` returning `access_token` with a one-hour integer `expires_in`. + pub fn ok(access_token: impl Into) -> Self { + MockTokenResponse::Ok { + access_token: access_token.into(), + expires_in: Some("3600".to_string()), + delay: Duration::ZERO, + } + } + + /// Overrides the `expires_in` fragment (see [`MockTokenResponse::Ok`]). + pub fn with_expires_in(mut self, raw: Option<&str>) -> Self { + if let MockTokenResponse::Ok { expires_in, .. } = &mut self { + *expires_in = raw.map(str::to_string); + } + self + } + + /// Delays a `200 OK` reply by `delay` (see [`MockTokenResponse::Ok`]). + pub fn with_delay(mut self, delay: Duration) -> Self { + if let MockTokenResponse::Ok { delay: d, .. } = &mut self { + *d = delay; + } + self + } + + /// An error reply with the given HTTP status and body. + pub fn error(status: u16, body: impl Into) -> Self { + MockTokenResponse::Error { + status, + body: body.into(), + } + } +} + +/// A mock Unity Catalog OAuth token endpoint bound to loopback. +/// +/// The returned handle shares state with the background server task: tests set +/// the scripted replies and read [`mint_count`](Self::mint_count) to assert how +/// many times the SDK actually hit the endpoint (i.e. minted rather than served +/// from cache). +pub struct MockOAuthServer { + /// Replies consumed in order; once drained, the server task falls back to a + /// defensive default so an unexpected extra mint fails an assertion cleanly + /// instead of hanging. + responses: Arc>>, + /// Number of token requests fully received (mints observed by the server). + mint_count: Arc, +} + +impl MockOAuthServer { + /// Replaces the scripted replies, consumed front to back by later mints. + pub async fn set_responses(&self, responses: Vec) { + let mut queue = self.responses.lock().await; + queue.clear(); + queue.extend(responses); + } + + /// Number of token requests the endpoint has received so far. + pub fn mint_count(&self) -> usize { + self.mint_count.load(Ordering::SeqCst) + } +} + +/// Starts the mock endpoint and returns the handle plus its base URL +/// (`http://127.0.0.1:`). Pass the base URL to `unity_catalog_url`; the +/// SDK appends `/oidc/v1/token`. +pub async fn start_mock_oauth_server() -> (MockOAuthServer, String) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind mock OAuth server"); + let base_url = format!("http://{}", listener.local_addr().expect("local_addr")); + + let responses = Arc::new(Mutex::new(VecDeque::new())); + let default_response = Arc::new(MockTokenResponse::ok("mock-default-token")); + let mint_count = Arc::new(AtomicUsize::new(0)); + + let server = MockOAuthServer { + responses: Arc::clone(&responses), + mint_count: Arc::clone(&mint_count), + }; + + tokio::spawn(async move { + loop { + let socket = match listener.accept().await { + Ok((socket, _)) => socket, + Err(e) => { + error!("mock OAuth accept error: {e}"); + return; + } + }; + let responses = Arc::clone(&responses); + let default_response = Arc::clone(&default_response); + let mint_count = Arc::clone(&mint_count); + tokio::spawn(handle_connection( + socket, + responses, + default_response, + mint_count, + )); + } + }); + + (server, base_url) +} + +/// Reads one HTTP/1.1 request to completion, counts it as a mint, and writes the +/// next scripted reply. +async fn handle_connection( + mut socket: TcpStream, + responses: Arc>>, + default_response: Arc, + mint_count: Arc, +) { + if !read_request(&mut socket).await { + return; + } + + // The request arrived in full: count it as a mint before replying. + mint_count.fetch_add(1, Ordering::SeqCst); + + let response = match responses.lock().await.pop_front() { + Some(response) => response, + None => (*default_response).clone(), + }; + + match response { + // Hold the connection open without replying so the caller's own timeout + // is what fires (used to model a hung token endpoint). + MockTokenResponse::Hang => std::future::pending::<()>().await, + MockTokenResponse::Ok { + access_token, + expires_in, + delay, + } => { + if !delay.is_zero() { + tokio::time::sleep(delay).await; + } + let body = token_body(&access_token, expires_in.as_deref()); + let _ = write_response(&mut socket, 200, "OK", &body).await; + } + MockTokenResponse::Error { status, body } => { + let _ = write_response(&mut socket, status, reason_phrase(status), &body).await; + } + } +} + +/// Reads until the header terminator, then drains `Content-Length` body bytes. +/// Returns `false` if the peer closed before a full request arrived. +async fn read_request(socket: &mut TcpStream) -> bool { + let mut buf = Vec::new(); + let mut chunk = [0u8; 1024]; + + let header_end = loop { + if let Some(pos) = find_subslice(&buf, b"\r\n\r\n") { + break pos + 4; + } + match socket.read(&mut chunk).await { + Ok(0) => return false, + Ok(n) => buf.extend_from_slice(&chunk[..n]), + Err(_) => return false, + } + }; + + let content_length = parse_content_length(&buf[..header_end]).unwrap_or(0); + let mut remaining = content_length.saturating_sub(buf.len() - header_end); + while remaining > 0 { + match socket.read(&mut chunk).await { + Ok(0) => break, + Ok(n) => remaining = remaining.saturating_sub(n), + Err(_) => break, + } + } + debug!("mock OAuth received a token request"); + true +} + +/// Builds the JSON token body, inserting `expires_in` verbatim when present. +fn token_body(access_token: &str, expires_in: Option<&str>) -> String { + match expires_in { + Some(raw) => format!( + r#"{{"access_token":"{access_token}","token_type":"Bearer","expires_in":{raw}}}"# + ), + None => format!(r#"{{"access_token":"{access_token}","token_type":"Bearer"}}"#), + } +} + +/// Writes a minimal HTTP/1.1 response and closes the connection. +async fn write_response( + socket: &mut TcpStream, + status: u16, + reason: &str, + body: &str, +) -> std::io::Result<()> { + let response = format!( + "HTTP/1.1 {status} {reason}\r\n\ + Content-Type: application/json\r\n\ + Content-Length: {}\r\n\ + Connection: close\r\n\ + \r\n\ + {body}", + body.len() + ); + socket.write_all(response.as_bytes()).await?; + socket.flush().await?; + let _ = socket.shutdown().await; + Ok(()) +} + +fn reason_phrase(status: u16) -> &'static str { + match status { + 400 => "Bad Request", + 401 => "Unauthorized", + 403 => "Forbidden", + 500 => "Internal Server Error", + 503 => "Service Unavailable", + _ => "Status", + } +} + +fn parse_content_length(headers: &[u8]) -> Option { + let text = String::from_utf8_lossy(headers); + for line in text.lines() { + if let Some((name, value)) = line.split_once(':') { + if name.trim().eq_ignore_ascii_case("content-length") { + return value.trim().parse().ok(); + } + } + } + None +} + +fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) +} diff --git a/rust/tests/src/oauth_token_tests.rs b/rust/tests/src/oauth_token_tests.rs new file mode 100644 index 00000000..3df14502 --- /dev/null +++ b/rust/tests/src/oauth_token_tests.rs @@ -0,0 +1,964 @@ +//! Integration tests for OAuth token caching. +//! +//! These exercise the real minting path end to end: the SDK's `reqwest` client +//! POSTs to a loopback mock Unity Catalog endpoint (`mock_oauth`), the response +//! flows through `DefaultTokenFactory` into the shared `TokenCache`, and back +//! out through `OAuthHeadersProvider`. +//! +//! Behavior is induced from the server side — the mock varies `expires_in`, the +//! HTTP status, and response timing — and asserted from the client side via the +//! returned token and `mock_oauth`'s mint counter (how many times the SDK hit +//! the endpoint rather than serving from cache). + +// The shared gRPC mock is only partially used here (stream creation and error +// injection), so silence dead-code warnings for its unused response variants. +#[allow(dead_code)] +mod mock_grpc; +mod mock_oauth; +mod utils; + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use databricks_zerobus_ingest_sdk::{ + HeadersProvider, NoTlsConfig, OAuthHeadersProvider, ZerobusError, ZerobusSdk, +}; + +use futures::future::join_all; +use mock_grpc::{start_mock_server, MockResponse}; +use mock_oauth::{start_mock_oauth_server, MockTokenResponse}; +use utils::setup_tracing; + +const TABLE_NAME: &str = "catalog.schema.orders"; +const OTHER_TABLE: &str = "catalog.schema.events"; +const CLIENT_ID: &str = "test-client-id"; +const CLIENT_SECRET: &str = "test-client-secret"; +const WORKSPACE_ID: &str = "test-workspace"; + +/// Builds an `OAuthHeadersProvider` with its own cache, pointed at the mock UC +/// endpoint. This standalone provider is enough for the pure token-cache tests +/// (no stream, no gRPC server). +fn oauth_provider(uc_url: String) -> OAuthHeadersProvider { + OAuthHeadersProvider::new( + CLIENT_ID.to_string(), + CLIENT_SECRET.to_string(), + TABLE_NAME.to_string(), + WORKSPACE_ID.to_string(), + uc_url, + ) +} + +/// Extracts the bearer token from a set of provider headers. +fn bearer_token(headers: &HashMap<&'static str, String>) -> String { + headers + .get("authorization") + .and_then(|value| value.strip_prefix("Bearer ")) + .expect("headers must carry a Bearer authorization") + .to_string() +} + +/// Token minting and caching through the real HTTP path, without a stream. +mod token_minting_and_caching_tests { + use super::*; + + #[tokio::test] + async fn mints_once_then_serves_from_cache() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + oauth + .set_responses(vec![MockTokenResponse::ok("token-1")]) + .await; + let provider = oauth_provider(uc_url); + + let first = bearer_token(&provider.get_headers().await.unwrap()); + let second = bearer_token(&provider.get_headers().await.unwrap()); + + assert_eq!(first, "token-1"); + assert_eq!( + second, "token-1", + "the second call must serve the cached token" + ); + assert_eq!( + oauth.mint_count(), + 1, + "a cached token must not be re-minted" + ); + } + + #[tokio::test] + async fn quoted_expires_in_is_cached() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // UC returns expires_in as a quoted string; caching must still kick in. + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1").with_expires_in(Some("\"3600\"")) + ]) + .await; + let provider = oauth_provider(uc_url); + + provider.get_headers().await.unwrap(); + provider.get_headers().await.unwrap(); + + assert_eq!( + oauth.mint_count(), + 1, + "a quoted expires_in must parse to a TTL and enable caching" + ); + } + + #[tokio::test] + async fn missing_expires_in_is_not_cached() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1").with_expires_in(None), + MockTokenResponse::ok("token-2"), + ]) + .await; + let provider = oauth_provider(uc_url); + + let first = bearer_token(&provider.get_headers().await.unwrap()); + let second = bearer_token(&provider.get_headers().await.unwrap()); + + assert_eq!(first, "token-1"); + assert_eq!( + second, "token-2", + "a token with no TTL is not cacheable, so the next call re-mints" + ); + assert_eq!(oauth.mint_count(), 2); + } + + #[tokio::test] + async fn server_error_surfaces_as_retryable() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + oauth + .set_responses(vec![MockTokenResponse::error(503, "upstream unavailable")]) + .await; + let provider = oauth_provider(uc_url); + + let err = provider + .get_headers() + .await + .expect_err("a 5xx response must fail the mint"); + assert!( + matches!(err, ZerobusError::TokenFetchError(_)), + "got {err:?}" + ); + assert!(err.is_retryable(), "a 5xx token error is retryable"); + } + + #[tokio::test] + async fn client_error_surfaces_as_non_retryable() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + oauth + .set_responses(vec![MockTokenResponse::error(401, "invalid_client")]) + .await; + let provider = oauth_provider(uc_url); + + let err = provider + .get_headers() + .await + .expect_err("a 4xx response must fail the mint"); + assert!( + matches!(err, ZerobusError::InvalidUCTokenError(_)), + "got {err:?}" + ); + assert!(!err.is_retryable(), "a 4xx token error is not retryable"); + } + + #[tokio::test] + async fn invalidate_forces_a_remint() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1"), + MockTokenResponse::ok("token-2"), + ]) + .await; + let provider = oauth_provider(uc_url); + + let first = bearer_token(&provider.get_headers().await.unwrap()); + provider.invalidate().await; + let second = bearer_token(&provider.get_headers().await.unwrap()); + + assert_eq!(first, "token-1"); + assert_eq!( + second, "token-2", + "invalidate must drop the rejected token so the next call re-mints" + ); + assert_eq!(oauth.mint_count(), 2); + } + + #[tokio::test] + async fn proactive_refresh_replaces_the_cached_token() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // token-1 is short-lived (60s < the 300s refresh window) so the second call + // proactively refreshes; token-2 is long-lived so the third call is a plain + // cache hit. This is the successful-refresh path (the hung/backoff tests only + // exercise a *failing* refresh). + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1").with_expires_in(Some("60")), + MockTokenResponse::ok("token-2"), + ]) + .await; + let provider = oauth_provider(uc_url); + + let first = bearer_token(&provider.get_headers().await.unwrap()); + let second = bearer_token(&provider.get_headers().await.unwrap()); + let third = bearer_token(&provider.get_headers().await.unwrap()); + + assert_eq!(first, "token-1", "cold miss mints token-1"); + assert_eq!( + second, "token-2", + "the in-window token is proactively refreshed to token-2" + ); + assert_eq!( + third, "token-2", + "the refreshed token is then cached and reused" + ); + assert_eq!( + oauth.mint_count(), + 2, + "a cold mint plus one refresh, and no mint on the third call" + ); + } + + #[tokio::test] + async fn dead_on_arrival_token_surfaces_retryable_error() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // A 1s token delivered ~1.2s late is already past its start-anchored expiry. + // On a cold miss (no token to fall back to) it must surface a retryable error + // rather than be cached. + oauth + .set_responses(vec![MockTokenResponse::ok("token-1") + .with_expires_in(Some("1")) + .with_delay(Duration::from_millis(1200))]) + .await; + let provider = oauth_provider(uc_url); + + let err = provider + .get_headers() + .await + .expect_err("a dead-on-arrival token must fail the cold miss"); + assert!( + matches!(err, ZerobusError::TokenFetchError(_)), + "got {err:?}" + ); + assert!( + err.is_retryable(), + "a dead-on-arrival token is a retryable fetch error" + ); + assert_eq!(oauth.mint_count(), 1); + } + + #[tokio::test] + async fn dead_on_arrival_refresh_falls_back_to_cached() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // Seed a still-valid token in its refresh window; the refresh then returns a + // dead-on-arrival token (1s TTL, ~1.2s late), so the cache must fall back to + // the still-valid seed rather than install and serve the DOA token. + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1").with_expires_in(Some("60")), + MockTokenResponse::ok("token-2") + .with_expires_in(Some("1")) + .with_delay(Duration::from_millis(1200)), + ]) + .await; + let provider = oauth_provider(uc_url); + + let first = bearer_token(&provider.get_headers().await.unwrap()); + let second = bearer_token(&provider.get_headers().await.unwrap()); + + assert_eq!(first, "token-1"); + assert_eq!( + second, "token-1", + "a dead-on-arrival refresh must fall back to the cached token, not serve the DOA one" + ); + assert_eq!( + oauth.mint_count(), + 2, + "the refresh was attempted, then discarded as dead on arrival" + ); + } + + #[tokio::test] + async fn unusable_token_is_rejected() { + setup_tracing(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // A token carrying a control character cannot be an HTTP header value, so the + // factory must reject it before it is ever cached. (The raw string embeds a + // JSON `\n` escape, so the response parses but the token contains a newline.) + oauth + .set_responses(vec![MockTokenResponse::ok(r"bad\ntoken")]) + .await; + let provider = oauth_provider(uc_url); + + let err = provider + .get_headers() + .await + .expect_err("an unusable token must fail the mint"); + assert!( + matches!(err, ZerobusError::InvalidUCTokenError(_)), + "got {err:?}" + ); + assert!( + !err.is_retryable(), + "an unusable token is a non-retryable error" + ); + } +} + +/// The whole chain: SDK stream creation minting from mock UC and connecting to +/// the mock gRPC data plane. +mod stream_creation_tests { + use super::*; + + /// Number of builds launched together in the concurrency tests. Bumping it + /// stresses single-flight / backoff harder; the asserted mint counts do not + /// depend on it (the per-key slot lock serializes the callers). + const CONCURRENT_BUILDS: usize = 3; + + fn build_sdk(grpc_url: String, uc_url: String) -> ZerobusSdk { + ZerobusSdk::builder() + .endpoint(grpc_url) + .unity_catalog_url(uc_url) + .tls_config(Arc::new(NoTlsConfig)) + .build() + .expect("SDK should build") + } + + /// One `CreateStream` success for each of `count` streams. + fn create_stream_responses(count: usize) -> Vec { + (0..count) + .map(|i| MockResponse::CreateStream { + stream_id: format!("s{i}"), + delay_ms: 0, + }) + .collect() + } + + #[tokio::test] + async fn shared_cache_reuses_token_across_streams() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + oauth + .set_responses(vec![MockTokenResponse::ok("token-1")]) + .await; + grpc.inject_responses( + TABLE_NAME, + vec![ + MockResponse::CreateStream { + stream_id: "s1".to_string(), + delay_ms: 0, + }, + MockResponse::CreateStream { + stream_id: "s2".to_string(), + delay_ms: 0, + }, + ], + ) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + let first = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!(first.is_ok(), "first stream: {:?}", first.err()); + + let second = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!(second.is_ok(), "second stream: {:?}", second.err()); + + assert_eq!( + oauth.mint_count(), + 1, + "both streams from one SDK must share a single cached token" + ); + } + + #[tokio::test] + async fn auth_rejection_invalidates_and_remints() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1"), + MockTokenResponse::ok("token-2"), + ]) + .await; + // The server rejects the first attempt's credential and accepts the retry's. + grpc.inject_responses( + TABLE_NAME, + vec![ + MockResponse::Error { + status: tonic::Status::unauthenticated("stale token"), + delay_ms: 0, + }, + MockResponse::CreateStream { + stream_id: "s1".to_string(), + delay_ms: 0, + }, + ], + ) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + let stream = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(true) + .recovery_retries(1) + .recovery_backoff_ms(0) + .build() + .await; + + assert!( + stream.is_ok(), + "the one-shot auth retry should succeed: {:?}", + stream.err() + ); + assert_eq!( + oauth.mint_count(), + 2, + "the auth rejection must invalidate the token and re-mint on the retry" + ); + } + + // Real (not virtual) time: the mint is real socket I/O, so a paused clock + // would auto-advance to the stream-creation timeout and fire it before the + // real mint could complete. This mirrors the repo's `test_timeouted_stream_creation`, + // which also bounds a real connection with a real millisecond budget. + #[tokio::test] + async fn hung_refresh_falls_back_to_cached_token() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // The first mint yields a short-lived token (60s < the 300s refresh + // buffer, so it is immediately in its refresh window); the proactive + // refresh triggered by the second stream then hangs and never replies. + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1").with_expires_in(Some("60")), + MockTokenResponse::Hang, + ]) + .await; + grpc.inject_responses( + TABLE_NAME, + vec![ + MockResponse::CreateStream { + stream_id: "s1".to_string(), + delay_ms: 0, + }, + MockResponse::CreateStream { + stream_id: "s2".to_string(), + delay_ms: 0, + }, + ], + ) + .await; + // refresh_timeout = recovery_timeout_ms / 2 = 500ms. The 60s cached token + // outlives that cap, so the bounded refresh applies and can fall back to + // it; the 1s outer budget leaves room for the fallback plus the connect. + let sdk = ZerobusSdk::builder() + .endpoint(grpc_url) + .unity_catalog_url(uc_url) + .tls_config(Arc::new(NoTlsConfig)) + .build() + .expect("SDK should build"); + + // First stream: a cold mint of the short-lived token. + let first = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .recovery_timeout_ms(1000) + .build() + .await; + assert!(first.is_ok(), "first stream: {:?}", first.err()); + assert_eq!(oauth.mint_count(), 1); + + // Second stream: the cached token is in its refresh window, so a + // proactive refresh fires and hangs. The 500ms cap must make the SDK + // fall back to the still-valid cached token rather than hang. + let second = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .recovery_timeout_ms(1000) + .build() + .await; + assert!( + second.is_ok(), + "a hung refresh must fall back to the cached token, not hang: {:?}", + second.err() + ); + assert_eq!( + oauth.mint_count(), + 2, + "the refresh must be attempted before the fallback" + ); + } + + #[tokio::test] + async fn retryable_token_error_is_retried_then_succeeds() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // The first mint hits a 5xx (retryable); the retry mints successfully. + oauth + .set_responses(vec![ + MockTokenResponse::error(503, "token endpoint unavailable"), + MockTokenResponse::ok("token-1"), + ]) + .await; + grpc.inject_responses( + TABLE_NAME, + vec![MockResponse::CreateStream { + stream_id: "s1".to_string(), + delay_ms: 0, + }], + ) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + let stream = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(true) + .recovery_retries(1) + .recovery_backoff_ms(0) + .build() + .await; + + assert!( + stream.is_ok(), + "a retryable token error should be retried into a successful mint: {:?}", + stream.err() + ); + assert_eq!( + oauth.mint_count(), + 2, + "the failed mint and its retry are both hits on the endpoint" + ); + } + + #[tokio::test] + async fn non_retryable_token_error_fails_without_retry() { + setup_tracing(); + // The mint fails before any connection is attempted, so the gRPC endpoint is + // never contacted and needs no scripted response (the server is kept only to + // give the SDK a valid endpoint). + let (_grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // A 4xx (bad credentials) is non-retryable and is not a server auth + // rejection, so the retry loop must not re-mint even with a budget. + oauth + .set_responses(vec![MockTokenResponse::error(401, "invalid_client")]) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + let stream = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(true) + .recovery_retries(3) + .recovery_backoff_ms(0) + .build() + .await; + + let err = stream + .err() + .expect("a 4xx mint error must fail stream creation"); + assert!( + matches!(err, ZerobusError::InvalidUCTokenError(_)), + "got {err:?}" + ); + assert_eq!( + oauth.mint_count(), + 1, + "a non-retryable token error must not be retried" + ); + } + + #[tokio::test] + async fn concurrent_stream_creation_mints_once() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // One token available; the concurrent creations must share it via + // single-flight rather than each minting one. + oauth + .set_responses(vec![MockTokenResponse::ok("token-1")]) + .await; + grpc.inject_responses(TABLE_NAME, create_stream_responses(CONCURRENT_BUILDS)) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + // Builds polled together on one task: whichever reaches the per-key slot + // first mints; the others wait on it and reuse the result. + let builds: Vec<_> = (0..CONCURRENT_BUILDS) + .map(|_| { + sdk.stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + }) + .collect(); + let results = join_all(builds).await; + + for (i, result) in results.iter().enumerate() { + assert!( + result.is_ok(), + "concurrent stream {i}: {:?}", + result.as_ref().err() + ); + } + assert_eq!( + oauth.mint_count(), + 1, + "concurrent creations must single-flight into one mint" + ); + } + + #[tokio::test] + async fn refresh_backoff_suppresses_a_concurrent_mint_stampede() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // Seed a short-lived token (already in its 300s refresh window), then fail + // the refresh. Only the slot-winner's refresh reaches the endpoint; the rest + // of the burst is suppressed by the post-fallback backoff. One 503 suffices: + // if the backoff failed to hold, the extra refreshers would fall through to + // the mock's default OK response and still push mint_count past 2. + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1").with_expires_in(Some("60")), + MockTokenResponse::error(503, "token endpoint down"), + ]) + .await; + // One CreateStream for the seed, plus one per concurrent build. + grpc.inject_responses(TABLE_NAME, create_stream_responses(1 + CONCURRENT_BUILDS)) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + // Seed the cache with the short-lived token. + let seed = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!(seed.is_ok(), "seed stream: {:?}", seed.err()); + assert_eq!(oauth.mint_count(), 1); + + // Concurrent burst: whichever build wins the slot refreshes (503 → falls + // back to the cached token and arms a ~5s backoff); the others wake inside + // the backoff window and serve the cached token without refreshing. + let builds: Vec<_> = (0..CONCURRENT_BUILDS) + .map(|_| { + sdk.stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + }) + .collect(); + let results = join_all(builds).await; + + for (i, result) in results.iter().enumerate() { + assert!( + result.is_ok(), + "burst stream {i}: {:?}", + result.as_ref().err() + ); + } + assert_eq!( + oauth.mint_count(), + 2, + "one cold mint plus exactly one failed refresh; backoff must suppress the rest of the burst" + ); + } + + #[tokio::test] + async fn repeated_auth_rejection_stops_after_one_remint() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // Every minted token is rejected. The initial-setup auth retry is one-shot, + // so the second rejection fails the build rather than re-minting again: the + // two scripted tokens are the initial mint and its single retry, and the two + // rejections are all the server ever sends. + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1"), + MockTokenResponse::ok("token-2"), + ]) + .await; + grpc.inject_responses( + TABLE_NAME, + vec![ + MockResponse::Error { + status: tonic::Status::unauthenticated("stale token"), + delay_ms: 0, + }, + MockResponse::Error { + status: tonic::Status::unauthenticated("stale token"), + delay_ms: 0, + }, + ], + ) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + // recovery_retries(3) leaves plenty of generic retry budget; the one-shot + // auth cap — not the budget — is what must stop the retries. + let stream = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(true) + .recovery_retries(3) + .recovery_backoff_ms(0) + .build() + .await; + + let err = stream + .err() + .expect("repeated auth rejection must fail stream creation"); + assert!( + !err.is_retryable(), + "the surfaced auth rejection is non-retryable: {err:?}" + ); + assert_eq!( + oauth.mint_count(), + 2, + "one-shot auth retry: the initial mint plus exactly one re-mint, then it gives up" + ); + } + + #[tokio::test] + async fn rejected_cached_token_forces_next_stream_to_remint() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-1"), + MockTokenResponse::ok("token-2"), + ]) + .await; + // Stream 1 succeeds; stream 2 (reusing the cached token) is rejected as if + // the token were revoked mid-lifetime; stream 3 succeeds with a fresh token. + grpc.inject_responses( + TABLE_NAME, + vec![ + MockResponse::CreateStream { + stream_id: "s1".to_string(), + delay_ms: 0, + }, + MockResponse::Error { + status: tonic::Status::unauthenticated("token revoked"), + delay_ms: 0, + }, + MockResponse::CreateStream { + stream_id: "s3".to_string(), + delay_ms: 0, + }, + ], + ) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + // 1) A successful create mints and caches token-1. + let first = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!(first.is_ok(), "first stream: {:?}", first.err()); + assert_eq!(oauth.mint_count(), 1); + + // 2) The next create reuses the cached token (no mint), but the server + // rejects it, so the connection invalidates the shared cache. With + // recovery off the stream fails without retrying. `mint_count` staying at + // 1 is what proves this stream reused the cached token rather than minting. + let second = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + let err = second.err().expect("the rejected stream must fail"); + assert!( + !err.is_retryable(), + "auth rejection is non-retryable: {err:?}" + ); + assert_eq!( + oauth.mint_count(), + 1, + "stream 2 must reuse the cached token, not mint a new one" + ); + + // 3) Because the reused token was invalidated, the next create must re-mint. + let third = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!( + third.is_ok(), + "third stream should re-mint and succeed: {:?}", + third.err() + ); + assert_eq!( + oauth.mint_count(), + 2, + "the rejection of the reused cached token must force the next stream to re-mint" + ); + } + + #[tokio::test] + async fn rejection_is_scoped_to_its_table() { + setup_tracing(); + let (grpc, grpc_url) = start_mock_server().await.unwrap(); + let (oauth, uc_url) = start_mock_oauth_server().await; + // Two tables share one SDK cache (distinct keys). A rejection on table A must + // invalidate only A's entry, leaving table B's cached token intact. + oauth + .set_responses(vec![ + MockTokenResponse::ok("token-a"), + MockTokenResponse::ok("token-b"), + ]) + .await; + grpc.inject_responses( + TABLE_NAME, + vec![ + MockResponse::CreateStream { + stream_id: "a1".to_string(), + delay_ms: 0, + }, + MockResponse::Error { + status: tonic::Status::unauthenticated("token revoked"), + delay_ms: 0, + }, + ], + ) + .await; + grpc.inject_responses( + OTHER_TABLE, + vec![ + MockResponse::CreateStream { + stream_id: "b1".to_string(), + delay_ms: 0, + }, + MockResponse::CreateStream { + stream_id: "b2".to_string(), + delay_ms: 0, + }, + ], + ) + .await; + let sdk = build_sdk(grpc_url, uc_url); + + // Seed one cached token per table. + let a1 = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!(a1.is_ok(), "table A stream 1: {:?}", a1.err()); + let b1 = sdk + .stream_builder() + .table(OTHER_TABLE) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!(b1.is_ok(), "table B stream 1: {:?}", b1.err()); + assert_eq!(oauth.mint_count(), 2, "one mint per table"); + + // A stream on table A reuses A's cached token, is rejected, and invalidates + // only A's cache entry. + let a2 = sdk + .stream_builder() + .table(TABLE_NAME) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!(a2.is_err(), "table A stream 2 must be rejected"); + + // Table B's cached token must be untouched, so a later B stream reuses it. + let b2 = sdk + .stream_builder() + .table(OTHER_TABLE) + .oauth(CLIENT_ID, CLIENT_SECRET) + .json() + .recovery(false) + .build() + .await; + assert!( + b2.is_ok(), + "table B stream 2 should reuse B's cached token: {:?}", + b2.err() + ); + assert_eq!( + oauth.mint_count(), + 2, + "table B must not re-mint: a rejection on table A is scoped to A's cache entry" + ); + } +}