Merge pull request #5006 from DataDog/appsec-security-sampling-with-inferred-route

Add a fallback to inferred route for AppSec Api security sampler

Author: y9v

lib/datadog/appsec/api_security/route_extractor.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative '../../tracing/contrib/rack/route_inference'
+
 module Datadog
   module AppSec
     module APISecurity
@@ -66,7 +68,7 @@ def self.route_pattern(request)
             #       to generic request path
             (pattern || request.path).delete_suffix(RAILS_FORMAT_SUFFIX)
           else
-            request.path
+            Tracing::Contrib::Rack::RouteInference.read_or_infer(request.env)
           end
         end
       end

lib/datadog/appsec/api_security/sampler.rb

@@ -43,6 +43,7 @@ def initialize(sample_delay)
 
         def sample?(request, response)
           return true if @sample_delay_seconds.zero?
+          return false if response.status == 404
 
           key = Zlib.crc32("#{request.request_method}#{RouteExtractor.route_pattern(request)}#{response.status}")
           current_timestamp = Core::Utils::Time.now.to_i

lib/datadog/tracing/contrib/rack/middlewares.rb

@@ -13,7 +13,7 @@
 require_relative 'header_collection'
 require_relative 'header_tagging'
 require_relative 'request_queue'
-require_relative 'route_from_path_inference'
+require_relative 'route_inference'
 require_relative 'trace_proxy_middleware'
 
 module Datadog
@@ -235,7 +235,7 @@ def set_route_and_endpoint_tags(trace:, request_span:, status:, env:)
               # we infer the route from the full request path.
               if last_script_name == '' && env['SCRIPT_NAME'] != '' &&
                   !Datadog.configuration.tracing.resource_renaming.always_simplified_endpoint &&
-                  (inferred_route = RouteFromPathInference.infer(env['PATH_INFO']))
+                  (inferred_route = RouteInference.infer(env['PATH_INFO']))
                 set_endpoint_tag(request_span, last_route + inferred_route)
               end
 
@@ -250,13 +250,8 @@ def set_route_and_endpoint_tags(trace:, request_span:, status:, env:)
 
             if Datadog.configuration.tracing.resource_renaming.always_simplified_endpoint ||
                 request_span.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE).nil?
-              # when http.route tag wasn't set or resource_renaming.always_simplified_endpoint is set to true,
-              # we will try to infer the path for the full request path,
-              # which is SCRIPT_NAME + PATH_INFO for mounted rack applications
-              full_path = env['SCRIPT_NAME'].to_s + env['PATH_INFO'].to_s
-
-              if (inferred_route = RouteFromPathInference.infer(full_path))
-                set_endpoint_tag(request_span, inferred_route)
+              if (inferred_route = RouteInference.read_or_infer(env))
+                set_endpoint_tag(request_span, inferred_route) if inferred_route
               end
             elsif !request_span.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ENDPOINT)
               set_endpoint_tag(request_span, request_span.get_tag(Tracing::Metadata::Ext::HTTP::TAG_ROUTE))

lib/datadog/tracing/contrib/rack/route_from_path_inference.rb renamed to lib/datadog/tracing/contrib/rack/route_inference.rb

@@ -6,7 +6,7 @@ module Contrib
       module Rack
         # This module provides logic for inferring HTTP route pattern
         # from an HTTP path.
-        module RouteFromPathInference
+        module RouteInference
           MAX_NUMBER_OF_SEGMENTS = 8
 
           INT_PARAM_REGEX = /\A[0-9]+\z/.freeze
@@ -15,12 +15,19 @@ module RouteFromPathInference
           HEX_ID_PARAM_REGEX = /\A(?=.*\d)[A-Fa-f0-9._-]{6,}\z/.freeze
           STRING_PARAM_REGEX = /\A.{20,}|.*[%&'()*+,:=@].*\z/.freeze
 
+          DATADOG_INFERRED_ROUTE_ENV_KEY = 'datadog.inferred_route'
+
           module_function
 
+          def read_or_infer(request_env)
+            request_env[DATADOG_INFERRED_ROUTE_ENV_KEY] ||=
+              infer(request_env['SCRIPT_NAME'].to_s + request_env['PATH_INFO'].to_s)
+          end
+
           def infer(path)
             segments = path.delete_prefix('/').split('/', MAX_NUMBER_OF_SEGMENTS + 1).first(MAX_NUMBER_OF_SEGMENTS)
 
-            segments.map! do |segment|
+            segments.map! do |segment| #: Array[String?]
               next if segment.empty?
 
               case segment
@@ -33,7 +40,7 @@ def infer(path)
               end
             end
 
-            segments.compact!
+            segments.compact! #: Array[String]
 
             "/#{segments.join("/")}"
           rescue

sig/datadog/appsec/api_security/route_extractor.rbs

@@ -16,7 +16,7 @@ module Datadog
 
         RAILS_FORMAT_SUFFIX: String
 
-        def self.route_pattern: (APISecurity::_Request request) -> String
+        def self.route_pattern: (APISecurity::_Request request) -> String?
       end
     end
   end

sig/datadog/tracing/contrib/rack/route_from_path_inference.rbs renamed to sig/datadog/tracing/contrib/rack/route_inference.rbs

@@ -2,7 +2,7 @@ module Datadog
   module Tracing
     module Contrib
       module Rack
-        module RouteFromPathInference
+        module RouteInference
           MAX_NUMBER_OF_SEGMENTS: ::Integer
 
           INT_PARAM_REGEX: ::Regexp
@@ -11,7 +11,11 @@ module Datadog
           HEX_ID_PARAM_REGEX: ::Regexp
           STRING_PARAM_REGEX: ::Regexp
 
-          def self.infer: (::String path) -> ::String?
+          DATADOG_INFERRED_ROUTE_ENV_KEY: ::String
+
+          def self?.read_or_infer: (Hash[::String, untyped] request_env) -> ::String?
+
+          def self?.infer: (::String path) -> ::String?
         end
       end
     end

spec/datadog/appsec/api_security/route_extractor_spec.rb

@@ -96,15 +96,22 @@
         before do
           allow(request).to receive(:env).and_return({
             'action_dispatch.routes' => route_set,
-            'action_dispatch.request.path_parameters' => {}
+            'action_dispatch.request.path_parameters' => {},
+            'PATH_INFO' => '/users/1'
           })
         end
 
         let(:router) { double('ActionDispatch::Routing::RouteSet::Router') }
         let(:route_set) { double('ActionDispatch::Routing::RouteSet', router: router) }
-        let(:request) { double('Rack::Request', env: {}, script_name: '', path: '/users/1') }
+        let(:request) { double('Rack::Request', script_name: '', path: '/users/1') }
 
-        it { expect(described_class.route_pattern(request)).to eq('/users/1') }
+        it { expect(described_class.route_pattern(request)).to eq('/users/{param:int}') }
+
+        it 'persists inferred route in the request env' do
+          expect { described_class.route_pattern(request) }
+            .to change { request.env[Datadog::Tracing::Contrib::Rack::RouteInference::DATADOG_INFERRED_ROUTE_ENV_KEY] }
+            .from(nil).to('/users/{param:int}')
+        end
       end
 
       context 'when route_uri_pattern is not set and request path_parameters is present' do
@@ -130,8 +137,8 @@
           let(:action_dispatch_request) { double('ActionDispatch::Request', env: {}, script_name: '', path: '/users/1') }
 
           it 'uses action dispatch request for route recognition' do
-            expect(router).to receive(:recognize).with(action_dispatch_request).and_return('/users/1')
-            expect(described_class.route_pattern(request)).to eq('/users/1')
+            expect(router).to receive(:recognize).with(action_dispatch_request).and_return('/users/:id(.:format)')
+            expect(described_class.route_pattern(request)).to eq('/users/:id')
           end
         end
 
@@ -143,16 +150,18 @@
           let(:request) { double('Rack::Request', env: {}, script_name: '', path: '/users/1', head?: false) }
 
           it 'uses action dispatch request for route recognition' do
-            expect(router).to receive(:recognize).with(request).and_return('/users/1')
-            expect(described_class.route_pattern(request)).to eq('/users/1')
+            expect(router).to receive(:recognize).with(request).and_return('/users/:id(.:format)')
+            expect(described_class.route_pattern(request)).to eq('/users/:id')
           end
         end
       end
 
       context 'when Rails router cannot recognize request' do
         before do
-          allow(request).to receive(:env).and_return({'action_dispatch.routes' => route_set})
-          allow(request).to receive(:path).and_return('/unmatched/route')
+          allow(request).to receive(:env).and_return({
+            'action_dispatch.routes' => route_set,
+            'PATH_INFO' => '/unmatched/route'
+          })
           allow(router).to receive(:recognize).with(request).and_return([])
         end
 
@@ -169,7 +178,7 @@
       end
 
       context 'when route has nested path' do
-        before { allow(request).to receive(:path).and_return('/some/other/path') }
+        before { allow(request).to receive(:env).and_return({'PATH_INFO' => '/some/other/path'}) }
 
         it { expect(described_class.route_pattern(request)).to eq('/some/other/path') }
       end

spec/datadog/appsec/api_security/sampler_spec.rb

@@ -74,6 +74,16 @@
       end
     end
 
+    context 'when response status is 404' do
+      let(:response) { double('Rack::Response', status: 404) }
+
+      it 'always returns false' do
+        3.times do
+          expect(sampler.sample?(request, response)).to be(false)
+        end
+      end
+    end
+
     context 'when sampling for the first time' do
       it { expect(sampler.sample?(request, response)).to be(true) }
     end

spec/datadog/tracing/contrib/rack/route_from_path_inference_spec.rb renamed to spec/datadog/tracing/contrib/rack/route_inference_spec.rb

@@ -1,8 +1,46 @@
 require 'datadog/tracing/contrib/support/spec_helper'
 
-require 'datadog/tracing/contrib/rack/route_from_path_inference'
+require 'datadog/tracing/contrib/rack/route_inference'
+
+RSpec.describe Datadog::Tracing::Contrib::Rack::RouteInference do
+  describe '.read_or_infer' do
+    context 'when inferred route was not yet persisted in request env' do
+      let(:env) do
+        {
+          'SCRIPT_NAME' => '/api',
+          'PATH_INFO' => '/users/1'
+        }
+      end
+
+      it 'returns inferred route' do
+        expect(described_class.read_or_infer(env)).to eq('/api/users/{param:int}')
+      end
+
+      it 'persists inferred route in request env' do
+        expect { described_class.read_or_infer(env) }.to change { env['datadog.inferred_route'] }
+          .from(nil).to('/api/users/{param:int}')
+      end
+    end
+
+    context 'when inferred route already persisted in request env' do
+      let(:env) do
+        {
+          'SCRIPT_NAME' => '/api',
+          'PATH_INFO' => '/users/1',
+          'datadog.inferred_route' => '/some_route'
+        }
+      end
+
+      it 'returns persisted inferred route' do
+        expect(described_class.read_or_infer(env)).to eq('/some_route')
+      end
+
+      it 'does not change inferred route value in request env' do
+        expect { described_class.read_or_infer(env) }.not_to change { env['datadog.inferred_route'] }
+      end
+    end
+  end
 
-RSpec.describe Datadog::Tracing::Contrib::Rack::RouteFromPathInference do
   describe '.infer' do
     it 'works with an empty path' do
       expect(described_class.infer('/')).to eq('/')