diff --git a/google/auth/compute_engine/credentials.py b/google/auth/compute_engine/credentials.py index 618fa5a2d..ed935c17c 100644 --- a/google/auth/compute_engine/credentials.py +++ b/google/auth/compute_engine/credentials.py @@ -28,6 +28,7 @@ from google.auth import exceptions from google.auth import iam from google.auth import jwt +from google.auth import metrics from google.auth.compute_engine import _metadata from google.oauth2 import _client @@ -94,6 +95,9 @@ def _retrieve_info(self, request): if self._scopes is None: self._scopes = info["scopes"] + def _metric_header_for_usage(self): + return metrics.CRED_TYPE_SA_MDS + def refresh(self, request): """Refresh the access token and scopes. diff --git a/google/auth/impersonated_credentials.py b/google/auth/impersonated_credentials.py index f978b64ef..c596222a4 100644 --- a/google/auth/impersonated_credentials.py +++ b/google/auth/impersonated_credentials.py @@ -37,6 +37,7 @@ from google.auth import credentials from google.auth import exceptions from google.auth import jwt +from google.auth import metrics _DEFAULT_TOKEN_LIFETIME_SECS = 3600 # 1 hour in seconds @@ -238,6 +239,9 @@ def __init__( self._quota_project_id = quota_project_id self._iam_endpoint_override = iam_endpoint_override + def _metric_header_for_usage(self): + return metrics.CRED_TYPE_SA_IMPERSONATE + @_helpers.copy_docstring(credentials.Credentials) def refresh(self, request): self._update_token(request) diff --git a/google/oauth2/credentials.py b/google/oauth2/credentials.py index 457db76ae..dd54edfb3 100644 --- a/google/oauth2/credentials.py +++ b/google/oauth2/credentials.py @@ -42,6 +42,7 @@ from google.auth import _helpers from google.auth import credentials from google.auth import exceptions +from google.auth import metrics from google.oauth2 import reauth _LOGGER = logging.getLogger(__name__) @@ -287,6 +288,9 @@ def with_token_uri(self, token_uri): enable_reauth_refresh=self._enable_reauth_refresh, ) + def _metric_header_for_usage(self): + return metrics.CRED_TYPE_USER + @_helpers.copy_docstring(credentials.Credentials) def refresh(self, request): scopes = self._scopes if self._scopes is not None else self._default_scopes diff --git a/google/oauth2/service_account.py b/google/oauth2/service_account.py index a0268970c..c80627398 100644 --- a/google/oauth2/service_account.py +++ b/google/oauth2/service_account.py @@ -78,6 +78,7 @@ from google.auth import credentials from google.auth import exceptions from google.auth import jwt +from google.auth import metrics from google.oauth2 import _client _DEFAULT_TOKEN_LIFETIME_SECS = 3600 # 1 hour in seconds @@ -400,6 +401,16 @@ def _make_authorization_grant_assertion(self): return token + def _use_self_signed_jwt(self): + # Since domain wide delegation doesn't work with self signed JWT. If + # subject exists, then we should not use self signed JWT. + return self._subject is None and self._jwt_credentials is not None + + def _metric_header_for_usage(self): + if self._use_self_signed_jwt(): + return metrics.CRED_TYPE_SA_JWT + return metrics.CRED_TYPE_SA_ASSERTION + @_helpers.copy_docstring(credentials.Credentials) def refresh(self, request): if ( @@ -414,9 +425,7 @@ def refresh(self, request): "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: + if self._use_self_signed_jwt(): self._jwt_credentials.refresh(request) self.token = self._jwt_credentials.token.decode() self.expiry = self._jwt_credentials.expiry diff --git a/tests/compute_engine/test_credentials.py b/tests/compute_engine/test_credentials.py index 6a2f8cc20..4234d98e4 100644 --- a/tests/compute_engine/test_credentials.py +++ b/tests/compute_engine/test_credentials.py @@ -185,6 +185,15 @@ def test_with_scopes(self): assert self.credentials._scopes == scopes + def test_token_usage_metrics(self): + self.credentials.token = "token" + self.credentials.expiry = None + + headers = {} + self.credentials.before_request(mock.Mock(), None, None, headers) + assert headers["authorization"] == "Bearer token" + assert headers["x-goog-api-client"] == "cred-type/mds" + class TestIDTokenCredentials(object): credentials = None diff --git a/tests/oauth2/test_credentials.py b/tests/oauth2/test_credentials.py index 5a63224fd..feb7e4311 100644 --- a/tests/oauth2/test_credentials.py +++ b/tests/oauth2/test_credentials.py @@ -69,6 +69,16 @@ def test_default_state(self): assert credentials.rapt_token == self.RAPT_TOKEN assert credentials.refresh_handler is None + def test_token_usage_metrics(self): + credentials = self.make_credentials() + credentials.token = "token" + credentials.expiry = None + + headers = {} + credentials.before_request(mock.Mock(), None, None, headers) + assert headers["authorization"] == "Bearer token" + assert headers["x-goog-api-client"] == "cred-type/u" + def test_refresh_handler_setter_and_getter(self): scopes = ["email", "profile"] original_refresh_handler = mock.Mock(return_value=("ACCESS_TOKEN_1", None)) diff --git a/tests/oauth2/test_service_account.py b/tests/oauth2/test_service_account.py index 8388acd06..36565b721 100644 --- a/tests/oauth2/test_service_account.py +++ b/tests/oauth2/test_service_account.py @@ -408,6 +408,37 @@ def test__create_self_signed_jwt_always_use_jwt_access(self, jwt): credentials._create_self_signed_jwt(None) jwt.from_signing_credentials.assert_not_called() + def test_token_usage_metrics_assertion(self): + credentials = service_account.Credentials( + SIGNER, + self.SERVICE_ACCOUNT_EMAIL, + self.TOKEN_URI, + always_use_jwt_access=False, + ) + credentials.token = "token" + credentials.expiry = None + + headers = {} + credentials.before_request(mock.Mock(), None, None, headers) + assert headers["authorization"] == "Bearer token" + assert headers["x-goog-api-client"] == "cred-type/sa" + + def test_token_usage_metrics_self_signed_jwt(self): + credentials = service_account.Credentials( + SIGNER, + self.SERVICE_ACCOUNT_EMAIL, + self.TOKEN_URI, + always_use_jwt_access=True, + ) + credentials._create_self_signed_jwt("foo.googleapis.com") + credentials.token = "token" + credentials.expiry = None + + headers = {} + credentials.before_request(mock.Mock(), None, None, headers) + assert headers["authorization"] == "Bearer token" + assert headers["x-goog-api-client"] == "cred-type/jwt" + @mock.patch("google.oauth2._client.jwt_grant", autospec=True) def test_refresh_success(self, jwt_grant): credentials = self.make_credentials() diff --git a/tests/test_impersonated_credentials.py b/tests/test_impersonated_credentials.py index 0e4dc08d7..8f2e1fdbe 100644 --- a/tests/test_impersonated_credentials.py +++ b/tests/test_impersonated_credentials.py @@ -162,6 +162,16 @@ def make_request( return request + def test_token_usage_metrics(self): + credentials = self.make_credentials() + credentials.token = "token" + credentials.expiry = None + + headers = {} + credentials.before_request(mock.Mock(), None, None, headers) + assert headers["authorization"] == "Bearer token" + assert headers["x-goog-api-client"] == "cred-type/imp" + @pytest.mark.parametrize("use_data_bytes", [True, False]) def test_refresh_success(self, use_data_bytes, mock_donor_credentials): credentials = self.make_credentials(lifetime=None)