IMDSv2 support for EC2 plugin

aws_xray_sdk/core/plugins/ec2_plugin.py

@@ -1,12 +1,14 @@
 import logging
 from future.standard_library import install_aliases
+from urllib.request import urlopen, Request
+
 install_aliases()
-from urllib.request import urlopen
 
 log = logging.getLogger(__name__)
 
 SERVICE_NAME = 'ec2'
 ORIGIN = 'AWS::EC2::Instance'
+IMDS_URL = 'http://169.254.169.254/latest/'
 
 
 def initialize():
@@ -17,16 +19,51 @@ def initialize():
     """
     global runtime_context
 
+    # Try the IMDSv2 endpoint for metadata
     try:
         runtime_context = {}
 
-        r = urlopen('http://169.254.169.254/latest/meta-data/instance-id', timeout=1)
-        runtime_context['instance_id'] = r.read().decode('utf-8')
+        # get session token with 60 seconds TTL to not have the token lying around for a long time
+        token = do_request(url=IMDS_URL + "api/token",
+                           headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"},
+                           method="PUT")
+
+        # get instance-id metadata
+        runtime_context['instance_id'] = do_request(url=IMDS_URL + "meta-data/instance-id",
+                                                    headers={"X-aws-ec2-metadata-token": token},
+                                                    method="GET")
+
+        # get availability-zone metadata
+        runtime_context['availability_zone'] = do_request(url=IMDS_URL + "meta-data/placement/availability-zone",
+                                                          headers={"X-aws-ec2-metadata-token": token},
+                                                          method="GET")
+
+    except Exception as e:
+        # Falling back to IMDSv1 endpoint
+        log.debug("failed to get ec2 instance metadata from IMDSv2 due to {}. Falling back to IMDSv1".format(e))
+
+        try:
+            runtime_context = {}
+
+            runtime_context['instance_id'] = do_request(url=IMDS_URL + "meta-data/instance-id")
+
+            runtime_context['availability_zone'] = do_request(url=IMDS_URL + "meta-data/placement/availability-zone-1")
+
+        except Exception as e:
+            runtime_context = None
+            log.debug("failed to get ec2 instance metadata from IMDSv1 due to {}".format(e))
+            log.warning("Failed to get ec2 instance metadata")
+
 
-        r = urlopen('http://169.254.169.254/latest/meta-data/placement/availability-zone',
-                    timeout=1)
-        runtime_context['availability_zone'] = r.read().decode('utf-8')
+def do_request(url, headers=None, method="GET"):
+    if headers is None:
+        headers = {}
 
-    except Exception:
-        runtime_context = None
-        log.warning("failed to get ec2 instance metadata.")
+    if url is None:
+        return None
+
+    req = Request(url=url)
+    req.headers = headers
+    req.method = method
+    res = urlopen(req, timeout=1)
+    return res.read().decode('utf-8')

tests/test_plugins.py

@@ -1,4 +1,5 @@
 from aws_xray_sdk.core.plugins.utils import get_plugin_modules
+from mock import patch
 
 supported_plugins = (
     'ec2_plugin',
@@ -13,3 +14,38 @@ def test_runtime_context_available():
     for plugin in plugins:
         plugin.initialize()
         assert hasattr(plugin, 'runtime_context')
+
+
+@patch('aws_xray_sdk.core.plugins.ec2_plugin.do_request')
+def test_ec2_plugin_imdsv2_success(mock_do_request):
+    mock_do_request.side_effect = ['token', 'i-0a1d026d92d4709cd', 'us-west-2b']
+
+    ec2_plugin = get_plugin_modules(('ec2_plugin',))[0]
+    ec2_plugin.initialize()
+    assert hasattr(ec2_plugin, 'runtime_context')
+    r_c = getattr(ec2_plugin, 'runtime_context')
+    assert r_c['instance_id'] == 'i-0a1d026d92d4709cd'
+    assert r_c['availability_zone'] == 'us-west-2b'
+
+
+@patch('aws_xray_sdk.core.plugins.ec2_plugin.do_request')
+def test_ec2_plugin_v2_fail_v1_success(mock_do_request):
+    mock_do_request.side_effect = [Exception("Boom!"), 'i-0a1d026d92d4709ab', 'us-west-2a']
+
+    ec2_plugin = get_plugin_modules(('ec2_plugin',))[0]
+    ec2_plugin.initialize()
+    assert hasattr(ec2_plugin, 'runtime_context')
+    r_c = getattr(ec2_plugin, 'runtime_context')
+    assert r_c['instance_id'] == 'i-0a1d026d92d4709ab'
+    assert r_c['availability_zone'] == 'us-west-2a'
+
+
+@patch('aws_xray_sdk.core.plugins.ec2_plugin.do_request')
+def test_ec2_plugin_v2_fail_v1_fail(mock_do_request):
+    mock_do_request.side_effect = [Exception("Boom v2!"), Exception("Boom v1!")]
+
+    ec2_plugin = get_plugin_modules(('ec2_plugin',))[0]
+    ec2_plugin.initialize()
+    assert hasattr(ec2_plugin, 'runtime_context')
+    r_c = getattr(ec2_plugin, 'runtime_context')
+    assert r_c is None

tox.ini

@@ -32,6 +32,7 @@ deps =
     testing.postgresql
     testing.mysqld
     webtest
+    mock
 
     # Python2 only deps
     py{27}: enum34