From f16a73553331b02810d0a8140849e0235bbcc519 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 12 Aug 2026 17:19:52 +0800 Subject: [PATCH 1/2] fix: deduplicate signer credential refreshes --- core/src/lib.rs | 7 +- core/src/signer.rs | 852 ++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 809 insertions(+), 50 deletions(-) diff --git a/core/src/lib.rs b/core/src/lib.rs index 54320926..3c2b9dc9 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -53,8 +53,11 @@ //! [`SigningCredential::is_valid`] controls whether a cached credential can be reused //! without refresh. [`SigningCredential::is_valid_at`] checks exact usability at the //! timestamp returned by [`SignRequest::required_valid_until`]. A refreshed credential -//! only needs to satisfy the exact operation requirement; provider errors are returned -//! without retrying internally or falling back to the old cached credential. +//! only needs to satisfy the exact operation requirement. Concurrent callers that need +//! the same shared-cache refresh use one provider invocation and receive its result. +//! Refresh failures are not cached after current waiters observe them, so a later call +//! can retry. Provider errors are returned without retrying internally or falling back +//! to the old cached credential. //! //! ## Example //! diff --git a/core/src/signer.rs b/core/src/signer.rs index 663c95d7..59b12bb7 100644 --- a/core/src/signer.rs +++ b/core/src/signer.rs @@ -23,11 +23,179 @@ use crate::Result; use crate::SignRequest; use crate::SignRequestDyn; use crate::SigningCredential; +use futures::channel::oneshot; use std::any::type_name; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, Mutex}; use std::time::Duration; +struct CredentialCache { + state: Mutex>, +} + +struct CredentialCacheState { + credential: Option, + refresh: Option>, + next_refresh_generation: u64, +} + +struct CredentialRefresh { + generation: u64, + waiters: Vec>>, +} + +#[derive(Clone)] +enum CredentialRefreshOutcome { + Loaded(K), + Failed(SharedRefreshError), +} + +#[derive(Clone)] +struct SharedRefreshError { + kind: crate::ErrorKind, + message: String, + context: Vec, + retryable: bool, +} + +impl SharedRefreshError { + fn capture(error: &Error) -> Self { + Self { + kind: error.kind(), + message: error.to_string(), + context: error.context().to_vec(), + retryable: error.is_retryable(), + } + } + + fn into_error(self) -> Error { + let mut error = Error::new(self.kind, self.message).set_retryable(self.retryable); + for context in self.context { + error = error.with_context(context); + } + error + } +} + +enum CredentialCacheAction { + Ready(K), + Wait(oneshot::Receiver>), + Refresh(CredentialRefreshGuard), +} + +impl CredentialCache { + fn new() -> Self { + Self { + state: Mutex::new(CredentialCacheState { + credential: None, + refresh: None, + next_refresh_generation: 0, + }), + } + } + + fn action( + self: &Arc, + cached_is_usable: impl FnOnce(&K) -> bool, + ) -> CredentialCacheAction { + let mut state = self.state.lock().expect("lock poisoned"); + if let Some(credential) = state.credential.as_ref() { + if cached_is_usable(credential) { + return CredentialCacheAction::Ready(credential.clone()); + } + } + + if let Some(refresh) = state.refresh.as_mut() { + let (sender, receiver) = oneshot::channel(); + refresh.waiters.push(sender); + return CredentialCacheAction::Wait(receiver); + } + + let generation = state.next_refresh_generation; + state.next_refresh_generation = state.next_refresh_generation.wrapping_add(1); + state.refresh = Some(CredentialRefresh { + generation, + waiters: Vec::new(), + }); + CredentialCacheAction::Refresh(CredentialRefreshGuard { + cache: self.clone(), + generation, + completed: false, + }) + } + + fn complete(&self, generation: u64, outcome: CredentialRefreshOutcome) { + let waiters = { + let mut state = self.state.lock().expect("lock poisoned"); + if state + .refresh + .as_ref() + .is_none_or(|refresh| refresh.generation != generation) + { + return; + } + + if let CredentialRefreshOutcome::Loaded(credential) = &outcome { + state.credential = Some(credential.clone()); + } + state + .refresh + .take() + .expect("matching refresh must exist") + .waiters + }; + + for waiter in waiters { + let _ = waiter.send(outcome.clone()); + } + } + + #[cfg(test)] + fn set_credential(&self, credential: K) { + self.state.lock().expect("lock poisoned").credential = Some(credential); + } +} + +struct CredentialRefreshGuard { + cache: Arc>, + generation: u64, + completed: bool, +} + +impl CredentialRefreshGuard { + fn complete_success(mut self, credential: K) { + self.completed = true; + self.cache.complete( + self.generation, + CredentialRefreshOutcome::Loaded(credential), + ); + } + + fn complete_failure(mut self, error: &Error) { + self.completed = true; + self.cache.complete( + self.generation, + CredentialRefreshOutcome::Failed(SharedRefreshError::capture(error)), + ); + } +} + +impl Drop for CredentialRefreshGuard { + fn drop(&mut self) { + if self.completed { + return; + } + + let error = Error::unexpected("credential refresh was cancelled") + .with_context(format!("credential_type: {}", type_name::())) + .set_retryable(true); + self.cache.complete( + self.generation, + CredentialRefreshOutcome::Failed(SharedRefreshError::capture(&error)), + ); + } +} + /// Loads credentials and atomically signs request heads. /// /// The service-specific [`SignRequest`] runs against a private candidate. Only the @@ -37,7 +205,7 @@ pub struct Signer { ctx: Context, loader: Arc>, builder: Arc>, - credential: Arc>>, + credential_cache: Arc>, } impl Debug for Signer { @@ -60,23 +228,25 @@ impl Signer { loader: Arc::new(loader), builder: Arc::new(builder), - credential: Arc::new(Mutex::new(None)), + credential_cache: Arc::new(CredentialCache::new()), } } - /// Replace the context while keeping credential provider and request signer. + /// Replace the context while keeping the credential provider, request signer, + /// and shared credential cache. pub fn with_context(mut self, ctx: Context) -> Self { self.ctx = ctx; self } - /// Replace the credential provider while keeping context and request signer. + /// Replace the credential provider while keeping context and request signer, + /// and create an isolated empty credential cache. pub fn with_credential_provider( mut self, provider: impl ProvideCredential, ) -> Self { self.loader = Arc::new(provider); - self.credential = Arc::new(Mutex::new(None)); // Clear cached credential + self.credential_cache = Arc::new(CredentialCache::new()); self } @@ -103,50 +273,18 @@ impl Signer { /// /// Cached credentials must be fresh according to [`SigningCredential::is_valid`] /// and usable through [`SignRequest::required_valid_until`]. A refreshed credential - /// only needs to satisfy the exact operation deadline. Provider errors are returned - /// without internal retry or fallback to the previous cached credential. + /// only needs to satisfy the exact operation deadline. Concurrent callers that need + /// the same cache refresh share one provider invocation and its result. A failed + /// refresh is returned to its current waiters but is not cached, so a later call can + /// retry. Provider errors are returned without internal retry or fallback to the + /// previous cached credential. Request signing runs after refresh coordination has + /// completed and remains concurrent. pub async fn sign( &self, req: &mut http::request::Parts, expires_in: Option, ) -> Result<()> { - let credential = self.credential.lock().expect("lock poisoned").clone(); - let credential = match credential { - Some(credential) - if credential.is_valid() - && credential.is_valid_at( - self.builder - .required_valid_until_dyn(&credential, expires_in), - ) => - { - credential - } - _ => { - let credential = self - .loader - .provide_credential_dyn(&self.ctx) - .await? - .ok_or_else(|| { - Error::credential_invalid("failed to load signing credential") - .with_context(format!("credential_type: {}", type_name::())) - })?; - - *self.credential.lock().expect("lock poisoned") = Some(credential.clone()); - - let required_until = self - .builder - .required_valid_until_dyn(&credential, expires_in); - if !credential.is_valid_at(required_until) { - return Err(Error::credential_invalid( - "refreshed signing credential expires before the requested operation deadline", - ) - .with_context(format!("credential_type: {}", type_name::())) - .with_context(format!("required_valid_until: {required_until}"))); - } - - credential - } - }; + let credential = self.credential(expires_in).await?; let mut candidate = req.clone(); self.builder @@ -157,6 +295,73 @@ impl Signer { req.headers = candidate.headers; Ok(()) } + + async fn credential(&self, expires_in: Option) -> Result { + let action = self.credential_cache.action(|credential| { + credential.is_valid() + && credential.is_valid_at( + self.builder + .required_valid_until_dyn(credential, expires_in), + ) + }); + + match action { + CredentialCacheAction::Ready(credential) => Ok(credential), + CredentialCacheAction::Wait(receiver) => match receiver.await { + Ok(CredentialRefreshOutcome::Loaded(credential)) => { + self.validate_refreshed_credential(credential, expires_in) + } + Ok(CredentialRefreshOutcome::Failed(error)) => Err(error.into_error()), + Err(_) => Err(Error::unexpected( + "credential refresh coordination ended without a result", + ) + .with_context(format!("credential_type: {}", type_name::())) + .set_retryable(true)), + }, + CredentialCacheAction::Refresh(refresh) => { + let result = self + .loader + .provide_credential_dyn(&self.ctx) + .await + .and_then(|credential| { + credential.ok_or_else(|| { + Error::credential_invalid("failed to load signing credential") + .with_context(format!("credential_type: {}", type_name::())) + }) + }); + + match result { + Ok(credential) => { + refresh.complete_success(credential.clone()); + self.validate_refreshed_credential(credential, expires_in) + } + Err(error) => { + refresh.complete_failure(&error); + Err(error) + } + } + } + } + } + + fn validate_refreshed_credential( + &self, + credential: K, + expires_in: Option, + ) -> Result { + let required_until = self + .builder + .required_valid_until_dyn(&credential, expires_in); + if !credential.is_valid_at(required_until) { + return Err(Error::credential_invalid( + "refreshed signing credential expires before the requested operation deadline", + ) + .with_context(format!("credential_type: {}", type_name::())) + .with_context(format!("required_valid_until: {required_until}"))); + } + + Ok(credential) + } } #[cfg(test)] @@ -164,6 +369,9 @@ mod tests { use super::*; use crate::time::Timestamp; use crate::{ErrorKind, ProvideCredential, SignRequest}; + use futures::channel::oneshot; + use futures::future::{join_all, pending}; + use futures::poll; use http::{HeaderValue, Method, Request, Version}; use std::collections::VecDeque; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -242,6 +450,13 @@ mod tests { } } + type ControlledResponse = Result>; + type ControlledProviderParts = ( + ControlledProvider, + Arc, + Vec>, + ); + #[derive(Debug)] struct SequenceProvider { responses: Mutex>>>, @@ -276,6 +491,190 @@ mod tests { } } + struct ControlledProvider { + responses: Mutex>>, + calls: Arc, + } + + impl Debug for ControlledProvider { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ControlledProvider").finish_non_exhaustive() + } + } + + impl ControlledProvider { + fn new(invocations: usize) -> ControlledProviderParts { + let calls = Arc::new(AtomicUsize::new(0)); + let (senders, receivers) = (0..invocations) + .map(|_| oneshot::channel()) + .unzip::<_, _, Vec<_>, VecDeque<_>>(); + ( + Self { + responses: Mutex::new(receivers), + calls: calls.clone(), + }, + calls, + senders, + ) + } + } + + impl ProvideCredential for ControlledProvider { + type Credential = ExpiringCredential; + + async fn provide_credential(&self, _ctx: &Context) -> Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + let response = self + .responses + .lock() + .expect("lock poisoned") + .pop_front() + .expect("controlled response must exist"); + response + .await + .map_err(|_| Error::unexpected("controlled response sender was dropped"))? + } + } + + struct ControlledRequestSigner { + started: Arc, + releases: Mutex>>, + } + + impl Debug for ControlledRequestSigner { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ControlledRequestSigner") + .finish_non_exhaustive() + } + } + + impl ControlledRequestSigner { + fn new(count: usize) -> (Self, Arc, Vec>) { + let started = Arc::new(AtomicUsize::new(0)); + let (senders, receivers) = (0..count) + .map(|_| oneshot::channel()) + .unzip::<_, _, Vec<_>, VecDeque<_>>(); + ( + Self { + started: started.clone(), + releases: Mutex::new(receivers), + }, + started, + senders, + ) + } + } + + impl SignRequest for ControlledRequestSigner { + type Credential = ExpiringCredential; + + fn required_valid_until( + &self, + credential: &Self::Credential, + _expires_in: Option, + ) -> Timestamp { + credential.required_until + } + + async fn sign_request( + &self, + _ctx: &Context, + req: &mut http::request::Parts, + credential: Option<&Self::Credential>, + _expires_in: Option, + ) -> Result<()> { + self.started.fetch_add(1, Ordering::SeqCst); + let release = self + .releases + .lock() + .expect("lock poisoned") + .pop_front() + .expect("signing release must exist"); + release + .await + .map_err(|_| Error::unexpected("signing release sender was dropped"))?; + req.headers.insert( + "x-credential-generation", + credential + .expect("credential must be present") + .generation + .to_string() + .parse()?, + ); + Ok(()) + } + } + + struct CancellationProvider { + calls: Arc, + credential: ExpiringCredential, + } + + impl Debug for CancellationProvider { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CancellationProvider") + .finish_non_exhaustive() + } + } + + impl ProvideCredential for CancellationProvider { + type Credential = ExpiringCredential; + + async fn provide_credential(&self, _ctx: &Context) -> Result> { + if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + pending::<()>().await; + } + Ok(Some(self.credential.clone())) + } + } + + const CREDENTIAL_SECRET: &str = "credential-secret-must-not-leak"; + + #[derive(Clone, Debug)] + struct SecretCredential { + secret: &'static str, + } + + impl SigningCredential for SecretCredential { + fn is_valid(&self) -> bool { + !self.secret.is_empty() + } + + fn is_valid_at(&self, _timestamp: Timestamp) -> bool { + false + } + } + + #[derive(Debug)] + struct SecretProvider; + + impl ProvideCredential for SecretProvider { + type Credential = SecretCredential; + + async fn provide_credential(&self, _ctx: &Context) -> Result> { + Ok(Some(SecretCredential { + secret: CREDENTIAL_SECRET, + })) + } + } + + #[derive(Debug)] + struct SecretRequestSigner; + + impl SignRequest for SecretRequestSigner { + type Credential = SecretCredential; + + async fn sign_request( + &self, + _ctx: &Context, + _req: &mut http::request::Parts, + _credential: Option<&Self::Credential>, + _expires_in: Option, + ) -> Result<()> { + Ok(()) + } + } + #[derive(Debug)] struct OperationSigner; @@ -377,6 +776,363 @@ mod tests { assert!(!parts.headers.contains_key("x-original")); } + #[test] + fn concurrent_cold_start_invokes_provider_once() { + futures::executor::block_on(async { + let base = Timestamp::from_second(500).expect("timestamp must be valid"); + let credential = ExpiringCredential { + generation: 1, + fresh: true, + expires_at: base + Duration::from_secs(30), + required_until: base + Duration::from_secs(10), + }; + let (provider, calls, mut responses) = ControlledProvider::new(1); + let signer = Signer::new(Context::new(), provider, OperationSigner); + let signers = (0..8).map(|_| signer.clone()).collect::>(); + let mut requests = (0..8).map(|_| request_parts()).collect::>(); + let mut batch = Box::pin(join_all( + requests + .iter_mut() + .zip(signers.iter()) + .map(|(request, signer)| signer.sign(request, None)), + )); + + assert!(poll!(&mut batch).is_pending()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + responses + .remove(0) + .send(Ok(Some(credential))) + .expect("controlled response must be received"); + for result in batch.await { + result.expect("concurrent cold-start signing must succeed"); + } + assert_eq!(calls.load(Ordering::SeqCst), 1); + }); + } + + #[test] + fn concurrent_stale_refresh_invokes_provider_once() { + futures::executor::block_on(async { + let base = Timestamp::from_second(600).expect("timestamp must be valid"); + let cached = ExpiringCredential { + generation: 1, + fresh: false, + expires_at: base + Duration::from_secs(30), + required_until: base + Duration::from_secs(10), + }; + let refreshed = ExpiringCredential { + generation: 2, + fresh: true, + expires_at: base + Duration::from_secs(30), + required_until: base + Duration::from_secs(10), + }; + let (provider, calls, mut responses) = ControlledProvider::new(1); + let signer = Signer::new(Context::new(), provider, OperationSigner); + signer.credential_cache.set_credential(cached); + let mut requests = (0..8).map(|_| request_parts()).collect::>(); + let mut batch = Box::pin(join_all( + requests + .iter_mut() + .map(|request| signer.sign(request, None)), + )); + + assert!(poll!(&mut batch).is_pending()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + responses + .remove(0) + .send(Ok(Some(refreshed))) + .expect("controlled response must be received"); + for result in batch.await { + result.expect("concurrent stale refresh must succeed"); + } + assert_eq!(calls.load(Ordering::SeqCst), 1); + for request in requests { + assert_eq!( + request.headers.get("x-credential-generation"), + Some(&HeaderValue::from_static("2")) + ); + } + }); + } + + #[test] + fn concurrent_refresh_failure_is_shared_and_later_call_retries() { + futures::executor::block_on(async { + let base = Timestamp::from_second(700).expect("timestamp must be valid"); + let refreshed = ExpiringCredential { + generation: 2, + fresh: true, + expires_at: base + Duration::from_secs(30), + required_until: base + Duration::from_secs(10), + }; + let (provider, calls, mut responses) = ControlledProvider::new(2); + let signer = Signer::new(Context::new(), provider, OperationSigner); + let mut requests = (0..8).map(|_| request_parts()).collect::>(); + let mut batch = Box::pin(join_all( + requests + .iter_mut() + .map(|request| signer.sign(request, None)), + )); + + assert!(poll!(&mut batch).is_pending()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + responses + .remove(0) + .send(Err(Error::rate_limited("injected refresh failure") + .with_context("refresh_generation: 1"))) + .expect("controlled failure must be received"); + for result in batch.await { + let error = result.expect_err("current waiter must receive refresh failure"); + assert_eq!(error.kind(), ErrorKind::RateLimited); + assert_eq!(error.to_string(), "injected refresh failure"); + assert_eq!(error.context(), &["refresh_generation: 1"]); + assert!(error.is_retryable()); + } + assert_eq!(calls.load(Ordering::SeqCst), 1); + + let mut retry_request = request_parts(); + let mut retry = Box::pin(signer.sign(&mut retry_request, None)); + assert!(poll!(&mut retry).is_pending()); + assert_eq!(calls.load(Ordering::SeqCst), 2); + responses + .remove(0) + .send(Ok(Some(refreshed))) + .expect("controlled recovery must be received"); + retry.await.expect("later caller must retry and recover"); + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!( + retry_request.headers.get("x-credential-generation"), + Some(&HeaderValue::from_static("2")) + ); + }); + } + + #[test] + fn concurrent_callers_apply_exact_refreshed_validity_check() { + futures::executor::block_on(async { + let base = Timestamp::from_second(800).expect("timestamp must be valid"); + let credential = ExpiringCredential { + generation: 1, + fresh: true, + expires_at: base + Duration::from_secs(10), + required_until: base + Duration::from_secs(10), + }; + let (provider, calls, mut responses) = ControlledProvider::new(1); + let signer = Signer::new(Context::new(), provider, OperationSigner); + let mut requests = (0..8).map(|_| request_parts()).collect::>(); + let mut batch = Box::pin(join_all( + requests + .iter_mut() + .map(|request| signer.sign(request, None)), + )); + + assert!(poll!(&mut batch).is_pending()); + responses + .remove(0) + .send(Ok(Some(credential))) + .expect("controlled response must be received"); + for result in batch.await { + let error = result.expect_err("exact deadline must reject the credential"); + assert_eq!(error.kind(), ErrorKind::CredentialInvalid); + assert!( + error + .to_string() + .contains("expires before the requested operation deadline") + ); + } + assert_eq!(calls.load(Ordering::SeqCst), 1); + }); + } + + #[test] + fn distinct_credential_caches_do_not_block_each_other() { + futures::executor::block_on(async { + let base = Timestamp::from_second(900).expect("timestamp must be valid"); + let credential_a = ExpiringCredential { + generation: 1, + fresh: true, + expires_at: base + Duration::from_secs(30), + required_until: base + Duration::from_secs(10), + }; + let credential_b = ExpiringCredential { + generation: 2, + fresh: true, + expires_at: base + Duration::from_secs(30), + required_until: base + Duration::from_secs(10), + }; + let (provider_a, calls_a, mut responses_a) = ControlledProvider::new(1); + let (provider_b, calls_b, mut responses_b) = ControlledProvider::new(1); + let signer_a = Signer::new(Context::new(), provider_a, OperationSigner); + let signer_b = Signer::new(Context::new(), provider_b, OperationSigner); + let mut request_a = request_parts(); + let mut request_b = request_parts(); + let mut future_a = Box::pin(signer_a.sign(&mut request_a, None)); + let mut future_b = Box::pin(signer_b.sign(&mut request_b, None)); + + assert!(poll!(&mut future_a).is_pending()); + assert!(poll!(&mut future_b).is_pending()); + assert_eq!(calls_a.load(Ordering::SeqCst), 1); + assert_eq!(calls_b.load(Ordering::SeqCst), 1); + + responses_b + .remove(0) + .send(Ok(Some(credential_b))) + .expect("second cache response must be received"); + future_b + .await + .expect("second cache must complete while first is blocked"); + assert!(poll!(&mut future_a).is_pending()); + + responses_a + .remove(0) + .send(Ok(Some(credential_a))) + .expect("first cache response must be received"); + future_a.await.expect("first cache must complete"); + }); + } + + #[test] + fn request_signing_remains_concurrent_after_refresh() { + futures::executor::block_on(async { + let base = Timestamp::from_second(950).expect("timestamp must be valid"); + let credential = ExpiringCredential { + generation: 1, + fresh: true, + expires_at: base + Duration::from_secs(30), + required_until: base + Duration::from_secs(10), + }; + let (provider, calls, mut responses) = ControlledProvider::new(1); + let (request_signer, started, releases) = ControlledRequestSigner::new(8); + let signer = Signer::new(Context::new(), provider, request_signer); + let mut requests = (0..8).map(|_| request_parts()).collect::>(); + let mut batch = Box::pin(join_all( + requests + .iter_mut() + .map(|request| signer.sign(request, None)), + )); + + assert!(poll!(&mut batch).is_pending()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(started.load(Ordering::SeqCst), 0); + responses + .remove(0) + .send(Ok(Some(credential))) + .expect("controlled response must be received"); + + assert!(poll!(&mut batch).is_pending()); + assert_eq!(started.load(Ordering::SeqCst), 8); + for release in releases { + release.send(()).expect("signing release must be received"); + } + for result in batch.await { + result.expect("concurrent request signing must succeed"); + } + }); + } + + #[test] + fn cache_sharing_and_reset_contract_is_preserved() { + let signer = Signer::new( + Context::new(), + StaticProvider, + MutatingSigner { fail: false }, + ); + let clone = signer.clone(); + let with_context = signer.clone().with_context(Context::new()); + let with_request_signer = signer + .clone() + .with_request_signer(MutatingSigner { fail: false }); + let with_provider = signer.clone().with_credential_provider(StaticProvider); + + assert!(Arc::ptr_eq( + &signer.credential_cache, + &clone.credential_cache + )); + assert!(Arc::ptr_eq( + &signer.credential_cache, + &with_context.credential_cache + )); + assert!(Arc::ptr_eq( + &signer.credential_cache, + &with_request_signer.credential_cache + )); + assert!(!Arc::ptr_eq( + &signer.credential_cache, + &with_provider.credential_cache + )); + } + + #[test] + fn cancelled_refresh_notifies_waiters_and_allows_retry() { + futures::executor::block_on(async { + let base = Timestamp::from_second(975).expect("timestamp must be valid"); + let calls = Arc::new(AtomicUsize::new(0)); + let provider = CancellationProvider { + calls: calls.clone(), + credential: ExpiringCredential { + generation: 2, + fresh: true, + expires_at: base + Duration::from_secs(30), + required_until: base + Duration::from_secs(10), + }, + }; + let signer = Signer::new(Context::new(), provider, OperationSigner); + let mut leader_request = request_parts(); + let mut waiter_request = request_parts(); + let mut leader = Box::pin(signer.sign(&mut leader_request, None)); + let mut waiter = Box::pin(signer.sign(&mut waiter_request, None)); + + assert!(poll!(&mut leader).is_pending()); + assert!(poll!(&mut waiter).is_pending()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + drop(leader); + + let error = waiter + .await + .expect_err("waiter must receive refresh cancellation"); + assert_eq!(error.kind(), ErrorKind::Unexpected); + assert_eq!(error.to_string(), "credential refresh was cancelled"); + assert!(error.is_retryable()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + let mut retry_request = request_parts(); + signer + .sign(&mut retry_request, None) + .await + .expect("later call must retry after cancellation"); + assert_eq!(calls.load(Ordering::SeqCst), 2); + }); + } + + #[test] + fn credential_values_are_redacted_from_debug_and_validation_errors() { + let signer = Signer::new(Context::new(), SecretProvider, SecretRequestSigner); + let mut request = request_parts(); + let error = futures::executor::block_on(signer.sign(&mut request, None)) + .expect_err("unusable credential must fail validation"); + + assert!(!format!("{signer:?}").contains(CREDENTIAL_SECRET)); + assert!(!format!("{error:?}").contains(CREDENTIAL_SECRET)); + assert!(!error.to_string().contains(CREDENTIAL_SECRET)); + } + + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn sign_future_remains_send_on_native_targets() { + fn assert_send(_future: T) {} + + let signer = Signer::new( + Context::new(), + StaticProvider, + MutatingSigner { fail: false }, + ); + let mut request = request_parts(); + assert_send(signer.sign(&mut request, None)); + } + #[test] fn refreshes_cached_credential_for_operation_requirement() { let base = Timestamp::from_second(1_000).expect("timestamp must be valid"); @@ -394,7 +1150,7 @@ mod tests { }; let (provider, calls) = SequenceProvider::new([Ok(Some(refreshed))]); let signer = Signer::new(Context::new(), provider, OperationSigner); - *signer.credential.lock().expect("lock poisoned") = Some(cached); + signer.credential_cache.set_credential(cached); let mut parts = request_parts(); futures::executor::block_on(signer.sign(&mut parts, None)) @@ -449,7 +1205,7 @@ mod tests { Ok(Some(refreshed)), ]); let signer = Signer::new(Context::new(), provider, OperationSigner); - *signer.credential.lock().expect("lock poisoned") = Some(cached); + signer.credential_cache.set_credential(cached); let mut parts = request_parts(); let original = parts.clone(); @@ -486,7 +1242,7 @@ mod tests { }; let (provider, calls) = SequenceProvider::new([Ok(None), Ok(Some(refreshed))]); let signer = Signer::new(Context::new(), provider, OperationSigner); - *signer.credential.lock().expect("lock poisoned") = Some(cached); + signer.credential_cache.set_credential(cached); let mut parts = request_parts(); let original = parts.clone(); @@ -514,7 +1270,7 @@ mod tests { StaticProvider, MutatingSigner { fail: false }, ); - *signer.credential.lock().expect("lock poisoned") = Some(TestCredential); + signer.credential_cache.set_credential(TestCredential); let debug = format!("{signer:?}"); assert!(debug.starts_with("Signer")); From cc43e5e0a6ef38952a40cd3a30e1c63a5689c695 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 14 Aug 2026 00:31:23 +0800 Subject: [PATCH 2/2] refactor: simplify signer credential refresh --- Cargo.toml | 1 + core/Cargo.toml | 1 + core/src/lib.rs | 10 +- core/src/signer.rs | 383 +++++++++++---------------------------------- 4 files changed, 101 insertions(+), 294 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 872fc0c5..895aab32 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,6 +57,7 @@ hmac = "0.13" http = "1" jiff = "0.2" log = "0.4" +mea = "0.6.4" pem = "4.0" percent-encoding = "2" p256 = { version = "0.14.0", default-features = false, features = ["ecdsa"] } diff --git a/core/Cargo.toml b/core/Cargo.toml index cc855828..6b88aa72 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -41,6 +41,7 @@ hmac = { workspace = true } http = { workspace = true } jiff = { workspace = true } log = { workspace = true } +mea = { workspace = true } percent-encoding = { workspace = true } sha1 = { workspace = true } sha2 = { workspace = true } diff --git a/core/src/lib.rs b/core/src/lib.rs index 3c2b9dc9..b43c566b 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -53,11 +53,11 @@ //! [`SigningCredential::is_valid`] controls whether a cached credential can be reused //! without refresh. [`SigningCredential::is_valid_at`] checks exact usability at the //! timestamp returned by [`SignRequest::required_valid_until`]. A refreshed credential -//! only needs to satisfy the exact operation requirement. Concurrent callers that need -//! the same shared-cache refresh use one provider invocation and receive its result. -//! Refresh failures are not cached after current waiters observe them, so a later call -//! can retry. Provider errors are returned without retrying internally or falling back -//! to the old cached credential. +//! only needs to satisfy the exact operation requirement. Credential refresh is +//! serialized per shared cache, so concurrent cold or stale callers reuse a successful +//! refresh. Refresh failures are not cached, so the next waiting or later caller can +//! retry. Provider errors are returned without retrying internally or falling back to +//! the old cached credential. //! //! ## Example //! diff --git a/core/src/signer.rs b/core/src/signer.rs index 59b12bb7..489159b1 100644 --- a/core/src/signer.rs +++ b/core/src/signer.rs @@ -23,179 +23,12 @@ use crate::Result; use crate::SignRequest; use crate::SignRequestDyn; use crate::SigningCredential; -use futures::channel::oneshot; +use mea::mutex::Mutex; use std::any::type_name; use std::fmt::{Debug, Formatter}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use std::time::Duration; -struct CredentialCache { - state: Mutex>, -} - -struct CredentialCacheState { - credential: Option, - refresh: Option>, - next_refresh_generation: u64, -} - -struct CredentialRefresh { - generation: u64, - waiters: Vec>>, -} - -#[derive(Clone)] -enum CredentialRefreshOutcome { - Loaded(K), - Failed(SharedRefreshError), -} - -#[derive(Clone)] -struct SharedRefreshError { - kind: crate::ErrorKind, - message: String, - context: Vec, - retryable: bool, -} - -impl SharedRefreshError { - fn capture(error: &Error) -> Self { - Self { - kind: error.kind(), - message: error.to_string(), - context: error.context().to_vec(), - retryable: error.is_retryable(), - } - } - - fn into_error(self) -> Error { - let mut error = Error::new(self.kind, self.message).set_retryable(self.retryable); - for context in self.context { - error = error.with_context(context); - } - error - } -} - -enum CredentialCacheAction { - Ready(K), - Wait(oneshot::Receiver>), - Refresh(CredentialRefreshGuard), -} - -impl CredentialCache { - fn new() -> Self { - Self { - state: Mutex::new(CredentialCacheState { - credential: None, - refresh: None, - next_refresh_generation: 0, - }), - } - } - - fn action( - self: &Arc, - cached_is_usable: impl FnOnce(&K) -> bool, - ) -> CredentialCacheAction { - let mut state = self.state.lock().expect("lock poisoned"); - if let Some(credential) = state.credential.as_ref() { - if cached_is_usable(credential) { - return CredentialCacheAction::Ready(credential.clone()); - } - } - - if let Some(refresh) = state.refresh.as_mut() { - let (sender, receiver) = oneshot::channel(); - refresh.waiters.push(sender); - return CredentialCacheAction::Wait(receiver); - } - - let generation = state.next_refresh_generation; - state.next_refresh_generation = state.next_refresh_generation.wrapping_add(1); - state.refresh = Some(CredentialRefresh { - generation, - waiters: Vec::new(), - }); - CredentialCacheAction::Refresh(CredentialRefreshGuard { - cache: self.clone(), - generation, - completed: false, - }) - } - - fn complete(&self, generation: u64, outcome: CredentialRefreshOutcome) { - let waiters = { - let mut state = self.state.lock().expect("lock poisoned"); - if state - .refresh - .as_ref() - .is_none_or(|refresh| refresh.generation != generation) - { - return; - } - - if let CredentialRefreshOutcome::Loaded(credential) = &outcome { - state.credential = Some(credential.clone()); - } - state - .refresh - .take() - .expect("matching refresh must exist") - .waiters - }; - - for waiter in waiters { - let _ = waiter.send(outcome.clone()); - } - } - - #[cfg(test)] - fn set_credential(&self, credential: K) { - self.state.lock().expect("lock poisoned").credential = Some(credential); - } -} - -struct CredentialRefreshGuard { - cache: Arc>, - generation: u64, - completed: bool, -} - -impl CredentialRefreshGuard { - fn complete_success(mut self, credential: K) { - self.completed = true; - self.cache.complete( - self.generation, - CredentialRefreshOutcome::Loaded(credential), - ); - } - - fn complete_failure(mut self, error: &Error) { - self.completed = true; - self.cache.complete( - self.generation, - CredentialRefreshOutcome::Failed(SharedRefreshError::capture(error)), - ); - } -} - -impl Drop for CredentialRefreshGuard { - fn drop(&mut self) { - if self.completed { - return; - } - - let error = Error::unexpected("credential refresh was cancelled") - .with_context(format!("credential_type: {}", type_name::())) - .set_retryable(true); - self.cache.complete( - self.generation, - CredentialRefreshOutcome::Failed(SharedRefreshError::capture(&error)), - ); - } -} - /// Loads credentials and atomically signs request heads. /// /// The service-specific [`SignRequest`] runs against a private candidate. Only the @@ -205,7 +38,7 @@ pub struct Signer { ctx: Context, loader: Arc>, builder: Arc>, - credential_cache: Arc>, + credential: Arc>>, } impl Debug for Signer { @@ -228,7 +61,7 @@ impl Signer { loader: Arc::new(loader), builder: Arc::new(builder), - credential_cache: Arc::new(CredentialCache::new()), + credential: Arc::new(Mutex::new(None)), } } @@ -246,7 +79,7 @@ impl Signer { provider: impl ProvideCredential, ) -> Self { self.loader = Arc::new(provider); - self.credential_cache = Arc::new(CredentialCache::new()); + self.credential = Arc::new(Mutex::new(None)); self } @@ -273,12 +106,11 @@ impl Signer { /// /// Cached credentials must be fresh according to [`SigningCredential::is_valid`] /// and usable through [`SignRequest::required_valid_until`]. A refreshed credential - /// only needs to satisfy the exact operation deadline. Concurrent callers that need - /// the same cache refresh share one provider invocation and its result. A failed - /// refresh is returned to its current waiters but is not cached, so a later call can - /// retry. Provider errors are returned without internal retry or fallback to the - /// previous cached credential. Request signing runs after refresh coordination has - /// completed and remains concurrent. + /// only needs to satisfy the exact operation deadline. Credential refresh is + /// serialized per shared cache. A failed refresh is not cached, so the next waiting + /// or later caller can retry. Provider errors are returned without internal retry or + /// fallback to the previous cached credential. Request signing runs after refresh + /// coordination has completed and remains concurrent. pub async fn sign( &self, req: &mut http::request::Parts, @@ -297,51 +129,31 @@ impl Signer { } async fn credential(&self, expires_in: Option) -> Result { - let action = self.credential_cache.action(|credential| { - credential.is_valid() + let mut cached = self.credential.lock().await; + if let Some(credential) = cached.as_ref() { + if credential.is_valid() && credential.is_valid_at( self.builder .required_valid_until_dyn(credential, expires_in), ) - }); - - match action { - CredentialCacheAction::Ready(credential) => Ok(credential), - CredentialCacheAction::Wait(receiver) => match receiver.await { - Ok(CredentialRefreshOutcome::Loaded(credential)) => { - self.validate_refreshed_credential(credential, expires_in) - } - Ok(CredentialRefreshOutcome::Failed(error)) => Err(error.into_error()), - Err(_) => Err(Error::unexpected( - "credential refresh coordination ended without a result", - ) - .with_context(format!("credential_type: {}", type_name::())) - .set_retryable(true)), - }, - CredentialCacheAction::Refresh(refresh) => { - let result = self - .loader - .provide_credential_dyn(&self.ctx) - .await - .and_then(|credential| { - credential.ok_or_else(|| { - Error::credential_invalid("failed to load signing credential") - .with_context(format!("credential_type: {}", type_name::())) - }) - }); - - match result { - Ok(credential) => { - refresh.complete_success(credential.clone()); - self.validate_refreshed_credential(credential, expires_in) - } - Err(error) => { - refresh.complete_failure(&error); - Err(error) - } - } + { + return Ok(credential.clone()); } } + + let credential = self + .loader + .provide_credential_dyn(&self.ctx) + .await? + .ok_or_else(|| { + Error::credential_invalid("failed to load signing credential") + .with_context(format!("credential_type: {}", type_name::())) + })?; + + *cached = Some(credential.clone()); + drop(cached); + + self.validate_refreshed_credential(credential, expires_in) } fn validate_refreshed_credential( @@ -374,6 +186,7 @@ mod tests { use futures::poll; use http::{HeaderValue, Method, Request, Version}; use std::collections::VecDeque; + use std::sync::Mutex as StdMutex; use std::sync::atomic::{AtomicUsize, Ordering}; #[derive(Clone, Debug)] @@ -459,7 +272,7 @@ mod tests { #[derive(Debug)] struct SequenceProvider { - responses: Mutex>>>, + responses: StdMutex>>>, calls: Arc, } @@ -470,7 +283,7 @@ mod tests { let calls = Arc::new(AtomicUsize::new(0)); ( Self { - responses: Mutex::new(responses.into_iter().collect()), + responses: StdMutex::new(responses.into_iter().collect()), calls: calls.clone(), }, calls, @@ -492,7 +305,7 @@ mod tests { } struct ControlledProvider { - responses: Mutex>>, + responses: StdMutex>>, calls: Arc, } @@ -510,7 +323,7 @@ mod tests { .unzip::<_, _, Vec<_>, VecDeque<_>>(); ( Self { - responses: Mutex::new(receivers), + responses: StdMutex::new(receivers), calls: calls.clone(), }, calls, @@ -538,7 +351,7 @@ mod tests { struct ControlledRequestSigner { started: Arc, - releases: Mutex>>, + releases: StdMutex>>, } impl Debug for ControlledRequestSigner { @@ -557,7 +370,7 @@ mod tests { ( Self { started: started.clone(), - releases: Mutex::new(receivers), + releases: StdMutex::new(receivers), }, started, senders, @@ -722,6 +535,13 @@ mod tests { parts } + fn set_cached_credential(signer: &Signer, credential: K) { + *signer + .credential + .try_lock() + .expect("credential cache must be unlocked") = Some(credential); + } + #[test] fn failure_leaves_entire_request_head_unchanged() { let signer = Signer::new( @@ -829,7 +649,7 @@ mod tests { }; let (provider, calls, mut responses) = ControlledProvider::new(1); let signer = Signer::new(Context::new(), provider, OperationSigner); - signer.credential_cache.set_credential(cached); + set_cached_credential(&signer, cached); let mut requests = (0..8).map(|_| request_parts()).collect::>(); let mut batch = Box::pin(join_all( requests @@ -858,7 +678,7 @@ mod tests { } #[test] - fn concurrent_refresh_failure_is_shared_and_later_call_retries() { + fn concurrent_refresh_failure_allows_waiter_retry() { futures::executor::block_on(async { let base = Timestamp::from_second(700).expect("timestamp must be valid"); let refreshed = ExpiringCredential { @@ -884,34 +704,38 @@ mod tests { .send(Err(Error::rate_limited("injected refresh failure") .with_context("refresh_generation: 1"))) .expect("controlled failure must be received"); - for result in batch.await { - let error = result.expect_err("current waiter must receive refresh failure"); - assert_eq!(error.kind(), ErrorKind::RateLimited); - assert_eq!(error.to_string(), "injected refresh failure"); - assert_eq!(error.context(), &["refresh_generation: 1"]); - assert!(error.is_retryable()); - } - assert_eq!(calls.load(Ordering::SeqCst), 1); - let mut retry_request = request_parts(); - let mut retry = Box::pin(signer.sign(&mut retry_request, None)); - assert!(poll!(&mut retry).is_pending()); + assert!(poll!(&mut batch).is_pending()); assert_eq!(calls.load(Ordering::SeqCst), 2); responses .remove(0) .send(Ok(Some(refreshed))) .expect("controlled recovery must be received"); - retry.await.expect("later caller must retry and recover"); + + let results = batch.await; + let errors = results + .iter() + .filter_map(|result| result.as_ref().err()) + .collect::>(); + assert_eq!(errors.len(), 1); + assert_eq!(errors[0].kind(), ErrorKind::RateLimited); + assert_eq!(errors[0].to_string(), "injected refresh failure"); + assert_eq!(errors[0].context(), &["refresh_generation: 1"]); + assert!(errors[0].is_retryable()); + assert_eq!(results.iter().filter(|result| result.is_ok()).count(), 7); + assert_eq!(calls.load(Ordering::SeqCst), 2); + + let mut later_request = request_parts(); + signer + .sign(&mut later_request, None) + .await + .expect("later caller must reuse the recovered credential"); assert_eq!(calls.load(Ordering::SeqCst), 2); - assert_eq!( - retry_request.headers.get("x-credential-generation"), - Some(&HeaderValue::from_static("2")) - ); }); } #[test] - fn concurrent_callers_apply_exact_refreshed_validity_check() { + fn refreshed_credential_is_checked_for_exact_operation_deadline() { futures::executor::block_on(async { let base = Timestamp::from_second(800).expect("timestamp must be valid"); let credential = ExpiringCredential { @@ -922,27 +746,23 @@ mod tests { }; let (provider, calls, mut responses) = ControlledProvider::new(1); let signer = Signer::new(Context::new(), provider, OperationSigner); - let mut requests = (0..8).map(|_| request_parts()).collect::>(); - let mut batch = Box::pin(join_all( - requests - .iter_mut() - .map(|request| signer.sign(request, None)), - )); + let mut request = request_parts(); + let mut signing = Box::pin(signer.sign(&mut request, None)); - assert!(poll!(&mut batch).is_pending()); + assert!(poll!(&mut signing).is_pending()); responses .remove(0) .send(Ok(Some(credential))) .expect("controlled response must be received"); - for result in batch.await { - let error = result.expect_err("exact deadline must reject the credential"); - assert_eq!(error.kind(), ErrorKind::CredentialInvalid); - assert!( - error - .to_string() - .contains("expires before the requested operation deadline") - ); - } + let error = signing + .await + .expect_err("exact deadline must reject the credential"); + assert_eq!(error.kind(), ErrorKind::CredentialInvalid); + assert!( + error + .to_string() + .contains("expires before the requested operation deadline") + ); assert_eq!(calls.load(Ordering::SeqCst), 1); }); } @@ -1047,26 +867,17 @@ mod tests { .with_request_signer(MutatingSigner { fail: false }); let with_provider = signer.clone().with_credential_provider(StaticProvider); + assert!(Arc::ptr_eq(&signer.credential, &clone.credential)); + assert!(Arc::ptr_eq(&signer.credential, &with_context.credential)); assert!(Arc::ptr_eq( - &signer.credential_cache, - &clone.credential_cache - )); - assert!(Arc::ptr_eq( - &signer.credential_cache, - &with_context.credential_cache - )); - assert!(Arc::ptr_eq( - &signer.credential_cache, - &with_request_signer.credential_cache - )); - assert!(!Arc::ptr_eq( - &signer.credential_cache, - &with_provider.credential_cache + &signer.credential, + &with_request_signer.credential )); + assert!(!Arc::ptr_eq(&signer.credential, &with_provider.credential)); } #[test] - fn cancelled_refresh_notifies_waiters_and_allows_retry() { + fn cancelled_refresh_releases_lock_and_waiter_retries() { futures::executor::block_on(async { let base = Timestamp::from_second(975).expect("timestamp must be valid"); let calls = Arc::new(AtomicUsize::new(0)); @@ -1090,20 +901,14 @@ mod tests { assert_eq!(calls.load(Ordering::SeqCst), 1); drop(leader); - let error = waiter + waiter .await - .expect_err("waiter must receive refresh cancellation"); - assert_eq!(error.kind(), ErrorKind::Unexpected); - assert_eq!(error.to_string(), "credential refresh was cancelled"); - assert!(error.is_retryable()); - assert_eq!(calls.load(Ordering::SeqCst), 1); - - let mut retry_request = request_parts(); - signer - .sign(&mut retry_request, None) - .await - .expect("later call must retry after cancellation"); + .expect("waiter must retry after leader cancellation"); assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!( + waiter_request.headers.get("x-credential-generation"), + Some(&HeaderValue::from_static("2")) + ); }); } @@ -1150,7 +955,7 @@ mod tests { }; let (provider, calls) = SequenceProvider::new([Ok(Some(refreshed))]); let signer = Signer::new(Context::new(), provider, OperationSigner); - signer.credential_cache.set_credential(cached); + set_cached_credential(&signer, cached); let mut parts = request_parts(); futures::executor::block_on(signer.sign(&mut parts, None)) @@ -1205,7 +1010,7 @@ mod tests { Ok(Some(refreshed)), ]); let signer = Signer::new(Context::new(), provider, OperationSigner); - signer.credential_cache.set_credential(cached); + set_cached_credential(&signer, cached); let mut parts = request_parts(); let original = parts.clone(); @@ -1242,7 +1047,7 @@ mod tests { }; let (provider, calls) = SequenceProvider::new([Ok(None), Ok(Some(refreshed))]); let signer = Signer::new(Context::new(), provider, OperationSigner); - signer.credential_cache.set_credential(cached); + set_cached_credential(&signer, cached); let mut parts = request_parts(); let original = parts.clone(); @@ -1270,7 +1075,7 @@ mod tests { StaticProvider, MutatingSigner { fail: false }, ); - signer.credential_cache.set_credential(TestCredential); + set_cached_credential(&signer, TestCredential); let debug = format!("{signer:?}"); assert!(debug.starts_with("Signer"));