feat: Add trust boundary support for external accounts. (#1809)

* feat: Add trust boundary support for external accounts.

* Add trust boundary support to external account authorized users.

* Fix lint issues

* Implement additional unit tests for external account authorized user.

* fix formatting issue

* Add a unit test with invalid audiance

* add missing unit tests

Author: nbayati

google/auth/_constants.py

@@ -0,0 +1,5 @@
+"""Shared constants."""
+
+_SERVICE_ACCOUNT_TRUST_BOUNDARY_LOOKUP_ENDPOINT = "https://iamcredentials.{universe_domain}/v1/projects/-/serviceAccounts/{service_account_email}/allowedLocations"
+_WORKFORCE_POOL_TRUST_BOUNDARY_LOOKUP_ENDPOINT = "https://iamcredentials.{universe_domain}/v1/locations/global/workforcePools/{pool_id}/allowedLocations"
+_WORKLOAD_IDENTITY_POOL_TRUST_BOUNDARY_LOOKUP_ENDPOINT = "https://iamcredentials.{universe_domain}/v1/projects/{project_number}/locations/global/workloadIdentityPools/{pool_id}/allowedLocations"

google/auth/external_account.py

@@ -36,6 +36,7 @@
 import json
 import re
 
+from google.auth import _constants
 from google.auth import _helpers
 from google.auth import credentials
 from google.auth import exceptions
@@ -81,6 +82,7 @@ class Credentials(
     credentials.Scoped,
     credentials.CredentialsWithQuotaProject,
     credentials.CredentialsWithTokenUri,
+    credentials.CredentialsWithTrustBoundary,
     metaclass=abc.ABCMeta,
 ):
     """Base class for all external account credentials.
@@ -173,10 +175,7 @@ def __init__(
         self._scopes = scopes
         self._default_scopes = default_scopes
         self._workforce_pool_user_project = workforce_pool_user_project
-        self._trust_boundary = {
-            "locations": [],
-            "encoded_locations": "0x0",
-        }  # expose a placeholder trust boundary value.
+        self._trust_boundary = trust_boundary
 
         if self._client_id:
             self._client_auth = utils.ClientAuthentication(
@@ -242,6 +241,7 @@ def _constructor_args(self):
             "scopes": self._scopes,
             "default_scopes": self._default_scopes,
             "universe_domain": self._universe_domain,
+            "trust_boundary": self._trust_boundary,
         }
         if not self.is_workforce_pool:
             args.pop("workforce_pool_user_project")
@@ -412,8 +412,23 @@ def get_project_id(self, request):
 
         return None
 
-    @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
+        """Refreshes the access token.
+
+        For impersonated credentials, this method will refresh the underlying
+        source credentials and the impersonated credentials. For non-impersonated
+        credentials, it will refresh the access token and the trust boundary.
+        """
+        self._refresh_token(request)
+        # If we are impersonating, the trust boundary is handled by the
+        # impersonated credentials object. We need to get it from there.
+        if self._service_account_impersonation_url:
+            self._trust_boundary = self._impersonated_credentials._trust_boundary
+        else:
+            # Otherwise, refresh the trust boundary for the external account.
+            self._refresh_trust_boundary(request)
+
+    def _refresh_token(self, request):
         scopes = self._scopes if self._scopes is not None else self._default_scopes
 
         # Inject client certificate into request.
@@ -463,6 +478,40 @@ def refresh(self, request):
 
             self.expiry = now + lifetime
 
+    def _build_trust_boundary_lookup_url(self):
+        """Builds and returns the URL for the trust boundary lookup API."""
+        url = None
+        # Try to parse as a workload identity pool.
+        # Audience format: //iam.googleapis.com/projects/PROJECT_NUMBER/locations/global/workloadIdentityPools/POOL_ID/providers/PROVIDER_ID
+        workload_match = re.search(
+            r"projects/([^/]+)/locations/global/workloadIdentityPools/([^/]+)",
+            self._audience,
+        )
+        if workload_match:
+            project_number, pool_id = workload_match.groups()
+            url = _constants._WORKLOAD_IDENTITY_POOL_TRUST_BOUNDARY_LOOKUP_ENDPOINT.format(
+                universe_domain=self._universe_domain,
+                project_number=project_number,
+                pool_id=pool_id,
+            )
+        else:
+            # If that fails, try to parse as a workforce pool.
+            # Audience format: //iam.googleapis.com/locations/global/workforcePools/POOL_ID/providers/PROVIDER_ID
+            workforce_match = re.search(
+                r"locations/[^/]+/workforcePools/([^/]+)", self._audience
+            )
+            if workforce_match:
+                pool_id = workforce_match.groups()[0]
+                url = _constants._WORKFORCE_POOL_TRUST_BOUNDARY_LOOKUP_ENDPOINT.format(
+                    universe_domain=self._universe_domain, pool_id=pool_id
+                )
+
+        if url:
+            return url
+        else:
+            # If both fail, the audience format is invalid.
+            raise exceptions.InvalidValue("Invalid audience format.")
+
     def _make_copy(self):
         kwargs = self._constructor_args()
         new_cred = self.__class__(**kwargs)
@@ -489,6 +538,12 @@ def with_universe_domain(self, universe_domain):
         cred._universe_domain = universe_domain
         return cred
 
+    @_helpers.copy_docstring(credentials.CredentialsWithTrustBoundary)
+    def with_trust_boundary(self, trust_boundary):
+        cred = self._make_copy()
+        cred._trust_boundary = trust_boundary
+        return cred
+
     def _should_initialize_impersonated_credentials(self):
         return (
             self._service_account_impersonation_url is not None
@@ -537,6 +592,7 @@ def _initialize_impersonated_credentials(self):
             lifetime=self._service_account_impersonation_options.get(
                 "token_lifetime_seconds"
             ),
+            trust_boundary=self._trust_boundary,
         )
 
     def _create_default_metrics_options(self):
@@ -623,6 +679,7 @@ def from_info(cls, info, **kwargs):
             universe_domain=info.get(
                 "universe_domain", credentials.DEFAULT_UNIVERSE_DOMAIN
             ),
+            trust_boundary=info.get("trust_boundary"),
             **kwargs
         )
 

google/auth/external_account_authorized_user.py

@@ -36,7 +36,9 @@
 import datetime
 import io
 import json
+import re
 
+from google.auth import _constants
 from google.auth import _helpers
 from google.auth import credentials
 from google.auth import exceptions
@@ -50,6 +52,7 @@ class Credentials(
     credentials.CredentialsWithQuotaProject,
     credentials.ReadOnlyScoped,
     credentials.CredentialsWithTokenUri,
+    credentials.CredentialsWithTrustBoundary,
 ):
     """Credentials for External Account Authorized Users.
 
@@ -83,6 +86,7 @@ def __init__(
         scopes=None,
         quota_project_id=None,
         universe_domain=credentials.DEFAULT_UNIVERSE_DOMAIN,
+        trust_boundary=None,
     ):
         """Instantiates a external account authorized user credentials object.
 
@@ -108,6 +112,7 @@ def __init__(
             create the credentials.
         universe_domain (Optional[str]): The universe domain. The default value
             is googleapis.com.
+        trust_boundary (Mapping[str,str]): A credential trust boundary.
 
         Returns:
             google.auth.external_account_authorized_user.Credentials: The
@@ -118,7 +123,7 @@ def __init__(
         self.token = token
         self.expiry = expiry
         self._audience = audience
-        self._refresh_token = refresh_token
+        self._refresh_token_val = refresh_token
         self._token_url = token_url
         self._token_info_url = token_info_url
         self._client_id = client_id
@@ -128,6 +133,7 @@ def __init__(
         self._scopes = scopes
         self._universe_domain = universe_domain or credentials.DEFAULT_UNIVERSE_DOMAIN
         self._cred_file_path = None
+        self._trust_boundary = trust_boundary
 
         if not self.valid and not self.can_refresh:
             raise exceptions.InvalidOperation(
@@ -164,7 +170,7 @@ def info(self):
     def constructor_args(self):
         return {
             "audience": self._audience,
-            "refresh_token": self._refresh_token,
+            "refresh_token": self._refresh_token_val,
             "token_url": self._token_url,
             "token_info_url": self._token_info_url,
             "client_id": self._client_id,
@@ -175,6 +181,7 @@ def constructor_args(self):
             "scopes": self._scopes,
             "quota_project_id": self._quota_project_id,
             "universe_domain": self._universe_domain,
+            "trust_boundary": self._trust_boundary,
         }
 
     @property
@@ -184,7 +191,7 @@ def scopes(self):
 
     @property
     def requires_scopes(self):
-        """ False: OAuth 2.0 credentials have their scopes set when
+        """False: OAuth 2.0 credentials have their scopes set when
         the initial token is requested and can not be changed."""
         return False
 
@@ -201,13 +208,13 @@ def client_secret(self):
     @property
     def audience(self):
         """Optional[str]: The STS audience which contains the resource name for the
-            workforce pool and the provider identifier in that pool."""
+        workforce pool and the provider identifier in that pool."""
         return self._audience
 
     @property
     def refresh_token(self):
         """Optional[str]: The OAuth 2.0 refresh token."""
-        return self._refresh_token
+        return self._refresh_token_val
 
     @property
     def token_url(self):
@@ -226,13 +233,18 @@ def revoke_url(self):
 
     @property
     def is_user(self):
-        """ True: This credential always represents a user."""
+        """True: This credential always represents a user."""
         return True
 
     @property
     def can_refresh(self):
         return all(
-            (self._refresh_token, self._token_url, self._client_id, self._client_secret)
+            (
+                self._refresh_token_val,
+                self._token_url,
+                self._client_id,
+                self._client_secret,
+            )
         )
 
     def get_project_id(self, request=None):
@@ -266,7 +278,7 @@ def to_json(self, strip=None):
         strip = strip if strip else []
         return json.dumps({k: v for (k, v) in self.info.items() if k not in strip})
 
-    def refresh(self, request):
+    def _refresh_token(self, request):
         """Refreshes the access token.
 
         Args:
@@ -285,18 +297,29 @@ def refresh(self, request):
             )
 
         now = _helpers.utcnow()
-        response_data = self._make_sts_request(request)
+        response_data = self._sts_client.refresh_token(request, self._refresh_token_val)
 
         self.token = response_data.get("access_token")
 
         lifetime = datetime.timedelta(seconds=response_data.get("expires_in"))
         self.expiry = now + lifetime
 
         if "refresh_token" in response_data:
-            self._refresh_token = response_data["refresh_token"]
+            self._refresh_token_val = response_data["refresh_token"]
+
+    def _build_trust_boundary_lookup_url(self):
+        """Builds and returns the URL for the trust boundary lookup API."""
+        # Audience format: //iam.googleapis.com/locations/global/workforcePools/POOL_ID/providers/PROVIDER_ID
+        match = re.search(r"locations/[^/]+/workforcePools/([^/]+)", self._audience)
+
+        if not match:
+            raise exceptions.InvalidValue("Invalid workforce pool audience format.")
+
+        pool_id = match.groups()[0]
 
-    def _make_sts_request(self, request):
-        return self._sts_client.refresh_token(request, self._refresh_token)
+        return _constants._WORKFORCE_POOL_TRUST_BOUNDARY_LOOKUP_ENDPOINT.format(
+            universe_domain=self._universe_domain, pool_id=pool_id
+        )
 
     @_helpers.copy_docstring(credentials.Credentials)
     def get_cred_info(self):
@@ -331,6 +354,12 @@ def with_universe_domain(self, universe_domain):
         cred._universe_domain = universe_domain
         return cred
 
+    @_helpers.copy_docstring(credentials.CredentialsWithTrustBoundary)
+    def with_trust_boundary(self, trust_boundary):
+        cred = self._make_copy()
+        cred._trust_boundary = trust_boundary
+        return cred
+
     @classmethod
     def from_info(cls, info, **kwargs):
         """Creates a Credentials instance from parsed external account info.
@@ -375,6 +404,7 @@ def from_info(cls, info, **kwargs):
             universe_domain=info.get(
                 "universe_domain", credentials.DEFAULT_UNIVERSE_DOMAIN
             ),
+            trust_boundary=info.get("trust_boundary"),
             **kwargs
         )
 

google/auth/impersonated_credentials.py

@@ -286,7 +286,7 @@ def _refresh_token(self, request):
             self._source_credentials.token_state == credentials.TokenState.STALE
             or self._source_credentials.token_state == credentials.TokenState.INVALID
         ):
-            self._source_credentials.refresh(request)
+            self._source_credentials._refresh_token(request)
 
         body = {
             "delegates": self._delegates,
@@ -526,13 +526,15 @@ def from_impersonated_service_account_info(cls, info, scopes=None):
         target_principal = impersonation_url[start_index + 1 : end_index]
         delegates = info.get("delegates")
         quota_project_id = info.get("quota_project_id")
+        trust_boundary = info.get("trust_boundary")
 
         return cls(
             source_credentials,
             target_principal,
             scopes,
             delegates,
             quota_project_id=quota_project_id,
+            trust_boundary=trust_boundary,
         )
 
 

google/oauth2/service_account.py

@@ -73,6 +73,7 @@
 import copy
 import datetime
 
+from google.auth import _constants
 from google.auth import _helpers
 from google.auth import _service_account_info
 from google.auth import credentials
@@ -84,9 +85,6 @@
 
 _DEFAULT_TOKEN_LIFETIME_SECS = 3600  # 1 hour in seconds
 _GOOGLE_OAUTH2_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token"
-_TRUST_BOUNDARY_LOOKUP_ENDPOINT = (
-    "https://iamcredentials.{}/v1/projects/-/serviceAccounts/{}/allowedLocations"
-)
 
 
 class Credentials(
@@ -520,8 +518,9 @@ def _build_trust_boundary_lookup_url(self):
             raise ValueError(
                 "Service account email is required to build the trust boundary lookup URL."
             )
-        return _TRUST_BOUNDARY_LOOKUP_ENDPOINT.format(
-            self._universe_domain, self._service_account_email
+        return _constants._SERVICE_ACCOUNT_TRUST_BOUNDARY_LOOKUP_ENDPOINT.format(
+            universe_domain=self._universe_domain,
+            service_account_email=self._service_account_email,
         )
 
     @_helpers.copy_docstring(credentials.Signing)