feat: universe domain support for service account (#1286)

* feat: universe domain support for service account

* update

* update

* update

* update token

Author: arithmetic1728

google/oauth2/service_account.py

@@ -76,10 +76,12 @@
 from google.auth import _helpers
 from google.auth import _service_account_info
 from google.auth import credentials
+from google.auth import exceptions
 from google.auth import jwt
 from google.oauth2 import _client
 
 _DEFAULT_TOKEN_LIFETIME_SECS = 3600  # 1 hour in seconds
+_DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
 _GOOGLE_OAUTH2_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token"
 
 
@@ -136,6 +138,7 @@ def __init__(
         quota_project_id=None,
         additional_claims=None,
         always_use_jwt_access=False,
+        universe_domain=_DEFAULT_UNIVERSE_DOMAIN,
     ):
         """
         Args:
@@ -156,6 +159,9 @@ def __init__(
                 the JWT assertion used in the authorization grant.
             always_use_jwt_access (Optional[bool]): Whether self signed JWT should
                 be always used.
+            universe_domain (str): The universe domain. The default
+                universe domain is googleapis.com. For default value self
+                signed jwt is used for token refresh.
 
         .. note:: Typically one of the helper constructors
             :meth:`from_service_account_file` or
@@ -173,6 +179,13 @@ def __init__(
         self._quota_project_id = quota_project_id
         self._token_uri = token_uri
         self._always_use_jwt_access = always_use_jwt_access
+        if not universe_domain:
+            self._universe_domain = _DEFAULT_UNIVERSE_DOMAIN
+        else:
+            self._universe_domain = universe_domain
+
+        if universe_domain != _DEFAULT_UNIVERSE_DOMAIN:
+            self._always_use_jwt_access = True
 
         self._jwt_credentials = None
 
@@ -202,6 +215,7 @@ def _from_signer_and_info(cls, signer, info, **kwargs):
             service_account_email=info["client_email"],
             token_uri=info["token_uri"],
             project_id=info.get("project_id"),
+            universe_domain=info.get("universe_domain", _DEFAULT_UNIVERSE_DOMAIN),
             **kwargs
         )
 
@@ -262,20 +276,28 @@ def requires_scopes(self):
         """
         return True if not self._scopes else False
 
-    @_helpers.copy_docstring(credentials.Scoped)
-    def with_scopes(self, scopes, default_scopes=None):
-        return self.__class__(
+    def _make_copy(self):
+        cred = self.__class__(
             self._signer,
             service_account_email=self._service_account_email,
-            scopes=scopes,
-            default_scopes=default_scopes,
+            scopes=copy.copy(self._scopes),
+            default_scopes=copy.copy(self._default_scopes),
             token_uri=self._token_uri,
             subject=self._subject,
             project_id=self._project_id,
             quota_project_id=self._quota_project_id,
             additional_claims=self._additional_claims.copy(),
             always_use_jwt_access=self._always_use_jwt_access,
+            universe_domain=self._universe_domain,
         )
+        return cred
+
+    @_helpers.copy_docstring(credentials.Scoped)
+    def with_scopes(self, scopes, default_scopes=None):
+        cred = self._make_copy()
+        cred._scopes = scopes
+        cred._default_scopes = default_scopes
+        return cred
 
     def with_always_use_jwt_access(self, always_use_jwt_access):
         """Create a copy of these credentials with the specified always_use_jwt_access value.
@@ -286,19 +308,20 @@ def with_always_use_jwt_access(self, always_use_jwt_access):
         Returns:
             google.auth.service_account.Credentials: A new credentials
                 instance.
+        Raises:
+            google.auth.exceptions.InvalidValue: If the universe domain is not
+                default and always_use_jwt_access is False.
         """
-        return self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            scopes=self._scopes,
-            default_scopes=self._default_scopes,
-            token_uri=self._token_uri,
-            subject=self._subject,
-            project_id=self._project_id,
-            quota_project_id=self._quota_project_id,
-            additional_claims=self._additional_claims.copy(),
-            always_use_jwt_access=always_use_jwt_access,
-        )
+        cred = self._make_copy()
+        if (
+            cred._universe_domain != _DEFAULT_UNIVERSE_DOMAIN
+            and not always_use_jwt_access
+        ):
+            raise exceptions.InvalidValue(
+                "always_use_jwt_access should be True for non-default universe domain"
+            )
+        cred._always_use_jwt_access = always_use_jwt_access
+        return cred
 
     def with_subject(self, subject):
         """Create a copy of these credentials with the specified subject.
@@ -310,18 +333,9 @@ def with_subject(self, subject):
             google.auth.service_account.Credentials: A new credentials
                 instance.
         """
-        return self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            scopes=self._scopes,
-            default_scopes=self._default_scopes,
-            token_uri=self._token_uri,
-            subject=subject,
-            project_id=self._project_id,
-            quota_project_id=self._quota_project_id,
-            additional_claims=self._additional_claims.copy(),
-            always_use_jwt_access=self._always_use_jwt_access,
-        )
+        cred = self._make_copy()
+        cred._subject = subject
+        return cred
 
     def with_claims(self, additional_claims):
         """Returns a copy of these credentials with modified claims.
@@ -337,51 +351,21 @@ def with_claims(self, additional_claims):
         """
         new_additional_claims = copy.deepcopy(self._additional_claims)
         new_additional_claims.update(additional_claims or {})
-
-        return self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            scopes=self._scopes,
-            default_scopes=self._default_scopes,
-            token_uri=self._token_uri,
-            subject=self._subject,
-            project_id=self._project_id,
-            quota_project_id=self._quota_project_id,
-            additional_claims=new_additional_claims,
-            always_use_jwt_access=self._always_use_jwt_access,
-        )
+        cred = self._make_copy()
+        cred._additional_claims = new_additional_claims
+        return cred
 
     @_helpers.copy_docstring(credentials.CredentialsWithQuotaProject)
     def with_quota_project(self, quota_project_id):
-
-        return self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            default_scopes=self._default_scopes,
-            scopes=self._scopes,
-            token_uri=self._token_uri,
-            subject=self._subject,
-            project_id=self._project_id,
-            quota_project_id=quota_project_id,
-            additional_claims=self._additional_claims.copy(),
-            always_use_jwt_access=self._always_use_jwt_access,
-        )
+        cred = self._make_copy()
+        cred._quota_project_id = quota_project_id
+        return cred
 
     @_helpers.copy_docstring(credentials.CredentialsWithTokenUri)
     def with_token_uri(self, token_uri):
-
-        return self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            default_scopes=self._default_scopes,
-            scopes=self._scopes,
-            token_uri=token_uri,
-            subject=self._subject,
-            project_id=self._project_id,
-            quota_project_id=self._quota_project_id,
-            additional_claims=self._additional_claims.copy(),
-            always_use_jwt_access=self._always_use_jwt_access,
-        )
+        cred = self._make_copy()
+        cred._token_uri = token_uri
+        return cred
 
     def _make_authorization_grant_assertion(self):
         """Create the OAuth 2.0 assertion.
@@ -418,6 +402,18 @@ def _make_authorization_grant_assertion(self):
 
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
+        if (
+            self._universe_domain != _DEFAULT_UNIVERSE_DOMAIN
+            and not self._jwt_credentials
+        ):
+            raise exceptions.RefreshError(
+                "self._jwt_credentials is missing for non-default universe domain"
+            )
+        if self._universe_domain != _DEFAULT_UNIVERSE_DOMAIN and self._subject:
+            raise exceptions.RefreshError(
+                "domain wide delegation is not supported for non-default universe domain"
+            )
+
         # Since domain wide delegation doesn't work with self signed JWT. If
         # subject exists, then we should not use self signed JWT.
         if self._subject is None and self._jwt_credentials is not None:
@@ -544,6 +540,7 @@ def __init__(
         target_audience,
         additional_claims=None,
         quota_project_id=None,
+        universe_domain=_DEFAULT_UNIVERSE_DOMAIN,
     ):
         """
         Args:
@@ -556,6 +553,11 @@ def __init__(
             additional_claims (Mapping[str, str]): Any additional claims for
                 the JWT assertion used in the authorization grant.
             quota_project_id (Optional[str]): The project ID used for quota and billing.
+            universe_domain (str): The universe domain. The default
+                universe domain is googleapis.com. For default value IAM ID
+                token endponint is used for token refresh. Note that
+                iam.serviceAccountTokenCreator role is required to use the IAM
+                endpoint.
         .. note:: Typically one of the helper constructors
             :meth:`from_service_account_file` or
             :meth:`from_service_account_info` are used instead of calling the
@@ -569,6 +571,14 @@ def __init__(
         self._quota_project_id = quota_project_id
         self._use_iam_endpoint = False
 
+        if not universe_domain:
+            self._universe_domain = _DEFAULT_UNIVERSE_DOMAIN
+        else:
+            self._universe_domain = universe_domain
+
+        if universe_domain != _DEFAULT_UNIVERSE_DOMAIN:
+            self._use_iam_endpoint = True
+
         if additional_claims is not None:
             self._additional_claims = additional_claims
         else:
@@ -592,6 +602,8 @@ def _from_signer_and_info(cls, signer, info, **kwargs):
         """
         kwargs.setdefault("service_account_email", info["client_email"])
         kwargs.setdefault("token_uri", info["token_uri"])
+        if "universe_domain" in info:
+            kwargs["universe_domain"] = info["universe_domain"]
         return cls(signer, **kwargs)
 
     @classmethod
@@ -632,6 +644,20 @@ def from_service_account_file(cls, filename, **kwargs):
         )
         return cls._from_signer_and_info(signer, info, **kwargs)
 
+    def _make_copy(self):
+        cred = self.__class__(
+            self._signer,
+            service_account_email=self._service_account_email,
+            token_uri=self._token_uri,
+            target_audience=self._target_audience,
+            additional_claims=self._additional_claims.copy(),
+            quota_project_id=self.quota_project_id,
+            universe_domain=self._universe_domain,
+        )
+        # _use_iam_endpoint is not exposed in the constructor
+        cred._use_iam_endpoint = self._use_iam_endpoint
+        return cred
+
     def with_target_audience(self, target_audience):
         """Create a copy of these credentials with the specified target
         audience.
@@ -644,14 +670,9 @@ def with_target_audience(self, target_audience):
             google.auth.service_account.IDTokenCredentials: A new credentials
                 instance.
         """
-        return self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            token_uri=self._token_uri,
-            target_audience=target_audience,
-            additional_claims=self._additional_claims.copy(),
-            quota_project_id=self.quota_project_id,
-        )
+        cred = self._make_copy()
+        cred._target_audience = target_audience
+        return cred
 
     def _with_use_iam_endpoint(self, use_iam_endpoint):
         """Create a copy of these credentials with the use_iam_endpoint value.
@@ -666,39 +687,29 @@ def _with_use_iam_endpoint(self, use_iam_endpoint):
         Returns:
             google.auth.service_account.IDTokenCredentials: A new credentials
                 instance.
+        Raises:
+            google.auth.exceptions.InvalidValue: If the universe domain is not
+                default and use_iam_endpoint is False.
         """
-        cred = self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            token_uri=self._token_uri,
-            target_audience=self._target_audience,
-            additional_claims=self._additional_claims.copy(),
-            quota_project_id=self.quota_project_id,
-        )
+        cred = self._make_copy()
+        if cred._universe_domain != _DEFAULT_UNIVERSE_DOMAIN and not use_iam_endpoint:
+            raise exceptions.InvalidValue(
+                "use_iam_endpoint should be True for non-default universe domain"
+            )
         cred._use_iam_endpoint = use_iam_endpoint
         return cred
 
     @_helpers.copy_docstring(credentials.CredentialsWithQuotaProject)
     def with_quota_project(self, quota_project_id):
-        return self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            token_uri=self._token_uri,
-            target_audience=self._target_audience,
-            additional_claims=self._additional_claims.copy(),
-            quota_project_id=quota_project_id,
-        )
+        cred = self._make_copy()
+        cred._quota_project_id = quota_project_id
+        return cred
 
     @_helpers.copy_docstring(credentials.CredentialsWithTokenUri)
     def with_token_uri(self, token_uri):
-        return self.__class__(
-            self._signer,
-            service_account_email=self._service_account_email,
-            token_uri=token_uri,
-            target_audience=self._target_audience,
-            additional_claims=self._additional_claims.copy(),
-            quota_project_id=self._quota_project_id,
-        )
+        cred = self._make_copy()
+        cred._token_uri = token_uri
+        return cred
 
     def _make_authorization_grant_assertion(self):
         """Create the OAuth 2.0 assertion.

system_tests/secrets.tar.enc

@@ -0,0 +1,15 @@
+{
+    "type": "service_account",
+    "universe_domain": "universe.foo",
+    "project_id": "example_project",
+    "private_key_id": "1",
+    "private_key": "-----BEGIN RSA PRIVATE KEY-----\nMIIEpAIBAAKCAQEA4ej0p7bQ7L/r4rVGUz9RN4VQWoej1Bg1mYWIDYslvKrk1gpj\n7wZgkdmM7oVK2OfgrSj/FCTkInKPqaCR0gD7K80q+mLBrN3PUkDrJQZpvRZIff3/\nxmVU1WeruQLFJjnFb2dqu0s/FY/2kWiJtBCakXvXEOb7zfbINuayL+MSsCGSdVYs\nSliS5qQpgyDap+8b5fpXZVJkq92hrcNtbkg7hCYUJczt8n9hcCTJCfUpApvaFQ18\npe+zpyl4+WzkP66I28hniMQyUlA1hBiskT7qiouq0m8IOodhv2fagSZKjOTTU2xk\nSBc//fy3ZpsL7WqgsZS7Q+0VRK8gKfqkxg5OYQIDAQABAoIBAQDGGHzQxGKX+ANk\nnQi53v/c6632dJKYXVJC+PDAz4+bzU800Y+n/bOYsWf/kCp94XcG4Lgsdd0Gx+Zq\nHD9CI1IcqqBRR2AFscsmmX6YzPLTuEKBGMW8twaYy3utlFxElMwoUEsrSWRcCA1y\nnHSDzTt871c7nxCXHxuZ6Nm/XCL7Bg8uidRTSC1sQrQyKgTPhtQdYrPQ4WZ1A4J9\nIisyDYmZodSNZe5P+LTJ6M1SCgH8KH9ZGIxv3diMwzNNpk3kxJc9yCnja4mjiGE2\nYCNusSycU5IhZwVeCTlhQGcNeV/skfg64xkiJE34c2y2ttFbdwBTPixStGaF09nU\nZ422D40BAoGBAPvVyRRsC3BF+qZdaSMFwI1yiXY7vQw5+JZh01tD28NuYdRFzjcJ\nvzT2n8LFpj5ZfZFvSMLMVEFVMgQvWnN0O6xdXvGov6qlRUSGaH9u+TCPNnIldjMP\nB8+xTwFMqI7uQr54wBB+Poq7dVRP+0oHb0NYAwUBXoEuvYo3c/nDoRcZAoGBAOWl\naLHjMv4CJbArzT8sPfic/8waSiLV9Ixs3Re5YREUTtnLq7LoymqB57UXJB3BNz/2\neCueuW71avlWlRtE/wXASj5jx6y5mIrlV4nZbVuyYff0QlcG+fgb6pcJQuO9DxMI\naqFGrWP3zye+LK87a6iR76dS9vRU+bHZpSVvGMKJAoGAFGt3TIKeQtJJyqeUWNSk\nklORNdcOMymYMIlqG+JatXQD1rR6ThgqOt8sgRyJqFCVT++YFMOAqXOBBLnaObZZ\nCFbh1fJ66BlSjoXff0W+SuOx5HuJJAa5+WtFHrPajwxeuRcNa8jwxUsB7n41wADu\nUqWWSRedVBg4Ijbw3nWwYDECgYB0pLew4z4bVuvdt+HgnJA9n0EuYowVdadpTEJg\nsoBjNHV4msLzdNqbjrAqgz6M/n8Ztg8D2PNHMNDNJPVHjJwcR7duSTA6w2p/4k28\nbvvk/45Ta3XmzlxZcZSOct3O31Cw0i2XDVc018IY5be8qendDYM08icNo7vQYkRH\n504kQQKBgQDjx60zpz8ozvm1XAj0wVhi7GwXe+5lTxiLi9Fxq721WDxPMiHDW2XL\nYXfFVy/9/GIMvEiGYdmarK1NW+VhWl1DC5xhDg0kvMfxplt4tynoq1uTsQTY31Mx\nBeF5CT/JuNYk3bEBF0H/Q3VGO1/ggVS+YezdFbLWIRoMnLj6XCFEGg==\n-----END RSA PRIVATE KEY-----\n",
+    "client_email": "[email protected]",
+    "client_id": "1234",
+    "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+    "token_uri": "https://oauth2.universe.foo/token",
+    "auth_provider_x509_cert_url": "https://www.universe.foo/oauth2/v1/certs",
+    "client_x509_cert_url": "https://www.universe.foo/robot/v1/metadata/x509/foo.iam.gserviceaccount.com"
+}
+  
+