From a0f1a331c7af8ab5ad0d40f58d20099d76164d3c Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 12 Jan 2026 17:35:14 +0000 Subject: [PATCH 1/2] impl(auth): add retry to impersonated idtoken provider --- .../src/credentials/idtoken/impersonated.rs | 202 +++++++++++++++--- 1 file changed, 176 insertions(+), 26 deletions(-) diff --git a/src/auth/src/credentials/idtoken/impersonated.rs b/src/auth/src/credentials/idtoken/impersonated.rs index e645f90d43..1462b48c6b 100644 --- a/src/auth/src/credentials/idtoken/impersonated.rs +++ b/src/auth/src/credentials/idtoken/impersonated.rs @@ -78,11 +78,15 @@ use crate::{ }, errors, headers_util::{self, ID_TOKEN_REQUEST_TYPE, metrics_header_value}, + retry::Builder as RetryTokenProviderBuilder, token::{CachedTokenProvider, Token, TokenProvider}, token_cache::TokenCache, }; use async_trait::async_trait; +use gax::backoff_policy::BackoffPolicyArg; use gax::error::CredentialsError; +use gax::retry_policy::RetryPolicyArg; +use gax::retry_throttler::RetryThrottlerArg; use http::{Extensions, HeaderMap}; use reqwest::Client; use serde_json::Value; @@ -114,6 +118,7 @@ pub struct Builder { pub(crate) include_email: Option, target_audience: String, service_account_impersonation_url: Option, + retry_builder: RetryTokenProviderBuilder, } impl Builder { @@ -131,6 +136,7 @@ impl Builder { include_email: None, target_audience: target_audience.into(), service_account_impersonation_url: None, + retry_builder: RetryTokenProviderBuilder::default(), } } @@ -167,6 +173,7 @@ impl Builder { "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{}:generateIdToken", target_principal.into() )), + retry_builder: RetryTokenProviderBuilder::default(), } } @@ -223,6 +230,83 @@ impl Builder { self } + /// Configure the retry policy for fetching tokens. + /// + /// The retry policy controls how to handle retries, and sets limits on + /// the number of attempts or the total time spent retrying. + /// + /// # Example + /// + /// ``` + /// # use google_cloud_auth::credentials::idtoken; + /// # use serde_json::json; + /// use gax::retry_policy::{AlwaysRetry, RetryPolicyExt}; + /// + /// let impersonated_credential = json!({ /* add details here */ }); + /// + /// let audience = "https://my-service.a.run.app"; + /// let credentials = idtoken::impersonated::Builder::new(audience, impersonated_credential) + /// .with_retry_policy(AlwaysRetry.with_attempt_limit(3)) + /// .build(); + /// ``` + pub fn with_retry_policy>(mut self, v: V) -> Self { + self.retry_builder = self.retry_builder.with_retry_policy(v.into()); + self + } + + /// Configure the retry backoff policy. + /// + /// The backoff policy controls how long to wait in between retry attempts. + /// + /// # Example + /// + /// ``` + /// # use google_cloud_auth::credentials::idtoken; + /// # use serde_json::json; + /// use gax::exponential_backoff::ExponentialBackoff; + /// + /// let impersonated_credential = json!({ /* add details here */ }); + /// + /// let audience = "https://my-service.a.run.app"; + /// let credentials = idtoken::impersonated::Builder::new(audience, impersonated_credential) + /// .with_backoff_policy(ExponentialBackoff::default()) + /// .build(); + /// ``` + pub fn with_backoff_policy>(mut self, v: V) -> Self { + self.retry_builder = self.retry_builder.with_backoff_policy(v.into()); + self + } + + /// Configure the retry throttler. + /// + /// Advanced applications may want to configure a retry throttler to + /// [Address Cascading Failures] and when [Handling Overload] conditions. + /// The authentication library throttles its retry loop, using a policy to + /// control the throttling algorithm. Use this method to fine tune or + /// customize the default retry throttler. + /// + /// [Handling Overload]: https://sre.google/sre-book/handling-overload/ + /// [Address Cascading Failures]: https://sre.google/sre-book/addressing-cascading-failures/ + /// + /// # Example + /// + /// ``` + /// # use google_cloud_auth::credentials::idtoken; + /// # use serde_json::json; + /// use gax::retry_throttler::AdaptiveThrottler; + /// + /// let impersonated_credential = json!({ /* add details here */ }); + /// + /// let audience = "https://my-service.a.run.app"; + /// let credentials = idtoken::impersonated::Builder::new(audience, impersonated_credential) + /// .with_retry_throttler(AdaptiveThrottler::default()) + /// .build(); + /// ``` + pub fn with_retry_throttler>(mut self, v: V) -> Self { + self.retry_builder = self.retry_builder.with_retry_throttler(v.into()); + self + } + /// Returns a [Credentials] instance with the configured settings. /// /// # Errors @@ -259,6 +343,7 @@ impl Builder { include_email: self.include_email, target_audience: self.target_audience, }; + let token_provider = self.retry_builder.build(token_provider); Ok(IDTokenCredentials { inner: Arc::new(ImpersonatedServiceAccount { token_provider: TokenCache::new(token_provider), @@ -385,6 +470,9 @@ struct GenerateIdTokenResponse { mod tests { use super::*; use crate::credentials::idtoken::tests::generate_test_id_token; + use crate::credentials::tests::{ + get_mock_auth_retry_policy, get_mock_backoff_policy, get_mock_retry_throttler, + }; use httptest::{Expectation, Server, matchers::*, responders::*}; use serde_json::json; @@ -556,7 +644,9 @@ mod tests { } #[tokio::test] - async fn test_impersonated_id_token_fail() -> TestResult { + async fn test_impersonated_id_token_metrics_header() -> TestResult { + let audience = "test-audience"; + let token_string = generate_test_id_token(audience); let server = Server::run(); server.expect( Expectation::matching(request::method_path("POST", "/token")).respond_with( @@ -567,12 +657,22 @@ mod tests { })), ), ); + server.expect( - Expectation::matching(request::method_path( - "POST", - "/v1/projects/-/serviceAccounts/test-principal:generateIdToken", - )) - .respond_with(status_code(500)), + Expectation::matching(all_of![ + request::method_path( + "POST", + "/v1/projects/-/serviceAccounts/test-principal:generateIdToken" + ), + request::headers(contains(("x-goog-api-client", matches("cred-type/imp")))), + request::headers(contains(( + "x-goog-api-client", + matches("auth-request-type/it") + ))) + ]) + .respond_with(json_encoded(json!({ + "token": token_string, + }))), ); let impersonated_credential = json!({ @@ -586,19 +686,20 @@ mod tests { "token_uri": server.url("/token").to_string() } }); - let creds = Builder::new("test-audience", impersonated_credential).build()?; + let creds = Builder::new(audience, impersonated_credential).build()?; - let err = creds.id_token().await.unwrap_err(); - assert!(err.is_transient()); + let token = creds.id_token().await?; + assert_eq!(token, token_string); Ok(()) } #[tokio::test] - async fn test_impersonated_id_token_metrics_header() -> TestResult { + async fn test_impersonated_id_token_retries_for_success() -> TestResult { let audience = "test-audience"; let token_string = generate_test_id_token(audience); - let server = Server::run(); + let mut server = Server::run(); + // Source credential token endpoint server.expect( Expectation::matching(request::method_path("POST", "/token")).respond_with( json_encoded(json!({ @@ -609,21 +710,18 @@ mod tests { ), ); + // Impersonation endpoint + let impersonation_path = "/v1/projects/-/serviceAccounts/test-principal:generateIdToken"; server.expect( - Expectation::matching(all_of![ - request::method_path( - "POST", - "/v1/projects/-/serviceAccounts/test-principal:generateIdToken" - ), - request::headers(contains(("x-goog-api-client", matches("cred-type/imp")))), - request::headers(contains(( - "x-goog-api-client", - matches("auth-request-type/it") - ))) - ]) - .respond_with(json_encoded(json!({ - "token": token_string, - }))), + Expectation::matching(request::method_path("POST", impersonation_path)) + .times(3) + .respond_with(cycle![ + status_code(503).body("try-again"), + status_code(503).body("try-again"), + json_encoded(json!({ + "token": token_string, + })), + ]), ); let impersonated_credential = json!({ @@ -637,11 +735,63 @@ mod tests { "token_uri": server.url("/token").to_string() } }); - let creds = Builder::new(audience, impersonated_credential).build()?; + + let creds = Builder::new(audience, impersonated_credential) + .with_retry_policy(get_mock_auth_retry_policy(3)) + .with_backoff_policy(get_mock_backoff_policy()) + .with_retry_throttler(get_mock_retry_throttler()) + .build()?; let token = creds.id_token().await?; assert_eq!(token, token_string); + server.verify_and_clear(); + Ok(()) + } + + #[tokio::test] + async fn test_impersonated_id_token_does_not_retry_on_non_transient_failures() -> TestResult { + let server = Server::run(); + // Source credential token endpoint + server.expect( + Expectation::matching(request::method_path("POST", "/token")).respond_with( + json_encoded(json!({ + "access_token": "test-user-account-token", + "expires_in": 3600, + "token_type": "Bearer", + })), + ), + ); + + // Impersonation endpoint + let impersonation_path = "/v1/projects/-/serviceAccounts/test-principal:generateIdToken"; + server.expect( + Expectation::matching(request::method_path("POST", impersonation_path)) + .times(1) + .respond_with(status_code(401)), + ); + + let impersonated_credential = json!({ + "type": "impersonated_service_account", + "service_account_impersonation_url": server.url("/v1/projects/-/serviceAccounts/test-principal:generateAccessToken").to_string(), + "source_credentials": { + "type": "authorized_user", + "client_id": "test-client-id", + "client_secret": "test-client-secret", + "refresh_token": "test-refresh-token", + "token_uri": server.url("/token").to_string() + } + }); + + let creds = Builder::new("test-audience", impersonated_credential) + .with_retry_policy(get_mock_auth_retry_policy(3)) + .with_backoff_policy(get_mock_backoff_policy()) + .with_retry_throttler(get_mock_retry_throttler()) + .build()?; + + let err = creds.id_token().await.unwrap_err(); + assert!(!err.is_transient()); + Ok(()) } } From 1eb1ad1d4e3272dc33e945b82396a90d8ec0b8e5 Mon Sep 17 00:00:00 2001 From: Alvaro Viebrantz Date: Mon, 12 Jan 2026 18:16:25 +0000 Subject: [PATCH 2/2] fix: httptest server doesnt need to be mut --- src/auth/src/credentials/idtoken/impersonated.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/auth/src/credentials/idtoken/impersonated.rs b/src/auth/src/credentials/idtoken/impersonated.rs index 1462b48c6b..fc9ef16482 100644 --- a/src/auth/src/credentials/idtoken/impersonated.rs +++ b/src/auth/src/credentials/idtoken/impersonated.rs @@ -698,7 +698,7 @@ mod tests { async fn test_impersonated_id_token_retries_for_success() -> TestResult { let audience = "test-audience"; let token_string = generate_test_id_token(audience); - let mut server = Server::run(); + let server = Server::run(); // Source credential token endpoint server.expect( Expectation::matching(request::method_path("POST", "/token")).respond_with( @@ -745,7 +745,6 @@ mod tests { let token = creds.id_token().await?; assert_eq!(token, token_string); - server.verify_and_clear(); Ok(()) }