From 9fc1b4a63a248e3d5183c1695cea37786c2e842c Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 21 May 2020 00:42:02 +0530 Subject: [PATCH 1/5] Implemented support for IMDSv2 and fallback to IMDSv1. Added unit tests for the same --- aws_xray_sdk/core/plugins/ec2_plugin.py | 53 +++++++++++++++++++++---- tests/test_plugins.py | 32 +++++++++++++++ 2 files changed, 77 insertions(+), 8 deletions(-) diff --git a/aws_xray_sdk/core/plugins/ec2_plugin.py b/aws_xray_sdk/core/plugins/ec2_plugin.py index 71d5f196..7f3c0fa4 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 + token = do_request(url=IMDS_URL + "api/token", + headers={"X-aws-ec2-metadata-token-ttl-seconds": "21600"}, + 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.warning("failed to get ec2 instance metadata from IMDSv2 due to {}. Falling back to IMDSv1".format(e)) + + try: + runtime_context = {} - r = urlopen('http://169.254.169.254/latest/meta-data/placement/availability-zone', - timeout=1) - runtime_context['availability_zone'] = r.read().decode('utf-8') + 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.warning("failed to get ec2 instance metadata from IMDSv1 due to {}".format(e)) + + +def do_request(url, headers=None, method="GET"): + if headers is None: + headers = {} + try: + 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') except Exception: - runtime_context = None - log.warning("failed to get ec2 instance metadata.") + raise diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 268a4fc6..e4ce4250 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 unittest.mock import patch supported_plugins = ( 'ec2_plugin', @@ -13,3 +14,34 @@ def test_runtime_context_available(): for plugin in plugins: plugin.initialize() assert hasattr(plugin, 'runtime_context') + + +@patch('urllib.request.urlopen') +def test_ec2_plugin(mock_urlopen): + # set up mock response + mock_urlopen.return_value.read.return_value.decode.side_effect = ['token', 'i-0a1d026d92d4709cd', 'us-west-2b', + Exception("Boom!"), 'i-0a1d026d92d4709ab', + 'us-west-2a', + Exception("Boom v2!"), Exception("Boom v1!")] + + ec2_plugin = get_plugin_modules(('ec2_plugin',)) + for plugin in ec2_plugin: + # for IMDSv2 success + plugin.initialize() + assert hasattr(plugin, 'runtime_context') + r_c = getattr(plugin, 'runtime_context') + assert r_c['instance_id'] == 'i-0a1d026d92d4709cd' + assert r_c['availability_zone'] == 'us-west-2b' + + # for IMDSv2 fail and IMDSv1 success + plugin.initialize() + assert hasattr(plugin, 'runtime_context') + r_c = getattr(plugin, 'runtime_context') + assert r_c['instance_id'] == 'i-0a1d026d92d4709ab' + assert r_c['availability_zone'] == 'us-west-2a' + + # for both failure + plugin.initialize() + assert hasattr(plugin, 'runtime_context') + r_c = getattr(plugin, 'runtime_context') + assert r_c is None From d085955820ca98477fc355e4dc32f7dd1d980e94 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 21 May 2020 01:33:04 +0530 Subject: [PATCH 2/5] Patching function instead of urlopen. Using mock since its compatible with python27 --- tests/test_plugins.py | 15 +++++++-------- tox.ini | 1 + 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/test_plugins.py b/tests/test_plugins.py index e4ce4250..55a15f6a 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -1,5 +1,5 @@ from aws_xray_sdk.core.plugins.utils import get_plugin_modules -from unittest.mock import patch +from mock import patch supported_plugins = ( 'ec2_plugin', @@ -16,13 +16,12 @@ def test_runtime_context_available(): assert hasattr(plugin, 'runtime_context') -@patch('urllib.request.urlopen') -def test_ec2_plugin(mock_urlopen): - # set up mock response - mock_urlopen.return_value.read.return_value.decode.side_effect = ['token', 'i-0a1d026d92d4709cd', 'us-west-2b', - Exception("Boom!"), 'i-0a1d026d92d4709ab', - 'us-west-2a', - Exception("Boom v2!"), Exception("Boom v1!")] +@patch('aws_xray_sdk.core.plugins.ec2_plugin.do_request') +def test_ec2_plugin(mock_do_request): + mock_do_request.side_effect = ['token', 'i-0a1d026d92d4709cd', 'us-west-2b', + Exception("Boom!"), 'i-0a1d026d92d4709ab', + 'us-west-2a', + Exception("Boom v2!"), Exception("Boom v1!")] ec2_plugin = get_plugin_modules(('ec2_plugin',)) for plugin in ec2_plugin: 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 From e614152b74995e6028ab1d2e3804cd12611f130e Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Thu, 21 May 2020 02:02:20 +0530 Subject: [PATCH 3/5] Changed the token ttl to 60 seconds to be safe --- aws_xray_sdk/core/plugins/ec2_plugin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aws_xray_sdk/core/plugins/ec2_plugin.py b/aws_xray_sdk/core/plugins/ec2_plugin.py index 7f3c0fa4..495ff9f7 100644 --- a/aws_xray_sdk/core/plugins/ec2_plugin.py +++ b/aws_xray_sdk/core/plugins/ec2_plugin.py @@ -23,9 +23,9 @@ def initialize(): try: runtime_context = {} - # get session token + # 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": "21600"}, + headers={"X-aws-ec2-metadata-token-ttl-seconds": "60"}, method="PUT") # get instance-id metadata From b79cdeb6dc92b27fd381028a0abcbe2974ae3e66 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Fri, 22 May 2020 01:30:07 +0530 Subject: [PATCH 4/5] Some refactoring --- aws_xray_sdk/core/plugins/ec2_plugin.py | 24 +++++----- tests/test_plugins.py | 58 ++++++++++++++----------- 2 files changed, 45 insertions(+), 37 deletions(-) diff --git a/aws_xray_sdk/core/plugins/ec2_plugin.py b/aws_xray_sdk/core/plugins/ec2_plugin.py index 495ff9f7..c5aea567 100644 --- a/aws_xray_sdk/core/plugins/ec2_plugin.py +++ b/aws_xray_sdk/core/plugins/ec2_plugin.py @@ -40,7 +40,7 @@ def initialize(): except Exception as e: # Falling back to IMDSv1 endpoint - log.warning("failed to get ec2 instance metadata from IMDSv2 due to {}. Falling back to IMDSv1".format(e)) + log.debug("failed to get ec2 instance metadata from IMDSv2 due to {}. Falling back to IMDSv1".format(e)) try: runtime_context = {} @@ -51,19 +51,19 @@ def initialize(): except Exception as e: runtime_context = None - log.warning("failed to get ec2 instance metadata from IMDSv1 due to {}".format(e)) + log.debug("failed to get ec2 instance metadata from IMDSv1 due to {}".format(e)) + log.warning("Failed to get ec2 instance metadata") def do_request(url, headers=None, method="GET"): if headers is None: headers = {} - try: - 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') - except Exception: - raise + + 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 55a15f6a..3a59610c 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -17,30 +17,38 @@ def test_runtime_context_available(): @patch('aws_xray_sdk.core.plugins.ec2_plugin.do_request') -def test_ec2_plugin(mock_do_request): - mock_do_request.side_effect = ['token', 'i-0a1d026d92d4709cd', 'us-west-2b', - Exception("Boom!"), 'i-0a1d026d92d4709ab', - 'us-west-2a', - Exception("Boom v2!"), Exception("Boom v1!")] - - ec2_plugin = get_plugin_modules(('ec2_plugin',)) - for plugin in ec2_plugin: - # for IMDSv2 success - plugin.initialize() - assert hasattr(plugin, 'runtime_context') - r_c = getattr(plugin, 'runtime_context') - assert r_c['instance_id'] == 'i-0a1d026d92d4709cd' - assert r_c['availability_zone'] == 'us-west-2b' +def test_ec2_plugin_imdsv2_success(mock_do_request): + mock_do_request.side_effect = ['token', 'i-0a1d026d92d4709cd', 'us-west-2b'] - # for IMDSv2 fail and IMDSv1 success - plugin.initialize() - assert hasattr(plugin, 'runtime_context') - r_c = getattr(plugin, 'runtime_context') - assert r_c['instance_id'] == 'i-0a1d026d92d4709ab' - assert r_c['availability_zone'] == 'us-west-2a' + ec2_plugin = get_plugin_modules(('ec2_plugin',))[0] + # for IMDSv2 success + 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' - # for both failure - plugin.initialize() - assert hasattr(plugin, 'runtime_context') - r_c = getattr(plugin, 'runtime_context') - assert r_c is None + +@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] + # for IMDSv2 success + 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] + # for IMDSv2 success + ec2_plugin.initialize() + assert hasattr(ec2_plugin, 'runtime_context') + r_c = getattr(ec2_plugin, 'runtime_context') + assert r_c is None From 09dec8ad5c7c3ed7acbb4c42094980072ebe9b19 Mon Sep 17 00:00:00 2001 From: Prashant Srivastava Date: Fri, 22 May 2020 01:33:55 +0530 Subject: [PATCH 5/5] Removed unnecessary comments from tests --- tests/test_plugins.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 3a59610c..ca3e73c3 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -21,7 +21,6 @@ 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] - # for IMDSv2 success ec2_plugin.initialize() assert hasattr(ec2_plugin, 'runtime_context') r_c = getattr(ec2_plugin, 'runtime_context') @@ -34,7 +33,6 @@ 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] - # for IMDSv2 success ec2_plugin.initialize() assert hasattr(ec2_plugin, 'runtime_context') r_c = getattr(ec2_plugin, 'runtime_context') @@ -47,7 +45,6 @@ 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] - # for IMDSv2 success ec2_plugin.initialize() assert hasattr(ec2_plugin, 'runtime_context') r_c = getattr(ec2_plugin, 'runtime_context')