Merge pull request #4468 from DataDog/appsec-56960-add-collection-of-rack-request-headers-to-identity

[APPSEC-56960] Add collection of request headers for identity.set_user

Author: Strech

lib/datadog/appsec/contrib/rack/ext.rb

@@ -6,6 +6,26 @@ module Contrib
       module Rack
         # Rack integration constants
         module Ext
+          IDENTITY_COLLECTABLE_REQUEST_HEADERS = [
+            'accept-encoding',
+            'accept-language',
+            'cf-connecting-ip',
+            'cf-connecting-ipv6',
+            'content-encoding',
+            'content-language',
+            'content-length',
+            'fastly-client-ip',
+            'forwarded',
+            'forwarded-for',
+            'host',
+            'true-client-ip',
+            'via',
+            'x-client-ip',
+            'x-cluster-client-ip',
+            'x-forwarded',
+            'x-forwarded-for',
+            'x-real-ip'
+          ].freeze
         end
       end
     end

lib/datadog/appsec/contrib/rack/gateway/watcher.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative '../ext'
 require_relative '../../../instrumentation/gateway'
 require_relative '../../../event'
 
@@ -110,6 +111,25 @@ def watch_request_body(gateway = Instrumentation.gateway)
                   stack.call(gateway_request.request)
                 end
               end
+
+              # NOTE: In the current state we unable to substibe twice to the same
+              #       event within the same group. Ideally this code should live
+              #       somewhere closer to identity related monitor.
+              # WARNING: The Gateway is a subject of refactoring
+              def watch_request_finish(gateway = Instrumentation.gateway)
+                gateway.watch('rack.request.finish', :appsec) do |stack, gateway_request|
+                  context = gateway_request.env[AppSec::Ext::CONTEXT_KEY]
+                  next stack.call(gateway_request.request) if context.span.nil? || !gateway.pushed?('identity.set_user')
+
+                  gateway_request.headers.each do |name, value|
+                    next unless Ext::IDENTITY_COLLECTABLE_REQUEST_HEADERS.include?(name)
+
+                    context.span["http.request.headers.#{name}"] = value
+                  end
+
+                  stack.call(gateway_request.request)
+                end
+              end
             end
           end
         end

lib/datadog/appsec/contrib/rack/request_middleware.rb

@@ -77,6 +77,8 @@ def call(env)
             gateway_response = nil
 
             interrupt_params = catch(::Datadog::AppSec::Ext::INTERRUPT) do
+              # TODO: This event should be renamed into `rack.request.start` to
+              #       reflect that it's the beginning of the request-cycle
               http_response, _gateway_request = Instrumentation.gateway.push('rack.request', gateway_request) do
                 @app.call(env)
               end
@@ -85,6 +87,7 @@ def call(env)
                 http_response[2], http_response[0], http_response[1], context: ctx
               )
 
+              Instrumentation.gateway.push('rack.request.finish', gateway_request)
               Instrumentation.gateway.push('rack.response', gateway_response)
 
               nil

sig/datadog/appsec/contrib/rack/ext.rbs

@@ -3,7 +3,7 @@ module Datadog
     module Contrib
       module Rack
         module Ext
-          APP: untyped
+          IDENTITY_COLLECTABLE_REQUEST_HEADERS: ::Array[::String]
         end
       end
     end

spec/datadog/appsec/contrib/integration/rack/request_headers_collection_for_identity_spec.rb

@@ -0,0 +1,113 @@
+# frozen_string_literal: true
+
+require 'datadog/tracing/contrib/support/spec_helper'
+require 'datadog/appsec/spec_helper'
+require 'rack/test'
+
+require 'datadog/tracing'
+require 'datadog/appsec'
+
+RSpec.describe 'Rack-request headers collection for identity.set_user' do
+  include Rack::Test::Methods
+
+  before do
+    Datadog.configure do |c|
+      c.tracing.enabled = true
+
+      c.appsec.enabled = true
+      c.appsec.instrument :rack
+
+      c.appsec.standalone.enabled = false
+      c.appsec.waf_timeout = 10_000_000 # in us
+      c.appsec.ip_passlist = []
+      c.appsec.ip_denylist = []
+      c.appsec.user_id_denylist = []
+      c.appsec.ruleset = :recommended
+      c.appsec.api_security.enabled = false
+      c.appsec.api_security.sample_rate = 0.0
+
+      c.remote.enabled = false
+    end
+
+    allow(Datadog::AppSec::Instrumentation).to receive(:gateway).and_return(gateway)
+
+    # NOTE: Don't reach the agent in any way
+    allow_any_instance_of(Datadog::Tracing::Transport::HTTP::Client).to receive(:send_request)
+    allow_any_instance_of(Datadog::Tracing::Transport::Traces::Transport).to receive(:native_events_supported?)
+      .and_return(true)
+
+    Datadog::AppSec::Contrib::Rack::Gateway::Watcher.watch_request_finish
+  end
+
+  after do
+    Datadog.configuration.reset!
+    Datadog.registry[:rack].reset_configuration!
+  end
+
+  let(:gateway) { Datadog::AppSec::Instrumentation::Gateway.new }
+
+  let(:http_service_entry_span) do
+    Datadog::Tracing::Transport::TraceFormatter.format!(trace)
+    spans.find { |s| s.name == 'rack.request' }
+  end
+
+  let(:app) do
+    stack = Rack::Builder.new do
+      use Datadog::Tracing::Contrib::Rack::TraceMiddleware
+      use Datadog::AppSec::Contrib::Rack::RequestMiddleware
+
+      map '/with-identity-set-user' do
+        run(
+          lambda do |_env|
+            Datadog::Kit::Identity.set_user(
+              Datadog::Tracing.active_trace, Datadog::Tracing.active_span, id: '42'
+            )
+
+            [200, { 'Content-Type' => 'text/html' }, ['OK']]
+          end
+        )
+      end
+
+      map '/without-identity-set-user' do
+        run ->(_env) { [200, { 'Content-Type' => 'text/html' }, ['OK']] }
+      end
+    end
+
+    stack.to_app
+  end
+
+  subject(:response) { last_response }
+
+  context 'when identity.set_user event was pushed' do
+    before do
+      headers = {
+        'HTTP_CF_RAY' => '230b030023ae2822-SJC',
+        'HTTP_CF_CONNECTING_IPV6' => '2001:db8:3333:4444:5555:6666:1.2.3.4'
+      }
+      get('/with-identity-set-user', {}, headers)
+    end
+
+    it 'collects identity related request headers' do
+      expect(response).to be_ok
+
+      expect(http_service_entry_span.tags['http.request.headers.cf-connecting-ipv6'])
+        .to eq('2001:db8:3333:4444:5555:6666:1.2.3.4')
+    end
+  end
+
+  context 'when identity.set_user event was not pushed' do
+    before do
+      headers = {
+        'HTTP_CF_RAY' => '230b030023ae2822-SJC',
+        'HTTP_CF_CONNECTING_IPV6' => '2001:db8:3333:4444:5555:6666:1.2.3.4'
+      }
+      get('/without-identity-set-user', {}, headers)
+    end
+
+    it 'does not collect identity related request headers' do
+      expect(response).to be_ok
+
+      expect(http_service_entry_span.tags).not_to have_key('http.request.headers.cf-connecting-ipv6')
+    end
+  end
+end