diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py index e5abf2097..b35775a23 100644 --- a/google/auth/compute_engine/_metadata.py +++ b/google/auth/compute_engine/_metadata.py @@ -26,15 +26,18 @@ from six.moves.urllib import parse as urlparse from google.auth import _helpers +from google.auth import environment_vars from google.auth import exceptions _LOGGER = logging.getLogger(__name__) -_METADATA_ROOT = 'http://metadata.google.internal/computeMetadata/v1/' +_METADATA_ROOT = 'http://{}/computeMetadata/v1/'.format( + os.getenv(environment_vars.GCE_METADATA_ROOT, 'metadata.google.internal')) # This is used to ping the metadata server, it avoids the cost of a DNS # lookup. -_METADATA_IP_ROOT = 'http://169.254.169.254' +_METADATA_IP_ROOT = 'http://{}'.format( + os.getenv(environment_vars.GCE_METADATA_IP, '169.254.169.254')) _METADATA_FLAVOR_HEADER = 'metadata-flavor' _METADATA_FLAVOR_VALUE = 'Google' _METADATA_HEADERS = {_METADATA_FLAVOR_HEADER: _METADATA_FLAVOR_VALUE} diff --git a/google/auth/environment_vars.py b/google/auth/environment_vars.py index b4ed2b28a..0110e6a3c 100644 --- a/google/auth/environment_vars.py +++ b/google/auth/environment_vars.py @@ -37,3 +37,13 @@ CLOUD_SDK_CONFIG_DIR = 'CLOUDSDK_CONFIG' """Environment variable defines the location of Google Cloud SDK's config files.""" + +# These two variables allow for customization of the addresses used when +# contacting the GCE metadata service. +GCE_METADATA_ROOT = 'GCE_METADATA_ROOT' +"""Environment variable providing an alternate hostname or host:port to be +used for GCE metadata requests.""" + +GCE_METADATA_IP = 'GCE_METADATA_IP' +"""Environment variable providing an alternate ip:port to be used for ip-only +GCE metadata requests.""" diff --git a/tests/compute_engine/test__metadata.py b/tests/compute_engine/test__metadata.py index 6a45c426d..0adf88252 100644 --- a/tests/compute_engine/test__metadata.py +++ b/tests/compute_engine/test__metadata.py @@ -14,12 +14,15 @@ import datetime import json +import os import mock import pytest from six.moves import http_client +from six.moves import reload_module from google.auth import _helpers +from google.auth import environment_vars from google.auth import exceptions from google.auth.compute_engine import _metadata @@ -67,6 +70,26 @@ def test_ping_failure_connection_failed(mock_request): assert not _metadata.ping(request_mock) +def test_ping_success_custom_root(mock_request): + request_mock = mock_request('', headers=_metadata._METADATA_HEADERS) + + fake_ip = '1.2.3.4' + os.environ[environment_vars.GCE_METADATA_IP] = fake_ip + reload_module(_metadata) + + try: + assert _metadata.ping(request_mock) + finally: + del os.environ[environment_vars.GCE_METADATA_IP] + reload_module(_metadata) + + request_mock.assert_called_once_with( + method='GET', + url='http://' + fake_ip, + headers=_metadata._METADATA_HEADERS, + timeout=_metadata._METADATA_DEFAULT_TIMEOUT) + + def test_get_success_json(mock_request): key, value = 'foo', 'bar' @@ -96,6 +119,26 @@ def test_get_success_text(mock_request): assert result == data +def test_get_success_custom_root(mock_request): + request_mock = mock_request( + '{}', headers={'content-type': 'application/json'}) + + fake_root = 'another.metadata.service' + os.environ[environment_vars.GCE_METADATA_ROOT] = fake_root + reload_module(_metadata) + + try: + _metadata.get(request_mock, PATH) + finally: + del os.environ[environment_vars.GCE_METADATA_ROOT] + reload_module(_metadata) + + request_mock.assert_called_once_with( + method='GET', + url='http://{}/computeMetadata/v1/{}'.format(fake_root, PATH), + headers=_metadata._METADATA_HEADERS) + + def test_get_failure(mock_request): request_mock = mock_request( 'Metadata error', status=http_client.NOT_FOUND)