google-cloud-auth 1.10.0

// Copyright 2025 Google LLC
//
// 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
//
//     https://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.

//! Obtain [OIDC ID tokens] using [impersonated service accounts].
//!
//! When the principal you are using doesn't have the permissions you need to
//! accomplish your task, or you want to use a service account in a development
//! environment, you can use service account impersonation. The typical principals
//! used to impersonate a service account are [User Account] or another [Service Account].
//!
//! The principal that is trying to impersonate a target service account should have
//! [Service Account Token Creator Role] on the target service account.
//!
//! ## Example: Creating impersonated credentials from a JSON object with target audience and sending ID Tokens.
//!
//! ```
//! # use google_cloud_auth::credentials::idtoken;
//! # use serde_json::json;
//! # use reqwest;
//! # async fn sample() -> anyhow::Result<()> {
//! let source_credentials = json!({
//!     "type": "authorized_user",
//!     "client_id": "test-client-id",
//!     "client_secret": "test-client-secret",
//!     "refresh_token": "test-refresh-token"
//! });
//!
//! let impersonated_credential = json!({
//!     "type": "impersonated_service_account",
//!     "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/test-principal:generateAccessToken",
//!     "source_credentials": source_credentials,
//! });
//!
//! let audience = "https://my-service.a.run.app";
//! let credentials = idtoken::impersonated::Builder::new(audience, impersonated_credential)
//!     .build()?;
//! let id_token = credentials.id_token().await?;
//!
//! // Make request with ID Token as Bearer Token.
//! let client = reqwest::Client::new();
//! let target_url = format!("{audience}/api/method");
//! client.get(target_url)
//!     .bearer_auth(id_token)
//!     .send()
//!     .await?;
//! # Ok(()) }
//! ```
//!
//! [Impersonated service accounts]: https://cloud.google.com/docs/authentication/use-service-account-impersonation
//! [OIDC ID tokens]: https://cloud.google.com/docs/authentication/token-types#identity-tokens
//! [User Account]: https://cloud.google.com/docs/authentication#user-accounts
//! [Service Account]: https://cloud.google.com/iam/docs/service-account-overview
//! [Service Account Token Creator Role]: https://cloud.google.com/docs/authentication/use-service-account-impersonation#required-roles

use crate::{
    BuildResult, Result,
    credentials::{
        CacheableResource, Credentials,
        idtoken::{
            IDTokenCredentials, dynamic::IDTokenCredentialsProvider, parse_id_token_from_str,
        },
        impersonated::{
            BuilderSource, IMPERSONATED_CREDENTIAL_TYPE, ImpersonationUrl, MSG,
            build_components_from_credentials, build_components_from_json,
        },
    },
    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 google_cloud_gax::backoff_policy::BackoffPolicyArg;
use google_cloud_gax::error::CredentialsError;
use google_cloud_gax::retry_policy::RetryPolicyArg;
use google_cloud_gax::retry_throttler::RetryThrottlerArg;
use http::{Extensions, HeaderMap};
use reqwest::Client;
use serde_json::Value;
use std::sync::Arc;

/// A builder for constructing Impersonated Service Account [IDTokenCredentials] instance.
///
/// # Example
/// ```
/// # use google_cloud_auth::credentials::idtoken;
/// # async fn sample() -> anyhow::Result<()> {
/// let impersonated_credential = serde_json::json!({
///     "type": "impersonated_service_account",
///     "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/test-principal:generateAccessToken",
///     "source_credentials": {
///         "type": "authorized_user",
///         "client_id": "test-client-id",
///         "client_secret": "test-client-secret",
///         "refresh_token": "test-refresh-token"
///     }
/// });
///
/// let audience = "https://my-service.a.run.app";
/// let credentials = idtoken::impersonated::Builder::new(audience, impersonated_credential).build()?;
/// # Ok(()) }
/// ```
pub struct Builder {
    source: BuilderSource,
    delegates: Option<Vec<String>>,
    pub(crate) include_email: Option<bool>,
    target_audience: String,
    service_account_impersonation_url: Option<ImpersonationUrl>,
    retry_builder: RetryTokenProviderBuilder,
}

impl Builder {
    /// Creates a new builder using `impersonated_service_account` JSON value.
    ///
    /// The `impersonated_service_account` JSON is typically generated using
    /// [application default login] with the [impersonation flag].
    ///
    /// [impersonation flag]: https://cloud.google.com/docs/authentication/use-service-account-impersonation#adc
    /// [application default login]: https://cloud.google.com/sdk/gcloud/reference/auth/application-default/login
    pub fn new<S: Into<String>>(target_audience: S, impersonated_credential: Value) -> Self {
        Self {
            source: BuilderSource::FromJson(impersonated_credential),
            delegates: None,
            include_email: None,
            target_audience: target_audience.into(),
            service_account_impersonation_url: None,
            retry_builder: RetryTokenProviderBuilder::default(),
        }
    }

    /// Creates a new builder with a source [Credentials] object, target principal and audience.
    /// Target principal is the email of the service account to impersonate.
    ///
    /// # Example
    /// ```
    /// # use google_cloud_auth::credentials::idtoken;
    /// # use google_cloud_auth::credentials::user_account;
    /// # use serde_json::json;
    /// #
    /// # fn example() -> anyhow::Result<()> {
    /// let source_credentials = user_account::Builder::new(json!({ /* add details here */ })).build()?;
    ///
    /// let audience = "https://my-service.a.run.app";
    /// let credentials = idtoken::impersonated::Builder::from_source_credentials(audience, "test-principal", source_credentials)
    ///     .build()?;
    /// # Ok(()) }
    /// // Now you can use credentials.id_token().await to fetch the token.
    /// ```
    pub fn from_source_credentials<SA: Into<String>, SP: Into<String>>(
        target_audience: SA,
        target_principal: SP,
        source_credentials: Credentials,
    ) -> Self {
        Self {
            source: BuilderSource::FromCredentials(source_credentials),
            delegates: None,
            include_email: None,
            target_audience: target_audience.into(),
            service_account_impersonation_url: Some(ImpersonationUrl::target_principal(
                target_principal.into(),
            )),
            retry_builder: RetryTokenProviderBuilder::default(),
        }
    }

    /// Should include email claims in the ID Token.
    ///
    /// # Example
    ///
    /// ```
    /// # use google_cloud_auth::credentials::idtoken;
    /// # use serde_json::json;
    /// 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_include_email()
    ///     .build();
    /// // Now you can use credentials.id_token().await to fetch the token.
    /// ```
    pub fn with_include_email(mut self) -> Self {
        self.include_email = Some(true);
        self
    }

    /// Sets the chain of delegates.
    ///
    /// # Example
    ///
    /// ```
    /// # use google_cloud_auth::credentials::idtoken;
    /// # use serde_json::json;
    /// 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_delegates(vec!["delegate1-sa@example.com", "delegate2-sa@example.com"])
    ///     .build();
    /// // Now you can use credentials.id_token().await to fetch the token.
    /// ```
    pub fn with_delegates<I, S>(mut self, delegates: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.delegates = Some(delegates.into_iter().map(|s| s.into()).collect());
        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 google_cloud_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<V: Into<RetryPolicyArg>>(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 google_cloud_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<V: Into<BackoffPolicyArg>>(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 google_cloud_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<V: Into<RetryThrottlerArg>>(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
    ///
    /// Returns an Error for one of the following cases:
    /// - If the `impersonated_service_account` provided to [`Builder::new`] cannot
    ///   be successfully deserialized into the expected format. This typically happens
    ///   if the JSON value is malformed or missing required fields. For more information,
    ///   see the guide on how to [use service account impersonation].
    /// - If the `impersonated_service_account` provided to [`Builder::new`] has a
    ///   `source_credentials` of `impersonated_service_account` type.
    ///
    /// [use service account impersonation]: https://cloud.google.com/docs/authentication/use-service-account-impersonation#adc
    pub fn build(self) -> BuildResult<IDTokenCredentials> {
        let components = match self.source {
            BuilderSource::FromJson(json) => build_components_from_json(json)?,
            BuilderSource::FromCredentials(source_credentials) => {
                build_components_from_credentials(
                    source_credentials,
                    self.service_account_impersonation_url,
                )?
            }
        };
        let token_provider = ImpersonatedTokenProvider {
            source_credentials: components.source_credentials,
            service_account_impersonation_url: components.service_account_impersonation_url,
            delegates: self.delegates.or(components.delegates),
            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),
            }),
        })
    }
}

#[derive(Debug)]
struct ImpersonatedServiceAccount<T>
where
    T: CachedTokenProvider,
{
    token_provider: T,
}

#[async_trait::async_trait]
impl<T> IDTokenCredentialsProvider for ImpersonatedServiceAccount<T>
where
    T: CachedTokenProvider,
{
    async fn id_token(&self) -> Result<String> {
        let cached_token = self.token_provider.token(Extensions::new()).await?;
        match cached_token {
            CacheableResource::New { data, .. } => Ok(data.token),
            CacheableResource::NotModified => {
                Err(CredentialsError::from_msg(false, "failed to fetch token"))
            }
        }
    }
}

#[derive(Debug)]
pub(crate) struct ImpersonatedTokenProvider {
    pub(crate) source_credentials: Credentials,
    pub(crate) service_account_impersonation_url: ImpersonationUrl,
    pub(crate) delegates: Option<Vec<String>>,
    pub(crate) target_audience: String,
    pub(crate) include_email: Option<bool>,
}

#[derive(serde::Serialize, serde::Deserialize, Debug, PartialEq)]
struct GenerateIdTokenRequest {
    #[serde(skip_serializing_if = "Option::is_none")]
    delegates: Option<Vec<String>>,
    audience: String,
    #[serde(skip_serializing_if = "Option::is_none", rename = "includeEmail")]
    include_email: Option<bool>,
}

async fn generate_id_token(
    source_headers: HeaderMap,
    delegates: Option<Vec<String>>,
    audience: String,
    include_email: Option<bool>,
    service_account_impersonation_url: &str,
) -> Result<Token> {
    let client = Client::new();

    let body = GenerateIdTokenRequest {
        audience,
        delegates,
        include_email,
    };

    let response = client
        .post(service_account_impersonation_url)
        .header("Content-Type", "application/json")
        .header(
            headers_util::X_GOOG_API_CLIENT,
            metrics_header_value(ID_TOKEN_REQUEST_TYPE, IMPERSONATED_CREDENTIAL_TYPE),
        )
        .headers(source_headers)
        .json(&body)
        .send()
        .await
        .map_err(|e| errors::from_http_error(e, MSG))?;

    if !response.status().is_success() {
        let err = errors::from_http_response(response, MSG).await;
        return Err(err);
    }

    let token_response = response
        .json::<GenerateIdTokenResponse>()
        .await
        .map_err(|e| {
            let retryable = !e.is_decode();
            CredentialsError::from_source(retryable, e)
        })?;

    parse_id_token_from_str(token_response.token)
}

#[async_trait]
impl TokenProvider for ImpersonatedTokenProvider {
    async fn token(&self) -> Result<Token> {
        let source_headers = self.source_credentials.headers(Extensions::new()).await?;
        let source_headers = match source_headers {
            CacheableResource::New { data, .. } => data,
            CacheableResource::NotModified => {
                unreachable!("requested source credentials without a caching etag")
            }
        };

        // We resolve the URL on every token call because fetching the universe domain
        // is async and must be done here rather than in the builder.
        // Since `token()` takes `&self`, we cannot mutate `self` to cache the URL
        // without using a lock for inner mutability, which is worse than just
        // building the URL string for each request.
        let url = self
            .service_account_impersonation_url
            .id_token_url(&self.source_credentials)
            .await;

        generate_id_token(
            source_headers,
            self.delegates.clone(),
            self.target_audience.clone(),
            self.include_email,
            &url,
        )
        .await
    }
}

#[derive(serde::Deserialize)]
struct GenerateIdTokenResponse {
    #[serde(rename = "token")]
    token: String,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::credentials::idtoken::tests::generate_test_id_token;
    use crate::credentials::tests::MockCredentials;
    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;

    type TestResult = anyhow::Result<()>;

    impl Builder {
        fn with_impersonation_endpoint(mut self, endpoint: &str) -> Self {
            self.service_account_impersonation_url = self
                .service_account_impersonation_url
                .map(|u| u.with_endpoint(endpoint));
            self
        }
    }

    #[tokio::test]
    async fn test_impersonated_service_account_id_token() -> 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(
                json_encoded(json!({
                    "access_token": "test-user-account-token",
                    "expires_in": 3600,
                    "token_type": "Bearer",
                })),
            ),
        );

        server.expect(
            Expectation::matching(all_of![
                request::method_path(
                    "POST",
                    "/v1/projects/-/serviceAccounts/test-principal:generateIdToken"
                ),
                request::headers(contains((
                    "authorization",
                    "Bearer test-user-account-token"
                ))),
                request::body(json_decoded(eq(json!({
                    "audience": audience,
                }))))
            ])
            .respond_with(json_encoded(json!({
                "token": token_string,
            }))),
        );

        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(audience, impersonated_credential.clone()).build()?;

        let token = creds.id_token().await?;
        assert_eq!(token, token_string);

        Ok(())
    }

    #[tokio::test]
    async fn test_impersonated_id_token_with_delegates_and_email() -> 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(
                json_encoded(json!({
                    "access_token": "test-user-account-token",
                    "expires_in": 3600,
                    "token_type": "Bearer",
                })),
            ),
        );

        server.expect(
            Expectation::matching(all_of![
                request::method_path(
                    "POST",
                    "/v1/projects/-/serviceAccounts/test-principal:generateIdToken"
                ),
                request::headers(contains((
                    "authorization",
                    "Bearer test-user-account-token"
                ))),
                request::body(json_decoded(eq(json!({
                    "audience": audience,
                    "delegates": ["delegate1", "delegate2"],
                    "includeEmail": true
                }))))
            ])
            .respond_with(json_encoded(json!({
                "token": token_string,
            }))),
        );

        let impersonated_credential = json!({
            "type": "impersonated_service_account",
            "service_account_impersonation_url": server.url("/v1/projects/-/serviceAccounts/test-principal:generateIdToken").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_delegates(vec!["delegate1", "delegate2"])
            .with_include_email()
            .build()?;

        let token = creds.id_token().await?;
        assert_eq!(token, token_string);

        Ok(())
    }

    #[tokio::test]
    async fn test_impersonated_id_token_from_source_credentials() -> 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(
                json_encoded(json!({
                    "access_token": "test-user-account-token",
                    "expires_in": 3600,
                    "token_type": "Bearer",
                })),
            ),
        );

        server.expect(
            Expectation::matching(all_of![
                request::method_path(
                    "POST",
                    "/v1/projects/-/serviceAccounts/test-principal:generateIdToken"
                ),
                request::headers(contains((
                    "authorization",
                    "Bearer test-user-account-token"
                ))),
                request::body(json_decoded(eq(json!({
                    "audience": audience,
                }))))
            ])
            .respond_with(json_encoded(json!({
                "token": token_string,
            }))),
        );

        let source_credentials = crate::credentials::user_account::Builder::new(json!({
            "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()
        }))
        .build()?;

        let endpoint = server.url("/").to_string();
        let endpoint = endpoint.trim_end_matches('/');
        let creds =
            Builder::from_source_credentials(audience, "test-principal", source_credentials)
                .with_impersonation_endpoint(endpoint)
                .build()?;

        let token = creds.id_token().await?;
        assert_eq!(token, token_string);

        Ok(())
    }

    #[tokio::test]
    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(
                json_encoded(json!({
                    "access_token": "test-user-account-token",
                    "expires_in": 3600,
                    "token_type": "Bearer",
                })),
            ),
        );

        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,
            }))),
        );

        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(audience, impersonated_credential).build()?;

        let token = creds.id_token().await?;
        assert_eq!(token, token_string);

        Ok(())
    }

    #[tokio::test]
    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();
        // 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(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!({
            "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(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);

        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(())
    }

    #[tokio::test]
    async fn test_impersonated_id_token_custom_universe_domain() -> TestResult {
        let audience = "test-audience";
        let token_string = generate_test_id_token(audience);
        let server = Server::run();
        let universe_domain = "my-custom-universe.com".to_string();

        server.expect(
            Expectation::matching(all_of![
                request::method_path(
                    "POST",
                    "/v1/projects/-/serviceAccounts/test-principal:generateIdToken"
                ),
                request::headers(contains((
                    "authorization",
                    "Bearer test-user-account-token"
                ))),
                request::body(json_decoded(eq(json!({
                    "audience": audience,
                }))))
            ])
            .respond_with(json_encoded(json!({
                "token": token_string,
            }))),
        );

        let universe_domain_clone = universe_domain.clone();
        let mut mock = MockCredentials::new();
        mock.expect_universe_domain()
            .returning(move || Some(universe_domain_clone.clone()));
        mock.expect_headers().returning(move |_| {
            let mut headers = HeaderMap::new();
            headers.insert(
                "authorization",
                "Bearer test-user-account-token".parse().unwrap(),
            );
            Ok(CacheableResource::New {
                entity_tag: Default::default(),
                data: headers,
            })
        });
        let source_credentials = Credentials::from(mock);

        let builder = Builder::from_source_credentials(
            audience,
            "test-principal",
            source_credentials.clone(),
        );

        // resolve url with universe domain
        let url = builder
            .service_account_impersonation_url
            .as_ref()
            .expect("url should be set from the with_target_principal call")
            .id_token_url(&source_credentials)
            .await;

        assert_eq!(
            url,
            format!(
                "https://iamcredentials.{universe_domain}/v1/projects/-/serviceAccounts/test-principal:generateIdToken"
            )
        );

        let endpoint = server.url("/").to_string();
        let endpoint = endpoint.trim_end_matches('/');

        let creds = builder.with_impersonation_endpoint(endpoint).build()?;

        let token = creds.id_token().await?;
        assert_eq!(token, token_string);

        Ok(())
    }
}