Merge pull request #293 from dhermes/fix-114

Use basic httplib call to check if on GCE.
This commit is contained in:
Nathaniel Manista
2015-08-31 15:37:42 -07:00
2 changed files with 101 additions and 95 deletions

View File

@@ -110,6 +110,11 @@ DEFAULT_ENV_NAME = 'UNKNOWN'
# If set to True _get_environment avoid GCE check (_detect_gce_environment)
NO_GCE_CHECK = os.environ.setdefault('NO_GCE_CHECK', 'False')
_SERVER_SOFTWARE = 'SERVER_SOFTWARE'
_GCE_METADATA_HOST = '169.254.169.254'
_METADATA_FLAVOR_HEADER = 'Metadata-Flavor'
_DESIRED_METADATA_FLAVOR = 'Google'
class SETTINGS(object):
"""Settings namespace for globally defined values."""
@@ -1053,34 +1058,34 @@ class AccessTokenCredentials(OAuth2Credentials):
self._do_revoke(http_request, self.access_token)
def _detect_gce_environment(urlopen=None):
def _detect_gce_environment():
"""Determine if the current environment is Compute Engine.
Args:
urlopen: Optional argument. Function used to open a connection to a
URL.
Returns:
Boolean indicating whether or not the current environment is Google
Compute Engine.
"""
urlopen = urlopen or urllib.request.urlopen
# Note: the explicit `timeout` below is a workaround. The underlying
# issue is that resolving an unknown host on some networks will take
# 20-30 seconds; making this timeout short fixes the issue, but
# 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".
# NOTE: The explicit ``timeout`` is a workaround. The underlying
# issue is that resolving an unknown host on some networks will take
# 20-30 seconds; making this timeout short fixes the issue, but
# 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".
connection = six.moves.http_client.HTTPConnection(
_GCE_METADATA_HOST, timeout=1)
try:
response = urlopen('http://169.254.169.254/', timeout=1)
return response.info().get('Metadata-Flavor', '') == 'Google'
except socket.timeout:
headers = {_METADATA_FLAVOR_HEADER: _DESIRED_METADATA_FLAVOR}
connection.request('GET', '/', headers=headers)
response = connection.getresponse()
if response.status == 200:
return (response.getheader(_METADATA_FLAVOR_HEADER) ==
_DESIRED_METADATA_FLAVOR)
except socket.error: # socket.timeout or socket.error(64, 'Host is down')
logger.info('Timeout attempting to reach GCE metadata service.')
return False
except urllib.error.URLError as e:
if isinstance(getattr(e, 'reason', None), socket.timeout):
logger.info('Timeout attempting to reach GCE metadata service.')
return False
finally:
connection.close()
def _in_gae_environment():
@@ -1094,7 +1099,7 @@ def _in_gae_environment():
try:
import google.appengine
server_software = os.environ.get('SERVER_SOFTWARE', '')
server_software = os.environ.get(_SERVER_SOFTWARE, '')
if server_software.startswith('Google App Engine/'):
SETTINGS.env_name = 'GAE_PRODUCTION'
return True
@@ -1107,20 +1112,16 @@ def _in_gae_environment():
return False
def _in_gce_environment(urlopen=None):
def _in_gce_environment():
"""Detect if the code is running in the Compute Engine environment.
Args:
urlopen: Optional argument. Function used to open a connection to a
URL.
Returns:
True if running in the GCE environment, False otherwise.
"""
if SETTINGS.env_name is not None:
return SETTINGS.env_name == 'GCE_PRODUCTION'
if NO_GCE_CHECK != 'True' and _detect_gce_environment(urlopen=urlopen):
if NO_GCE_CHECK != 'True' and _detect_gce_environment():
SETTINGS.env_name = 'GCE_PRODUCTION'
return True
return False

View File

@@ -225,27 +225,18 @@ class GoogleCredentialsTests(unittest.TestCase):
def test_environment_check_gae_production(self):
with mock_module_import('google.appengine'):
os.environ['SERVER_SOFTWARE'] = 'Google App Engine/XYZ'
self.assertTrue(_in_gae_environment())
self.assertFalse(_in_gce_environment())
self._environment_check_gce_helper(
server_software='Google App Engine/XYZ')
def test_environment_check_gae_local(self):
with mock_module_import('google.appengine'):
os.environ['SERVER_SOFTWARE'] = 'Development/XYZ'
self.assertTrue(_in_gae_environment())
self.assertFalse(_in_gce_environment())
self._environment_check_gce_helper(
server_software='Development/XYZ')
def test_environment_check_fastpath(self):
os.environ['SERVER_SOFTWARE'] = 'Development/XYZ'
with mock_module_import('google.appengine'):
with mock.patch.object(urllib.request, 'urlopen',
return_value=MockResponse({}),
autospec=True) as urlopen:
self.assertTrue(_in_gae_environment())
self.assertFalse(_in_gce_environment())
# We already know are in GAE, so we shouldn't actually do
# the urlopen.
self.assertFalse(urlopen.called)
self._environment_check_gce_helper(
server_software='Development/XYZ')
def test_environment_caching(self):
os.environ['SERVER_SOFTWARE'] = 'Development/XYZ'
@@ -256,68 +247,82 @@ class GoogleCredentialsTests(unittest.TestCase):
# is cached.
self.assertTrue(_in_gae_environment())
def test_environment_check_gae_module_on_gce(self):
with mock_module_import('google.appengine'):
os.environ['SERVER_SOFTWARE'] = ''
response = MockResponse({'Metadata-Flavor': 'Google'})
with mock.patch.object(urllib.request, 'urlopen',
return_value=response,
autospec=True) as urlopen:
self.assertFalse(_in_gae_environment())
self.assertTrue(_in_gce_environment())
urlopen.assert_called_once_with(
'http://169.254.169.254/', timeout=1)
def _environment_check_gce_helper(self, status_ok=True, socket_error=False,
server_software=''):
response = mock.MagicMock()
if status_ok:
response.status = 200
response.getheader = mock.MagicMock(
name='getheader',
return_value=client._DESIRED_METADATA_FLAVOR)
else:
response.status = 404
def test_environment_check_gae_module_unknown(self):
with mock_module_import('google.appengine'):
os.environ['SERVER_SOFTWARE'] = ''
with mock.patch.object(urllib.request, 'urlopen',
return_value=MockResponse({}),
autospec=True) as urlopen:
self.assertFalse(_in_gae_environment())
self.assertFalse(_in_gce_environment())
urlopen.assert_called_once_with(
'http://169.254.169.254/', timeout=1)
connection = mock.MagicMock()
connection.getresponse = mock.MagicMock(name='getresponse',
return_value=response)
if socket_error:
connection.getresponse.side_effect = socket.error()
with mock.patch('oauth2client.client.os') as os_module:
os_module.environ = {client._SERVER_SOFTWARE: server_software}
with mock.patch('oauth2client.client.six') as six_module:
http_client = six_module.moves.http_client
http_client.HTTPConnection = mock.MagicMock(
name='HTTPConnection', return_value=connection)
if server_software == '':
self.assertFalse(_in_gae_environment())
else:
self.assertTrue(_in_gae_environment())
if status_ok and not socket_error and server_software == '':
self.assertTrue(_in_gce_environment())
else:
self.assertFalse(_in_gce_environment())
if server_software == '':
http_client.HTTPConnection.assert_called_once_with(
client._GCE_METADATA_HOST, timeout=1)
connection.getresponse.assert_called_once_with()
# Remaining calls are not "getresponse"
headers = {
client._METADATA_FLAVOR_HEADER: (
client._DESIRED_METADATA_FLAVOR),
}
self.assertEqual(connection.method_calls, [
mock.call.request('GET', '/',
headers=headers),
mock.call.close(),
])
self.assertEqual(response.method_calls, [])
if status_ok and not socket_error:
response.getheader.assert_called_once_with(
client._METADATA_FLAVOR_HEADER)
else:
self.assertEqual(http_client.HTTPConnection.mock_calls, [])
self.assertEqual(connection.getresponse.mock_calls, [])
# Remaining calls are not "getresponse"
self.assertEqual(connection.method_calls, [])
self.assertEqual(response.method_calls, [])
self.assertEqual(response.getheader.mock_calls, [])
def test_environment_check_gce_production(self):
os.environ['SERVER_SOFTWARE'] = ''
response = MockResponse({'Metadata-Flavor': 'Google'})
with mock.patch.object(urllib.request, 'urlopen',
return_value=response,
autospec=True) as urlopen:
self.assertFalse(_in_gae_environment())
self.assertTrue(_in_gce_environment())
urlopen.assert_called_once_with(
'http://169.254.169.254/', timeout=1)
self._environment_check_gce_helper(status_ok=True)
def test_environment_check_gce_prod_with_working_gae_imports(self):
with mock_module_import('google.appengine'):
self._environment_check_gce_helper(status_ok=True)
def test_environment_check_gce_timeout(self):
os.environ['SERVER_SOFTWARE'] = ''
response = MockResponse({'Metadata-Flavor': 'Google'})
with mock.patch.object(urllib.request, 'urlopen',
return_value=response,
autospec=True) as urlopen:
urlopen.side_effect = socket.timeout()
self.assertFalse(_in_gce_environment())
urlopen.assert_called_once_with(
'http://169.254.169.254/', timeout=1)
self._environment_check_gce_helper(socket_error=True)
with mock.patch.object(urllib.request, 'urlopen',
return_value=response,
autospec=True) as urlopen:
urlopen.side_effect = urllib.error.URLError(socket.timeout())
self.assertFalse(_in_gce_environment())
urlopen.assert_called_once_with(
'http://169.254.169.254/', timeout=1)
def test_environ_check_gae_module_unknown(self):
with mock_module_import('google.appengine'):
self._environment_check_gce_helper(status_ok=False)
def test_environment_check_unknown(self):
os.environ['SERVER_SOFTWARE'] = ''
with mock.patch.object(urllib.request, 'urlopen',
return_value=MockResponse({}),
autospec=True) as urlopen:
self.assertFalse(_in_gce_environment())
self.assertFalse(_in_gae_environment())
urlopen.assert_called_once_with(
'http://169.254.169.254/', timeout=1)
self._environment_check_gce_helper(status_ok=False)
def test_get_environment_variable_file(self):
environment_variable_file = datafile(