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

@@ -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
 

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

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

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 = {