feat: adding universe domain support for downscroped credentials

google/auth/credentials.py

@@ -24,6 +24,8 @@
 from google.auth import metrics
 from google.auth._refresh_worker import RefreshThreadManager
 
+DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
+
 
 class Credentials(metaclass=abc.ABCMeta):
     """Base class for all credentials.

google/auth/downscoped.py

@@ -63,7 +63,7 @@
 # The token exchange requested_token_type. This is always an access_token.
 _STS_REQUESTED_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"
 # The STS token URL used to exchanged a short lived access token for a downscoped one.
-_STS_TOKEN_URL = "https://sts.googleapis.com/v1/token"
+_STS_TOKEN_URL_PATTERN = "https://sts.{}/v1/token"
 # The subject token type to use when exchanging a short lived access token for a
 # downscoped token.
 _STS_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"
@@ -437,7 +437,11 @@ class Credentials(credentials.CredentialsWithQuotaProject):
     """
 
     def __init__(
-        self, source_credentials, credential_access_boundary, quota_project_id=None
+        self,
+        source_credentials,
+        credential_access_boundary,
+        quota_project_id=None,
+        universe_domain=credentials.DEFAULT_UNIVERSE_DOMAIN,
     ):
         """Instantiates a downscoped credentials object using the provided source
         credentials and credential access boundary rules.
@@ -467,7 +471,10 @@ def __init__(
         self._source_credentials = source_credentials
         self._credential_access_boundary = credential_access_boundary
         self._quota_project_id = quota_project_id
-        self._sts_client = sts.Client(_STS_TOKEN_URL)
+        self._universe_domain = universe_domain or credentials.DEFAULT_UNIVERSE_DOMAIN
+        self._sts_client = sts.Client(
+            _STS_TOKEN_URL_PATTERN.format(self.universe_domain)
+        )
 
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):

tests/test_downscoped.py

@@ -25,6 +25,7 @@
 from google.auth import downscoped
 from google.auth import exceptions
 from google.auth import transport
+from google.auth.credentials import DEFAULT_UNIVERSE_DOMAIN
 from google.auth.credentials import TokenState
 
 
@@ -447,7 +448,11 @@ def test_to_json(self):
 
 class TestCredentials(object):
     @staticmethod
-    def make_credentials(source_credentials=SourceCredentials(), quota_project_id=None):
+    def make_credentials(
+        source_credentials=SourceCredentials(),
+        quota_project_id=None,
+        universe_domain=None,
+    ):
         availability_condition = make_availability_condition(
             EXPRESSION, TITLE, DESCRIPTION
         )
@@ -458,7 +463,10 @@ def make_credentials(source_credentials=SourceCredentials(), quota_project_id=No
         credential_access_boundary = make_credential_access_boundary(rules)
 
         return downscoped.Credentials(
-            source_credentials, credential_access_boundary, quota_project_id
+            source_credentials,
+            credential_access_boundary,
+            quota_project_id,
+            universe_domain,
         )
 
     @staticmethod
@@ -496,6 +504,20 @@ def test_default_state(self):
         assert not credentials.expired
         # No quota project ID set.
         assert not credentials.quota_project_id
+        assert credentials.universe_domain == DEFAULT_UNIVERSE_DOMAIN
+
+    def test_create_with_customized_universe_domain(self):
+        test_universe_domain = "foo.com"
+        credentials = self.make_credentials(universe_domain=test_universe_domain)
+        # No token acquired yet.
+        assert not credentials.token
+        assert not credentials.valid
+        # Expiration hasn't been set yet.
+        assert not credentials.expiry
+        assert not credentials.expired
+        # No quota project ID set.
+        assert not credentials.quota_project_id
+        assert credentials.universe_domain == test_universe_domain
 
     def test_with_quota_project(self):
         credentials = self.make_credentials()