ec2: Add support for AWS IMDS v2 (session-oriented) (#55)

* ec2: Add support for AWS IMDS v2 (session-oriented)

AWS now supports a new version of fetching Instance Metadata[1].

Update cloud-init's ec2 utility functions and update ec2 derived
datasources accordingly.  For DataSourceEc2 (versus ec2-look-alikes)
cloud-init will issue the PUT request to obtain an API token for
the maximum lifetime and then all subsequent interactions with the
IMDS will include the token in the header.

If the API token endpoint is unreachable on Ec2 platform, log a
warning and fallback to using IMDS v1 and which does not use
session tokens when communicating with the Instance metadata
service. 

We handle read errors, typically seen if the IMDS is beyond one 
etwork hop (IMDSv2 responses have a ttl=1), by setting the api token
to a disabled value and then using IMDSv1 paths.

To support token-based headers, ec2_utils functions were updated
to support custom headers_cb and exception_cb callback functions
so Ec2 could store, or refresh API tokens in the event of token
becoming stale.

[1] https://docs.aws.amazon.com/AWSEC2/latest/ \
UserGuide/ec2-instance-metadata.html \
#instance-metadata-v2-how-it-works

2 files changed, 21 insertions, 6 deletions

diff --git a/tests/unittests/test_datasource/test_cloudstack.py b/tests/unittests/test_datasource/test_cloudstack.py
index d6d2d6b2..83c2f753 100644
--- a/tests/unittests/test_datasource/test_cloudstack.py
+++ b/tests/unittests/test_datasource/test_cloudstack.py
@@ -10,6 +10,9 @@ from cloudinit.tests.helpers import CiTestCase, ExitStack, mock
 import os
 import time
 
+MOD_PATH = 'cloudinit.sources.DataSourceCloudStack'
+DS_PATH = MOD_PATH + '.DataSourceCloudStack'
+
 
 class TestCloudStackPasswordFetching(CiTestCase):
 
@@ -17,7 +20,7 @@ class TestCloudStackPasswordFetching(CiTestCase):
         super(TestCloudStackPasswordFetching, self).setUp()
         self.patches = ExitStack()
         self.addCleanup(self.patches.close)
-        mod_name = 'cloudinit.sources.DataSourceCloudStack'
+        mod_name = MOD_PATH
         self.patches.enter_context(mock.patch('{0}.ec2'.format(mod_name)))
         self.patches.enter_context(mock.patch('{0}.uhelp'.format(mod_name)))
         default_gw = "192.201.20.0"
@@ -56,7 +59,9 @@ class TestCloudStackPasswordFetching(CiTestCase):
         ds.get_data()
         self.assertEqual({}, ds.get_config_obj())
 
-    def test_password_sets_password(self):
+    @mock.patch(DS_PATH + '.wait_for_metadata_service')
+    def test_password_sets_password(self, m_wait):
+        m_wait.return_value = True
         password = 'SekritSquirrel'
         self._set_password_server_response(password)
         ds = DataSourceCloudStack(
@@ -64,7 +69,9 @@ class TestCloudStackPasswordFetching(CiTestCase):
         ds.get_data()
         self.assertEqual(password, ds.get_config_obj()['password'])
 
-    def test_bad_request_doesnt_stop_ds_from_working(self):
+    @mock.patch(DS_PATH + '.wait_for_metadata_service')
+    def test_bad_request_doesnt_stop_ds_from_working(self, m_wait):
+        m_wait.return_value = True
         self._set_password_server_response('bad_request')
         ds = DataSourceCloudStack(
             {}, None, helpers.Paths({'run_dir': self.tmp}))
@@ -79,7 +86,9 @@ class TestCloudStackPasswordFetching(CiTestCase):
                     request_types.append(arg.split()[1])
         self.assertEqual(expected_request_types, request_types)
 
-    def test_valid_response_means_password_marked_as_saved(self):
+    @mock.patch(DS_PATH + '.wait_for_metadata_service')
+    def test_valid_response_means_password_marked_as_saved(self, m_wait):
+        m_wait.return_value = True
         password = 'SekritSquirrel'
         subp = self._set_password_server_response(password)
         ds = DataSourceCloudStack(
@@ -92,7 +101,9 @@ class TestCloudStackPasswordFetching(CiTestCase):
         subp = self._set_password_server_response(response_string)
         ds = DataSourceCloudStack(
             {}, None, helpers.Paths({'run_dir': self.tmp}))
-        ds.get_data()
+        with mock.patch(DS_PATH + '.wait_for_metadata_service') as m_wait:
+            m_wait.return_value = True
+            ds.get_data()
         self.assertRequestTypesSent(subp, ['send_my_password'])
 
     def test_password_not_saved_if_empty(self):
diff --git a/tests/unittests/test_datasource/test_ec2.py b/tests/unittests/test_datasource/test_ec2.py
index 5e1dd777..34a089f2 100644
--- a/tests/unittests/test_datasource/test_ec2.py
+++ b/tests/unittests/test_datasource/test_ec2.py
@@ -191,7 +191,9 @@ def register_mock_metaserver(base_url, data):
             register(base_url, 'not found', status=404)
 
     def myreg(*argc, **kwargs):
-        return httpretty.register_uri(httpretty.GET, *argc, **kwargs)
+        url = argc[0]
+        method = httpretty.PUT if ec2.API_TOKEN_ROUTE in url else httpretty.GET
+        return httpretty.register_uri(method, *argc, **kwargs)
 
     register_helper(myreg, base_url, data)
 
@@ -237,6 +239,8 @@ class TestEc2(test_helpers.HttprettyTestCase):
         if md:
             all_versions = (
                 [ds.min_metadata_version] + ds.extended_metadata_versions)
+            token_url = self.data_url('latest', data_item='api/token')
+            register_mock_metaserver(token_url, 'API-TOKEN')
             for version in all_versions:
                 metadata_url = self.data_url(version) + '/'
                 if version == md_version: