feat: service account is able to use a private token endpoint

google/oauth2/service_account.py

@@ -125,6 +125,7 @@ def __init__(
         signer,
         service_account_email,
         token_uri,
+        private_token_uri=None,
         scopes=None,
         default_scopes=None,
         subject=None,
@@ -142,6 +143,8 @@ def __init__(
             default_scopes (Sequence[str]): Default scopes passed by a
                 Google client library. Use 'scopes' for user-defined scopes.
             token_uri (str): The OAuth 2.0 Token URI.
+            private_token_uri (str): The Token URI in a Private Service Connect
+                network.
             subject (str): For domain-wide delegation, the email address of the
                 user to for which to request delegated access.
             project_id  (str): Project ID associated with the service account
@@ -168,6 +171,7 @@ def __init__(
         self._project_id = project_id
         self._quota_project_id = quota_project_id
         self._token_uri = token_uri
+        self._private_token_uri = private_token_uri
         self._always_use_jwt_access = always_use_jwt_access
 
         self._jwt_credentials = None
@@ -266,6 +270,7 @@ def with_scopes(self, scopes, default_scopes=None):
             scopes=scopes,
             default_scopes=default_scopes,
             token_uri=self._token_uri,
+            private_token_uri=self._private_token_uri,
             subject=self._subject,
             project_id=self._project_id,
             quota_project_id=self._quota_project_id,
@@ -289,6 +294,7 @@ def with_always_use_jwt_access(self, always_use_jwt_access):
             scopes=self._scopes,
             default_scopes=self._default_scopes,
             token_uri=self._token_uri,
+            private_token_uri=self._private_token_uri,
             subject=self._subject,
             project_id=self._project_id,
             quota_project_id=self._quota_project_id,
@@ -312,6 +318,7 @@ def with_subject(self, subject):
             scopes=self._scopes,
             default_scopes=self._default_scopes,
             token_uri=self._token_uri,
+            private_token_uri=self._private_token_uri,
             subject=subject,
             project_id=self._project_id,
             quota_project_id=self._quota_project_id,
@@ -340,6 +347,7 @@ def with_claims(self, additional_claims):
             scopes=self._scopes,
             default_scopes=self._default_scopes,
             token_uri=self._token_uri,
+            private_token_uri=self._private_token_uri,
             subject=self._subject,
             project_id=self._project_id,
             quota_project_id=self._quota_project_id,
@@ -356,6 +364,7 @@ def with_quota_project(self, quota_project_id):
             default_scopes=self._default_scopes,
             scopes=self._scopes,
             token_uri=self._token_uri,
+            private_token_uri=self._private_token_uri,
             subject=self._subject,
             project_id=self._project_id,
             quota_project_id=quota_project_id,
@@ -404,9 +413,8 @@ def refresh(self, request):
             self.expiry = self._jwt_credentials.expiry
         else:
             assertion = self._make_authorization_grant_assertion()
-            access_token, expiry, _ = _client.jwt_grant(
-                request, self._token_uri, assertion
-            )
+            token_uri = self._private_token_uri or self._token_uri
+            access_token, expiry, _ = _client.jwt_grant(request, token_uri, assertion)
             self.token = access_token
             self.expiry = expiry
 
@@ -503,6 +511,7 @@ def __init__(
         service_account_email,
         token_uri,
         target_audience,
+        private_token_uri=None,
         additional_claims=None,
         quota_project_id=None,
     ):
@@ -514,6 +523,8 @@ def __init__(
             target_audience (str): The intended audience for these credentials,
                 used when requesting the ID Token. The ID Token's ``aud`` claim
                 will be set to this string.
+            private_token_uri (str): The Token URI in a Private Service Connect
+                network.
             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.
@@ -526,6 +537,7 @@ def __init__(
         self._signer = signer
         self._service_account_email = service_account_email
         self._token_uri = token_uri
+        self._private_token_uri = private_token_uri
         self._target_audience = target_audience
         self._quota_project_id = quota_project_id
 
@@ -609,6 +621,7 @@ def with_target_audience(self, target_audience):
             service_account_email=self._service_account_email,
             token_uri=self._token_uri,
             target_audience=target_audience,
+            private_token_uri=self._private_token_uri,
             additional_claims=self._additional_claims.copy(),
             quota_project_id=self.quota_project_id,
         )
@@ -620,6 +633,7 @@ def with_quota_project(self, quota_project_id):
             service_account_email=self._service_account_email,
             token_uri=self._token_uri,
             target_audience=self._target_audience,
+            private_token_uri=self._private_token_uri,
             additional_claims=self._additional_claims.copy(),
             quota_project_id=quota_project_id,
         )
@@ -658,8 +672,9 @@ def _make_authorization_grant_assertion(self):
     @_helpers.copy_docstring(credentials.Credentials)
     def refresh(self, request):
         assertion = self._make_authorization_grant_assertion()
+        token_uri = self._private_token_uri or self._token_uri
         access_token, expiry, _ = _client.id_token_jwt_grant(
-            request, self._token_uri, assertion
+            request, token_uri, assertion
         )
         self.token = access_token
         self.expiry = expiry

tests/oauth2/test_service_account.py

@@ -47,6 +47,7 @@
 class TestCredentials(object):
     SERVICE_ACCOUNT_EMAIL = "service-account@example.com"
     TOKEN_URI = "https://example.com/oauth2/token"
+    PRIVATE_TOKEN_URI = "https://another-example.com/oauth2/token"
 
     @classmethod
     def make_credentials(cls):
@@ -185,6 +186,14 @@ def test__make_authorization_grant_assertion_subject(self):
         payload = jwt.decode(token, PUBLIC_CERT_BYTES)
         assert payload["sub"] == subject
 
+    def test__make_authorization_grant_assertion_private_token_uri(self):
+        credentials = self.make_credentials()
+        credentials._private_token_uri = self.PRIVATE_TOKEN_URI
+
+        token = credentials._make_authorization_grant_assertion()
+        payload = jwt.decode(token, PUBLIC_CERT_BYTES)
+        assert payload["aud"] == credentials._token_uri
+
     def test_apply_with_quota_project_id(self):
         credentials = service_account.Credentials(
             SIGNER,
@@ -322,6 +331,30 @@ def test_refresh_success(self, jwt_grant):
         # expired)
         assert credentials.valid
 
+    @mock.patch("google.oauth2._client.jwt_grant", autospec=True)
+    def test_refresh_private_token_uri(self, jwt_grant):
+        credentials = self.make_credentials()
+        credentials._private_token_uri = self.PRIVATE_TOKEN_URI
+        request = mock.create_autospec(transport.Request, instance=True)
+        token = "token"
+        jwt_grant.return_value = (
+            token,
+            _helpers.utcnow() + datetime.timedelta(seconds=500),
+            {},
+        )
+
+        # Refresh credentials
+        credentials.refresh(request)
+
+        # Check jwt grant call.
+        assert jwt_grant.called
+
+        _, token_uri, assertion = jwt_grant.call_args[0]
+        assert token_uri == credentials._private_token_uri
+
+        payload = jwt.decode(assertion, PUBLIC_CERT_BYTES)
+        assert payload["aud"] == credentials._token_uri
+
     @mock.patch("google.oauth2._client.jwt_grant", autospec=True)
     def test_before_request_refreshes(self, jwt_grant):
         credentials = self.make_credentials()
@@ -375,6 +408,7 @@ def test_refresh_with_jwt_credentials(self, make_jwt):
 class TestIDTokenCredentials(object):
     SERVICE_ACCOUNT_EMAIL = "service-account@example.com"
     TOKEN_URI = "https://example.com/oauth2/token"
+    PRIVATE_TOKEN_URI = "https://another-example.com/oauth2/token"
     TARGET_AUDIENCE = "https://example.com"
 
     @classmethod
@@ -443,6 +477,13 @@ def test__make_authorization_grant_assertion(self):
         assert payload["aud"] == self.TOKEN_URI
         assert payload["target_audience"] == self.TARGET_AUDIENCE
 
+    def test__make_authorization_grant_assertion_private_token_uri(self):
+        credentials = self.make_credentials()
+        credentials._private_token_uri = self.PRIVATE_TOKEN_URI
+        token = credentials._make_authorization_grant_assertion()
+        payload = jwt.decode(token, PUBLIC_CERT_BYTES)
+        assert payload["aud"] == self.TOKEN_URI
+
     @mock.patch("google.oauth2._client.id_token_jwt_grant", autospec=True)
     def test_refresh_success(self, id_token_jwt_grant):
         credentials = self.make_credentials()
@@ -474,6 +515,29 @@ def test_refresh_success(self, id_token_jwt_grant):
         # expired)
         assert credentials.valid
 
+    @mock.patch("google.oauth2._client.id_token_jwt_grant", autospec=True)
+    def test_refresh_private_token_uri(self, id_token_jwt_grant):
+        credentials = self.make_credentials()
+        credentials._private_token_uri = self.PRIVATE_TOKEN_URI
+        token = "token"
+        id_token_jwt_grant.return_value = (
+            token,
+            _helpers.utcnow() + datetime.timedelta(seconds=500),
+            {},
+        )
+        request = mock.create_autospec(transport.Request, instance=True)
+
+        # Refresh credentials
+        credentials.refresh(request)
+
+        # Check jwt grant call.
+        assert id_token_jwt_grant.called
+
+        _, token_uri, assertion = id_token_jwt_grant.call_args[0]
+        assert token_uri == credentials._private_token_uri
+        payload = jwt.decode(assertion, PUBLIC_CERT_BYTES)
+        payload["aud"] = credentials._token_uri
+
     @mock.patch("google.oauth2._client.id_token_jwt_grant", autospec=True)
     def test_before_request_refreshes(self, id_token_jwt_grant):
         credentials = self.make_credentials()