Implement additional unit tests for the trust boundary

Author: nbayati

google/auth/credentials.py

@@ -25,12 +25,12 @@
 from google.auth import exceptions
 from google.auth import metrics
 from google.auth._credentials_base import _BaseCredentials
+from google.auth._default import _LOGGER
 from google.auth._refresh_worker import RefreshThreadManager
 
 DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
 NO_OP_TRUST_BOUNDARY_LOCATIONS: "typing.Tuple[str]" = ()
 NO_OP_TRUST_BOUNDARY_ENCODED_LOCATIONS = "0x0"
-TRUST_BOUNDARY_ENV_VAR = "GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED"
 
 
 class Credentials(_BaseCredentials):
@@ -310,7 +310,16 @@ def apply(self, headers, token=None):
         """Apply the token to the authentication header."""
         super().apply(headers, token)
         if self._trust_boundary is not None:
-            headers["x-allowed-locations"] = self._trust_boundary["encodedLocations"]
+            if (
+                self._trust_boundary["encodedLocations"]
+                == NO_OP_TRUST_BOUNDARY_ENCODED_LOCATIONS
+            ):
+                # STS expects an empty string if the trust boundary value is no-op.
+                headers["x-allowed-locations"] = ""
+            else:
+                headers["x-allowed-locations"] = self._trust_boundary[
+                    "encodedLocations"
+                ]
 
     def _refresh_trust_boundary(self, request):
         """Triggers a refresh of the trust boundary and updates the cache if necessary.
@@ -333,7 +342,11 @@ def _refresh_trust_boundary(self, request):
             # If the call to the lookup API failed, check if there is a trust boundary
             # already cached. If there is, do nothing. If not, then throw the error.
             if self._trust_boundary is None:
-                raise (error)
+                raise error
+            if _helpers.is_logging_enabled(_LOGGER):
+                _LOGGER.debug(
+                    "Using cached trust boundary due to refresh error: %s", error
+                )
             return
         else:
             self._trust_boundary = new_trust_boundary
@@ -353,9 +366,12 @@ def _lookup_trust_boundary(self, request):
                 retrieved.
         """
         from google.oauth2 import _client
-    
-         # Verify the trust boundary feature flag is enabled.
-        if os.getenv(TRUST_BOUNDARY_ENV_VAR, "").lower() != "true":
+
+        # Verify the trust boundary feature flag is enabled.
+        if (
+            os.getenv(environment_vars.GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED, "").lower()
+            != "true"
+        ):
             # Skip the lookup and return early if it's not explicitly enabled.
             return
 
@@ -364,6 +380,8 @@ def _lookup_trust_boundary(self, request):
             return
 
         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)
 
     @abc.abstractmethod
@@ -378,22 +396,9 @@ def _build_trust_boundary_lookup_url(self):
             str: The URL for the trust boundary lookup endpoint, or None
                  if lookup should be skipped (e.g., for non-applicable universe domains).
         """
-        raise NotImplementedError("_build_trust_boundary_lookup_url must be implemented")
-
-    @staticmethod
-    def _parse_trust_boundary(trust_boundary_string: str):
-        try:
-            trust_boundary = json.loads(trust_boundary_string)
-            if (
-                "locations" not in trust_boundary
-                or "encodedLocations" not in trust_boundary
-            ):
-                raise exceptions.MalformedError
-            return trust_boundary
-        except Exception:
-            raise exceptions.MalformedError(
-                "Cannot parse trust boundary {}".format(trust_boundary_string)
-            )
+        raise NotImplementedError(
+            "_build_trust_boundary_lookup_url must be implemented"
+        )
 
     def _has_no_op_trust_boundary(self):
         # A no-op trust boundary is indicated by encodedLocations being "0x0".
@@ -490,8 +495,7 @@ def default_scopes(self):
 
     @abc.abstractproperty
     def requires_scopes(self):
-        """True if these credentials require scopes to obtain an access token.
-        """
+        """True if these credentials require scopes to obtain an access token."""
         return False
 
     def has_scopes(self, scopes):

google/auth/environment_vars.py

@@ -82,3 +82,7 @@
 AWS_SESSION_TOKEN = "AWS_SESSION_TOKEN"
 AWS_REGION = "AWS_REGION"
 AWS_DEFAULT_REGION = "AWS_DEFAULT_REGION"
+
+GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED = "GOOGLE_AUTH_TRUST_BOUNDARY_ENABLED"
+"""Environment variable controlling whether to enable trust boundary feature.
+The default value is false. Users have to explicitly set this value to true."""

google/oauth2/_client.py

@@ -256,15 +256,17 @@ def _token_endpoint_request(
             an error.
     """
 
-    response_status_ok, response_data, retryable_error = _token_endpoint_request_no_throw(
-        request,
-        token_uri,
-        body,
-        access_token=access_token,
-        use_json=use_json,
-        can_retry=can_retry,
-        headers=headers,
-        **kwargs
+    response_status_ok, response_data, retryable_error = (
+        _token_endpoint_request_no_throw(
+            request,
+            token_uri,
+            body,
+            access_token=access_token,
+            use_json=use_json,
+            can_retry=can_retry,
+            headers=headers,
+            **kwargs
+        )
     )
     if not response_status_ok:
         _handle_error_response(response_data, retryable_error)
@@ -509,7 +511,7 @@ def refresh_grant(
 
 
 def lookup_trust_boundary(request, url, access_token):
-    """ Implements the global lookup of a credential trust boundary.
+    """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
     pools and workforce pools implementation may have trust boundaries configured.
@@ -549,7 +551,7 @@ def lookup_trust_boundary(request, url, access_token):
     """
 
     response_data = _lookup_trust_boundary_request(request, url, access_token, True)
-    #In case of no-op response, the "locations" list may or may not be present as an empty list.
+    # 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(
             "Invalid trust boundary info: {}".format(response_data)
@@ -583,8 +585,10 @@ def _lookup_trust_boundary_request(
         google.auth.exceptions.RefreshError: If the token endpoint returned
             an error.
     """
-    response_status_ok, response_data, retryable_error = _lookup_trust_boundary_request_no_throw(
-        request, url, access_token=access_token, can_retry=can_retry, **kwargs
+    response_status_ok, response_data, retryable_error = (
+        _lookup_trust_boundary_request_no_throw(
+            request, url, access_token=access_token, can_retry=can_retry, **kwargs
+        )
     )
     if not response_status_ok:
         _handle_error_response(response_data, retryable_error)

tests/oauth2/test__client.py

@@ -650,6 +650,7 @@ def test_lookup_trust_boundary():
     assert response["encodedLocations"] == "0x80080000000000"
     assert response["locations"] == ["us-central1", "us-east1"]
 
+
 def test_lookup_trust_boundary_no_op_response_without_locations():
     response_data = {"encodedLocations": "0x0"}
 
@@ -665,6 +666,7 @@ def test_lookup_trust_boundary_no_op_response_without_locations():
     assert response["encodedLocations"] == "0x0"
     assert "locations" not in response
 
+
 def test_lookup_trust_boundary_no_op_response():
     response_data = {"locations": [], "encodedLocations": "0x0"}