impl(auth): add access boundary support for SA

[alvarowolfx]

Towards #4186

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 95.05%. Comparing base (1482475) to head (7940ddf).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4657      +/-   ##
==========================================
+ Coverage   95.00%   95.05%   +0.05%     
==========================================
  Files         196      196              
  Lines        7525     7546      +21     
==========================================
+ Hits         7149     7173      +24     
+ Misses        376      373       -3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[coryan]

I assume several credential types will need one of these and will need to changes to add the new header?

Could this be done with a decorator that wraps the existing credential type?

[alvarowolfx]

I was experimenting with the decorator idea and I'm still on the fence.

Pros:

• Indeed it's more reusable and encapsulates all the access boundary logic.
• We can hide the AccessBoundary struct and only provide the decorator

Cons:

• In the decorator, we have to open the CacheableResourceHeader, mutate and return it. This makes the maybe_access_boundary method useless in the AuthHeaderBuilder and some of the recents changes to it to make it more extensible.
• IMHO the code is less clear that a given Credential type has support for Access Boundaries, we have to find the CredentialWithAccessBoundary decorator instead of a having an AccessBoundary module and explicitly adding to the headers.
• If we introduce more decorators, code might look weird with something like CredentialsWithYetAnotherField::new(CredentialsWithAnotherData::new(CredentialsWIthAccessBoundaries::new())), unless I'm missing a better pattern to use for decorators in Rust.
• The iam_endpoint_override is also used on Credential types that has support for remote Signers (via IAM API), so it's not just for Access Boundaries.
• For Service Account specifically, the change introduced some extra clone() because now the TokenProvider is not being constructed together with AccessBoundary and I can't reuse the same parsed ServiceAccountKey.

diff --git a/src/auth/src/access_boundary.rs b/src/auth/src/access_boundary.rs
index eb4f5cf63..b3c3e4bfe 100644
--- a/src/auth/src/access_boundary.rs
+++ b/src/auth/src/access_boundary.rs
@@ -12,12 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::credentials::CacheableResource;
+use crate::constants::TRUST_BOUNDARY_HEADER;
+use crate::credentials::{CacheableResource, CredentialsProvider, dynamic};
 use crate::errors::CredentialsError;
 use crate::headers_util::AuthHeadersBuilder;
 use crate::mds::client::Client as MDSClient;
 use crate::token::CachedTokenProvider;
-use http::Extensions;
+use crate::{Result, errors};
+use http::{Extensions, HeaderMap, HeaderValue};
 use reqwest::Client;
 use std::clone::Clone;
 use std::fmt::Debug;
@@ -57,9 +59,9 @@ impl BoundaryValue {
 }
 
 impl AccessBoundary {
-    pub(crate) fn new<T>(token_provider: T, url: String) -> Self
+    pub(crate) fn new<T>(token_provider: Arc<T>, url: String) -> Self
     where
-        T: CachedTokenProvider + 'static,
+        T: dynamic::CredentialsProvider + 'static,
     {
         let (tx_header, rx_header) = watch::channel(None);
 
@@ -76,7 +78,7 @@ impl AccessBoundary {
 
     pub(crate) fn new_for_mds<T>(token_provider: T, mds_client: MDSClient) -> Self
     where
-        T: CachedTokenProvider + 'static,
+        T: dynamic::CredentialsProvider + 'static,
     {
         let (tx_header, rx_header) = watch::channel(None);
 
@@ -122,29 +124,88 @@ impl AccessBoundary {
     }
 }
 
+/// A decorator for [crate::credentials::CredentialsProvider] with access boundary information.
+#[derive(Clone, Debug)]
+pub(crate) struct CredentialsWithAccessBoundary<T>
+where
+    T: dynamic::CredentialsProvider + 'static,
+{
+    credentials: Arc<T>,
+    access_boundary: Arc<AccessBoundary>,
+}
+
+impl<T> CredentialsWithAccessBoundary<T>
+where
+    T: dynamic::CredentialsProvider + 'static,
+{
+    pub(crate) fn new(credentials: T, access_boundary_url: String) -> Self {
+        let credentials = Arc::new(credentials);
+        let access_boundary = Arc::new(AccessBoundary::new(
+            credentials.clone(),
+            access_boundary_url,
+        ));
+        Self {
+            credentials,
+            access_boundary,
+        }
+    }
+}
+
+/// Decorates [Credentials] with access boundary information.
+impl<T> CredentialsProvider for CredentialsWithAccessBoundary<T>
+where
+    T: dynamic::CredentialsProvider + 'static,
+{
+    async fn headers(&self, extensions: Extensions) -> Result<CacheableResource<HeaderMap>> {
+        let cached_headers = self.credentials.headers(extensions).await?;
+        let (entity_tag, mut headers) = match cached_headers {
+            CacheableResource::New { entity_tag, data } => (entity_tag, data),
+            CacheableResource::NotModified => {
+                unreachable!("requested access boundary without a caching etag")
+            }
+        };
+
+        if let Some(value) = self.access_boundary.header_value() {
+            headers.insert(
+                TRUST_BOUNDARY_HEADER,
+                HeaderValue::from_str(&value).map_err(errors::non_retryable)?,
+            );
+        }
+
+        Ok(CacheableResource::New {
+            entity_tag,
+            data: headers,
+        })
+    }
+
+    async fn universe_domain(&self) -> Option<String> {
+        self.credentials.universe_domain().await
+    }
+}
+
 // internal trait for testability and avoid dependency on reqwest
 // which causes issues with tokio::time::advance and tokio::task::yield_now
 #[async_trait::async_trait]
 pub(crate) trait AccessBoundaryProvider: std::fmt::Debug + Send + Sync {
-    async fn fetch_access_boundary(&self) -> Result<Option<String>, CredentialsError>;
+    async fn fetch_access_boundary(&self) -> Result<Option<String>>;
 }
 
 // default implementation that uses IAM Access Boundaries API
 #[derive(Debug)]
 struct IAMAccessBoundaryProvider<T>
 where
-    T: CachedTokenProvider + 'static,
+    T: dynamic::CredentialsProvider + 'static,
 {
-    token_provider: T,
+    token_provider: Arc<T>,
     url: String,
 }
 
 #[async_trait::async_trait]
 impl<T> AccessBoundaryProvider for IAMAccessBoundaryProvider<T>
 where
-    T: CachedTokenProvider + 'static,
+    T: dynamic::CredentialsProvider + 'static,
 {
-    async fn fetch_access_boundary(&self) -> Result<Option<String>, CredentialsError> {
+    async fn fetch_access_boundary(&self) -> Result<Option<String>> {
         fetch_access_boundary(&self.token_provider, &self.url).await
     }
 }
@@ -157,15 +218,11 @@ struct AllowedLocationsResponse {
     encoded_locations: String,
 }
 
-async fn fetch_access_boundary<T>(
-    token_provider: &T,
-    url: &str,
-) -> Result<Option<String>, CredentialsError>
+async fn fetch_access_boundary<T>(token_provider: &T, url: &str) -> Result<Option<String>>
 where
-    T: CachedTokenProvider + 'static,
+    T: dynamic::CredentialsProvider + 'static,
 {
-    let token = token_provider.token(Extensions::new()).await?;
-    let headers = AuthHeadersBuilder::new(&token).build()?;
+    let headers = token_provider.headers(Extensions::new()).await?;
     let headers = match headers {
         CacheableResource::New { data, .. } => data,
         CacheableResource::NotModified => {
@@ -212,11 +269,11 @@ where
 }
 
 async fn refresh_task_mds<T>(
-    token_provider: T,
+    token_provider: Arc<T>,
     mds_client: MDSClient,
     tx_header: watch::Sender<Option<BoundaryValue>>,
 ) where
-    T: CachedTokenProvider + 'static,
+    T: dynamic::CredentialsProvider + 'static,
 {
     let mut provider = IAMAccessBoundaryProvider {
         token_provider,
diff --git a/src/auth/src/credentials/service_account.rs b/src/auth/src/credentials/service_account.rs
index 6e2d08801..530a5541b 100644
--- a/src/auth/src/credentials/service_account.rs
+++ b/src/auth/src/credentials/service_account.rs
@@ -72,7 +72,7 @@
 
 pub(crate) mod jws;
 
-use crate::access_boundary::AccessBoundary;
+use crate::access_boundary::CredentialsWithAccessBoundary;
 use crate::build_errors::Error as BuilderError;
 use crate::constants::DEFAULT_SCOPE;
 use crate::credentials::dynamic::{AccessTokenCredentialsProvider, CredentialsProvider};
@@ -297,7 +297,17 @@ impl Builder {
     ///
     /// [creating service account keys]: https://cloud.google.com/iam/docs/keys-create-delete#creating
     pub fn build(self) -> BuildResult<Credentials> {
-        Ok(self.build_access_token_credentials()?.into())
+        let iam_endpoint = self.iam_endpoint_override.clone();
+        let service_account_key =
+            serde_json::from_value::<ServiceAccountKey>(self.service_account_key.clone())
+                .map_err(BuilderError::parsing)?;
+        let client_email = service_account_key.client_email.clone();
+        let access_boundary_url = crate::access_boundary::service_account_lookup_url(
+            &client_email,
+            iam_endpoint.as_deref(),
+        );
+        let creds = self.build_access_token_credentials()?;
+        Ok(CredentialsWithAccessBoundary::new(creds, access_boundary_url).into())
     }
 
     /// Returns an [AccessTokenCredentials] instance with the configured settings.
@@ -338,24 +348,13 @@ impl Builder {
     /// [service account keys]: https://cloud.google.com/iam/docs/keys-create-delete#creating
     pub fn build_access_token_credentials(self) -> BuildResult<AccessTokenCredentials> {
         let quota_project_id = self.quota_project_id.clone();
-        let iam_endpoint = self.iam_endpoint_override.clone();
         let token_provider = self.build_token_provider()?;
-        let client_email = token_provider.service_account_key.client_email.clone();
 
         let token_provider = TokenCache::new(token_provider);
-        let access_boundary_url = crate::access_boundary::service_account_lookup_url(
-            &client_email,
-            iam_endpoint.as_deref(),
-        );
-        let access_boundary = Arc::new(AccessBoundary::new(
-            token_provider.clone(),
-            access_boundary_url,
-        ));
         Ok(AccessTokenCredentials {
             inner: Arc::new(ServiceAccountCredentials {
                 quota_project_id,
                 token_provider,
-                access_boundary,
             }),
         })
     }
@@ -486,7 +485,6 @@ where
 {
     token_provider: T,
     quota_project_id: Option<String>,
-    access_boundary: Arc<AccessBoundary>,
 }
 
 #[derive(Debug)]
@@ -598,11 +596,9 @@ where
 {
     async fn headers(&self, extensions: Extensions) -> Result<CacheableResource<HeaderMap>> {
         let token = self.token_provider.token(extensions).await?;
-        let access_boundary = self.access_boundary.header_value();
 
         AuthHeadersBuilder::new(&token)
             .maybe_quota_project_id(self.quota_project_id.as_deref())
-            .maybe_access_boundary(access_boundary.as_deref())
             .build()
     }
 }

[coryan]

This is getting too long for a review. Maybe time to create a doc.

In the decorator, we have to open the CacheableResourceHeader, mutate and return it. This makes the maybe_access_boundary method useless in the AuthHeaderBuilder and some of the recents changes to it to make it more extensible.

I don't understand... How is the calling code any different?

        CacheableResource::NotModified => {
            unreachable!("requested access boundary without a caching etag")
        }

That should be

        CacheableResource::NotModified => return Ok(CacheableResource::NotModified),

You are telling the caller that the value they have cached is still valid, no need to recompute it.