diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py index c1fecd3b7..cef574bcc 100644 --- a/google/auth/compute_engine/_metadata.py +++ b/google/auth/compute_engine/_metadata.py @@ -96,10 +96,7 @@ def _get_metadata_ip_root(use_mtls: bool): # Timeout in seconds to wait for the GCE metadata server when detecting the # GCE environment. -try: - _METADATA_DEFAULT_TIMEOUT = int(os.getenv(environment_vars.GCE_METADATA_TIMEOUT, 3)) -except ValueError: # pragma: NO COVER - _METADATA_DEFAULT_TIMEOUT = 3 +_METADATA_PING_DEFAULT_TIMEOUT = 3 # The number of tries to perform when waiting for the GCE metadata server # when detecting the GCE environment. @@ -209,6 +206,16 @@ def ping( # could lead to false negatives in the event that we are on GCE, but # the metadata resolution was particularly slow. The latter case is # "unlikely". + + if timeout is None: + try: + timeout = float(os.getenv( + environment_vars.GCE_METADATA_TIMEOUT, + str(_METADATA_PING_DEFAULT_TIMEOUT))) + except ValueError: + timeout = _METADATA_PING_DEFAULT_TIMEOUT + + retries = 0 headers = _METADATA_HEADERS.copy() headers[metrics.API_CLIENT_HEADER] = metrics.mds_ping() diff --git a/google/auth/environment_vars.py b/google/auth/environment_vars.py index ca041ca16..d4623f8ea 100644 --- a/google/auth/environment_vars.py +++ b/google/auth/environment_vars.py @@ -94,6 +94,9 @@ Used to distinguish between GAE gen1 and GAE gen2+. """ +GCE_METADATA_TIMEOUT = "GCE_METADATA_TIMEOUT" +"""Environment variable for setting timeouts in seconds for metadata queries.""" + # AWS environment variables used with AWS workload identity pools to retrieve # AWS security credentials and the AWS region needed to create a serialized # signed requests to the AWS STS GetCalledIdentity API that can be exchanged diff --git a/tests/compute_engine/test__metadata.py b/tests/compute_engine/test__metadata.py index 7bd7bc6ed..18c6c794c 100644 --- a/tests/compute_engine/test__metadata.py +++ b/tests/compute_engine/test__metadata.py @@ -143,9 +143,46 @@ def test_ping_success(mock_metrics_header_value): request.assert_called_once_with( method="GET", - url="http://169.254.169.254", - headers=MDS_PING_REQUEST_HEADER, - timeout=_metadata._METADATA_DEFAULT_TIMEOUT, + url=_metadata._METADATA_IP_ROOT, + headers=_metadata._METADATA_HEADER, + timeout=_metadata._METADATA_PING_DEFAULT_TIMEOUT, + ) + +@mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE) +def test_ping_success_with_gce_metadata_timeout(mock_metrics_header_value): + request = make_request("", headers=_metadata._METADATA_HEADERS) + gce_metadata_timeout = .5 + os.environ[ + environment_vars.GCE_METADATA_TIMEOUT] = str(gce_metadata_timeout) + + try: + assert _metadata.ping(request) + finally: + del os.environ[environment_vars.GCE_METADATA_TIMEOUT] + + request.assert_called_once_with( + method="GET", + url=_metadata._METADATA_IP_ROOT, + headers=_metadata._METADATA_HEADER, + timeout=gce_metadata_timeout, + ) + +@mock.patch("google.auth.metrics.mds_ping", return_value=MDS_PING_METRICS_HEADER_VALUE) +def test_ping_success_with_invalid_gce_metadata_timeout(mock_metrics_header_value): + request = make_request("", headers=_metadata._METADATA_HEADERS) + os.environ[ + environment_vars.GCE_METADATA_TIMEOUT] = "Not a valid float value!" + + try: + assert _metadata.ping(request) + finally: + del os.environ[environment_vars.GCE_METADATA_TIMEOUT] + + request.assert_called_once_with( + method="GET", + url=_metadata._METADATA_IP_ROOT, + headers=_metadata._METADATA_HEADERS, + timeout=_metadata._METADATA_PING_DEFAULT_TIMEOUT, # Fallback value. ) @@ -197,7 +234,7 @@ def test_ping_success_custom_root(mock_metrics_header_value): method="GET", url="http://" + fake_ip, headers=MDS_PING_REQUEST_HEADER, - timeout=_metadata._METADATA_DEFAULT_TIMEOUT, + timeout=_metadata._METADATA_PING_DEFAULT_TIMEOUT, )