From 36e8b01e82ab0ea3a7bf98e385561747896016f8 Mon Sep 17 00:00:00 2001 From: Danny Hermes Date: Thu, 27 Aug 2015 20:43:38 -0700 Subject: [PATCH] Using basic httplib call to check if on GCE. This avoids the potential issue of urllib detecting a proxy and changing the actual request underneath our code. The basic httplib library does not have any proxy detection, so this is not even possible there. Fixes #114. --- oauth2client/client.py | 53 +++++++-------- tests/test_client.py | 147 +++++++++++++++++++++-------------------- 2 files changed, 103 insertions(+), 97 deletions(-) diff --git a/oauth2client/client.py b/oauth2client/client.py index 3582990ec..8485c57eb 100644 --- a/oauth2client/client.py +++ b/oauth2client/client.py @@ -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.""" @@ -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(): @@ -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 diff --git a/tests/test_client.py b/tests/test_client.py index d59a55202..3a9bd26eb 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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' @@ -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(