diff --git a/aws_xray_sdk/core/plugins/ec2_plugin.py b/aws_xray_sdk/core/plugins/ec2_plugin.py index 71d5f196..c5aea567 100644 --- a/aws_xray_sdk/core/plugins/ec2_plugin.py +++ b/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') diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 268a4fc6..ca3e73c3 100644 --- a/tests/test_plugins.py +++ b/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 diff --git a/tox.ini b/tox.ini index 8e2d496a..66e34730 100644 --- a/tox.ini +++ b/tox.ini @@ -32,6 +32,7 @@ deps = testing.postgresql testing.mysqld webtest + mock # Python2 only deps py{27}: enum34