Skip to content
This repository has been archived by the owner on Nov 5, 2019. It is now read-only.

Commit

Permalink
Merge pull request #293 from dhermes/fix-114
Browse files Browse the repository at this point in the history 
Use basic httplib call to check if on GCE.
  • Loading branch information
@nathanielmanistaatgoogle
nathanielmanistaatgoogle committed Aug 31, 2015
2 parents 6e3c0c5 + 36e8b01 commit 5dda032
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 97 deletions.
53 changes: 27 additions & 26 deletions oauth2client/client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,11 @@
# 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."""
Expand Down Expand Up @@ -1053,34 +1058,34 @@ def _revoke(self, http_request):
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():
Expand All @@ -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
Expand All @@ -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
Expand Down
147 changes: 76 additions & 71 deletions tests/test_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -225,27 +225,18 @@ def test_create_scoped(self):

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'
Expand All @@ -256,68 +247,82 @@ def test_environment_caching(self):
# 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 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)
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

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)

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)
self._environment_check_gce_helper(socket_error=True)

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(
Expand Down

0 comments on commit 5dda032

Please sign in to comment.