google_cloud_auth/credentials/

mds.rs

// Copyright 2024 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.

//! [Metadata Service] Credentials type.
//!
//! Google Cloud environments such as [Google Compute Engine (GCE)][gce-link],
//! [Google Kubernetes Engine (GKE)][gke-link], or [Cloud Run] provide a metadata service.
//! This is a local service to the VM (or pod) which (as the name implies) provides
//! metadata information about the VM. The service also provides access
//! tokens associated with the [default service account] for the corresponding
//! VM.
//!
//! The default host name of the metadata service is `metadata.google.internal`.
//! If you would like to use a different hostname, you can set it using the
//! `GCE_METADATA_HOST` environment variable.
//!
//! You can use this access token to securely authenticate with Google Cloud,
//! without having to download secrets or other credentials. The types in this
//! module allow you to retrieve these access tokens, and can be used with
//! the Google Cloud client libraries for Rust.
//!
//! ## Example: Creating credentials with a custom quota project
//!
//! ```
//! # use google_cloud_auth::credentials::mds::Builder;
//! # use google_cloud_auth::credentials::Credentials;
//! # use http::Extensions;
//! # async fn sample() -> anyhow::Result<()> {
//! let credentials: Credentials = Builder::default()
//!     .with_quota_project_id("my-quota-project")
//!     .build()?;
//! let headers = credentials.headers(Extensions::new()).await?;
//! println!("Headers: {headers:?}");
//! # Ok(()) }
//! ```
//!
//! ## Example: Creating credentials with custom retry behavior
//!
//! ```
//! # use google_cloud_auth::credentials::mds::Builder;
//! # use google_cloud_auth::credentials::Credentials;
//! # use http::Extensions;
//! # use std::time::Duration;
//! # async fn sample() -> anyhow::Result<()> {
//! use google_cloud_gax::retry_policy::{AlwaysRetry, RetryPolicyExt};
//! use google_cloud_gax::exponential_backoff::ExponentialBackoff;
//! let backoff = ExponentialBackoff::default();
//! let credentials: Credentials = Builder::default()
//!     .with_retry_policy(AlwaysRetry.with_attempt_limit(3))
//!     .with_backoff_policy(backoff)
//!     .build()?;
//! let headers = credentials.headers(Extensions::new()).await?;
//! println!("Headers: {headers:?}");
//! # Ok(()) }
//! ```
//!
//! [Application Default Credentials]: https://cloud.google.com/docs/authentication/application-default-credentials
//! [Cloud Run]: https://cloud.google.com/run
//! [default service account]: https://cloud.google.com/iam/docs/service-account-types#default
//! [gce-link]: https://cloud.google.com/products/compute
//! [gke-link]: https://cloud.google.com/kubernetes-engine
//! [Metadata Service]: https://cloud.google.com/compute/docs/metadata/overview

use crate::access_boundary::CredentialsWithAccessBoundary;
use crate::credentials::dynamic::{AccessTokenCredentialsProvider, CredentialsProvider};
use crate::credentials::{AccessToken, AccessTokenCredentials, CacheableResource, Credentials};
use crate::headers_util::AuthHeadersBuilder;
use crate::mds::client::Client as MDSClient;
use crate::retry::{Builder as RetryTokenProviderBuilder, TokenProviderWithRetry};
use crate::token::{CachedTokenProvider, Token, TokenProvider};
use crate::token_cache::TokenCache;
use crate::{BuildResult, Result};
use async_trait::async_trait;
use google_cloud_gax::backoff_policy::{BackoffPolicy, BackoffPolicyArg};
use google_cloud_gax::error::CredentialsError;
use google_cloud_gax::retry_policy::{RetryPolicy, RetryPolicyArg};
use google_cloud_gax::retry_throttler::{RetryThrottlerArg, SharedRetryThrottler};
use http::{Extensions, HeaderMap};
use std::default::Default;
use std::sync::{Arc, OnceLock};

// TODO(#2235) - Improve this message by talking about retries when really running with MDS
const MDS_NOT_FOUND_ERROR: &str = concat!(
    "Could not fetch an auth token to authenticate with Google Cloud. ",
    "The most common reason for this problem is that you are not running in a Google Cloud Environment ",
    "and you have not configured local credentials for development and testing. ",
    "To setup local credentials, run `gcloud auth application-default login`. ",
    "More information on how to authenticate client libraries can be found at https://cloud.google.com/docs/authentication/client-libraries"
);

#[derive(Debug)]
struct MDSCredentials<T>
where
    T: CachedTokenProvider,
{
    quota_project_id: Option<String>,
    universe_domain_override: Option<String>,
    universe_domain: OnceLock<Option<String>>,
    token_provider: T,
    mds_client: MDSClient,
    retry_policy: Arc<dyn RetryPolicy>,
    backoff_policy: Arc<dyn BackoffPolicy>,
    retry_throttler: SharedRetryThrottler,
}

/// Creates [Credentials] instances backed by the [Metadata Service].
///
/// While the Google Cloud client libraries for Rust default to credentials
/// backed by the metadata service, some applications may need to:
/// * Customize the metadata service credentials in some way
/// * Bypass the [Application Default Credentials] lookup and only
///   use the metadata server credentials
/// * Use the credentials directly outside the client libraries
///
/// [Application Default Credentials]: https://cloud.google.com/docs/authentication/application-default-credentials
/// [Metadata Service]: https://cloud.google.com/compute/docs/metadata/overview
#[derive(Debug)]
pub struct Builder {
    endpoint: Option<String>,
    quota_project_id: Option<String>,
    universe_domain: Option<String>,
    scopes: Option<Vec<String>>,
    created_by_adc: bool,
    retry_builder: RetryTokenProviderBuilder,
    iam_endpoint_override: Option<String>,
    is_access_boundary_enabled: bool,
}

impl Default for Builder {
    fn default() -> Self {
        Self {
            endpoint: None,
            quota_project_id: None,
            universe_domain: None,
            scopes: None,
            created_by_adc: false,
            retry_builder: RetryTokenProviderBuilder::default(),
            iam_endpoint_override: None,
            is_access_boundary_enabled: true,
        }
    }
}

impl Builder {
    /// Sets the endpoint for this credentials.
    ///
    /// A trailing slash is significant, so specify the base URL without a trailing
    /// slash. If not set, the credentials use `http://metadata.google.internal`.
    ///
    /// # Example
    /// ```
    /// # use google_cloud_auth::credentials::mds::Builder;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let credentials = Builder::default()
    ///     .with_endpoint("https://metadata.google.foobar")
    ///     .build()?;
    /// # Ok(()) }
    /// ```
    pub fn with_endpoint<S: Into<String>>(mut self, endpoint: S) -> Self {
        self.endpoint = Some(endpoint.into());
        self
    }

    /// Set the [quota project] for this credentials.
    ///
    /// In some services, you can use a service account in
    /// one project for authentication and authorization, and charge
    /// the usage to a different project. This may require that the
    /// service account has `serviceusage.services.use` permissions on the quota project.
    ///
    /// [quota project]: https://cloud.google.com/docs/quotas/quota-project
    pub fn with_quota_project_id<S: Into<String>>(mut self, quota_project_id: S) -> Self {
        self.quota_project_id = Some(quota_project_id.into());
        self
    }

    /// Sets the Google Cloud universe domain for these credentials.
    ///
    /// The universe domain is the default service domain for a given Cloud universe.
    ///
    /// # Example
    /// ```
    /// # use google_cloud_auth::credentials::mds::Builder;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let credentials = Builder::default()
    ///     .with_universe_domain("googleapis.com")
    ///     .build()?;
    /// # Ok(()) }
    /// ```
    ///
    /// [universe domain]: https://cloud.google.com/docs/authentication/universe-domain
    pub fn with_universe_domain<S: Into<String>>(mut self, universe_domain: S) -> Self {
        self.universe_domain = Some(universe_domain.into());
        self
    }

    /// Sets the [scopes] for this credentials.
    ///
    /// Metadata server issues tokens based on the requested scopes.
    /// If no scopes are specified, the credentials defaults to all
    /// scopes configured for the [default service account] on the instance.
    ///
    /// [default service account]: https://cloud.google.com/iam/docs/service-account-types#default
    /// [scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
    pub fn with_scopes<I, S>(mut self, scopes: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.scopes = Some(scopes.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.
    ///
    /// ```
    /// # use google_cloud_auth::credentials::mds::Builder;
    /// # async fn sample() -> anyhow::Result<()> {
    /// use google_cloud_gax::retry_policy::{AlwaysRetry, RetryPolicyExt};
    /// let credentials = Builder::default()
    ///     .with_retry_policy(AlwaysRetry.with_attempt_limit(3))
    ///     .build()?;
    /// # Ok(()) }
    /// ```
    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.
    ///
    /// ```
    /// # use google_cloud_auth::credentials::mds::Builder;
    /// # use std::time::Duration;
    /// # async fn sample() -> anyhow::Result<()> {
    /// use google_cloud_gax::exponential_backoff::ExponentialBackoff;
    /// let policy = ExponentialBackoff::default();
    /// let credentials = Builder::default()
    ///     .with_backoff_policy(policy)
    ///     .build()?;
    /// # Ok(()) }
    /// ```
    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/
    ///
    /// ```
    /// # use google_cloud_auth::credentials::mds::Builder;
    /// # async fn sample() -> anyhow::Result<()> {
    /// use google_cloud_gax::retry_throttler::AdaptiveThrottler;
    /// let credentials = Builder::default()
    ///     .with_retry_throttler(AdaptiveThrottler::default())
    ///     .build()?;
    /// # Ok(()) }
    /// ```
    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
    }

    // This method is used to build mds credentials from ADC
    pub(crate) fn from_adc() -> Self {
        Self {
            created_by_adc: true,
            ..Default::default()
        }
    }

    fn build_token_provider(self) -> TokenProviderWithRetry<MDSAccessTokenProvider> {
        let tp = MDSAccessTokenProvider::builder()
            .endpoint(self.endpoint)
            .maybe_scopes(self.scopes)
            .created_by_adc(self.created_by_adc)
            .build();
        self.retry_builder.build(tp)
    }

    /// Returns a [Credentials] instance with the configured settings.
    pub fn build(self) -> BuildResult<Credentials> {
        Ok(self.build_credentials()?.into())
    }

    /// Returns an [AccessTokenCredentials] instance with the configured settings.
    ///
    /// # Example
    ///
    /// ```
    /// # use google_cloud_auth::credentials::mds::Builder;
    /// # use google_cloud_auth::credentials::{AccessTokenCredentials, AccessTokenCredentialsProvider};
    /// # async fn sample() -> anyhow::Result<()> {
    /// let credentials: AccessTokenCredentials = Builder::default()
    ///     .with_quota_project_id("my-quota-project")
    ///     .build_access_token_credentials()?;
    /// let access_token = credentials.access_token().await?;
    /// println!("Token: {}", access_token.token);
    /// # Ok(()) }
    /// ```
    pub fn build_access_token_credentials(self) -> BuildResult<AccessTokenCredentials> {
        Ok(self.build_credentials()?.into())
    }

    fn build_credentials(
        self,
    ) -> BuildResult<CredentialsWithAccessBoundary<MDSCredentials<TokenCache>>> {
        let iam_endpoint = self.iam_endpoint_override.clone();
        let is_access_boundary_enabled = self.is_access_boundary_enabled;
        let mds_client = MDSClient::new(self.endpoint.clone());
        let retry_builder = self.retry_builder.clone();
        let (backoff_policy, retry_throttler, retry_policy) = retry_builder.resolve();
        let mdsc = MDSCredentials {
            quota_project_id: self.quota_project_id.clone(),
            universe_domain_override: self.universe_domain.clone(),
            universe_domain: OnceLock::new(),
            token_provider: TokenCache::new(self.build_token_provider()),
            mds_client: mds_client.clone(),
            backoff_policy,
            retry_throttler,
            retry_policy,
        };
        if !is_access_boundary_enabled {
            return Ok(CredentialsWithAccessBoundary::new_no_op(mdsc));
        }
        Ok(CredentialsWithAccessBoundary::new_for_mds(
            mdsc,
            mds_client,
            iam_endpoint,
        ))
    }

    /// Returns a [crate::signer::Signer] instance with the configured settings.
    ///
    /// The returned [crate::signer::Signer] uses the [IAM signBlob API] to sign content. This API
    /// requires a network request for each signing operation.
    ///
    /// # Example
    ///
    /// ```
    /// # use google_cloud_auth::credentials::mds::Builder;
    /// # use google_cloud_auth::signer::Signer;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let signer: Signer = Builder::default().build_signer()?;
    /// # Ok(()) }
    /// ```
    ///
    /// [IAM signBlob API]: https://cloud.google.com/iam/docs/reference/credentials/rest/v1/projects.serviceAccounts/signBlob
    pub fn build_signer(self) -> BuildResult<crate::signer::Signer> {
        let client = MDSClient::new(self.endpoint.clone());
        let iam_endpoint = self.iam_endpoint_override.clone();
        let credentials = self.build()?;
        let signing_provider = crate::signer::mds::MDSSigner::new(client, credentials);
        let signing_provider = iam_endpoint
            .iter()
            .fold(signing_provider, |signing_provider, endpoint| {
                signing_provider.with_iam_endpoint_override(endpoint)
            });
        Ok(crate::signer::Signer {
            inner: Arc::new(signing_provider),
        })
    }
}

#[async_trait::async_trait]
impl<T> CredentialsProvider for MDSCredentials<T>
where
    T: CachedTokenProvider,
{
    async fn headers(&self, extensions: Extensions) -> Result<CacheableResource<HeaderMap>> {
        let token = self.token_provider.token(extensions).await?;

        AuthHeadersBuilder::new(&token)
            .maybe_quota_project_id(self.quota_project_id.as_deref())
            .build()
    }

    async fn universe_domain(&self) -> Option<String> {
        if let Some(ud) = &self.universe_domain_override {
            return Some(ud.clone());
        }
        if let Some(ud) = self.universe_domain.get() {
            return ud.clone();
        }

        // No overrides and no cache. Try to fetch from MDS.
        let response = self
            .mds_client
            .universe_domain()
            .with_backoff_policy(self.backoff_policy.clone().into())
            .with_retry_policy(self.retry_policy.clone().into())
            .with_retry_throttler(self.retry_throttler.clone().into())
            .send()
            .await;
        match response {
            Ok(universe_domain) => {
                let _ = self.universe_domain.set(Some(universe_domain.clone()));
                Some(universe_domain)
            }
            Err(e) => {
                if !e.is_transient() {
                    // Only cache None if the error is permanent (e.g., 404 on GDU)
                    let _ = self.universe_domain.set(None);
                }
                // Return None but do not cache it if it's transient,
                // allowing subsequent calls to retry or try again.
                None
            }
        }
    }
}

#[async_trait::async_trait]
impl<T> AccessTokenCredentialsProvider for MDSCredentials<T>
where
    T: CachedTokenProvider,
{
    async fn access_token(&self) -> Result<AccessToken> {
        let token = self.token_provider.token(Extensions::new()).await?;
        token.into()
    }
}

#[derive(Debug, Default)]
struct MDSAccessTokenProviderBuilder {
    scopes: Option<Vec<String>>,
    endpoint: Option<String>,
    created_by_adc: bool,
}

impl MDSAccessTokenProviderBuilder {
    fn build(self) -> MDSAccessTokenProvider {
        MDSAccessTokenProvider {
            client: MDSClient::new(self.endpoint),
            scopes: self.scopes,
            created_by_adc: self.created_by_adc,
        }
    }

    fn maybe_scopes(mut self, v: Option<Vec<String>>) -> Self {
        self.scopes = v;
        self
    }

    fn endpoint<T>(mut self, v: Option<T>) -> Self
    where
        T: Into<String>,
    {
        self.endpoint = v.map(Into::into);
        self
    }

    fn created_by_adc(mut self, v: bool) -> Self {
        self.created_by_adc = v;
        self
    }
}

#[derive(Debug, Clone)]
struct MDSAccessTokenProvider {
    scopes: Option<Vec<String>>,
    client: MDSClient,
    created_by_adc: bool,
}

impl MDSAccessTokenProvider {
    fn builder() -> MDSAccessTokenProviderBuilder {
        MDSAccessTokenProviderBuilder::default()
    }

    // During ADC, if no credentials are found in the well-known location and the GOOGLE_APPLICATION_CREDENTIALS
    // environment variable is not set, we default to MDS credentials without checking if the code is really
    // running in an environment with MDS. To help users who got to this state because of lack of credentials
    // setup on their machines, we provide a detailed error message to them talking about local setup and other
    // auth mechanisms available to them.
    // If the endpoint is overridden, even if ADC was used to create the MDS credentials, we do not give a detailed
    // error message because they deliberately wanted to use an MDS.
    fn error_message(&self) -> &str {
        if self.use_adc_message() {
            MDS_NOT_FOUND_ERROR
        } else {
            "failed to fetch token"
        }
    }

    fn use_adc_message(&self) -> bool {
        self.created_by_adc && self.client.is_default_endpoint
    }
}

#[async_trait]
impl TokenProvider for MDSAccessTokenProvider {
    async fn token(&self) -> Result<Token> {
        self.client
            .access_token(self.scopes.clone())
            .send()
            .await
            .map_err(|e| CredentialsError::new(e.is_transient(), self.error_message(), e))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::credentials::QUOTA_PROJECT_KEY;
    use crate::credentials::tests::{
        find_source_error, get_headers_from_cache, get_mock_auth_retry_policy,
        get_mock_backoff_policy, get_mock_retry_throttler, get_token_from_headers,
        get_token_type_from_headers,
    };
    use crate::errors;
    use crate::errors::CredentialsError;
    use crate::mds::client::MDSTokenResponse;
    use crate::mds::{
        GCE_METADATA_HOST_ENV_VAR, MDS_DEFAULT_URI, MDS_UNIVERSE_DOMAIN_URI, METADATA_ROOT,
    };
    use crate::token::tests::MockTokenProvider;
    use crate::token_cache::TokenCache;
    use base64::{Engine, prelude::BASE64_STANDARD};
    use http::HeaderValue;
    use http::header::AUTHORIZATION;
    use httptest::cycle;
    use httptest::matchers::{all_of, contains, request, url_decoded};
    use httptest::responders::{json_encoded, status_code};
    use httptest::{Expectation, Server};
    use reqwest::StatusCode;
    use scoped_env::ScopedEnv;
    use serde_json::json;
    use serial_test::{parallel, serial};
    use std::error::Error;
    use std::time::Duration;
    use test_case::test_case;
    use tokio::time::Instant;
    use url::Url;

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

    impl Builder {
        fn maybe_iam_endpoint_override(mut self, iam_endpoint_override: Option<String>) -> Self {
            self.iam_endpoint_override = iam_endpoint_override;
            self
        }

        fn without_access_boundary(mut self) -> Self {
            self.is_access_boundary_enabled = false;
            self
        }
    }

    #[tokio::test]
    #[parallel]
    async fn test_mds_retries_on_transient_failures() -> TestResult {
        let mut server = Server::run();
        server.expect(
            Expectation::matching(request::path(format!("{MDS_DEFAULT_URI}/token")))
                .times(3)
                .respond_with(status_code(503)),
        );

        let provider = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .with_retry_policy(get_mock_auth_retry_policy(3))
            .with_backoff_policy(get_mock_backoff_policy())
            .with_retry_throttler(get_mock_retry_throttler())
            .build_token_provider();

        let err = provider.token().await.unwrap_err();
        assert!(err.is_transient(), "{err:?}");
        server.verify_and_clear();
        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn test_mds_does_not_retry_on_non_transient_failures() -> TestResult {
        let mut server = Server::run();
        server.expect(
            Expectation::matching(request::path(format!("{MDS_DEFAULT_URI}/token")))
                .times(1)
                .respond_with(status_code(401)),
        );

        let provider = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .with_retry_policy(get_mock_auth_retry_policy(1))
            .with_backoff_policy(get_mock_backoff_policy())
            .with_retry_throttler(get_mock_retry_throttler())
            .build_token_provider();

        let err = provider.token().await.unwrap_err();
        assert!(!err.is_transient());
        server.verify_and_clear();
        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn test_mds_retries_for_success() -> TestResult {
        let mut server = Server::run();
        let response = MDSTokenResponse {
            access_token: "test-access-token".to_string(),
            expires_in: Some(3600),
            token_type: "test-token-type".to_string(),
        };

        server.expect(
            Expectation::matching(request::path(format!("{MDS_DEFAULT_URI}/token")))
                .times(3)
                .respond_with(cycle![
                    status_code(503).body("try-again"),
                    status_code(503).body("try-again"),
                    status_code(200)
                        .append_header("Content-Type", "application/json")
                        .body(serde_json::to_string(&response).unwrap()),
                ]),
        );

        let provider = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .with_retry_policy(get_mock_auth_retry_policy(3))
            .with_backoff_policy(get_mock_backoff_policy())
            .with_retry_throttler(get_mock_retry_throttler())
            .build_token_provider();

        let token = provider.token().await?;
        assert_eq!(token.token, "test-access-token");

        server.verify_and_clear();
        Ok(())
    }

    #[test]
    #[parallel]
    fn validate_default_endpoint_urls() {
        let default_endpoint_address = Url::parse(&format!("{METADATA_ROOT}{MDS_DEFAULT_URI}"));
        assert!(
            default_endpoint_address.is_ok(),
            "{default_endpoint_address:?}"
        );

        let token_endpoint_address = Url::parse(&format!("{METADATA_ROOT}{MDS_DEFAULT_URI}/token"));
        assert!(token_endpoint_address.is_ok(), "{token_endpoint_address:?}");
    }

    #[tokio::test]
    #[parallel]
    async fn headers_success() -> TestResult {
        let token = Token {
            token: "test-token".to_string(),
            token_type: "Bearer".to_string(),
            expires_at: None,
            metadata: None,
        };

        let mut mock = MockTokenProvider::new();
        mock.expect_token().times(1).return_once(|| Ok(token));

        let mdsc = MDSCredentials {
            quota_project_id: None,
            token_provider: TokenCache::new(mock),
            universe_domain_override: None,
            universe_domain: OnceLock::new(),
            mds_client: MDSClient::new(None),
            backoff_policy: Arc::new(get_mock_backoff_policy()),
            retry_throttler: Arc::new(std::sync::Mutex::new(get_mock_retry_throttler())),
            retry_policy: Arc::new(get_mock_auth_retry_policy(1)),
        };

        let mut extensions = Extensions::new();
        let cached_headers = mdsc.headers(extensions.clone()).await.unwrap();
        let (headers, entity_tag) = match cached_headers {
            CacheableResource::New { entity_tag, data } => (data, entity_tag),
            CacheableResource::NotModified => unreachable!("expecting new headers"),
        };
        let token = headers.get(AUTHORIZATION).unwrap();
        assert_eq!(headers.len(), 1, "{headers:?}");
        assert_eq!(token, HeaderValue::from_static("Bearer test-token"));
        assert!(token.is_sensitive());

        extensions.insert(entity_tag);

        let cached_headers = mdsc.headers(extensions).await?;

        match cached_headers {
            CacheableResource::New { .. } => unreachable!("expecting new headers"),
            CacheableResource::NotModified => CacheableResource::<HeaderMap>::NotModified,
        };
        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn access_token_success() -> TestResult {
        let server = Server::run();
        let response = MDSTokenResponse {
            access_token: "test-access-token".to_string(),
            expires_in: Some(3600),
            token_type: "Bearer".to_string(),
        };
        server.expect(
            Expectation::matching(all_of![request::path(format!("{MDS_DEFAULT_URI}/token")),])
                .respond_with(json_encoded(response)),
        );

        let creds = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .without_access_boundary()
            .build_access_token_credentials()
            .unwrap();

        let access_token = creds.access_token().await.unwrap();
        assert_eq!(access_token.token, "test-access-token");

        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn headers_failure() {
        let mut mock = MockTokenProvider::new();
        mock.expect_token()
            .times(1)
            .return_once(|| Err(errors::non_retryable_from_str("fail")));

        let mdsc = MDSCredentials {
            quota_project_id: None,
            token_provider: TokenCache::new(mock),
            universe_domain_override: None,
            universe_domain: OnceLock::new(),
            mds_client: MDSClient::new(None),
            backoff_policy: Arc::new(get_mock_backoff_policy()),
            retry_throttler: Arc::new(std::sync::Mutex::new(get_mock_retry_throttler())),
            retry_policy: Arc::new(get_mock_auth_retry_policy(1)),
        };
        let result = mdsc.headers(Extensions::new()).await;
        assert!(result.is_err(), "{result:?}");
    }

    #[test]
    #[parallel]
    fn error_message_with_adc() {
        let provider = MDSAccessTokenProvider::builder()
            .created_by_adc(true)
            .build();

        let want = MDS_NOT_FOUND_ERROR;
        let got = provider.error_message();
        assert!(got.contains(want), "{got}, {provider:?}");
    }

    #[test_case(false, false)]
    #[test_case(false, true)]
    #[test_case(true, true)]
    fn error_message_without_adc(adc: bool, overridden: bool) {
        let endpoint = if overridden {
            Some("http://127.0.0.1")
        } else {
            None
        };
        let provider = MDSAccessTokenProvider::builder()
            .endpoint(endpoint)
            .created_by_adc(adc)
            .build();

        let not_want = MDS_NOT_FOUND_ERROR;
        let got = provider.error_message();
        assert!(!got.contains(not_want), "{got}, {provider:?}");
    }

    #[tokio::test]
    #[serial]
    async fn adc_no_mds() -> TestResult {
        let Err(err) = Builder::from_adc().build_token_provider().token().await else {
            // The environment has an MDS, skip the test.
            return Ok(());
        };

        let original_err = err.source().unwrap();
        assert!(
            original_err.to_string().contains("application-default"),
            "display={err}, debug={err:?}"
        );

        Ok(())
    }

    #[tokio::test]
    #[serial]
    async fn adc_overridden_mds() -> TestResult {
        let _e = ScopedEnv::set(GCE_METADATA_HOST_ENV_VAR, "metadata.overridden");

        let err = Builder::from_adc()
            .build_token_provider()
            .token()
            .await
            .unwrap_err();

        let _e = ScopedEnv::remove(GCE_METADATA_HOST_ENV_VAR);

        let original_err = find_source_error::<CredentialsError>(&err).unwrap();
        assert!(original_err.is_transient());
        assert!(
            !original_err.to_string().contains("application-default"),
            "display={err}, debug={err:?}"
        );
        let source = find_source_error::<reqwest::Error>(&err);
        assert!(matches!(source, Some(e) if e.status().is_none()), "{err:?}");

        Ok(())
    }

    #[tokio::test]
    #[serial]
    async fn builder_no_mds() -> TestResult {
        let Err(e) = Builder::default().build_token_provider().token().await else {
            // The environment has an MDS, skip the test.
            return Ok(());
        };

        let original_err = find_source_error::<CredentialsError>(&e).unwrap();
        assert!(
            !format!("{:?}", original_err.source()).contains("application-default"),
            "{e:?}"
        );

        Ok(())
    }

    #[tokio::test]
    #[serial]
    async fn test_gce_metadata_host_env_var() -> TestResult {
        let server = Server::run();
        let scopes = ["scope1", "scope2"];
        let response = MDSTokenResponse {
            access_token: "test-access-token".to_string(),
            expires_in: Some(3600),
            token_type: "test-token-type".to_string(),
        };
        server.expect(
            Expectation::matching(all_of![
                request::path(format!("{MDS_DEFAULT_URI}/token")),
                request::query(url_decoded(contains(("scopes", scopes.join(",")))))
            ])
            .respond_with(json_encoded(response)),
        );

        let addr = server.addr().to_string();
        let _e = ScopedEnv::set(GCE_METADATA_HOST_ENV_VAR, &addr);
        let mdsc = Builder::default()
            .with_scopes(["scope1", "scope2"])
            .without_access_boundary()
            .build()
            .unwrap();
        let headers = mdsc.headers(Extensions::new()).await.unwrap();
        let _e = ScopedEnv::remove(GCE_METADATA_HOST_ENV_VAR);

        assert_eq!(
            get_token_from_headers(headers).unwrap(),
            "test-access-token"
        );
        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn headers_success_with_quota_project() -> TestResult {
        let server = Server::run();
        let scopes = ["scope1", "scope2"];
        let response = MDSTokenResponse {
            access_token: "test-access-token".to_string(),
            expires_in: Some(3600),
            token_type: "test-token-type".to_string(),
        };
        server.expect(
            Expectation::matching(all_of![
                request::path(format!("{MDS_DEFAULT_URI}/token")),
                request::query(url_decoded(contains(("scopes", scopes.join(",")))))
            ])
            .respond_with(json_encoded(response)),
        );

        let mdsc = Builder::default()
            .with_scopes(["scope1", "scope2"])
            .with_endpoint(format!("http://{}", server.addr()))
            .with_quota_project_id("test-project")
            .without_access_boundary()
            .build()?;

        let headers = get_headers_from_cache(mdsc.headers(Extensions::new()).await.unwrap())?;
        let token = headers.get(AUTHORIZATION).unwrap();
        let quota_project = headers.get(QUOTA_PROJECT_KEY).unwrap();

        assert_eq!(headers.len(), 2, "{headers:?}");
        assert_eq!(
            token,
            HeaderValue::from_static("test-token-type test-access-token")
        );
        assert!(token.is_sensitive());
        assert_eq!(quota_project, HeaderValue::from_static("test-project"));
        assert!(!quota_project.is_sensitive());

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[parallel]
    async fn token_caching() -> TestResult {
        let server = Server::run();
        let scopes = vec!["scope1".to_string()];
        let response = MDSTokenResponse {
            access_token: "test-access-token".to_string(),
            expires_in: Some(3600),
            token_type: "test-token-type".to_string(),
        };
        server.expect(
            Expectation::matching(all_of![
                request::path(format!("{MDS_DEFAULT_URI}/token")),
                request::query(url_decoded(contains(("scopes", scopes.join(",")))))
            ])
            .times(1)
            .respond_with(json_encoded(response)),
        );

        let mdsc = Builder::default()
            .with_scopes(scopes)
            .with_endpoint(format!("http://{}", server.addr()))
            .without_access_boundary()
            .build()?;
        let headers = mdsc.headers(Extensions::new()).await?;
        assert_eq!(
            get_token_from_headers(headers).unwrap(),
            "test-access-token"
        );
        let headers = mdsc.headers(Extensions::new()).await?;
        assert_eq!(
            get_token_from_headers(headers).unwrap(),
            "test-access-token"
        );

        Ok(())
    }

    #[tokio::test(start_paused = true)]
    #[parallel]
    async fn token_provider_full() -> TestResult {
        let server = Server::run();
        let scopes = vec!["scope1".to_string()];
        let response = MDSTokenResponse {
            access_token: "test-access-token".to_string(),
            expires_in: Some(3600),
            token_type: "test-token-type".to_string(),
        };
        server.expect(
            Expectation::matching(all_of![
                request::path(format!("{MDS_DEFAULT_URI}/token")),
                request::query(url_decoded(contains(("scopes", scopes.join(",")))))
            ])
            .respond_with(json_encoded(response)),
        );

        let token = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .with_scopes(scopes)
            .build_token_provider()
            .token()
            .await?;

        let now = tokio::time::Instant::now();
        assert_eq!(token.token, "test-access-token");
        assert_eq!(token.token_type, "test-token-type");
        assert!(
            token
                .expires_at
                .is_some_and(|d| d >= now + Duration::from_secs(3600))
        );

        Ok(())
    }

    #[tokio::test(start_paused = true)]
    #[parallel]
    async fn token_provider_full_no_scopes() -> TestResult {
        let server = Server::run();
        let response = MDSTokenResponse {
            access_token: "test-access-token".to_string(),
            expires_in: Some(3600),
            token_type: "test-token-type".to_string(),
        };
        server.expect(
            Expectation::matching(request::path(format!("{MDS_DEFAULT_URI}/token")))
                .respond_with(json_encoded(response)),
        );

        let token = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .build_token_provider()
            .token()
            .await?;

        let now = Instant::now();
        assert_eq!(token.token, "test-access-token");
        assert_eq!(token.token_type, "test-token-type");
        assert!(
            token
                .expires_at
                .is_some_and(|d| d == now + Duration::from_secs(3600))
        );

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[parallel]
    async fn credential_provider_full() -> TestResult {
        let server = Server::run();
        let scopes = vec!["scope1".to_string()];
        let response = MDSTokenResponse {
            access_token: "test-access-token".to_string(),
            expires_in: None,
            token_type: "test-token-type".to_string(),
        };
        server.expect(
            Expectation::matching(all_of![
                request::path(format!("{MDS_DEFAULT_URI}/token")),
                request::query(url_decoded(contains(("scopes", scopes.join(",")))))
            ])
            .respond_with(json_encoded(response)),
        );

        let mdsc = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .with_scopes(scopes)
            .without_access_boundary()
            .build()?;
        let headers = mdsc.headers(Extensions::new()).await?;
        assert_eq!(
            get_token_from_headers(headers.clone()).unwrap(),
            "test-access-token"
        );
        assert_eq!(
            get_token_type_from_headers(headers).unwrap(),
            "test-token-type"
        );

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[parallel]
    async fn credentials_headers_retryable_error() -> TestResult {
        let server = Server::run();
        let scopes = vec!["scope1".to_string()];
        server.expect(
            Expectation::matching(all_of![
                request::path(format!("{MDS_DEFAULT_URI}/token")),
                request::query(url_decoded(contains(("scopes", scopes.join(",")))))
            ])
            .respond_with(status_code(503)),
        );

        let mdsc = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .with_scopes(scopes)
            .without_access_boundary()
            .build()?;
        let err = mdsc.headers(Extensions::new()).await.unwrap_err();
        let original_err = find_source_error::<CredentialsError>(&err).unwrap();
        assert!(original_err.is_transient());
        let source = find_source_error::<google_cloud_gax::error::Error>(&err);
        assert!(
            matches!(source, Some(e) if e.http_status_code() == Some(StatusCode::SERVICE_UNAVAILABLE.into())),
            "{err:?}"
        );

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[parallel]
    async fn credentials_headers_nonretryable_error() -> TestResult {
        let server = Server::run();
        let scopes = vec!["scope1".to_string()];
        server.expect(
            Expectation::matching(all_of![
                request::path(format!("{MDS_DEFAULT_URI}/token")),
                request::query(url_decoded(contains(("scopes", scopes.join(",")))))
            ])
            .respond_with(status_code(401)),
        );

        let mdsc = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .with_scopes(scopes)
            .without_access_boundary()
            .build()?;

        let err = mdsc.headers(Extensions::new()).await.unwrap_err();
        let original_err = find_source_error::<CredentialsError>(&err).unwrap();
        assert!(!original_err.is_transient());
        let source = find_source_error::<google_cloud_gax::error::Error>(&err);
        assert!(
            matches!(source, Some(e) if e.http_status_code() == Some(StatusCode::UNAUTHORIZED.into())),
            "{err:?}"
        );

        Ok(())
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    #[parallel]
    async fn credentials_headers_malformed_response_is_nonretryable() -> TestResult {
        let server = Server::run();
        let scopes = vec!["scope1".to_string()];
        server.expect(
            Expectation::matching(all_of![
                request::path(format!("{MDS_DEFAULT_URI}/token")),
                request::query(url_decoded(contains(("scopes", scopes.join(",")))))
            ])
            .respond_with(json_encoded("bad json")),
        );

        let mdsc = Builder::default()
            .with_endpoint(format!("http://{}", server.addr()))
            .with_scopes(scopes)
            .without_access_boundary()
            .build()?;

        let e = mdsc.headers(Extensions::new()).await.err().unwrap();
        assert!(!e.is_transient());

        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn get_default_universe_domain() -> TestResult {
        let server = Server::run();
        server.expect(
            Expectation::matching(all_of![request::path(MDS_UNIVERSE_DOMAIN_URI),])
                .respond_with(status_code(404)),
        );

        let mut mock = MockTokenProvider::new();
        mock.expect_token()
            .returning(|| Err(crate::errors::non_retryable_from_str("fail")));

        let creds = MDSCredentials {
            quota_project_id: None,
            universe_domain_override: None,
            universe_domain: OnceLock::new(),
            token_provider: TokenCache::new(mock),
            mds_client: crate::mds::client::Client::new(Some(format!("http://{}", server.addr()))),
            backoff_policy: Arc::new(get_mock_backoff_policy()),
            retry_throttler: Arc::new(std::sync::Mutex::new(get_mock_retry_throttler())),
            retry_policy: Arc::new(get_mock_auth_retry_policy(1)),
        };

        let universe_domain = creds.universe_domain().await;
        assert!(universe_domain.is_none());
        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn get_universe_domain_override() -> TestResult {
        let creds = Builder::default()
            .with_universe_domain("my-universe-domain.com")
            .without_access_boundary()
            .build()?;
        let universe_domain = creds.universe_domain().await;
        assert_eq!(universe_domain.as_deref(), Some("my-universe-domain.com"));
        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn get_universe_domain_from_mds() -> TestResult {
        let server = Server::run();
        server.expect(
            Expectation::matching(all_of![request::path(MDS_UNIVERSE_DOMAIN_URI),])
                .respond_with(status_code(200).body("my-universe-domain.com")),
        );

        let mut mock = MockTokenProvider::new();
        mock.expect_token()
            .returning(|| Err(crate::errors::non_retryable_from_str("fail")));

        let creds = MDSCredentials {
            quota_project_id: None,
            universe_domain_override: None,
            universe_domain: OnceLock::new(),
            token_provider: TokenCache::new(mock),
            mds_client: crate::mds::client::Client::new(Some(format!("http://{}", server.addr()))),
            backoff_policy: Arc::new(get_mock_backoff_policy()),
            retry_throttler: Arc::new(std::sync::Mutex::new(get_mock_retry_throttler())),
            retry_policy: Arc::new(get_mock_auth_retry_policy(1)),
        };
        let universe_domain = creds.universe_domain().await;
        assert_eq!(universe_domain.as_deref(), Some("my-universe-domain.com"));
        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn get_universe_domain_retries_on_transient_failures() -> TestResult {
        let server = Server::run();
        server.expect(
            Expectation::matching(all_of![request::path(MDS_UNIVERSE_DOMAIN_URI),])
                .times(3)
                .respond_with(cycle![
                    status_code(503).body("transient error"),
                    status_code(503).body("transient error"),
                    status_code(200).body("my-universe-domain.com"),
                ]),
        );

        let mut mock = MockTokenProvider::new();
        mock.expect_token()
            .returning(|| Err(crate::errors::non_retryable_from_str("fail")));

        let creds = MDSCredentials {
            quota_project_id: None,
            universe_domain_override: None,
            universe_domain: OnceLock::new(),
            token_provider: TokenCache::new(mock),
            mds_client: crate::mds::client::Client::new(Some(format!("http://{}", server.addr()))),
            backoff_policy: Arc::new(get_mock_backoff_policy()),
            retry_throttler: Arc::new(std::sync::Mutex::new(get_mock_retry_throttler())),
            retry_policy: Arc::new(get_mock_auth_retry_policy(3)),
        };

        let universe_domain = creds.universe_domain().await;
        assert_eq!(universe_domain.as_deref(), Some("my-universe-domain.com"));

        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn get_universe_domain_caching() -> TestResult {
        let server = Server::run();
        server.expect(
            Expectation::matching(all_of![request::path(MDS_UNIVERSE_DOMAIN_URI),])
                .times(2)
                .respond_with(cycle![
                    status_code(503).body("transient error"),
                    status_code(200).body("my-universe-domain.com"),
                ]),
        );

        let mut mock = MockTokenProvider::new();
        mock.expect_token()
            .returning(|| Err(crate::errors::non_retryable_from_str("fail")));

        let creds = MDSCredentials {
            quota_project_id: None,
            universe_domain_override: None,
            universe_domain: OnceLock::new(),
            token_provider: TokenCache::new(mock),
            mds_client: crate::mds::client::Client::new(Some(format!("http://{}", server.addr()))),
            backoff_policy: Arc::new(get_mock_backoff_policy()),
            retry_throttler: Arc::new(std::sync::Mutex::new(get_mock_retry_throttler())),
            retry_policy: Arc::new(get_mock_auth_retry_policy(1)),
        };

        let universe_domain = creds.universe_domain().await;
        assert_eq!(universe_domain, None);

        let universe_domain = creds.universe_domain().await;
        assert_eq!(universe_domain.as_deref(), Some("my-universe-domain.com"));

        let universe_domain = creds.universe_domain().await;
        assert_eq!(universe_domain.as_deref(), Some("my-universe-domain.com"));

        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn get_universe_domain_caching_permanent_error() -> TestResult {
        let server = Server::run();
        server.expect(
            Expectation::matching(all_of![request::path(MDS_UNIVERSE_DOMAIN_URI),])
                .times(1)
                .respond_with(status_code(404).body("permanent error")),
        );

        let mut mock = MockTokenProvider::new();
        mock.expect_token()
            .returning(|| Err(crate::errors::non_retryable_from_str("fail")));

        let creds = MDSCredentials {
            quota_project_id: None,
            universe_domain_override: None,
            universe_domain: OnceLock::new(),
            token_provider: TokenCache::new(mock),
            mds_client: crate::mds::client::Client::new(Some(format!("http://{}", server.addr()))),
            backoff_policy: Arc::new(get_mock_backoff_policy()),
            retry_throttler: Arc::new(std::sync::Mutex::new(get_mock_retry_throttler())),
            retry_policy: Arc::new(get_mock_auth_retry_policy(1)),
        };

        let universe_domain = creds.universe_domain().await;
        assert_eq!(universe_domain, None);

        let universe_domain = creds.universe_domain().await;
        assert_eq!(universe_domain, None);

        Ok(())
    }

    #[tokio::test]
    #[parallel]
    async fn get_mds_signer() -> TestResult {
        let server = Server::run();
        server.expect(
            Expectation::matching(all_of![request::path(format!("{MDS_DEFAULT_URI}/token")),])
                .respond_with(json_encoded(MDSTokenResponse {
                    access_token: "test-access-token".to_string(),
                    expires_in: None,
                    token_type: "Bearer".to_string(),
                })),
        );
        server.expect(
            Expectation::matching(all_of![request::path(format!("{MDS_DEFAULT_URI}/email")),])
                .respond_with(status_code(200).body("test-client-email")),
        );
        server.expect(
            Expectation::matching(all_of![
                request::method_path(
                    "POST",
                    "/v1/projects/-/serviceAccounts/test-client-email:signBlob"
                ),
                request::headers(contains(("authorization", "Bearer test-access-token"))),
            ])
            .respond_with(json_encoded(json!({
                "signedBlob": BASE64_STANDARD.encode("signed_blob"),
            }))),
        );

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

        let signer = Builder::default()
            .with_endpoint(&endpoint)
            .maybe_iam_endpoint_override(Some(endpoint))
            .without_access_boundary()
            .build_signer()?;

        let client_email = signer.client_email().await?;
        assert_eq!(client_email, "test-client-email");

        let signature = signer.sign(b"test").await?;
        assert_eq!(signature.as_ref(), b"signed_blob");

        Ok(())
    }

    #[tokio::test]
    #[parallel]
    #[cfg(google_cloud_unstable_trusted_boundaries)]
    async fn e2e_access_boundary() -> TestResult {
        use crate::credentials::tests::get_access_boundary_from_headers;
        use crate::mds::MDS_UNIVERSE_DOMAIN_URI;

        let server = Server::run();
        server.expect(
            Expectation::matching(all_of![request::path(format!("{MDS_DEFAULT_URI}/token")),])
                .respond_with(json_encoded(MDSTokenResponse {
                    access_token: "test-access-token".to_string(),
                    expires_in: None,
                    token_type: "Bearer".to_string(),
                })),
        );
        server.expect(
            Expectation::matching(all_of![request::path(format!("{MDS_DEFAULT_URI}/email")),])
                .respond_with(status_code(200).body("test-client-email")),
        );
        server.expect(
            Expectation::matching(all_of![request::path(MDS_UNIVERSE_DOMAIN_URI),])
                .respond_with(status_code(404)),
        );
        server.expect(
            Expectation::matching(all_of![
                request::method_path(
                    "GET",
                    "/v1/projects/-/serviceAccounts/test-client-email/allowedLocations"
                ),
                request::headers(contains(("authorization", "Bearer test-access-token"))),
            ])
            .respond_with(json_encoded(json!({
                "locations": ["us-central1", "us-east1"],
                "encodedLocations": "0x1234"
            }))),
        );

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

        let creds = Builder::default()
            .with_endpoint(&endpoint)
            .maybe_iam_endpoint_override(Some(endpoint))
            .build_credentials()?;

        // let the access boundary background thread update
        creds.wait_for_boundary().await;

        let headers = creds.headers(Extensions::new()).await?;
        let token = get_token_from_headers(headers.clone());
        let access_boundary = get_access_boundary_from_headers(headers);
        assert!(token.is_some(), "should have some token: {token:?}");
        assert_eq!(
            access_boundary.as_deref(),
            Some("0x1234"),
            "should be 0x1234 but found: {access_boundary:?}"
        );

        Ok(())
    }
}