Allow customizing the GCE metadata service address via an env var. (#148)

The goal here is to make it possible for a user of a binary that depends on
this library (eg the google cloud SDK) to be able to customize where it looks
for the GCE metadata service. (An adventurous user can already customize the
GCE metadata service location via the existing global vars in this library.)

Author: craigcitro

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}

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."""

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)