Add x-allowed-location header to all IAM requests.

Author: nbayati

google/auth/credentials.py

@@ -306,17 +306,19 @@ def with_trust_boundary(self, trust_boundary):
         """
         raise NotImplementedError("This credential does not support trust boundaries.")
 
-    def apply(self, headers, token=None):
-        """Apply the token to the authentication header."""
-        super().apply(headers, token)
+    def _get_trust_boundary_header(self):
         if self._trust_boundary is not None:
             if self._has_no_op_trust_boundary():
                 # STS expects an empty string if the trust boundary value is no-op.
-                headers["x-allowed-locations"] = ""
+                return {"x-allowed-locations": ""}
             else:
-                headers["x-allowed-locations"] = self._trust_boundary[
-                    "encodedLocations"
-                ]
+                return {"x-allowed-locations": self._trust_boundary["encodedLocations"]}
+        return {}
+
+    def apply(self, headers, token=None):
+        """Apply the token to the authentication header."""
+        super().apply(headers, token)
+        headers.update(self._get_trust_boundary_header())
 
     def _refresh_trust_boundary(self, request):
         """Triggers a refresh of the trust boundary and updates the cache if necessary.
@@ -378,7 +380,10 @@ def _lookup_trust_boundary(self, request):
         url = self._build_trust_boundary_lookup_url()
         if not url:
             raise exceptions.InvalidValue("Failed to build trust boundary lookup URL.")
-        return _client.lookup_trust_boundary(request, url, self.token)
+
+        headers = {}
+        self.apply(headers)
+        return _client.lookup_trust_boundary(request, url, headers=headers)
 
     @abc.abstractmethod
     def _build_trust_boundary_lookup_url(self):

google/auth/impersonated_credentials.py

@@ -607,6 +607,7 @@ def refresh(self, request):
             "Content-Type": "application/json",
             metrics.API_CLIENT_HEADER: metrics.token_request_id_token_impersonate(),
         }
+        headers.update(self._target_credentials._get_trust_boundary_header())
 
         authed_session = AuthorizedSession(
             self._target_credentials._source_credentials, auth_request=request

google/oauth2/_client.py

@@ -327,6 +327,7 @@ def call_iam_generate_id_token_endpoint(
     signer_email,
     audience,
     access_token,
+    headers=None,
     universe_domain=credentials.DEFAULT_UNIVERSE_DOMAIN,
 ):
     """Call iam.generateIdToken endpoint to get ID token.
@@ -339,6 +340,9 @@ def call_iam_generate_id_token_endpoint(
             generateIdToken endpoint.
         audience (str): The audience for the ID token.
         access_token (str): The access token used to call the IAM endpoint.
+        headers (Optional[Mapping[str, str]]): The headers for the request.
+        universe_domain (str): The universe domain for the request. The
+            default is ``googleapis.com``.
 
     Returns:
         Tuple[str, datetime]: The ID token and expiration.
@@ -353,6 +357,7 @@ def call_iam_generate_id_token_endpoint(
         body,
         access_token=access_token,
         use_json=True,
+        headers=headers,
     )
 
     try:
@@ -510,7 +515,7 @@ def refresh_grant(
     return _handle_refresh_grant_response(response_data, refresh_token)
 
 
-def lookup_trust_boundary(request, url, access_token):
+def lookup_trust_boundary(request, url, headers=None):
     """Implements the global lookup of a credential trust boundary.
     For the lookup, we send a request to the global lookup endpoint and then
     parse the response. Service account credentials, workload identity
@@ -519,15 +524,7 @@ def lookup_trust_boundary(request, url, access_token):
         request (google.auth.transport.Request): A callable used to make
             HTTP requests.
         url (str): The trust boundary lookup url.
-        access_token (Optional(str)): The access token needed to make the request
         headers (Optional[Mapping[str, str]]): The headers for the request.
-        kwargs: Additional arguments passed on to the request method. The
-            kwargs will be passed to `requests.request` method, see:
-            https://docs.python-requests.org/en/latest/api/#requests.request.
-            For example, you can use `cert=("cert_pem_path", "key_pem_path")`
-            to set up client side SSL certificate, and use
-            `verify="ca_bundle_path"` to set up the CA certificates for sever
-            side SSL certificate verification.
     Returns:
         Mapping[str,list|str]: A dictionary containing
             "locations" as a list of allowed locations as strings and
@@ -550,7 +547,7 @@ def lookup_trust_boundary(request, url, access_token):
         exceptions.MalformedError: If the response is not in a valid format.
     """
 
-    response_data = _lookup_trust_boundary_request(request, url, access_token, True)
+    response_data = _lookup_trust_boundary_request(request, url, headers=headers)
     # In case of no-op response, the "locations" list may or may not be present as an empty list.
     if "encodedLocations" not in response_data:
         raise exceptions.MalformedError(
@@ -560,16 +557,16 @@ def lookup_trust_boundary(request, url, access_token):
 
 
 def _lookup_trust_boundary_request(
-    request, url, access_token, can_retry=True, **kwargs
+    request, url, can_retry=True, headers=None, **kwargs
 ):
     """Makes a request to the trust boundary lookup endpoint.
 
     Args:
         request (google.auth.transport.Request): A callable used to make
             HTTP requests.
         url (str): The trust boundary lookup url.
-        access_token (Optional(str)): The access token needed to make the request
         can_retry (bool): Enable or disable request retry behavior. Defaults to true.
+        headers (Optional[Mapping[str, str]]): The headers for the request.
         kwargs: Additional arguments passed on to the request method. The
             kwargs will be passed to `requests.request` method, see:
             https://docs.python-requests.org/en/latest/api/#requests.request.
@@ -587,7 +584,7 @@ def _lookup_trust_boundary_request(
     """
     response_status_ok, response_data, retryable_error = (
         _lookup_trust_boundary_request_no_throw(
-            request, url, access_token=access_token, can_retry=can_retry, **kwargs
+            request, url, can_retry, headers, **kwargs
         )
     )
     if not response_status_ok:
@@ -596,7 +593,7 @@ def _lookup_trust_boundary_request(
 
 
 def _lookup_trust_boundary_request_no_throw(
-    request, url, access_token=None, can_retry=True, **kwargs
+    request, url, can_retry=True, headers=None, **kwargs
 ):
     """Makes a request to the trust boundary lookup endpoint. This
         function doesn't throw on response errors.
@@ -605,8 +602,8 @@ def _lookup_trust_boundary_request_no_throw(
         request (google.auth.transport.Request): A callable used to make
             HTTP requests.
         url (str): The trust boundary lookup url.
-        access_token (Optional(str)): The access token needed to make the request
         can_retry (bool): Enable or disable request retry behavior. Defaults to true.
+        headers (Optional[Mapping[str, str]]): The headers for the request.
         kwargs: Additional arguments passed on to the request method. The
             kwargs will be passed to `requests.request` method, see:
             https://docs.python-requests.org/en/latest/api/#requests.request.
@@ -622,14 +619,12 @@ def _lookup_trust_boundary_request_no_throw(
           is retryable.
     """
 
-    headers_to_use = {"Authorization": "Bearer {}".format(access_token)}
-
     response_data = {}
     retryable_error = False
 
     retries = _exponential_backoff.ExponentialBackoff()
     for _ in retries:
-        response = request(method="GET", url=url, headers=headers_to_use, **kwargs)
+        response = request(method="GET", url=url, headers=headers, **kwargs)
         response_body = (
             response.data.decode("utf-8")
             if hasattr(response.data, "decode")

google/oauth2/service_account.py

@@ -554,6 +554,7 @@ class IDTokenCredentials(
     credentials.Signing,
     credentials.CredentialsWithQuotaProject,
     credentials.CredentialsWithTokenUri,
+    credentials.CredentialsWithTrustBoundary,
 ):
     """Open ID Connect ID Token-based service account credentials.
 
@@ -608,6 +609,7 @@ def __init__(
         additional_claims=None,
         quota_project_id=None,
         universe_domain=credentials.DEFAULT_UNIVERSE_DOMAIN,
+        trust_boundary=None,
     ):
         """
         Args:
@@ -625,6 +627,8 @@ def __init__(
                 token endponint is used for token refresh. Note that
                 iam.serviceAccountTokenCreator role is required to use the IAM
                 endpoint.
+            trust_boundary (Mapping[str,str]): A credential trust boundary.
+
         .. note:: Typically one of the helper constructors
             :meth:`from_service_account_file` or
             :meth:`from_service_account_info` are used instead of calling the
@@ -637,6 +641,7 @@ def __init__(
         self._target_audience = target_audience
         self._quota_project_id = quota_project_id
         self._use_iam_endpoint = False
+        self._trust_boundary = trust_boundary
 
         if not universe_domain:
             self._universe_domain = credentials.DEFAULT_UNIVERSE_DOMAIN
@@ -674,6 +679,8 @@ def _from_signer_and_info(cls, signer, info, **kwargs):
         kwargs.setdefault("token_uri", info["token_uri"])
         if "universe_domain" in info:
             kwargs["universe_domain"] = info["universe_domain"]
+        if "trust_boundary" in info:
+            kwargs["trust_boundary"] = info["trust_boundary"]
         return cls(signer, **kwargs)
 
     @classmethod
@@ -723,6 +730,7 @@ def _make_copy(self):
             additional_claims=self._additional_claims.copy(),
             quota_project_id=self.quota_project_id,
             universe_domain=self._universe_domain,
+            trust_boundary=self._trust_boundary,
         )
         # _use_iam_endpoint is not exposed in the constructor
         cred._use_iam_endpoint = self._use_iam_endpoint
@@ -784,6 +792,22 @@ def with_token_uri(self, token_uri):
         cred._token_uri = token_uri
         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 _build_trust_boundary_lookup_url(self):
+        """Builds and returns the URL for the trust boundary lookup API.
+
+        This is not used by IDTokenCredentials as it does not perform trust
+        boundary lookups. It is defined here to satisfy the abstract base class.
+        """
+        raise NotImplementedError(
+            "_build_trust_boundary_lookup_url is not implemented for IDTokenCredentials."
+        )
+
     def _make_authorization_grant_assertion(self):
         """Create the OAuth 2.0 assertion.
 
@@ -840,13 +864,17 @@ def _refresh_with_iam_endpoint(self, request):
             additional_claims={"scope": "https://www.googleapis.com/auth/iam"},
         )
         jwt_credentials.refresh(request)
+
+        headers = self._get_trust_boundary_header()
+
         self.token, self.expiry = _client.call_iam_generate_id_token_endpoint(
             request,
             self._iam_id_token_endpoint,
             self.signer_email,
             self._target_audience,
             jwt_credentials.token.decode(),
-            self._universe_domain,
+            headers=headers,
+            universe_domain=self._universe_domain,
         )
 
     @_helpers.copy_docstring(credentials.Credentials)