From 3148a6854b69ce1319b1f6cc884840641a4fb2a6 Mon Sep 17 00:00:00 2001 From: srprash <50466688+srprash@users.noreply.github.com> Date: Wed, 10 Jun 2020 18:17:29 +0530 Subject: [PATCH] Added support for IMDSv2 (#48) * Added support for IMDSv2. using net/http for fetching metadata. Updated tests * Refactored to more modular code. Added 1 second timeout and a retry attempt --- lib/aws-xray-sdk/facets/net_http.rb | 10 +++- lib/aws-xray-sdk/plugins/ec2.rb | 85 ++++++++++++++++++++++++----- test/aws-xray-sdk/tc_plugin.rb | 63 +++++++++++++++++---- test/aws-xray-sdk/tc_recorder.rb | 12 ++-- 4 files changed, 138 insertions(+), 32 deletions(-) diff --git a/lib/aws-xray-sdk/facets/net_http.rb b/lib/aws-xray-sdk/facets/net_http.rb index 1dd5803..8966de2 100644 --- a/lib/aws-xray-sdk/facets/net_http.rb +++ b/lib/aws-xray-sdk/facets/net_http.rb @@ -30,9 +30,15 @@ def xray_sampling_request?(req) req.path && (req.path == ('/GetSamplingRules') || req.path == ('/SamplingTargets')) end + # Instance Metadata Service provides endpoint 169.254.169.254 to + # provide EC2 metadata + def ec2_metadata_request?(req) + req.uri && req.uri.hostname == '169.254.169.254' + end + def request(req, body = nil, &block) - # Do not trace requests to xray or aws lambda runtime - if xray_sampling_request?(req) || lambda_runtime_request? + # Do not trace requests to xray or aws lambda runtime or ec2 metadata endpoint + if xray_sampling_request?(req) || lambda_runtime_request? || ec2_metadata_request?(req) return super end diff --git a/lib/aws-xray-sdk/plugins/ec2.rb b/lib/aws-xray-sdk/plugins/ec2.rb index bfb6cb2..1d995f5 100644 --- a/lib/aws-xray-sdk/plugins/ec2.rb +++ b/lib/aws-xray-sdk/plugins/ec2.rb @@ -1,36 +1,91 @@ -require 'open-uri' +require 'net/http' +require 'json' require 'aws-xray-sdk/logger' module XRay module Plugins - # A plugin that gets the EC2 instance-id and AZ if running on an EC2 instance. + # A plugin that gets the EC2 instance_id, availabiity_zone, instance_type, and ami_id if running on an EC2 instance. module EC2 include Logging ORIGIN = 'AWS::EC2::Instance'.freeze + # http://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html#instancedata-data-retrieval - ID_ADDR = 'http://169.254.169.254/latest/meta-data/instance-id'.freeze - AZ_ADDR = 'http://169.254.169.254/latest/meta-data/placement/availability-zone'.freeze + METADATA_BASE_URL = 'http://169.254.169.254/latest'.freeze def self.aws - @@aws ||= begin - instance_id = open(ID_ADDR, open_timeout: 1).read - az = open(AZ_ADDR, open_timeout: 1).read - { - ec2: { - instance_id: instance_id, - availability_zone: az - } + @@aws = {} + token = get_token + ec2_metadata = get_metadata(token) + @@aws = { + ec2: ec2_metadata + } + end + + + private # private methods + + def self.get_token + token_uri = URI(METADATA_BASE_URL + '/api/token') + + req = Net::HTTP::Put.new(token_uri) + req['X-aws-ec2-metadata-token-ttl-seconds'] = '60' + begin + return do_request(req) + rescue StandardError => e + Logging.logger.warn %(can not get the IMDSv2 token due to: #{e.message}.) + '' + end + end + + def self.get_metadata(token) + metadata_uri = URI(METADATA_BASE_URL + '/dynamic/instance-identity/document') + + req = Net::HTTP::Get.new(metadata_uri) + if token != '' + req['X-aws-ec2-metadata-token'] = token + end + + begin + metadata_json = do_request(req) + return parse_metadata(metadata_json) + rescue StandardError => e + Logging.logger.warn %(can not get the ec2 instance metadata due to: #{e.message}.) + {} + end + end + + def self.parse_metadata(json_str) + metadata = {} + data = JSON(json_str) + metadata['instance_id'] = data['instanceId'] + metadata['availability_zone'] = data['availabilityZone'] + metadata['instance_type'] = data['instanceType'] + metadata['ami_id'] = data['imageId'] + + metadata + end + + def self.do_request(request) + begin + response = Net::HTTP.start(request.uri.hostname, read_timeout: 1) { |http| + http.request(request) } + + if response.code == '200' + return response.body + else + raise(StandardError.new('Unsuccessful response::' + response.code + '::' + response.message)) + end rescue StandardError => e - # Two attempts in total to get EC2 metadata + # Two attempts in total to complete the request successfully @retries ||= 0 if @retries < 1 @retries += 1 retry else - @@aws = {} - Logging.logger.warn %(can not get the ec2 instance metadata due to: #{e.message}.) + Logging.logger.warn %(Failed to complete request due to: #{e.message}.) + raise e end end end diff --git a/test/aws-xray-sdk/tc_plugin.rb b/test/aws-xray-sdk/tc_plugin.rb index 0e349a6..a9becb3 100644 --- a/test/aws-xray-sdk/tc_plugin.rb +++ b/test/aws-xray-sdk/tc_plugin.rb @@ -16,27 +16,68 @@ def test_origin_all_set # when running on any machine. def test_get_runtime_context XRay::Plugins::ElasticBeanstalk.aws - stub_request(:any, XRay::Plugins::EC2::ID_ADDR).to_raise(StandardError) - stub_request(:any, XRay::Plugins::EC2::AZ_ADDR).to_raise(StandardError) + stub_request(:any, XRay::Plugins::EC2::METADATA_BASE_URL + '/api/token') + .to_raise(StandardError) + stub_request(:any, XRay::Plugins::EC2::METADATA_BASE_URL + '/dynamic/instance-identity/document') + .to_raise(StandardError) XRay::Plugins::EC2.aws XRay::Plugins::ECS.aws WebMock.reset! end - def test_mocked_ec2_metadata - instance_id = "abc" - az = "us-east-1a" - stub_request(:any, XRay::Plugins::EC2::ID_ADDR) - .to_return(body: instance_id, status: 200) - stub_request(:any, XRay::Plugins::EC2::AZ_ADDR) - .to_return(body: az, status: 200) + def test_ec2_metadata_v2_successful + dummy_json = '{\"availabilityZone\" : \"us-east-2a\", \"imageId\" : \"ami-03cca83dd001d4666\", + \"instanceId\" : \"i-07a181803de94c666\", \"instanceType\" : \"t3.xlarge\"}' + + stub_request(:put, XRay::Plugins::EC2::METADATA_BASE_URL + '/api/token') + .to_return(status: 200, body: 'some_token', headers: {}) + + stub_request(:get, XRay::Plugins::EC2::METADATA_BASE_URL + '/dynamic/instance-identity/document') + .to_return(status: 200, body: dummy_json, headers: {}) + expected = { ec2: { - instance_id: instance_id, - avaliablity_zone: az + instance_id: 'i-07a181803de94c666', + availability_zone: 'us-east-2a', + instance_type: 't3.xlarge', + ami_id: 'ami-03cca83dd001d4666' } } assert expected, XRay::Plugins::EC2.aws WebMock.reset! end + + def test_ec2_metadata_v1_successful + dummy_json = '{\"availabilityZone\" : \"cn-north-1a\", \"imageId\" : \"ami-03cca83dd001d4111\", + \"instanceId\" : \"i-07a181803de94c111\", \"instanceType\" : \"t2.xlarge\"}' + + stub_request(:put, XRay::Plugins::EC2::METADATA_BASE_URL + '/api/token') + .to_raise(StandardError) + + stub_request(:get, XRay::Plugins::EC2::METADATA_BASE_URL + '/dynamic/instance-identity/document') + .to_return(status: 200, body: dummy_json, headers: {}) + + expected = { + ec2: { + instance_id: 'i-07a181803de94c111', + availability_zone: 'cn-north-1a', + instance_type: 't2.xlarge', + ami_id: 'ami-03cca83dd001d4111' + } + } + assert expected, XRay::Plugins::EC2.aws + WebMock.reset! + end + + def test_ec2_metadata_fail + stub_request(:put, XRay::Plugins::EC2::METADATA_BASE_URL + '/api/token') + .to_raise(StandardError) + + stub_request(:get, XRay::Plugins::EC2::METADATA_BASE_URL + '/dynamic/instance-identity/document') + .to_raise(StandardError) + + expected = {} + assert expected, XRay::Plugins::EC2.aws + WebMock.reset! + end end diff --git a/test/aws-xray-sdk/tc_recorder.rb b/test/aws-xray-sdk/tc_recorder.rb index 3ece24f..d6b349b 100644 --- a/test/aws-xray-sdk/tc_recorder.rb +++ b/test/aws-xray-sdk/tc_recorder.rb @@ -159,10 +159,14 @@ def test_xray_metadata end def test_plugins_runtime_context - stub_request(:any, XRay::Plugins::EC2::ID_ADDR) - .to_return(body: 'some_id', status: 200) - stub_request(:any, XRay::Plugins::EC2::AZ_ADDR) - .to_return(body: 'some_az', status: 200) + dummy_json = '{\"availabilityZone\" : \"us-east-2a\", \"imageId\" : \"ami-03cca83dd001d4666\", + \"instanceId\" : \"i-07a181803de94c666\", \"instanceType\" : \"t3.xlarge\"}' + + stub_request(:put, 'http://169.254.169.254/latest/api/token') + .to_return(status: 200, body: 'some_token', headers: {}) + + stub_request(:get, 'http://169.254.169.254/latest/dynamic/instance-identity/document') + .to_return(status: 200, body: dummy_json, headers: {}) recorder = XRay::Recorder.new config = {