From be3989f567dd5eb1f5d1a8462cb909a5b9f9f8b0 Mon Sep 17 00:00:00 2001 From: Siva Date: Thu, 13 Aug 2026 01:00:25 -0700 Subject: [PATCH 1/3] fix(storage): re-read the Azure federated token file on every exchange `WorkloadIdentityCredential` reads the file pointed to by `AZURE_FEDERATED_TOKEN_FILE` once, when the credential is built, and reuses that client assertion for the lifetime of the process. Kubernetes rotates the projected service account token roughly hourly, while the access token obtained from Entra ID lives ~24h, so no token exchange is attempted for a day. By then the assertion the credential holds is long expired: Entra rejects the exchange with `AADSTS700024`, `azure_core` classifies the 401 as non-retryable, and the indexing pipeline dies with no path to recovery short of a restart. In practice only the indexer is affected. `StorageResolver::resolve()` builds a fresh `Storage` per call, so the searcher and the janitor re-read the token file on each operation, while the indexer holds one `Storage` for the life of the pipeline. Upgrading `azure_identity` is not an option here: `azure_storage_blobs` was never published past 0.21.0 and requires `azure_core ^0.21`, while `azure_identity >= 0.22` implements a different `TokenCredential` trait, so a bump means porting the whole Azure backend to the rewritten SDK. Add `RefreshingWorkloadIdentityCredential`, which runs the same federated credentials flow but reads the assertion from disk on every exchange and refreshes shortly before the access token expires rather than after. It is used only when all three workload identity environment variables are set; otherwise `from_uri()` falls through to the stock `azure_identity` credential chain, so access key and non-Azure paths are unchanged. The cached token is held in an `ArcSwapOption` rather than behind a lock, so no lock is held across the token exchange; refreshes are rare and idempotent, so a race between two callers is harmless. Cache entries record the scopes they were minted for, so a token is never reused for scopes it was not issued against. Tests inject a fake token endpoint through the `HttpClient` seam that `federated_credentials_flow::perform` already exposes. It accepts only the assertion the token file currently holds and rejects anything else with `AADSTS700024`, mirroring the real endpoint. One test reproduces the failure against the stock `WorkloadIdentityCredential`, another shows the new credential picks up the rotated assertion and recovers. Closes #6672 --- quickwit/Cargo.lock | 2 + quickwit/quickwit-storage/Cargo.toml | 4 + .../src/object_storage/azure_blob_storage.rs | 6 + .../object_storage/azure_workload_identity.rs | 556 ++++++++++++++++++ .../src/object_storage/mod.rs | 2 + 5 files changed, 570 insertions(+) create mode 100644 quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs diff --git a/quickwit/Cargo.lock b/quickwit/Cargo.lock index cd19a875db9..06da97094f8 100644 --- a/quickwit/Cargo.lock +++ b/quickwit/Cargo.lock @@ -9397,6 +9397,7 @@ version = "0.9.0" dependencies = [ "ahash", "anyhow", + "arc-swap", "async-trait", "aws-config", "aws-credential-types", @@ -9438,6 +9439,7 @@ dependencies = [ "tantivy", "tempfile", "thiserror 2.0.18", + "time", "tokio", "tokio-rustls 0.26.4", "tokio-stream", diff --git a/quickwit/quickwit-storage/Cargo.toml b/quickwit/quickwit-storage/Cargo.toml index cda13362a79..3ae8f5132f6 100644 --- a/quickwit/quickwit-storage/Cargo.toml +++ b/quickwit/quickwit-storage/Cargo.toml @@ -48,6 +48,8 @@ azure_core = { workspace = true, optional = true } azure_identity = { workspace = true, optional = true } azure_storage = { workspace = true, optional = true } azure_storage_blobs = { workspace = true, optional = true } +arc-swap = { workspace = true, optional = true } +time = { workspace = true, optional = true } quickwit-aws = { workspace = true } quickwit-common = { workspace = true } @@ -81,6 +83,8 @@ quickwit-common = { workspace = true, features = ["testsuite"] } [features] azure = [ + "dep:arc-swap", + "dep:time", "azure_core", "azure_identity", "azure_identity/enable_reqwest_rustls", diff --git a/quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs b/quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs index cbbf6ea8ded..247786b0102 100644 --- a/quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs +++ b/quickwit/quickwit-storage/src/object_storage/azure_blob_storage.rs @@ -44,6 +44,7 @@ use tokio_util::compat::FuturesAsyncReadCompatExt; use tokio_util::io::StreamReader; use tracing::{info, instrument, warn}; +use super::azure_workload_identity::RefreshingWorkloadIdentityCredential; use crate::debouncer::DebouncedStorage; use crate::metrics::object_storage_get_slice_in_flight_guards; use crate::stable_deref_bytes::into_owned_bytes; @@ -184,6 +185,11 @@ impl AzureBlobStorage { azure_storage_config.resolve_access_key() { StorageCredentials::access_key(storage_account_name.clone(), access_key) + } else if let Some(credential) = RefreshingWorkloadIdentityCredential::from_env() { + // `azure_identity` 0.21 caches the federated token file contents for the process + // lifetime, which breaks workload identity ~24h after startup. See the module docs + // of `azure_workload_identity`. + StorageCredentials::token_credential(Arc::new(credential)) } else if let Ok(credential) = azure_identity::create_credential() { StorageCredentials::token_credential(credential) } else { diff --git a/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs b/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs new file mode 100644 index 00000000000..b9e34f540ef --- /dev/null +++ b/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs @@ -0,0 +1,556 @@ +// Copyright 2021-Present Datadog, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Workload identity credential that re-reads the federated token file on every exchange. +//! +//! `azure_identity` 0.21's `WorkloadIdentityCredential` reads the file pointed to by +//! `AZURE_FEDERATED_TOKEN_FILE` once, when the credential is built, and reuses that client +//! assertion for the lifetime of the process. Kubernetes rotates the projected service account +//! token roughly hourly, while the access token obtained from Entra ID lives ~24h, so no token +//! exchange is even attempted for a day. By then the assertion on hand is long expired: Entra +//! rejects the exchange with `AADSTS700024`, the SDK classifies that as non-retryable, and the +//! indexing pipeline dies with no self-healing. +//! +//! This credential performs the same token exchange but reads the assertion from disk each time, +//! and refreshes before the access token expires rather than after. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; + +use arc_swap::ArcSwapOption; +use async_trait::async_trait; +use azure_core::auth::{AccessToken, TokenCredential}; +use azure_core::error::{Error as AzureCoreError, ErrorKind, ResultExt}; +use azure_core::{HttpClient, Url}; +use azure_identity::{TokenCredentialOptions, federated_credentials_flow}; +use time::OffsetDateTime; +use tracing::{debug, info}; + +const AZURE_TENANT_ID_ENV_VAR: &str = "AZURE_TENANT_ID"; +const AZURE_CLIENT_ID_ENV_VAR: &str = "AZURE_CLIENT_ID"; +const AZURE_FEDERATED_TOKEN_FILE_ENV_VAR: &str = "AZURE_FEDERATED_TOKEN_FILE"; + +/// Refresh the access token this long before it actually expires, so that a failed refresh is +/// retried while the token in hand is still usable. +const EXPIRATION_MARGIN: Duration = Duration::from_secs(5 * 60); + +/// An access token together with the scopes it was minted for. +/// +/// A token is only valid for the scopes it was requested with, so the scopes are part of the +/// cache entry: a request for different scopes must trigger a new exchange rather than reuse +/// whatever happens to be cached. +#[derive(Debug)] +struct CachedToken { + scopes: Vec, + access_token: AccessToken, +} + +impl CachedToken { + fn matches(&self, scopes: &[&str]) -> bool { + self.scopes.len() == scopes.len() + && self + .scopes + .iter() + .zip(scopes) + .all(|(cached_scope, scope)| cached_scope == scope) + } +} + +/// A [`TokenCredential`] for Azure AD Workload Identity that reads the federated token file on +/// every token exchange rather than caching its contents for the lifetime of the process. +/// +/// The cached access token is stored in an [`ArcSwapOption`] rather than behind a lock: refreshes +/// are rare (roughly daily) and idempotent, so two callers racing to refresh is harmless and +/// preferable to holding a lock across the token exchange. +/// +/// Only the most recently requested scopes are cached. Quickwit's Azure backend always requests +/// the same storage scope, so this single slot never thrashes in practice. +#[derive(Debug)] +pub(super) struct RefreshingWorkloadIdentityCredential { + http_client: Arc, + authority_host: Url, + tenant_id: String, + client_id: String, + token_file_path: PathBuf, + /// How long before actual expiry a cached token is treated as expiring. + expiration_margin: Duration, + cached_token_opt: ArcSwapOption, +} + +impl RefreshingWorkloadIdentityCredential { + fn new( + http_client: Arc, + authority_host: Url, + tenant_id: String, + client_id: String, + token_file_path: PathBuf, + expiration_margin: Duration, + ) -> Self { + Self { + http_client, + authority_host, + tenant_id, + client_id, + token_file_path, + expiration_margin, + cached_token_opt: ArcSwapOption::empty(), + } + } + + /// Builds a credential from the standard workload identity environment variables. + /// + /// Returns `None` when the process is not configured for workload identity, in which case the + /// caller should fall back to the default `azure_identity` credential chain. + pub(super) fn from_env() -> Option { + let tenant_id = non_empty_env_var(AZURE_TENANT_ID_ENV_VAR)?; + let client_id = non_empty_env_var(AZURE_CLIENT_ID_ENV_VAR)?; + let token_file_path = non_empty_env_var(AZURE_FEDERATED_TOKEN_FILE_ENV_VAR)?; + let options = TokenCredentialOptions::default(); + let authority_host = options.authority_host().ok()?; + + info!( + token_file_path = %token_file_path, + "using azure workload identity credential with per-exchange token file reads" + ); + Some(Self::new( + options.http_client(), + authority_host, + tenant_id, + client_id, + PathBuf::from(token_file_path), + EXPIRATION_MARGIN, + )) + } + + /// Reads the client assertion from disk and exchanges it for an access token. + async fn exchange_token(&self, scopes: &[&str]) -> azure_core::Result { + let assertion = read_token_file(&self.token_file_path)?; + let login_response = federated_credentials_flow::perform( + self.http_client.clone(), + &self.client_id, + assertion.trim(), + scopes, + &self.tenant_id, + &self.authority_host, + ) + .await + .context(ErrorKind::Credential, "federated token exchange failed")?; + + let expires_on = OffsetDateTime::now_utc() + Duration::from_secs(login_response.expires_in); + debug!(%expires_on, "obtained azure access token from federated credentials flow"); + Ok(AccessToken::new( + login_response.access_token().clone(), + expires_on, + )) + } +} + +#[async_trait] +impl TokenCredential for RefreshingWorkloadIdentityCredential { + async fn get_token(&self, scopes: &[&str]) -> azure_core::Result { + if let Some(cached_token) = self.cached_token_opt.load_full() + && cached_token.matches(scopes) + && !is_expiring(&cached_token.access_token, self.expiration_margin) + { + return Ok(cached_token.access_token.clone()); + } + let access_token = self.exchange_token(scopes).await?; + self.cached_token_opt.store(Some(Arc::new(CachedToken { + scopes: scopes.iter().map(|scope| scope.to_string()).collect(), + access_token: access_token.clone(), + }))); + Ok(access_token) + } + + async fn clear_cache(&self) -> azure_core::Result<()> { + self.cached_token_opt.store(None); + Ok(()) + } +} + +fn is_expiring(access_token: &AccessToken, expiration_margin: Duration) -> bool { + access_token.expires_on - expiration_margin <= OffsetDateTime::now_utc() +} + +fn non_empty_env_var(key: &str) -> Option { + let value = std::env::var(key).ok()?; + if value.trim().is_empty() { + return None; + } + Some(value) +} + +fn read_token_file(token_file_path: &Path) -> azure_core::Result { + std::fs::read_to_string(token_file_path).map_err(|error| { + AzureCoreError::full( + ErrorKind::Credential, + error, + format!( + "failed to read azure federated token file `{}`", + token_file_path.display() + ), + ) + }) +} + +#[cfg(test)] +mod tests { + use std::io::Write; + use std::sync::Mutex; + + use azure_core::auth::Secret; + use azure_core::headers::Headers; + use azure_core::{Request as AzureRequest, Response as AzureResponse, StatusCode}; + use azure_identity::WorkloadIdentityCredential; + use bytes::Bytes; + use tempfile::NamedTempFile; + + use super::*; + + /// Verbatim shape of what Entra returns when a federated client assertion has expired. This + /// is the response that kills the indexer at T+24h: `azure_core` classifies it as + /// non-retryable, so the pipeline dies rather than backing off and retrying. + const AADSTS700024_BODY: &str = r#"{"error":"invalid_client","error_description":"AADSTS700024: Client assertion is not within its valid time range. Current time: 2026-08-13T00:00:00Z, assertion valid from 2026-08-11T23:00:00Z, expiry time of assertion 2026-08-12T00:00:00Z.","error_codes":[700024]}"#; + + /// A stand-in for the Entra token endpoint. + /// + /// It accepts exactly one client assertion — the one the projected service account token file + /// currently holds — and rejects anything else with `AADSTS700024`, which is precisely how the + /// real endpoint treats a stale assertion. Every assertion presented is recorded, so a test + /// can assert on what the credential actually sent rather than only on the outcome. + #[derive(Debug)] + struct FakeEntra { + accepted_assertion: Mutex, + presented_assertions: Mutex>, + expires_in_secs: u64, + } + + impl FakeEntra { + fn new(accepted_assertion: &str, expires_in_secs: u64) -> Arc { + Arc::new(Self { + accepted_assertion: Mutex::new(accepted_assertion.to_string()), + presented_assertions: Mutex::new(Vec::new()), + expires_in_secs, + }) + } + + /// Simulates the previous assertion expiring: only `assertion` is accepted from now on. + fn rotate_accepted_assertion(&self, assertion: &str) { + *self.accepted_assertion.lock().unwrap() = assertion.to_string(); + } + + fn presented_assertions(&self) -> Vec { + self.presented_assertions.lock().unwrap().clone() + } + + fn exchange_count(&self) -> usize { + self.presented_assertions.lock().unwrap().len() + } + } + + fn azure_response(status: StatusCode, body: &str) -> AzureResponse { + let bytes = Bytes::from(body.to_string()); + AzureResponse::new( + status, + Headers::new(), + Box::pin(futures::stream::iter(std::iter::once(Ok(bytes)))), + ) + } + + /// Extracts a field from an `application/x-www-form-urlencoded` body. Test assertions are + /// restricted to characters that survive form encoding unchanged, so no decoding is needed. + fn form_field(body: &str, key: &str) -> Option { + for pair in body.split('&') { + let (pair_key, pair_value) = pair.split_once('=')?; + if pair_key == key { + return Some(pair_value.to_string()); + } + } + None + } + + #[async_trait] + impl HttpClient for FakeEntra { + async fn execute_request( + &self, + request: &AzureRequest, + ) -> azure_core::Result { + let azure_core::Body::Bytes(body_bytes) = request.body() else { + panic!("expected a bytes body on the token request"); + }; + let body = std::str::from_utf8(body_bytes).unwrap(); + let assertion = form_field(body, "client_assertion") + .expect("token request must carry a client_assertion"); + self.presented_assertions + .lock() + .unwrap() + .push(assertion.clone()); + + if assertion != *self.accepted_assertion.lock().unwrap() { + return Ok(azure_response(StatusCode::Unauthorized, AADSTS700024_BODY)); + } + let body = format!( + r#"{{"token_type":"Bearer","expires_in":{},"ext_expires_in":{},"access_token":"access-token-for-{}"}}"#, + self.expires_in_secs, self.expires_in_secs, assertion + ); + Ok(azure_response(StatusCode::Ok, &body)) + } + } + + fn token_file_containing(assertion: &str) -> NamedTempFile { + let mut token_file = NamedTempFile::new().unwrap(); + token_file.write_all(assertion.as_bytes()).unwrap(); + token_file.flush().unwrap(); + token_file + } + + /// Simulates kubelet rotating the projected service account token in place. + fn rotate_token_file(token_file: &NamedTempFile, assertion: &str) { + std::fs::write(token_file.path(), assertion.as_bytes()).unwrap(); + } + + fn authority_host() -> Url { + Url::parse("https://login.microsoftonline.com").unwrap() + } + + const STORAGE_SCOPE: &str = "https://storage.azure.com/.default"; + + fn refreshing_credential( + fake_entra: Arc, + token_file: &NamedTempFile, + expiration_margin: Duration, + ) -> RefreshingWorkloadIdentityCredential { + RefreshingWorkloadIdentityCredential::new( + fake_entra, + authority_host(), + "test-tenant-id".to_string(), + "test-client-id".to_string(), + token_file.path().to_path_buf(), + expiration_margin, + ) + } + + fn access_token_expiring_in(duration: Duration) -> AccessToken { + AccessToken::new( + Secret::new("token".to_string()), + OffsetDateTime::now_utc() + duration, + ) + } + + #[test] + fn test_is_expiring_within_margin() { + assert!(is_expiring( + &access_token_expiring_in(Duration::from_secs(60)), + EXPIRATION_MARGIN + )); + assert!(is_expiring( + &access_token_expiring_in(EXPIRATION_MARGIN), + EXPIRATION_MARGIN + )); + } + + #[test] + fn test_is_not_expiring_outside_margin() { + assert!(!is_expiring( + &access_token_expiring_in(EXPIRATION_MARGIN + Duration::from_secs(60)), + EXPIRATION_MARGIN + )); + assert!(!is_expiring( + &access_token_expiring_in(Duration::from_secs(24 * 3600)), + EXPIRATION_MARGIN + )); + } + + #[test] + fn test_read_token_file_picks_up_rotations() { + let token_file = token_file_containing("first-assertion"); + assert_eq!( + read_token_file(token_file.path()).unwrap(), + "first-assertion" + ); + rotate_token_file(&token_file, "second-assertion"); + assert_eq!( + read_token_file(token_file.path()).unwrap(), + "second-assertion" + ); + } + + /// Sanity check that the fake endpoint has teeth: if it accepted anything, the two scenario + /// tests below would pass even with a credential that never re-reads the file. + #[tokio::test] + async fn test_fake_entra_rejects_a_stale_assertion() { + let fake_entra = FakeEntra::new("assertion-hour-24", 86_400); + let token_file = token_file_containing("assertion-hour-0"); + let credential = + refreshing_credential(fake_entra.clone(), &token_file, Duration::from_secs(5 * 60)); + + let error = credential.get_token(&[STORAGE_SCOPE]).await.unwrap_err(); + assert!(matches!(error.kind(), ErrorKind::Credential)); + assert_eq!(fake_entra.presented_assertions(), vec!["assertion-hour-0"]); + } + + /// Reproduces AE-16269 against the real upstream credential. + /// + /// `WorkloadIdentityCredential` is built while the token file holds the hour-0 assertion. + /// Kubelet then rotates the file, and Entra stops honouring the old assertion — exactly the + /// state of the world at T+24h. The credential still presents the assertion it captured at + /// construction and is rejected with `AADSTS700024`, with no path to recovery short of a + /// process restart. + #[tokio::test] + async fn test_upstream_credential_ignores_the_rotated_token_file() { + let fake_entra = FakeEntra::new("assertion-hour-0", 86_400); + let token_file = token_file_containing("assertion-hour-0"); + + // `WorkloadIdentityCredential::create` reads the file exactly here, and never again. + let assertion_at_startup = std::fs::read_to_string(token_file.path()).unwrap(); + let upstream_credential = WorkloadIdentityCredential::new( + fake_entra.clone(), + authority_host(), + "test-tenant-id".to_string(), + "test-client-id".to_string(), + assertion_at_startup, + ); + upstream_credential + .get_token(&[STORAGE_SCOPE]) + .await + .expect("the first exchange succeeds while the startup assertion is still valid"); + + // 24h later: kubelet has rotated the file many times over and Entra has aged out the + // assertion the credential is holding. + rotate_token_file(&token_file, "assertion-hour-24"); + fake_entra.rotate_accepted_assertion("assertion-hour-24"); + upstream_credential.clear_cache().await.unwrap(); + + let error = upstream_credential + .get_token(&[STORAGE_SCOPE]) + .await + .unwrap_err(); + assert!(matches!(error.kind(), ErrorKind::Credential)); + // Note that `AADSTS700024` does not survive into the error: `azure_core` keeps the status + // and drops the response body, which is a large part of why this failure was hard to + // diagnose from indexer logs in the first place. + assert!( + format!("{error:?}").contains("Unauthorized"), + "expected the stale assertion to be rejected, got: {error:?}" + ); + assert_eq!( + fake_entra.presented_assertions(), + vec!["assertion-hour-0", "assertion-hour-0"], + "upstream re-presents the assertion it captured at startup, never re-reading the file" + ); + } + + /// The same scenario, with the fix: the rotated assertion is picked up and the exchange + /// succeeds, so the indexer survives past T+24h without a restart. + #[tokio::test] + async fn test_refreshing_credential_survives_token_rotation() { + let fake_entra = FakeEntra::new("assertion-hour-0", 86_400); + let token_file = token_file_containing("assertion-hour-0"); + // A margin wider than the token lifetime makes every cached token count as expiring, which + // is how we reach the refresh path deterministically instead of waiting 24h. + let credential = refreshing_credential( + fake_entra.clone(), + &token_file, + Duration::from_secs(48 * 3600), + ); + + credential + .get_token(&[STORAGE_SCOPE]) + .await + .expect("the first exchange succeeds"); + + rotate_token_file(&token_file, "assertion-hour-24"); + fake_entra.rotate_accepted_assertion("assertion-hour-24"); + + let access_token = credential + .get_token(&[STORAGE_SCOPE]) + .await + .expect("the refresh picks up the rotated assertion and succeeds"); + assert_eq!( + access_token.token.secret(), + "access-token-for-assertion-hour-24" + ); + assert_eq!( + fake_entra.presented_assertions(), + vec!["assertion-hour-0", "assertion-hour-24"], + "the credential re-reads the token file on every exchange" + ); + } + + /// The fix must not turn one exchange per day into one per request. + #[tokio::test] + async fn test_token_is_cached_until_it_approaches_expiry() { + let fake_entra = FakeEntra::new("assertion-hour-0", 86_400); + let token_file = token_file_containing("assertion-hour-0"); + let credential = + refreshing_credential(fake_entra.clone(), &token_file, Duration::from_secs(5 * 60)); + + for _ in 0..10 { + credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + } + assert_eq!(fake_entra.exchange_count(), 1); + + // A different scope must not reuse the cached token. + credential + .get_token(&["https://storage.azure.us/.default"]) + .await + .unwrap(); + assert_eq!(fake_entra.exchange_count(), 2); + } + + /// A failed exchange must leave the credential able to recover on the next call, rather than + /// wedging the way the upstream one does. + #[tokio::test] + async fn test_failed_exchange_recovers_on_next_call() { + let fake_entra = FakeEntra::new("assertion-hour-24", 86_400); + let token_file = token_file_containing("assertion-hour-0"); + let credential = + refreshing_credential(fake_entra.clone(), &token_file, Duration::from_secs(5 * 60)); + + credential + .get_token(&[STORAGE_SCOPE]) + .await + .expect_err("the stale assertion is rejected"); + + // Kubelet rotates the file; no restart, no intervention. + rotate_token_file(&token_file, "assertion-hour-24"); + credential + .get_token(&[STORAGE_SCOPE]) + .await + .expect("the credential self-heals once the file is rotated"); + } + + #[test] + fn test_cached_token_matches_scopes() { + let cached_token = CachedToken { + scopes: vec!["https://storage.azure.com/.default".to_string()], + access_token: access_token_expiring_in(Duration::from_secs(3600)), + }; + assert!(cached_token.matches(&["https://storage.azure.com/.default"])); + // A token minted for the public cloud must not be reused for a sovereign cloud scope. + assert!(!cached_token.matches(&["https://storage.azure.us/.default"])); + assert!(!cached_token.matches(&[])); + assert!(!cached_token.matches(&[ + "https://storage.azure.com/.default", + "https://storage.azure.us/.default", + ])); + } + + #[test] + fn test_read_token_file_missing() { + let error = read_token_file(Path::new("/does/not/exist")).unwrap_err(); + assert!(matches!(error.kind(), ErrorKind::Credential)); + } +} diff --git a/quickwit/quickwit-storage/src/object_storage/mod.rs b/quickwit/quickwit-storage/src/object_storage/mod.rs index e914c107291..ab327497bbd 100644 --- a/quickwit/quickwit-storage/src/object_storage/mod.rs +++ b/quickwit/quickwit-storage/src/object_storage/mod.rs @@ -26,4 +26,6 @@ mod s3_compatible_storage_resolver; #[cfg(feature = "azure")] mod azure_blob_storage; #[cfg(feature = "azure")] +mod azure_workload_identity; +#[cfg(feature = "azure")] pub use self::azure_blob_storage::{AzureBlobStorage, AzureBlobStorageFactory}; From 51f1fcb6b094323b9c2e93ebb1ae1da936a20c7c Mon Sep 17 00:00:00 2001 From: Siva Date: Thu, 13 Aug 2026 03:20:50 -0700 Subject: [PATCH 2/3] fix(storage): narrow when the refreshing credential is used, and survive a failed refresh Addresses review feedback on #6687. Defer to the stock credential chain in two cases where it would not have built a file-backed workload identity credential either, even though a projected token file is present: - an explicit `AZURE_CREDENTIAL_KIND` selecting another provider. The workload identity environment variables are injected automatically by the webhook, so they are routinely present even when an operator deliberately configured a different credential. Only `environment` and `workloadidentity` resolve to a workload identity credential upstream, so only those two are stood in for, using the same case and space normalization `SpecificAzureCredential` applies. - an inline `AZURE_FEDERATED_TOKEN`, which upstream prefers over the file. Reading the file instead could authenticate as a different subject, and an inline assertion has nothing to re-read. Also stop failing storage operations when a refresh fails inside the expiration margin. The margin exists precisely so a failed refresh can be retried while the token in hand is still usable, but the error was propagated regardless, so a transient hiccup near expiry would fail requests for up to five minutes even though a valid token was cached. Serve the cached token until it genuinely expires, and surface the error only once there is nothing left to fall back on. --- .../object_storage/azure_workload_identity.rs | 154 +++++++++++++++++- 1 file changed, 146 insertions(+), 8 deletions(-) diff --git a/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs b/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs index b9e34f540ef..e0f013bb93b 100644 --- a/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs +++ b/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs @@ -36,11 +36,18 @@ use azure_core::error::{Error as AzureCoreError, ErrorKind, ResultExt}; use azure_core::{HttpClient, Url}; use azure_identity::{TokenCredentialOptions, federated_credentials_flow}; use time::OffsetDateTime; -use tracing::{debug, info}; +use tracing::{debug, info, warn}; const AZURE_TENANT_ID_ENV_VAR: &str = "AZURE_TENANT_ID"; const AZURE_CLIENT_ID_ENV_VAR: &str = "AZURE_CLIENT_ID"; const AZURE_FEDERATED_TOKEN_FILE_ENV_VAR: &str = "AZURE_FEDERATED_TOKEN_FILE"; +const AZURE_FEDERATED_TOKEN_ENV_VAR: &str = "AZURE_FEDERATED_TOKEN"; +const AZURE_CREDENTIAL_KIND_ENV_VAR: &str = "AZURE_CREDENTIAL_KIND"; + +/// The only two `AZURE_CREDENTIAL_KIND` values that resolve to a workload identity credential +/// upstream, and therefore the only two this credential may stand in for. +const CREDENTIAL_KIND_ENVIRONMENT: &str = "environment"; +const CREDENTIAL_KIND_WORKLOAD_IDENTITY: &str = "workloadidentity"; /// Refresh the access token this long before it actually expires, so that a failed refresh is /// retried while the token in hand is still usable. @@ -114,6 +121,12 @@ impl RefreshingWorkloadIdentityCredential { /// Returns `None` when the process is not configured for workload identity, in which case the /// caller should fall back to the default `azure_identity` credential chain. pub(super) fn from_env() -> Option { + if !handles_env( + non_empty_env_var(AZURE_CREDENTIAL_KIND_ENV_VAR).as_deref(), + non_empty_env_var(AZURE_FEDERATED_TOKEN_ENV_VAR).as_deref(), + ) { + return None; + } let tenant_id = non_empty_env_var(AZURE_TENANT_ID_ENV_VAR)?; let client_id = non_empty_env_var(AZURE_CLIENT_ID_ENV_VAR)?; let token_file_path = non_empty_env_var(AZURE_FEDERATED_TOKEN_FILE_ENV_VAR)?; @@ -160,18 +173,38 @@ impl RefreshingWorkloadIdentityCredential { #[async_trait] impl TokenCredential for RefreshingWorkloadIdentityCredential { async fn get_token(&self, scopes: &[&str]) -> azure_core::Result { - if let Some(cached_token) = self.cached_token_opt.load_full() + let cached_token_opt = self.cached_token_opt.load_full(); + if let Some(cached_token) = &cached_token_opt && cached_token.matches(scopes) && !is_expiring(&cached_token.access_token, self.expiration_margin) { return Ok(cached_token.access_token.clone()); } - let access_token = self.exchange_token(scopes).await?; - self.cached_token_opt.store(Some(Arc::new(CachedToken { - scopes: scopes.iter().map(|scope| scope.to_string()).collect(), - access_token: access_token.clone(), - }))); - Ok(access_token) + let exchange_result = self.exchange_token(scopes).await; + let error = match exchange_result { + Ok(access_token) => { + self.cached_token_opt.store(Some(Arc::new(CachedToken { + scopes: scopes.iter().map(|scope| scope.to_string()).collect(), + access_token: access_token.clone(), + }))); + return Ok(access_token); + } + Err(error) => error, + }; + // The refresh margin exists so that a failed refresh can be retried while the token in + // hand is still usable. Keep serving the cached token until it genuinely expires rather + // than failing storage operations for a transient hiccup near expiry. + if let Some(cached_token) = &cached_token_opt + && cached_token.matches(scopes) + && !is_expiring(&cached_token.access_token, Duration::ZERO) + { + warn!( + %error, + "azure token refresh failed, reusing the cached token until it expires" + ); + return Ok(cached_token.access_token.clone()); + } + Err(error) } async fn clear_cache(&self) -> azure_core::Result<()> { @@ -180,6 +213,35 @@ impl TokenCredential for RefreshingWorkloadIdentityCredential { } } +/// Whether this credential should stand in for the stock `azure_identity` chain, given the +/// relevant environment. +/// +/// Two cases are deliberately left to the stock chain even though a projected token file is +/// present, because in both of them the stock chain would not have built a file-backed workload +/// identity credential either: +/// +/// - an explicit `AZURE_CREDENTIAL_KIND` selecting some other provider. The workload identity +/// environment variables are injected automatically by the workload identity webhook, so they are +/// routinely present even when an operator has deliberately configured a different credential. +/// - an inline `AZURE_FEDERATED_TOKEN`, which upstream prefers over the file. Honouring the file +/// instead could authenticate as a different subject, and an inline assertion has nothing to +/// re-read anyway. +fn handles_env( + credential_kind_opt: Option<&str>, + inline_federated_token_opt: Option<&str>, +) -> bool { + if let Some(credential_kind) = credential_kind_opt { + // Mirrors the normalization `SpecificAzureCredential::create` applies. + let credential_kind = credential_kind.replace(' ', "").to_lowercase(); + if credential_kind != CREDENTIAL_KIND_ENVIRONMENT + && credential_kind != CREDENTIAL_KIND_WORKLOAD_IDENTITY + { + return false; + } + } + inline_federated_token_opt.is_none() +} + fn is_expiring(access_token: &AccessToken, expiration_margin: Duration) -> bool { access_token.expires_on - expiration_margin <= OffsetDateTime::now_utc() } @@ -510,6 +572,82 @@ mod tests { assert_eq!(fake_entra.exchange_count(), 2); } + #[test] + fn test_handles_env_defers_to_an_explicit_credential_kind() { + // The workload identity variables are injected automatically by the webhook, so they are + // present + // even when an operator has explicitly selected a different provider. + assert!(!handles_env(Some("azurecli"), None)); + assert!(!handles_env(Some("virtualmachine"), None)); + assert!(!handles_env(Some("appservice"), None)); + assert!(!handles_env(Some("clientsecret"), None)); + + // These two do resolve to a workload identity credential upstream. + assert!(handles_env(Some("environment"), None)); + assert!(handles_env(Some("workloadidentity"), None)); + // Upstream normalizes case and spaces before matching, so we must too. + assert!(handles_env(Some(" Workload Identity "), None)); + assert!(!handles_env(Some(" Azure CLI "), None)); + + assert!(handles_env(None, None)); + } + + #[test] + fn test_handles_env_defers_to_an_inline_assertion() { + // Upstream prefers `AZURE_FEDERATED_TOKEN` over the file, and an inline assertion has + // nothing to re-read. + assert!(!handles_env(None, Some("inline-assertion"))); + assert!(!handles_env( + Some("workloadidentity"), + Some("inline-assertion") + )); + } + + /// A transient failure near expiry must not fail storage operations while the cached token is + /// still valid — that is the entire point of refreshing ahead of expiry. + #[tokio::test] + async fn test_refresh_failure_falls_back_to_the_still_valid_cached_token() { + let fake_entra = FakeEntra::new("assertion-hour-0", 3600); + let token_file = token_file_containing("assertion-hour-0"); + // Wider than the token lifetime, so the token is always inside the refresh margin while + // remaining a long way from actually expiring. + let credential = refreshing_credential( + fake_entra.clone(), + &token_file, + Duration::from_secs(48 * 3600), + ); + + let first_token = credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + assert_eq!(fake_entra.exchange_count(), 1); + + // Entra starts rejecting: a hiccup, a revoked assertion, anything transient. + fake_entra.rotate_accepted_assertion("some-other-assertion"); + + let second_token = credential + .get_token(&[STORAGE_SCOPE]) + .await + .expect("the still-valid cached token is served rather than failing the request"); + assert_eq!(second_token.token.secret(), first_token.token.secret()); + assert_eq!(fake_entra.exchange_count(), 2, "the refresh was attempted"); + } + + /// Once the cached token has genuinely expired there is nothing to fall back on, so the error + /// must surface rather than handing out an unusable token. + #[tokio::test] + async fn test_refresh_failure_surfaces_once_the_cached_token_expires() { + // `expires_in: 0` means the token is already expired when it is cached. + let fake_entra = FakeEntra::new("assertion-hour-0", 0); + let token_file = token_file_containing("assertion-hour-0"); + let credential = + refreshing_credential(fake_entra.clone(), &token_file, Duration::from_secs(5 * 60)); + + credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + fake_entra.rotate_accepted_assertion("some-other-assertion"); + + let error = credential.get_token(&[STORAGE_SCOPE]).await.unwrap_err(); + assert!(matches!(error.kind(), ErrorKind::Credential)); + } + /// A failed exchange must leave the credential able to recover on the next call, rather than /// wedging the way the upstream one does. #[tokio::test] From 8c4b33346a7bb3fa02fa03c85c05bb6231ec4af2 Mon Sep 17 00:00:00 2001 From: Siva Date: Thu, 13 Aug 2026 04:35:35 -0700 Subject: [PATCH 3/3] fix(storage): bound token exchange attempts and guard the expiry computation `get_token` runs on every storage request and the Azure backend keeps up to `max_concurrent_uploads` requests in flight, so the refresh path was reachable by a hundred callers at once. Two consequences: every in-flight request would exchange simultaneously when the token entered the refresh margin, and if the token endpoint were failing, every request would keep re-exchanging for the whole margin. The second is the dangerous one -- it turns a brief outage into a retry storm against a service that throttles, which can extend the outage it is reacting to. Widening the margin to five minutes made this materially worse than the 20s window the stock token cache uses. A caller now claims the right to exchange with a compare-and-swap on the last attempt timestamp, and callers that already hold a usable token keep using it rather than piling on. This bounds exchanges to one per cooldown whether the endpoint is healthy or failing, and it also bounds the warning log to one per cooldown instead of one per request. Callers with nothing usable in hand never consult the cooldown, so a cold start behind a brief outage is never held back. A wall clock that steps backwards is treated as the cooldown having elapsed, rather than blocking refreshes until the clock catches up. Separately, clamp the lifetime taken from `expires_in`. `OffsetDateTime + Duration` panics on overflow, so a malformed response could have taken the process down. Clamping rather than saturating also means a nonsensical lifetime produces an earlier refresh instead of a token that is never refreshed at all. --- .../object_storage/azure_workload_identity.rs | 237 +++++++++++++++++- 1 file changed, 228 insertions(+), 9 deletions(-) diff --git a/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs b/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs index e0f013bb93b..cdbe53fc8a7 100644 --- a/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs +++ b/quickwit/quickwit-storage/src/object_storage/azure_workload_identity.rs @@ -27,6 +27,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; +use std::sync::atomic::{AtomicI64, Ordering}; use std::time::Duration; use arc_swap::ArcSwapOption; @@ -53,6 +54,23 @@ const CREDENTIAL_KIND_WORKLOAD_IDENTITY: &str = "workloadidentity"; /// retried while the token in hand is still usable. const EXPIRATION_MARGIN: Duration = Duration::from_secs(5 * 60); +/// Minimum spacing between token exchange attempts while a usable token is still cached. +/// +/// `get_token` is called on every storage request and the Azure backend runs up to +/// `max_concurrent_uploads` requests at a time, so without this every in-flight request would +/// exchange at once when the token enters the refresh margin, and would keep re-exchanging on +/// every request for the whole margin if the token endpoint were failing. That turns a brief +/// outage into a retry storm against a service that throttles. +const REFRESH_RETRY_COOLDOWN: Duration = Duration::from_secs(30); + +/// Sentinel for "no exchange has been attempted yet", so the first request never waits. +const NO_EXCHANGE_ATTEMPTED: i64 = i64::MIN; + +/// Upper bound on a token lifetime taken from `expires_in`. Entra issues storage tokens with a +/// ~24h lifetime; this only exists so that a malformed response cannot overflow the expiry +/// computation, which would panic. +const MAX_TOKEN_LIFETIME: Duration = Duration::from_secs(48 * 3600); + /// An access token together with the scopes it was minted for. /// /// A token is only valid for the scopes it was requested with, so the scopes are part of the @@ -93,6 +111,10 @@ pub(super) struct RefreshingWorkloadIdentityCredential { token_file_path: PathBuf, /// How long before actual expiry a cached token is treated as expiring. expiration_margin: Duration, + /// Minimum spacing between exchange attempts while a usable token is still cached. + refresh_retry_cooldown: Duration, + /// Unix timestamp of the last claimed exchange attempt, or [`NO_EXCHANGE_ATTEMPTED`]. + last_exchange_attempt_unix_secs: AtomicI64, cached_token_opt: ArcSwapOption, } @@ -104,6 +126,7 @@ impl RefreshingWorkloadIdentityCredential { client_id: String, token_file_path: PathBuf, expiration_margin: Duration, + refresh_retry_cooldown: Duration, ) -> Self { Self { http_client, @@ -112,10 +135,49 @@ impl RefreshingWorkloadIdentityCredential { client_id, token_file_path, expiration_margin, + refresh_retry_cooldown, + last_exchange_attempt_unix_secs: AtomicI64::new(NO_EXCHANGE_ATTEMPTED), cached_token_opt: ArcSwapOption::empty(), } } + /// Claims the right to attempt a token exchange. + /// + /// Returns `false` when another caller has claimed one within the cooldown, in which case a + /// caller holding a still-usable token should keep using it rather than piling on. Callers + /// with nothing usable in hand must exchange regardless and do not consult this. + fn try_claim_exchange(&self) -> bool { + let now_unix_secs = OffsetDateTime::now_utc().unix_timestamp(); + let cooldown_secs = self.refresh_retry_cooldown.as_secs() as i64; + loop { + let last_attempt_unix_secs = + self.last_exchange_attempt_unix_secs.load(Ordering::Relaxed); + // A negative elapsed time means the wall clock stepped backwards, in which case the + // cooldown is treated as elapsed rather than blocking refreshes until the clock + // catches up. + let elapsed_secs = now_unix_secs.saturating_sub(last_attempt_unix_secs); + if last_attempt_unix_secs != NO_EXCHANGE_ATTEMPTED + && (0..cooldown_secs).contains(&elapsed_secs) + { + return false; + } + // Only the caller that wins this swap performs the exchange; the rest fall back to + // the token they already hold. + if self + .last_exchange_attempt_unix_secs + .compare_exchange_weak( + last_attempt_unix_secs, + now_unix_secs, + Ordering::AcqRel, + Ordering::Relaxed, + ) + .is_ok() + { + return true; + } + } + } + /// Builds a credential from the standard workload identity environment variables. /// /// Returns `None` when the process is not configured for workload identity, in which case the @@ -144,6 +206,7 @@ impl RefreshingWorkloadIdentityCredential { client_id, PathBuf::from(token_file_path), EXPIRATION_MARGIN, + REFRESH_RETRY_COOLDOWN, )) } @@ -161,7 +224,11 @@ impl RefreshingWorkloadIdentityCredential { .await .context(ErrorKind::Credential, "federated token exchange failed")?; - let expires_on = OffsetDateTime::now_utc() + Duration::from_secs(login_response.expires_in); + // `OffsetDateTime + Duration` panics on overflow, so a malformed `expires_in` must not + // reach it. Clamping rather than saturating also means a nonsensical lifetime results in + // an earlier refresh instead of a token that is never refreshed at all. + let expires_in = Duration::from_secs(login_response.expires_in).min(MAX_TOKEN_LIFETIME); + let expires_on = OffsetDateTime::now_utc() + expires_in; debug!(%expires_on, "obtained azure access token from federated credentials flow"); Ok(AccessToken::new( login_response.access_token().clone(), @@ -173,15 +240,29 @@ impl RefreshingWorkloadIdentityCredential { #[async_trait] impl TokenCredential for RefreshingWorkloadIdentityCredential { async fn get_token(&self, scopes: &[&str]) -> azure_core::Result { - let cached_token_opt = self.cached_token_opt.load_full(); + let cached_token_opt = match self.cached_token_opt.load_full() { + Some(cached_token) if cached_token.matches(scopes) => Some(cached_token), + _ => None, + }; if let Some(cached_token) = &cached_token_opt - && cached_token.matches(scopes) && !is_expiring(&cached_token.access_token, self.expiration_margin) { return Ok(cached_token.access_token.clone()); } - let exchange_result = self.exchange_token(scopes).await; - let error = match exchange_result { + // Nothing cached for these scopes, or the cached token is inside the refresh margin. A + // token that has not actually expired yet remains a valid fallback. + let usable_cached_token_opt = match cached_token_opt { + Some(cached_token) if !is_expiring(&cached_token.access_token, Duration::ZERO) => { + Some(cached_token) + } + _ => None, + }; + if let Some(cached_token) = &usable_cached_token_opt + && !self.try_claim_exchange() + { + return Ok(cached_token.access_token.clone()); + } + let error = match self.exchange_token(scopes).await { Ok(access_token) => { self.cached_token_opt.store(Some(Arc::new(CachedToken { scopes: scopes.iter().map(|scope| scope.to_string()).collect(), @@ -194,10 +275,7 @@ impl TokenCredential for RefreshingWorkloadIdentityCredential { // The refresh margin exists so that a failed refresh can be retried while the token in // hand is still usable. Keep serving the cached token until it genuinely expires rather // than failing storage operations for a transient hiccup near expiry. - if let Some(cached_token) = &cached_token_opt - && cached_token.matches(scopes) - && !is_expiring(&cached_token.access_token, Duration::ZERO) - { + if let Some(cached_token) = &usable_cached_token_opt { warn!( %error, "azure token refresh failed, reusing the cached token until it expires" @@ -393,6 +471,22 @@ mod tests { fake_entra: Arc, token_file: &NamedTempFile, expiration_margin: Duration, + ) -> RefreshingWorkloadIdentityCredential { + // No cooldown by default, so tests exercise the exchange path on every call unless they + // are specifically testing the cooldown. + refreshing_credential_with_cooldown( + fake_entra, + token_file, + expiration_margin, + Duration::ZERO, + ) + } + + fn refreshing_credential_with_cooldown( + fake_entra: Arc, + token_file: &NamedTempFile, + expiration_margin: Duration, + refresh_retry_cooldown: Duration, ) -> RefreshingWorkloadIdentityCredential { RefreshingWorkloadIdentityCredential::new( fake_entra, @@ -401,6 +495,7 @@ mod tests { "test-client-id".to_string(), token_file.path().to_path_buf(), expiration_margin, + refresh_retry_cooldown, ) } @@ -648,6 +743,130 @@ mod tests { assert!(matches!(error.kind(), ErrorKind::Credential)); } + /// `get_token` runs on every storage request, so a failing token endpoint must not be retried + /// once per request for the whole refresh margin. + #[tokio::test] + async fn test_failing_refresh_is_not_retried_on_every_request() { + let fake_entra = FakeEntra::new("assertion-hour-0", 3600); + let token_file = token_file_containing("assertion-hour-0"); + let credential = refreshing_credential_with_cooldown( + fake_entra.clone(), + &token_file, + // Always inside the refresh margin, but a long way from actually expiring. + Duration::from_secs(48 * 3600), + Duration::from_secs(30), + ); + + credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + assert_eq!(fake_entra.exchange_count(), 1); + + // The token endpoint starts failing. + fake_entra.rotate_accepted_assertion("some-other-assertion"); + for _ in 0..50 { + credential + .get_token(&[STORAGE_SCOPE]) + .await + .expect("the cached token is still valid, so requests keep succeeding"); + } + assert_eq!( + fake_entra.exchange_count(), + 2, + "only one retry should have been attempted within the cooldown" + ); + } + + /// The same burst without a cooldown hammers the endpoint, which is what the cooldown exists + /// to prevent. + #[tokio::test] + async fn test_without_a_cooldown_every_request_retries() { + let fake_entra = FakeEntra::new("assertion-hour-0", 3600); + let token_file = token_file_containing("assertion-hour-0"); + let credential = refreshing_credential_with_cooldown( + fake_entra.clone(), + &token_file, + Duration::from_secs(48 * 3600), + Duration::ZERO, + ); + + credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + fake_entra.rotate_accepted_assertion("some-other-assertion"); + for _ in 0..50 { + credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + } + assert_eq!(fake_entra.exchange_count(), 51); + } + + /// A caller with nothing usable in hand must never be held back by the cooldown, otherwise a + /// cold start behind a brief outage would be stuck serving errors it could have recovered + /// from. + #[tokio::test] + async fn test_cooldown_never_blocks_a_caller_without_a_usable_token() { + // `expires_in: 0` means every issued token is already expired, so there is never a usable + // fallback and every call must be free to exchange. + let fake_entra = FakeEntra::new("assertion-hour-0", 0); + let token_file = token_file_containing("assertion-hour-0"); + let credential = refreshing_credential_with_cooldown( + fake_entra.clone(), + &token_file, + Duration::from_secs(5 * 60), + Duration::from_secs(3600), + ); + + for _ in 0..5 { + credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + } + assert_eq!(fake_entra.exchange_count(), 5); + } + + /// Concurrent requests entering the refresh margin together must not all exchange at once. + #[tokio::test] + async fn test_concurrent_requests_do_not_stampede_the_token_endpoint() { + let fake_entra = FakeEntra::new("assertion-hour-0", 3600); + let token_file = token_file_containing("assertion-hour-0"); + let credential = Arc::new(refreshing_credential_with_cooldown( + fake_entra.clone(), + &token_file, + Duration::from_secs(48 * 3600), + Duration::from_secs(30), + )); + + credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + assert_eq!(fake_entra.exchange_count(), 1); + + // Mirrors `max_concurrent_uploads` worth of in-flight storage requests. + let mut handles = Vec::with_capacity(100); + for _ in 0..100 { + let credential = credential.clone(); + handles.push(tokio::spawn(async move { + credential.get_token(&[STORAGE_SCOPE]).await.unwrap() + })); + } + for handle in handles { + handle.await.unwrap(); + } + assert_eq!( + fake_entra.exchange_count(), + 2, + "a single caller should win the refresh, the rest reuse the cached token" + ); + } + + /// A malformed `expires_in` must not overflow the expiry computation, which would panic and + /// take the indexer down. + #[tokio::test] + async fn test_absurd_expires_in_is_clamped_rather_than_overflowing() { + let fake_entra = FakeEntra::new("assertion-hour-0", u64::MAX); + let token_file = token_file_containing("assertion-hour-0"); + let credential = + refreshing_credential(fake_entra.clone(), &token_file, Duration::from_secs(5 * 60)); + + let access_token = credential.get_token(&[STORAGE_SCOPE]).await.unwrap(); + assert!( + access_token.expires_on <= OffsetDateTime::now_utc() + MAX_TOKEN_LIFETIME, + "the token lifetime should be clamped to a sane bound" + ); + } + /// A failed exchange must leave the credential able to recover on the next call, rather than /// wedging the way the upstream one does. #[tokio::test]