Add retry and timeout as params to retrieving metadata service

botocore/credentials.py

@@ -30,6 +30,15 @@
 
 
 logger = logging.getLogger(__name__)
+DEFAULT_METADATA_SERVICE_TIMEOUT = 1
+METADATA_SECURITY_CREDENTIALS_URL = (
+    'http://169.254.169.254/latest/meta-data/iam/security-credentials/'
+)
+
+
+class _RetriesExceededError(Exception):
+    """Internal exception used when the number of retries are exceeded."""
+    pass
 
 
 class Credentials(object):
@@ -53,34 +62,54 @@ def __init__(self, access_key=None, secret_key=None, token=None):
         self.profiles = []
 
 
-def _search_md(url='http://169.254.169.254/latest/meta-data/iam/security-credentials/'):
+def retrieve_iam_role_credentials(url=METADATA_SECURITY_CREDENTIALS_URL,
+                                  timeout=None, num_retries=0):
+    if timeout is None:
+        timeout = DEFAULT_METADATA_SERVICE_TIMEOUT
     d = {}
     try:
-        r = requests.get(url, timeout=.1)
-        if r.status_code == 200 and r.content:
+        r = _get_request(url, timeout, num_retries)
+        if r.content:
             fields = r.content.decode('utf-8').split('\n')
             for field in fields:
                 if field.endswith('/'):
-                    d[field[0:-1]] = _search_md(url + field)
+                    d[field[0:-1]] = retrieve_iam_role_credentials(
+                        url + field, timeout, num_retries)
                 else:
-                    val = requests.get(url + field).content.decode('utf-8')
+                    val = _get_request(
+                        url + field,
+                        timeout=timeout,
+                        num_retries=num_retries).content.decode('utf-8')
                     if val[0] == '{':
                         val = json.loads(val)
-                    else:
-                        p = val.find('\n')
-                        if p > 0:
-                            val = r.content.decode('utf-8').split('\n')
                     d[field] = val
-    except (requests.Timeout, requests.ConnectionError):
-        pass
+        else:
+            logger.debug("Metadata service returned non 200 status code "
+                         "of %s for url: %s, content body: %s",
+                         r.status_code, url, r.content)
+    except _RetriesExceededError:
+        logger.debug("Max number of retries exceeded (%s) when "
+                     "attempting to retrieve data from metadata service.",
+                     num_retries)
     return d
 
 
+def _get_request(url, timeout, num_retries):
+    for i in range(num_retries + 1):
+        try:
+            response = requests.get(url, timeout=timeout)
+        except (requests.Timeout, requests.ConnectionError) as e:
+            logger.debug("Caught exception wil trying to retrieve credentials "
+                        "from metadata service: %s", e, exc_info=True)
+        else:
+            if response.status_code == 200:
+                return response
+    raise _RetriesExceededError()
+
+
 def search_iam_role(**kwargs):
     credentials = None
-    metadata = kwargs.get('metadata', None)
-    if metadata is None:
-        metadata = _search_md()
+    metadata = retrieve_iam_role_credentials()
     if metadata:
         for role_name in metadata:
             credentials = Credentials(metadata[role_name]['AccessKeyId'],
@@ -183,11 +212,10 @@ def search_boto_config(**kwargs):
                        ('iam-role', search_iam_role))
 
 
-def get_credentials(session, metadata=None):
+def get_credentials(session):
     credentials = None
     for cred_method, cred_fn in _credential_methods:
-        credentials = cred_fn(session=session,
-                              metadata=metadata)
+        credentials = cred_fn(session=session)
         if credentials:
             break
     return credentials

botocore/session.py

@@ -310,24 +310,17 @@ def set_credentials(self, access_key, secret_key, token=None):
                                                              token)
         self._credentials.method = 'explicit'
 
-    def get_credentials(self, metadata=None):
+    def get_credentials(self):
         """
         Return the :class:`botocore.credential.Credential` object
         associated with this session.  If the credentials have not
         yet been loaded, this will attempt to load them.  If they
         have already been loaded, this will return the cached
         credentials.
 
-        :type metadata: dict
-        :param metadata: This parameter allows you to pass in
-            EC2 instance metadata containing IAM Role credentials.
-            This metadata will be used rather than retrieving the
-            metadata from the metadata service.  This is mainly used
-            for unit testing.
         """
         if self._credentials is None:
-            self._credentials = botocore.credentials.get_credentials(self,
-                                                                     metadata)
+            self._credentials = botocore.credentials.get_credentials(self)
         return self._credentials
 
     def user_agent(self):

tests/unit/test_credentials.py

@@ -29,6 +29,8 @@
 
 import botocore.session
 import botocore.exceptions
+from botocore import credentials
+from botocore.vendored.requests import ConnectionError
 
 # Passed to session to keep it from finding default config file
 TESTENVVARS = {'config_file': (None, 'AWS_CONFIG_FILE', None)}
@@ -76,14 +78,11 @@ def test_credentials_file(self):
         self.assertEqual(credentials.secret_key, 'bar')
         self.assertEqual(credentials.method, 'credentials-file')
 
-    def test_bad_file(self):
+    @mock.patch('botocore.vendored.requests.get')
+    def test_bad_file(self, get):
         self.environ['AWS_CREDENTIAL_FILE'] = path('no_aws_credentials')
         credentials = self.session.get_credentials()
-        # There can be two cases, either credentials are None,
-        # or if we're able to get an IAM role the credentials method
-        # is not 'credentials-file'
-        if credentials is not None:
-            self.assertNotEqual(credentials.method, 'credentials-file')
+        self.assertIsNone(credentials)
 
 
 class ConfigTest(BaseEnvVar):
@@ -134,24 +133,27 @@ class IamRoleTest(BaseEnvVar):
     def setUp(self):
         super(IamRoleTest, self).setUp()
         self.session = botocore.session.get_session(env_vars=TESTENVVARS)
-
-    def test_iam_role(self):
         self.environ['BOTO_CONFIG'] = ''
-        credentials = self.session.get_credentials(metadata=metadata)
+
+    @mock.patch('botocore.credentials.retrieve_iam_role_credentials')
+    def test_iam_role(self, retriever):
+        retriever.return_value = metadata
+        credentials = self.session.get_credentials()
         self.assertEqual(credentials.method, 'iam-role')
         self.assertEqual(credentials.access_key, 'foo')
         self.assertEqual(credentials.secret_key, 'bar')
 
-    def test_empty_boto_config_is_ignored(self):
+    @mock.patch('botocore.credentials.retrieve_iam_role_credentials')
+    def test_empty_boto_config_is_ignored(self, retriever):
+        retriever.return_value = metadata
         self.environ['BOTO_CONFIG'] = path('boto_config_empty')
-        credentials = self.session.get_credentials(metadata=metadata)
+        credentials = self.session.get_credentials()
         self.assertEqual(credentials.method, 'iam-role')
         self.assertEqual(credentials.access_key, 'foo')
         self.assertEqual(credentials.secret_key, 'bar')
 
     @mock.patch('botocore.vendored.requests.get')
     def test_get_credentials_with_metadata_mock(self, get):
-        self.environ['BOTO_CONFIG'] = ''
         first = mock.Mock()
         first.status_code = 200
         first.content = 'foobar'.encode('utf-8')
@@ -164,6 +166,49 @@ def test_get_credentials_with_metadata_mock(self, get):
         credentials = self.session.get_credentials()
         self.assertEqual(credentials.method, 'iam-role')
 
+    @mock.patch('botocore.vendored.requests.get')
+    def test_timeout_argument_forwarded_to_requests(self, get):
+        first = mock.Mock()
+        first.status_code = 200
+        first.content = 'foobar'.encode('utf-8')
+
+        second = mock.Mock()
+        second.status_code = 200
+        second.content = json.dumps(metadata['foobar']).encode('utf-8')
+        get.side_effect = [first, second]
+
+        credentials.retrieve_iam_role_credentials(timeout=10)
+        self.assertEqual(get.call_args[1]['timeout'], 10)
+
+    @mock.patch('botocore.vendored.requests.get')
+    def test_request_timeout_occurs(self, get):
+        first = mock.Mock()
+        first.side_effect = ConnectionError
+
+        d = credentials.retrieve_iam_role_credentials(timeout=10)
+        self.assertEqual(d, {})
+
+    @mock.patch('botocore.vendored.requests.get')
+    def test_retry_errors(self, get):
+        # First attempt we get a connection error.
+        first = mock.Mock()
+        first.side_effect = ConnectionError
+
+        # Next attempt we get a response with the foobar key.
+        second = mock.Mock()
+        second.status_code = 200
+        second.content = 'foobar'.encode('utf-8')
+
+        # Next attempt we get a response with the foobar creds.
+        third = mock.Mock()
+        third.status_code = 200
+        third.content = json.dumps(metadata['foobar']).encode('utf-8')
+        get.side_effect = [first, second, third]
+
+        retrieved = credentials.retrieve_iam_role_credentials(
+            num_retries=1)
+        self.assertEqual(retrieved['foobar']['AccessKeyId'], 'foo')
+
 
 if __name__ == "__main__":
     unittest.main()