From d0435acd341cd076053d06cd982c2cb9a40a9426 Mon Sep 17 00:00:00 2001 From: Sijun Liu Date: Wed, 24 May 2023 23:16:06 -0700 Subject: [PATCH] feat: add metrics (part 3) --- google/auth/compute_engine/_metadata.py | 30 +++++++++--- google/auth/compute_engine/credentials.py | 7 ++- google/auth/impersonated_credentials.py | 10 +++- google/oauth2/_client.py | 33 ++++++++++--- google/oauth2/reauth.py | 18 +++++-- tests/compute_engine/test__metadata.py | 59 ++++++++++++++++++----- tests/compute_engine/test_credentials.py | 30 ++++++++++-- tests/oauth2/test__client.py | 59 ++++++++++++++++++++--- tests/oauth2/test_reauth.py | 32 ++++++++++-- tests/test_aws.py | 22 ++++++++- tests/test_external_account.py | 54 ++++++++++++++++++--- tests/test_identity_pool.py | 10 +++- tests/test_impersonated_credentials.py | 47 +++++++++++++++++- 13 files changed, 355 insertions(+), 56 deletions(-) diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py index 87e79a9e9..c5b9d5a2b 100644 --- a/google/auth/compute_engine/_metadata.py +++ b/google/auth/compute_engine/_metadata.py @@ -29,6 +29,7 @@ from google.auth import _helpers from google.auth import environment_vars from google.auth import exceptions +from google.auth import metrics _LOGGER = logging.getLogger(__name__) @@ -121,13 +122,13 @@ def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3): # the metadata resolution was particularly slow. The latter case is # "unlikely". retries = 0 + headers = _METADATA_HEADERS.copy() + headers[metrics.API_CLIENT_HEADER] = metrics.mds_ping() + while retries < retry_count: try: response = request( - url=_METADATA_IP_ROOT, - method="GET", - headers=_METADATA_HEADERS, - timeout=timeout, + url=_METADATA_IP_ROOT, method="GET", headers=headers, timeout=timeout ) metadata_flavor = response.headers.get(_METADATA_FLAVOR_HEADER) @@ -150,7 +151,13 @@ def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT, retry_count=3): def get( - request, path, root=_METADATA_ROOT, params=None, recursive=False, retry_count=5 + request, + path, + root=_METADATA_ROOT, + params=None, + recursive=False, + retry_count=5, + headers=None, ): """Fetch a resource from the metadata server. @@ -167,6 +174,7 @@ def get( details. retry_count (int): How many times to attempt connecting to metadata server using above timeout. + headers (Optional[Mapping[str, str]]): Headers for the request. Returns: Union[Mapping, str]: If the metadata server returns JSON, a mapping of @@ -180,6 +188,10 @@ def get( base_url = urlparse.urljoin(root, path) query_params = {} if params is None else params + headers_to_use = _METADATA_HEADERS.copy() + if headers: + headers_to_use.update(headers) + if recursive: query_params["recursive"] = "true" @@ -188,7 +200,7 @@ def get( retries = 0 while retries < retry_count: try: - response = request(url=url, method="GET", headers=_METADATA_HEADERS) + response = request(url=url, method="GET", headers=headers_to_use) break except exceptions.TransportError as e: @@ -300,8 +312,12 @@ def get_service_account_token(request, service_account="default", scopes=None): else: params = None + metrics_header = { + metrics.API_CLIENT_HEADER: metrics.token_request_access_token_mds() + } + path = "instance/service-accounts/{0}/token".format(service_account) - token_json = get(request, path, params=params) + token_json = get(request, path, params=params, headers=metrics_header) token_expiry = _helpers.utcnow() + datetime.timedelta( seconds=token_json["expires_in"] ) diff --git a/google/auth/compute_engine/credentials.py b/google/auth/compute_engine/credentials.py index ed935c17c..30ffb162b 100644 --- a/google/auth/compute_engine/credentials.py +++ b/google/auth/compute_engine/credentials.py @@ -378,7 +378,12 @@ def _call_metadata_identity_endpoint(self, request): try: path = "instance/service-accounts/default/identity" params = {"audience": self._target_audience, "format": "full"} - id_token = _metadata.get(request, path, params=params) + metrics_header = { + metrics.API_CLIENT_HEADER: metrics.token_request_id_token_mds() + } + id_token = _metadata.get( + request, path, params=params, headers=metrics_header + ) except exceptions.TransportError as caught_exc: new_exc = exceptions.RefreshError(caught_exc) six.raise_from(new_exc, caught_exc) diff --git a/google/auth/impersonated_credentials.py b/google/auth/impersonated_credentials.py index c596222a4..ddafc08ee 100644 --- a/google/auth/impersonated_credentials.py +++ b/google/auth/impersonated_credentials.py @@ -265,7 +265,10 @@ def _update_token(self, request): "lifetime": str(self._lifetime) + "s", } - headers = {"Content-Type": "application/json"} + headers = { + "Content-Type": "application/json", + metrics.API_CLIENT_HEADER: metrics.token_request_access_token_impersonate(), + } # Apply the source credentials authentication info. self._source_credentials.apply(headers) @@ -426,7 +429,10 @@ def refresh(self, request): "includeEmail": self._include_email, } - headers = {"Content-Type": "application/json"} + headers = { + "Content-Type": "application/json", + metrics.API_CLIENT_HEADER: metrics.token_request_id_token_impersonate(), + } authed_session = AuthorizedSession( self._target_credentials._source_credentials, auth_request=request diff --git a/google/oauth2/_client.py b/google/oauth2/_client.py index 03f6e8f03..e2c9509a9 100644 --- a/google/oauth2/_client.py +++ b/google/oauth2/_client.py @@ -34,6 +34,7 @@ from google.auth import _helpers from google.auth import exceptions from google.auth import jwt +from google.auth import metrics from google.auth import transport _URLENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded" @@ -146,6 +147,7 @@ def _token_endpoint_request_no_throw( access_token=None, use_json=False, can_retry=True, + headers=None, **kwargs ): """Makes a request to the OAuth 2.0 authorization server's token endpoint. @@ -161,6 +163,7 @@ def _token_endpoint_request_no_throw( use_json (Optional(bool)): Use urlencoded format or json format for the content type. The default value is False. can_retry (bool): Enable or disable request retry behavior. + headers (Optional[Mapping[str, str]]): The headers for the request. kwargs: Additional arguments passed on to the request method. The kwargs will be passed to `requests.request` method, see: https://docs.python-requests.org/en/latest/api/#requests.request. @@ -176,18 +179,21 @@ def _token_endpoint_request_no_throw( is retryable. """ if use_json: - headers = {"Content-Type": _JSON_CONTENT_TYPE} + headers_to_use = {"Content-Type": _JSON_CONTENT_TYPE} body = json.dumps(body).encode("utf-8") else: - headers = {"Content-Type": _URLENCODED_CONTENT_TYPE} + headers_to_use = {"Content-Type": _URLENCODED_CONTENT_TYPE} body = urllib.parse.urlencode(body).encode("utf-8") if access_token: - headers["Authorization"] = "Bearer {}".format(access_token) + headers_to_use["Authorization"] = "Bearer {}".format(access_token) + + if headers: + headers_to_use.update(headers) def _perform_request(): response = request( - method="POST", url=token_uri, headers=headers, body=body, **kwargs + method="POST", url=token_uri, headers=headers_to_use, body=body, **kwargs ) response_body = ( response.data.decode("utf-8") @@ -231,6 +237,7 @@ def _token_endpoint_request( access_token=None, use_json=False, can_retry=True, + headers=None, **kwargs ): """Makes a request to the OAuth 2.0 authorization server's token endpoint. @@ -245,6 +252,7 @@ def _token_endpoint_request( use_json (Optional(bool)): Use urlencoded format or json format for the content type. The default value is False. can_retry (bool): Enable or disable request retry behavior. + headers (Optional[Mapping[str, str]]): The headers for the request. kwargs: Additional arguments passed on to the request method. The kwargs will be passed to `requests.request` method, see: https://docs.python-requests.org/en/latest/api/#requests.request. @@ -268,6 +276,7 @@ def _token_endpoint_request( access_token=access_token, use_json=use_json, can_retry=can_retry, + headers=headers, **kwargs ) if not response_status_ok: @@ -301,7 +310,13 @@ def jwt_grant(request, token_uri, assertion, can_retry=True): body = {"assertion": assertion, "grant_type": _JWT_GRANT_TYPE} response_data = _token_endpoint_request( - request, token_uri, body, can_retry=can_retry + request, + token_uri, + body, + can_retry=can_retry, + headers={ + metrics.API_CLIENT_HEADER: metrics.token_request_access_token_sa_assertion() + }, ) try: @@ -384,7 +399,13 @@ def id_token_jwt_grant(request, token_uri, assertion, can_retry=True): body = {"assertion": assertion, "grant_type": _JWT_GRANT_TYPE} response_data = _token_endpoint_request( - request, token_uri, body, can_retry=can_retry + request, + token_uri, + body, + can_retry=can_retry, + headers={ + metrics.API_CLIENT_HEADER: metrics.token_request_id_token_sa_assertion() + }, ) try: diff --git a/google/oauth2/reauth.py b/google/oauth2/reauth.py index 84fd0be2b..114c47c48 100644 --- a/google/oauth2/reauth.py +++ b/google/oauth2/reauth.py @@ -37,6 +37,7 @@ from six.moves import range from google.auth import exceptions +from google.auth import metrics from google.oauth2 import _client from google.oauth2 import challenges @@ -94,9 +95,15 @@ def _get_challenges( body = {"supportedChallengeTypes": supported_challenge_types} if requested_scopes: body["oauthScopesForDomainPolicyLookup"] = requested_scopes + metrics_header = {metrics.API_CLIENT_HEADER: metrics.reauth_start()} return _client._token_endpoint_request( - request, _REAUTH_API + ":start", body, access_token=access_token, use_json=True + request, + _REAUTH_API + ":start", + body, + access_token=access_token, + use_json=True, + headers=metrics_header, ) @@ -123,6 +130,7 @@ def _send_challenge_result( "action": "RESPOND", "proposalResponse": client_input, } + metrics_header = {metrics.API_CLIENT_HEADER: metrics.reauth_continue()} return _client._token_endpoint_request( request, @@ -130,6 +138,7 @@ def _send_challenge_result( body, access_token=access_token, use_json=True, + headers=metrics_header, ) @@ -320,9 +329,10 @@ def refresh_grant( body["scope"] = " ".join(scopes) if rapt_token: body["rapt"] = rapt_token + metrics_header = {metrics.API_CLIENT_HEADER: metrics.token_request_user()} response_status_ok, response_data, retryable_error = _client._token_endpoint_request_no_throw( - request, token_uri, body + request, token_uri, body, headers=metrics_header ) if ( not response_status_ok @@ -345,7 +355,9 @@ def refresh_grant( response_status_ok, response_data, retryable_error, - ) = _client._token_endpoint_request_no_throw(request, token_uri, body) + ) = _client._token_endpoint_request_no_throw( + request, token_uri, body, headers=metrics_header + ) if not response_status_ok: _client._handle_error_response(response_data, retryable_error) diff --git a/tests/compute_engine/test__metadata.py b/tests/compute_engine/test__metadata.py index 7f3386276..f543426d3 100644 --- a/tests/compute_engine/test__metadata.py +++ b/tests/compute_engine/test__metadata.py @@ -38,6 +38,15 @@ DATA_DIR, "smbios_product_name_non_google" ) +ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/mds" +) +MDS_PING_METRICS_HEADER_VALUE = "gl-python/3.7 auth/1.1 auth-request-type/mds" +MDS_PING_REQUEST_HEADER = { + "metadata-flavor": "Google", + "x-goog-api-client": MDS_PING_METRICS_HEADER_VALUE, +} + def make_request(data, status=http_client.OK, headers=None, retry=False): response = mock.create_autospec(transport.Response, instance=True) @@ -87,7 +96,8 @@ def test_is_on_gce_linux_success(): assert _metadata.is_on_gce(request) -def test_ping_success(): +@mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE) +def test_ping_success(mock_metrics_header_value): request = make_request("", headers=_metadata._METADATA_HEADERS) assert _metadata.ping(request) @@ -95,12 +105,13 @@ def test_ping_success(): request.assert_called_once_with( method="GET", url=_metadata._METADATA_IP_ROOT, - headers=_metadata._METADATA_HEADERS, + headers=MDS_PING_REQUEST_HEADER, timeout=_metadata._METADATA_DEFAULT_TIMEOUT, ) -def test_ping_success_retry(): +@mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE) +def test_ping_success_retry(mock_metrics_header_value): request = make_request("", headers=_metadata._METADATA_HEADERS, retry=True) assert _metadata.ping(request) @@ -108,7 +119,7 @@ def test_ping_success_retry(): request.assert_called_with( method="GET", url=_metadata._METADATA_IP_ROOT, - headers=_metadata._METADATA_HEADERS, + headers=MDS_PING_REQUEST_HEADER, timeout=_metadata._METADATA_DEFAULT_TIMEOUT, ) assert request.call_count == 2 @@ -127,7 +138,8 @@ def test_ping_failure_connection_failed(): assert not _metadata.ping(request) -def test_ping_success_custom_root(): +@mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE) +def test_ping_success_custom_root(mock_metrics_header_value): request = make_request("", headers=_metadata._METADATA_HEADERS) fake_ip = "1.2.3.4" @@ -143,7 +155,7 @@ def test_ping_success_custom_root(): request.assert_called_once_with( method="GET", url="http://" + fake_ip, - headers=_metadata._METADATA_HEADERS, + headers=MDS_PING_REQUEST_HEADER, timeout=_metadata._METADATA_DEFAULT_TIMEOUT, ) @@ -341,8 +353,12 @@ def test_get_project_id(): assert project_id == project +@mock.patch( + "google.auth.metrics.token_request_access_token_mds", + return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, +) @mock.patch("google.auth._helpers.utcnow", return_value=datetime.datetime.min) -def test_get_service_account_token(utcnow): +def test_get_service_account_token(utcnow, mock_metrics_header_value): ttl = 500 request = make_request( json.dumps({"access_token": "token", "expires_in": ttl}), @@ -354,14 +370,21 @@ def test_get_service_account_token(utcnow): request.assert_called_once_with( method="GET", url=_metadata._METADATA_ROOT + PATH + "/token", - headers=_metadata._METADATA_HEADERS, + headers={ + "metadata-flavor": "Google", + "x-goog-api-client": ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, + }, ) assert token == "token" assert expiry == utcnow() + datetime.timedelta(seconds=ttl) +@mock.patch( + "google.auth.metrics.token_request_access_token_mds", + return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, +) @mock.patch("google.auth._helpers.utcnow", return_value=datetime.datetime.min) -def test_get_service_account_token_with_scopes_list(utcnow): +def test_get_service_account_token_with_scopes_list(utcnow, mock_metrics_header_value): ttl = 500 request = make_request( json.dumps({"access_token": "token", "expires_in": ttl}), @@ -373,14 +396,23 @@ def test_get_service_account_token_with_scopes_list(utcnow): request.assert_called_once_with( method="GET", url=_metadata._METADATA_ROOT + PATH + "/token" + "?scopes=foo%2Cbar", - headers=_metadata._METADATA_HEADERS, + headers={ + "metadata-flavor": "Google", + "x-goog-api-client": ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, + }, ) assert token == "token" assert expiry == utcnow() + datetime.timedelta(seconds=ttl) +@mock.patch( + "google.auth.metrics.token_request_access_token_mds", + return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, +) @mock.patch("google.auth._helpers.utcnow", return_value=datetime.datetime.min) -def test_get_service_account_token_with_scopes_string(utcnow): +def test_get_service_account_token_with_scopes_string( + utcnow, mock_metrics_header_value +): ttl = 500 request = make_request( json.dumps({"access_token": "token", "expires_in": ttl}), @@ -392,7 +424,10 @@ def test_get_service_account_token_with_scopes_string(utcnow): request.assert_called_once_with( method="GET", url=_metadata._METADATA_ROOT + PATH + "/token" + "?scopes=foo%2Cbar", - headers=_metadata._METADATA_HEADERS, + headers={ + "metadata-flavor": "Google", + "x-goog-api-client": ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, + }, ) assert token == "token" assert expiry == utcnow() + datetime.timedelta(seconds=ttl) diff --git a/tests/compute_engine/test_credentials.py b/tests/compute_engine/test_credentials.py index 4234d98e4..009e42848 100644 --- a/tests/compute_engine/test_credentials.py +++ b/tests/compute_engine/test_credentials.py @@ -43,6 +43,13 @@ b"bsxbLa6Fp0SYeYwO8ifEnkRvasVpc1WTQqfRB2JCj5pTBDzJpIpFCMmnQ" ) +ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/mds" +) +ID_TOKEN_REQUEST_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/mds" +) + class TestCredentials(object): credentials = None @@ -96,12 +103,16 @@ def test_refresh_success(self, get, utcnow): # expired) assert self.credentials.valid + @mock.patch( + "google.auth.metrics.token_request_access_token_mds", + return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, + ) @mock.patch( "google.auth._helpers.utcnow", return_value=datetime.datetime.min + _helpers.REFRESH_THRESHOLD, ) @mock.patch("google.auth.compute_engine._metadata.get", autospec=True) - def test_refresh_success_with_scopes(self, get, utcnow): + def test_refresh_success_with_scopes(self, get, utcnow, mock_metrics_header_value): get.side_effect = [ { # First request is for sevice account info. @@ -133,7 +144,10 @@ def test_refresh_success_with_scopes(self, get, utcnow): assert self.credentials.valid kwargs = get.call_args[1] - assert kwargs == {"params": {"scopes": "three,four"}} + assert kwargs["params"] == {"scopes": "three,four"} + assert kwargs["headers"] == { + "x-goog-api-client": ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE + } @mock.patch("google.auth.compute_engine._metadata.get", autospec=True) def test_refresh_error(self, get): @@ -732,11 +746,17 @@ def test_sign_bytes(self, sign, get): # The JWT token signature is 'signature' encoded in base 64: assert signature == b"signature" + @mock.patch( + "google.auth.metrics.token_request_id_token_mds", + return_value=ID_TOKEN_REQUEST_METRICS_HEADER_VALUE, + ) @mock.patch( "google.auth.compute_engine._metadata.get_service_account_info", autospec=True ) @mock.patch("google.auth.compute_engine._metadata.get", autospec=True) - def test_get_id_token_from_metadata(self, get, get_service_account_info): + def test_get_id_token_from_metadata( + self, get, get_service_account_info, mock_metrics_header_value + ): get.return_value = SAMPLE_ID_TOKEN get_service_account_info.return_value = {"email": "foo@example.com"} @@ -745,6 +765,10 @@ def test_get_id_token_from_metadata(self, get, get_service_account_info): ) cred.refresh(request=mock.Mock()) + assert get.call_args.kwargs["headers"] == { + "x-goog-api-client": ID_TOKEN_REQUEST_METRICS_HEADER_VALUE + } + assert cred.token == SAMPLE_ID_TOKEN assert cred.expiry == datetime.datetime.fromtimestamp(SAMPLE_ID_TOKEN_EXP) assert cred._use_metadata_identity_endpoint diff --git a/tests/oauth2/test__client.py b/tests/oauth2/test__client.py index b34d2eb35..9450fa1fd 100644 --- a/tests/oauth2/test__client.py +++ b/tests/oauth2/test__client.py @@ -46,6 +46,13 @@ " https://www.googleapis.com/auth/logging.write" ) +ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/sa" +) +ID_TOKEN_REQUEST_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/sa" +) + @pytest.mark.parametrize("retryable", [True, False]) def test__handle_error_response(retryable): @@ -483,47 +490,83 @@ def test_refresh_grant_no_access_token(): assert not excinfo.value.retryable +@mock.patch( + "google.auth.metrics.token_request_access_token_sa_assertion", + return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, +) @mock.patch("google.oauth2._client._parse_expiry", return_value=None) @mock.patch.object(_client, "_token_endpoint_request", autospec=True) -def test_jwt_grant_retry_default(mock_token_endpoint_request, mock_expiry): +def test_jwt_grant_retry_default( + mock_token_endpoint_request, mock_expiry, mock_metrics_header_value +): _client.jwt_grant(mock.Mock(), mock.Mock(), mock.Mock()) mock_token_endpoint_request.assert_called_with( - mock.ANY, mock.ANY, mock.ANY, can_retry=True + mock.ANY, + mock.ANY, + mock.ANY, + can_retry=True, + headers={"x-goog-api-client": ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE}, ) @pytest.mark.parametrize("can_retry", [True, False]) +@mock.patch( + "google.auth.metrics.token_request_access_token_sa_assertion", + return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, +) @mock.patch("google.oauth2._client._parse_expiry", return_value=None) @mock.patch.object(_client, "_token_endpoint_request", autospec=True) def test_jwt_grant_retry_with_retry( - mock_token_endpoint_request, mock_expiry, can_retry + mock_token_endpoint_request, mock_expiry, mock_metrics_header_value, can_retry ): _client.jwt_grant(mock.Mock(), mock.Mock(), mock.Mock(), can_retry=can_retry) mock_token_endpoint_request.assert_called_with( - mock.ANY, mock.ANY, mock.ANY, can_retry=can_retry + mock.ANY, + mock.ANY, + mock.ANY, + can_retry=can_retry, + headers={"x-goog-api-client": ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE}, ) +@mock.patch( + "google.auth.metrics.token_request_id_token_sa_assertion", + return_value=ID_TOKEN_REQUEST_METRICS_HEADER_VALUE, +) @mock.patch("google.auth.jwt.decode", return_value={"exp": 0}) @mock.patch.object(_client, "_token_endpoint_request", autospec=True) -def test_id_token_jwt_grant_retry_default(mock_token_endpoint_request, mock_jwt_decode): +def test_id_token_jwt_grant_retry_default( + mock_token_endpoint_request, mock_jwt_decode, mock_metrics_header_value +): _client.id_token_jwt_grant(mock.Mock(), mock.Mock(), mock.Mock()) mock_token_endpoint_request.assert_called_with( - mock.ANY, mock.ANY, mock.ANY, can_retry=True + mock.ANY, + mock.ANY, + mock.ANY, + can_retry=True, + headers={"x-goog-api-client": ID_TOKEN_REQUEST_METRICS_HEADER_VALUE}, ) @pytest.mark.parametrize("can_retry", [True, False]) +@mock.patch( + "google.auth.metrics.token_request_id_token_sa_assertion", + return_value=ID_TOKEN_REQUEST_METRICS_HEADER_VALUE, +) @mock.patch("google.auth.jwt.decode", return_value={"exp": 0}) @mock.patch.object(_client, "_token_endpoint_request", autospec=True) def test_id_token_jwt_grant_retry_with_retry( - mock_token_endpoint_request, mock_jwt_decode, can_retry + mock_token_endpoint_request, mock_jwt_decode, mock_metrics_header_value, can_retry ): _client.id_token_jwt_grant( mock.Mock(), mock.Mock(), mock.Mock(), can_retry=can_retry ) mock_token_endpoint_request.assert_called_with( - mock.ANY, mock.ANY, mock.ANY, can_retry=can_retry + mock.ANY, + mock.ANY, + mock.ANY, + can_retry=can_retry, + headers={"x-goog-api-client": ID_TOKEN_REQUEST_METRICS_HEADER_VALUE}, ) diff --git a/tests/oauth2/test_reauth.py b/tests/oauth2/test_reauth.py index 47aa8fa95..54f59422d 100644 --- a/tests/oauth2/test_reauth.py +++ b/tests/oauth2/test_reauth.py @@ -40,6 +40,12 @@ "encodedProofOfReauthToken": "new_rapt_token", } +REAUTH_START_METRICS_HEADER_VALUE = "gl-python/3.7 auth/1.1 auth-request-type/re-start" +REAUTH_CONTINUE_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/re-cont" +) +TOKEN_REQUEST_METRICS_HEADER_VALUE = "gl-python/3.7 auth/1.1 cred-type/u" + class MockChallenge(object): def __init__(self, name, locally_eligible, challenge_input): @@ -56,7 +62,10 @@ def test_is_interactive(): assert reauth.is_interactive() -def test__get_challenges(): +@mock.patch( + "google.auth.metrics.reauth_start", return_value=REAUTH_START_METRICS_HEADER_VALUE +) +def test__get_challenges(mock_metrics_header_value): with mock.patch( "google.oauth2._client._token_endpoint_request" ) as mock_token_endpoint_request: @@ -67,10 +76,14 @@ def test__get_challenges(): {"supportedChallengeTypes": ["SAML"]}, access_token="token", use_json=True, + headers={"x-goog-api-client": REAUTH_START_METRICS_HEADER_VALUE}, ) -def test__get_challenges_with_scopes(): +@mock.patch( + "google.auth.metrics.reauth_start", return_value=REAUTH_START_METRICS_HEADER_VALUE +) +def test__get_challenges_with_scopes(mock_metrics_header_value): with mock.patch( "google.oauth2._client._token_endpoint_request" ) as mock_token_endpoint_request: @@ -86,10 +99,15 @@ def test__get_challenges_with_scopes(): }, access_token="token", use_json=True, + headers={"x-goog-api-client": REAUTH_START_METRICS_HEADER_VALUE}, ) -def test__send_challenge_result(): +@mock.patch( + "google.auth.metrics.reauth_continue", + return_value=REAUTH_CONTINUE_METRICS_HEADER_VALUE, +) +def test__send_challenge_result(mock_metrics_header_value): with mock.patch( "google.oauth2._client._token_endpoint_request" ) as mock_token_endpoint_request: @@ -107,6 +125,7 @@ def test__send_challenge_result(): }, access_token="token", use_json=True, + headers={"x-goog-api-client": REAUTH_CONTINUE_METRICS_HEADER_VALUE}, ) @@ -270,7 +289,11 @@ def test_get_rapt_token(): ) -def test_refresh_grant_failed(): +@mock.patch( + "google.auth.metrics.token_request_user", + return_value=TOKEN_REQUEST_METRICS_HEADER_VALUE, +) +def test_refresh_grant_failed(mock_metrics_header_value): with mock.patch( "google.oauth2._client._token_endpoint_request_no_throw" ) as mock_token_request: @@ -299,6 +322,7 @@ def test_refresh_grant_failed(): "scope": "foo bar", "rapt": "rapt_token", }, + headers={"x-goog-api-client": TOKEN_REQUEST_METRICS_HEADER_VALUE}, ) diff --git a/tests/test_aws.py b/tests/test_aws.py index d50a8f4e1..45397d34f 100644 --- a/tests/test_aws.py +++ b/tests/test_aws.py @@ -28,6 +28,10 @@ from google.auth import transport +IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/imp" +) + CLIENT_ID = "username" CLIENT_SECRET = "password" # Base64 encoding of "username:password". @@ -1901,8 +1905,14 @@ def test_refresh_success_without_impersonation_use_default_scopes(self, utcnow): assert credentials.scopes is None assert credentials.default_scopes == SCOPES + @mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, + ) @mock.patch("google.auth._helpers.utcnow") - def test_refresh_success_with_impersonation_ignore_default_scopes(self, utcnow): + def test_refresh_success_with_impersonation_ignore_default_scopes( + self, utcnow, mock_metrics_header_value + ): utcnow.return_value = datetime.datetime.strptime( self.AWS_SIGNATURE_TIME, "%Y-%m-%dT%H:%M:%SZ" ) @@ -1937,6 +1947,7 @@ def test_refresh_success_with_impersonation_ignore_default_scopes(self, utcnow): "Content-Type": "application/json", "authorization": "Bearer {}".format(self.SUCCESS_RESPONSE["access_token"]), "x-goog-user-project": QUOTA_PROJECT_ID, + "x-goog-api-client": IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, } impersonation_request_data = { "delegates": None, @@ -1985,8 +1996,14 @@ def test_refresh_success_with_impersonation_ignore_default_scopes(self, utcnow): assert credentials.scopes == SCOPES assert credentials.default_scopes == ["ignored"] + @mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, + ) @mock.patch("google.auth._helpers.utcnow") - def test_refresh_success_with_impersonation_use_default_scopes(self, utcnow): + def test_refresh_success_with_impersonation_use_default_scopes( + self, utcnow, mock_metrics_header_value + ): utcnow.return_value = datetime.datetime.strptime( self.AWS_SIGNATURE_TIME, "%Y-%m-%dT%H:%M:%SZ" ) @@ -2021,6 +2038,7 @@ def test_refresh_success_with_impersonation_use_default_scopes(self, utcnow): "Content-Type": "application/json", "authorization": "Bearer {}".format(self.SUCCESS_RESPONSE["access_token"]), "x-goog-user-project": QUOTA_PROJECT_ID, + "x-goog-api-client": IMPERSONATE_ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, } impersonation_request_data = { "delegates": None, diff --git a/tests/test_external_account.py b/tests/test_external_account.py index 598c3760c..0e8017eee 100644 --- a/tests/test_external_account.py +++ b/tests/test_external_account.py @@ -26,6 +26,8 @@ from google.auth import transport +METRICS_HEADER_VALUE = "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/imp" + CLIENT_ID = "username" CLIENT_SECRET = "password" # Base64 encoding of "username:password" @@ -751,7 +753,13 @@ def test_refresh_workforce_with_client_auth_and_no_workforce_project_success( assert not credentials.expired assert credentials.token == response["access_token"] - def test_refresh_impersonation_without_client_auth_success(self): + @mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=METRICS_HEADER_VALUE, + ) + def test_refresh_impersonation_without_client_auth_success( + self, mock_metrics_header_value + ): # Simulate service account access token expires in 2800 seconds. expire_time = ( _helpers.utcnow().replace(microsecond=0) + datetime.timedelta(seconds=2800) @@ -776,6 +784,7 @@ def test_refresh_impersonation_without_client_auth_success(self): impersonation_headers = { "Content-Type": "application/json", "authorization": "Bearer {}".format(token_response["access_token"]), + "x-goog-api-client": METRICS_HEADER_VALUE, } impersonation_request_data = { "delegates": None, @@ -815,7 +824,13 @@ def test_refresh_impersonation_without_client_auth_success(self): assert not credentials.expired assert credentials.token == impersonation_response["accessToken"] - def test_refresh_workforce_impersonation_without_client_auth_success(self): + @mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=METRICS_HEADER_VALUE, + ) + def test_refresh_workforce_impersonation_without_client_auth_success( + self, mock_metrics_header_value + ): # Simulate service account access token expires in 2800 seconds. expire_time = ( _helpers.utcnow().replace(microsecond=0) + datetime.timedelta(seconds=2800) @@ -843,6 +858,7 @@ def test_refresh_workforce_impersonation_without_client_auth_success(self): impersonation_headers = { "Content-Type": "application/json", "authorization": "Bearer {}".format(token_response["access_token"]), + "x-goog-api-client": METRICS_HEADER_VALUE, } impersonation_request_data = { "delegates": None, @@ -1000,7 +1016,13 @@ def test_refresh_with_client_auth_success(self): assert not credentials.expired assert credentials.token == self.SUCCESS_RESPONSE["access_token"] - def test_refresh_impersonation_with_client_auth_success_ignore_default_scopes(self): + @mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=METRICS_HEADER_VALUE, + ) + def test_refresh_impersonation_with_client_auth_success_ignore_default_scopes( + self, mock_metrics_header_value + ): # Simulate service account access token expires in 2800 seconds. expire_time = ( _helpers.utcnow().replace(microsecond=0) + datetime.timedelta(seconds=2800) @@ -1028,6 +1050,7 @@ def test_refresh_impersonation_with_client_auth_success_ignore_default_scopes(se impersonation_headers = { "Content-Type": "application/json", "authorization": "Bearer {}".format(token_response["access_token"]), + "x-goog-api-client": METRICS_HEADER_VALUE, } impersonation_request_data = { "delegates": None, @@ -1071,7 +1094,13 @@ def test_refresh_impersonation_with_client_auth_success_ignore_default_scopes(se assert not credentials.expired assert credentials.token == impersonation_response["accessToken"] - def test_refresh_impersonation_with_client_auth_success_use_default_scopes(self): + @mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=METRICS_HEADER_VALUE, + ) + def test_refresh_impersonation_with_client_auth_success_use_default_scopes( + self, mock_metrics_header_value + ): # Simulate service account access token expires in 2800 seconds. expire_time = ( _helpers.utcnow().replace(microsecond=0) + datetime.timedelta(seconds=2800) @@ -1099,6 +1128,7 @@ def test_refresh_impersonation_with_client_auth_success_use_default_scopes(self) impersonation_headers = { "Content-Type": "application/json", "authorization": "Bearer {}".format(token_response["access_token"]), + "x-goog-api-client": METRICS_HEADER_VALUE, } impersonation_request_data = { "delegates": None, @@ -1488,7 +1518,13 @@ def test_project_id_without_scopes(self): assert credentials.get_project_id(None) is None - def test_get_project_id_cloud_resource_manager_success(self): + @mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=METRICS_HEADER_VALUE, + ) + def test_get_project_id_cloud_resource_manager_success( + self, mock_metrics_header_value + ): # STS token exchange request/response. token_response = self.SUCCESS_RESPONSE.copy() token_headers = {"Content-Type": "application/x-www-form-urlencoded"} @@ -1513,6 +1549,7 @@ def test_get_project_id_cloud_resource_manager_success(self): "Content-Type": "application/json", "x-goog-user-project": self.QUOTA_PROJECT_ID, "authorization": "Bearer {}".format(token_response["access_token"]), + "x-goog-api-client": METRICS_HEADER_VALUE, } impersonation_request_data = { "delegates": None, @@ -1638,7 +1675,11 @@ def test_workforce_pool_get_project_id_cloud_resource_manager_success(self): # No additional requests. assert len(request.call_args_list) == 2 - def test_refresh_impersonation_with_lifetime(self): + @mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=METRICS_HEADER_VALUE, + ) + def test_refresh_impersonation_with_lifetime(self, mock_metrics_header_value): # Simulate service account access token expires in 2800 seconds. expire_time = ( _helpers.utcnow().replace(microsecond=0) + datetime.timedelta(seconds=2800) @@ -1663,6 +1704,7 @@ def test_refresh_impersonation_with_lifetime(self): impersonation_headers = { "Content-Type": "application/json", "authorization": "Bearer {}".format(token_response["access_token"]), + "x-goog-api-client": METRICS_HEADER_VALUE, } impersonation_request_data = { "delegates": None, diff --git a/tests/test_identity_pool.py b/tests/test_identity_pool.py index 6c1c6623c..561af35ac 100644 --- a/tests/test_identity_pool.py +++ b/tests/test_identity_pool.py @@ -286,6 +286,9 @@ def assert_underlying_credentials_refresh( json.dumps({"userProject": workforce_pool_user_project}) ) + metrics_header_value = ( + "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/imp" + ) if service_account_impersonation_url: # Service account impersonation request/response. expire_time = ( @@ -299,6 +302,7 @@ def assert_underlying_credentials_refresh( impersonation_headers = { "Content-Type": "application/json", "authorization": "Bearer {}".format(token_response["access_token"]), + "x-goog-api-client": metrics_header_value, } impersonation_request_data = { "delegates": None, @@ -321,7 +325,11 @@ def assert_underlying_credentials_refresh( request = cls.make_mock_request(*[el for req in requests for el in req]) - credentials.refresh(request) + with mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=metrics_header_value, + ): + credentials.refresh(request) assert len(request.call_args_list) == len(requests) if credential_data: diff --git a/tests/test_impersonated_credentials.py b/tests/test_impersonated_credentials.py index 8f2e1fdbe..dc091fe61 100644 --- a/tests/test_impersonated_credentials.py +++ b/tests/test_impersonated_credentials.py @@ -55,6 +55,13 @@ SIGNER = crypt.RSASigner.from_string(PRIVATE_KEY_BYTES, "1") TOKEN_URI = "https://example.com/oauth2/token" +ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/at cred-type/imp" +) +ID_TOKEN_REQUEST_METRICS_HEADER_VALUE = ( + "gl-python/3.7 auth/1.1 auth-request-type/it cred-type/imp" +) + @pytest.fixture def mock_donor_credentials(): @@ -188,10 +195,18 @@ def test_refresh_success(self, use_data_bytes, mock_donor_credentials): use_data_bytes=use_data_bytes, ) - credentials.refresh(request) + with mock.patch( + "google.auth.metrics.token_request_access_token_impersonate", + return_value=ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE, + ): + credentials.refresh(request) assert credentials.valid assert not credentials.expired + assert ( + request.call_args.kwargs["headers"]["x-goog-api-client"] + == ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE + ) @pytest.mark.parametrize("use_data_bytes", [True, False]) def test_refresh_success_iam_endpoint_override( @@ -454,6 +469,36 @@ def test_id_token_success( assert id_creds.token == ID_TOKEN_DATA assert id_creds.expiry == datetime.datetime.fromtimestamp(ID_TOKEN_EXPIRY) + def test_id_token_metrics(self, mock_donor_credentials): + credentials = self.make_credentials(lifetime=None) + credentials.token = "token" + credentials.expiry = None + target_audience = "https://foo.bar" + + id_creds = impersonated_credentials.IDTokenCredentials( + credentials, target_audience=target_audience + ) + + with mock.patch( + "google.auth.metrics.token_request_id_token_impersonate", + return_value=ID_TOKEN_REQUEST_METRICS_HEADER_VALUE, + ): + with mock.patch( + "google.auth.transport.requests.AuthorizedSession.post", autospec=True + ) as mock_post: + data = {"token": ID_TOKEN_DATA} + mock_post.return_value = MockResponse(data, http_client.OK) + id_creds.refresh(None) + + assert id_creds.token == ID_TOKEN_DATA + assert id_creds.expiry == datetime.datetime.fromtimestamp( + ID_TOKEN_EXPIRY + ) + assert ( + mock_post.call_args.kwargs["headers"]["x-goog-api-client"] + == ID_TOKEN_REQUEST_METRICS_HEADER_VALUE + ) + def test_id_token_from_credential( self, mock_donor_credentials, mock_authorizedsession_idtoken ):