Merge pull request #2964 from newrelic/resurrect_aws_accoun_id_decode

Resurrect aws account id decoding

lib/new_relic/agent/aws.rb

@@ -5,13 +5,64 @@
 module NewRelic
   module Agent
     module Aws
-      def self.create_arn(service, resource, region)
-        return unless NewRelic::Agent.config[:'cloud.aws.account_id']
+      CHARACTERS = %w[A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 2 3 4 5 6 7].freeze
+      HEX_MASK = '7fffffffff80'
 
-        "arn:aws:#{service}:#{region}:#{NewRelic::Agent.config[:'cloud.aws.account_id']}:#{resource}"
+      def self.create_arn(service, resource, region, account_id)
+        # if any of the values are nil, we can't create an ARN
+        return unless service && resource && region && account_id
+
+        "arn:aws:#{service}:#{region}:#{account_id}:#{resource}"
       rescue => e
         NewRelic::Agent.logger.warn("Failed to create ARN: #{e}")
       end
+
+      def self.get_account_id(config)
+        # if it is set in the agent config, use that first
+        return NewRelic::Agent.config[:'cloud.aws.account_id'] if NewRelic::Agent.config[:'cloud.aws.account_id']
+
+        access_key_id = config.credentials.credentials.access_key_id if config&.credentials&.credentials&.respond_to?(:access_key_id)
+        return unless access_key_id
+
+        NewRelic::Agent::Aws.convert_access_key_to_account_id(access_key_id)
+      rescue => e
+        NewRelic::Agent.logger.debug("Failed to create account id: #{e}")
+      end
+
+      def self.convert_access_key_to_account_id(access_key)
+        decoded_key = Integer(decode_to_hex(access_key[4..-1]), 16)
+        mask = Integer(HEX_MASK, 16)
+        (decoded_key & mask) >> 7
+      end
+
+      def self.decode_to_hex(access_key)
+        bytes = access_key.delete('=').each_char.map { |c| CHARACTERS.index(c) }
+
+        bytes.each_slice(8).map do |section|
+          convert_section(section)
+        end.flatten[0...6].join
+      end
+
+      def self.convert_section(section)
+        buffer = 0
+        section.each do |chunk|
+          buffer = (buffer << 5) + chunk
+        end
+
+        chunk_count = (section.length * 5.0 / 8.0).floor
+
+        if section.length < 8
+          buffer >>= (5 - (chunk_count * 8)) % 5
+        end
+
+        decoded = []
+        chunk_count.times do |i|
+          shift = 8 * (chunk_count - 1 - i)
+          decoded << ((buffer >> shift) & 255).to_s(16)
+        end
+
+        decoded
+      end
     end
   end
 end

lib/new_relic/agent/instrumentation/aws_sdk_lambda/instrumentation.rb

@@ -63,10 +63,9 @@ def process_function_error(response)
     def generate_segment(action, options = {})
       function = function_name(options)
       region = aws_region
-      account_id = aws_account_id
       arn = aws_arn(function, region)
       segment = NewRelic::Agent::Tracer.start_segment(name: "Lambda/#{action}/#{function}")
-      segment.add_agent_attribute('cloud.account.id', account_id)
+      segment.add_agent_attribute('cloud.account.id', nr_account_id)
       segment.add_agent_attribute('cloud.platform', CLOUD_PLATFORM)
       segment.add_agent_attribute('cloud.region', region)
       segment.add_agent_attribute('cloud.resource_id', arn) if arn
@@ -77,18 +76,18 @@ def function_name(options = {})
       (options.fetch(:function_name, nil) if options.respond_to?(:fetch)) || NewRelic::UNKNOWN
     end
 
-    def aws_account_id
-      return unless self.respond_to?(:config)
-
-      config&.account_id || NewRelic::Agent.config[:'cloud.aws.account_id']
-    end
-
     def aws_region
       config&.region if self.respond_to?(:config)
     end
 
     def aws_arn(function, region)
-      NewRelic::Agent::Aws.create_arn(AWS_SERVICE, "function:#{function}", region)
+      NewRelic::Agent::Aws.create_arn(AWS_SERVICE, "function:#{function}", region, nr_account_id)
+    end
+
+    def nr_account_id
+      return @nr_account_id if defined?(@nr_account_id)
+
+      @nr_account_id = NewRelic::Agent::Aws.get_account_id(config)
     end
   end
 end

lib/new_relic/agent/instrumentation/dynamodb/instrumentation.rb

@@ -49,10 +49,16 @@ def build_request_with_new_relic(*args)
       @nr_captured_request = yield
     end
 
+    def nr_account_id
+      return @nr_account_id if defined?(@nr_account_id)
+
+      @nr_account_id = NewRelic::Agent::Aws.get_account_id(config)
+    end
+
     def get_arn(params)
       return unless params[:table_name]
 
-      NewRelic::Agent::Aws.create_arn(PRODUCT.downcase, "table/#{params[:table_name]}", config&.region)
+      NewRelic::Agent::Aws.create_arn(PRODUCT.downcase, "table/#{params[:table_name]}", config&.region, nr_account_id)
     end
   end
 end

test/new_relic/agent/aws_test.rb

@@ -11,15 +11,28 @@ def test_create_arn
     account_id = '123456789'
     resource = 'test/test-resource'
     expected = 'arn:aws:test-service:us-test-region-1:123456789:test/test-resource'
+    arn = NewRelic::Agent::Aws.create_arn(service, resource, region, account_id)
 
-    with_config('cloud.aws.account_id': account_id) do
-      arn = NewRelic::Agent::Aws.create_arn(service, resource, region)
+    assert_equal expected, arn
+  end
 
-      assert_equal expected, arn
-    end
+  def test_get_account_id_decodes_access_key
+    config = mock
+    mock_credentials = mock
+    mock_credentials.stubs(:credentials).returns(mock_credentials)
+    mock_credentials.stubs(:access_key_id).returns('AKIAIOSFODNN7EXAMPLE') # this is a fake access key id from aws docs
+    config.stubs(:credentials).returns(mock_credentials)
+
+    account_id = NewRelic::Agent::Aws.get_account_id(config)
+
+    assert_equal 36315003739, account_id
   end
 
-  def test_doesnt_create_arn_no_account_id
-    assert_nil NewRelic::Agent::Aws.create_arn('service', 'resource', 'region')
+  def test_get_account_id_uses_config
+    config = mock
+
+    with_config(:'cloud.aws.account_id' => '123456789') do
+      assert_equal '123456789', NewRelic::Agent::Aws.get_account_id(config)
+    end
   end
 end