Merge pull request #4699 from DataDog/appsec-57738-add-api-security-sampler

[APPSEC-57738] Add API Security sampler for `rack`, `rails`, `sinatra`, `grape` contribs

Author: Strech

.github/workflows/system-tests.yml

@@ -22,7 +22,7 @@ permissions: {}
 env:
   REGISTRY: ghcr.io
   REPO: ghcr.io/datadog/dd-trace-rb
-  SYSTEM_TESTS_REF: 76f6755f072c2f50f1c956c4eb31cad9bb0e534c # Automated: Can be updated by .github/workflows/update-system-tests.yml
+  SYSTEM_TESTS_REF: appsec-57872-enable-api-security # Automated: This reference is automatically updated.
 
 jobs:
   changes:

.rubocop.yml

@@ -35,6 +35,9 @@ AllCops:
     - lib/datadog/core/**/*
     - lib/datadog/core.rb
     - spec/datadog/core/**/*
+    - lib/datadog/appsec/**/*
+    - lib/datadog/appsec.rb
+    - spec/datadog/appsec/**/*
   NewCops: disable # Don't allow new cops to be enabled implicitly.
   SuggestExtensions: false # Stop pushing suggestions constantly.
 

lib/datadog/appsec/api_security.rb

@@ -1,9 +1,23 @@
 # frozen_string_literal: true
 
+require_relative 'api_security/sampler'
+
 module Datadog
   module AppSec
     # A namespace for API Security features.
     module APISecurity
+      def self.enabled?
+        Datadog.configuration.appsec.api_security.enabled?
+      end
+
+      def self.sample?(request, response)
+        Sampler.thread_local.sample?(request, response)
+      end
+
+      def self.sample_trace?(trace)
+        # NOTE: Reads as "if trace is priority sampled or if in standalone mode"
+        trace&.priority_sampled? || !Datadog.configuration.apm.tracing.enabled
+      end
     end
   end
 end

lib/datadog/appsec/api_security/lru_cache.rb

@@ -30,6 +30,14 @@ def [](key)
           end
         end
 
+        def store(key, value)
+          return @store[key] = value if @store.delete(key)
+
+          # NOTE: evict the oldest entry if store reached the maximum allowed size
+          @store.shift if @store.size >= @max_size
+          @store[key] = value
+        end
+
         # NOTE: If the key exists, it's moved to the end of the list and
         #       if does not, the given block will be executed and the result
         #       will be stored (which will add it to the end of the list).
@@ -40,8 +48,7 @@ def fetch_or_store(key)
 
           # NOTE: evict the oldest entry if store reached the maximum allowed size
           @store.shift if @store.size >= @max_size
-
-          @store[key] ||= yield
+          @store[key] = yield
         end
       end
     end

lib/datadog/appsec/api_security/route_extractor.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AppSec
+    module APISecurity
+      # This is a helper module to extract the route pattern from the Rack::Request.
+      module RouteExtractor
+        SINATRA_ROUTE_KEY = 'sinatra.route'
+        SINATRA_ROUTE_SEPARATOR = ' '
+        GRAPE_ROUTE_KEY = 'grape.routing_args'
+        RAILS_ROUTE_KEY = 'action_dispatch.route_uri_pattern'
+        RAILS_ROUTES_KEY = 'action_dispatch.routes'
+        RAILS_FORMAT_SUFFIX = '(.:format)'
+
+        # HACK: We rely on the fact that each contrib will modify `request.env`
+        #       and store information sufficient to compute the canonical
+        #       route (ex: `/users/:id`).
+        #
+        #       When contribs like Sinatra or Grape are used, they could be mounted
+        #       into the Rails app, hence you can see the use of the `script_name`
+        #       that will contain the path prefix of the mounted app.
+        #
+        #       Rack
+        #         does not support named arguments, so we have to use `path`
+        #       Sinatra
+        #         uses `sinatra.route` with a string like "GET /users/:id"
+        #       Grape
+        #         uses `grape.routing_args` with a hash with a `:route_info` key
+        #         that contains a `Grape::Router::Route` object that contains
+        #         `Grape::Router::Pattern` object with an `origin` method
+        #       Rails < 7.1 (slow path)
+        #         uses `action_dispatch.routes` to store `ActionDispatch::Routing::RouteSet`
+        #         which can recognize requests
+        #       Rails > 7.1 (fast path)
+        #         uses `action_dispatch.route_uri_pattern` with a string like
+        #         "/users/:id(.:format)"
+        #
+        # WARNING: This method works only *after* the request has been routed.
+        def self.route_pattern(request)
+          if request.env.key?(GRAPE_ROUTE_KEY)
+            pattern = request.env[GRAPE_ROUTE_KEY][:route_info]&.pattern&.origin
+            "#{request.script_name}#{pattern}"
+          elsif request.env.key?(SINATRA_ROUTE_KEY)
+            pattern = request.env[SINATRA_ROUTE_KEY].split(SINATRA_ROUTE_SEPARATOR, 2)[1]
+            "#{request.script_name}#{pattern}"
+          elsif request.env.key?(RAILS_ROUTE_KEY)
+            request.env[RAILS_ROUTE_KEY].delete_suffix(RAILS_FORMAT_SUFFIX)
+          elsif request.env.key?(RAILS_ROUTES_KEY)
+            pattern = request.env[RAILS_ROUTES_KEY].router
+              .recognize(request) { |route, _| break route.path.spec.to_s }
+
+            # NOTE: If rails is unable to recognize request it returns empty Array
+            pattern = nil if pattern&.empty?
+
+            # NOTE: If rails can't recognize the request, we are going to fallback
+            #       to generic request path
+            (pattern || request.path).delete_suffix(RAILS_FORMAT_SUFFIX)
+          else
+            request.path
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/appsec/api_security/sampler.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require 'zlib'
+require_relative 'lru_cache'
+require_relative 'route_extractor'
+require_relative '../../core/utils/time'
+
+module Datadog
+  module AppSec
+    module APISecurity
+      # A thread-local sampler for API security based on defined delay between
+      # samples with caching capability.
+      class Sampler
+        THREAD_KEY = :datadog_appsec_api_security_sampler
+        MAX_CACHE_SIZE = 4096
+
+        class << self
+          def thread_local
+            sampler = Thread.current.thread_variable_get(THREAD_KEY)
+            return sampler unless sampler.nil?
+
+            Thread.current.thread_variable_set(THREAD_KEY, new(sample_delay))
+          end
+
+          # @api private
+          def reset!
+            Thread.current.thread_variable_set(THREAD_KEY, nil)
+          end
+
+          private
+
+          def sample_delay
+            Datadog.configuration.appsec.api_security.sample_delay
+          end
+        end
+
+        def initialize(sample_delay)
+          raise ArgumentError, 'sample_delay must be an Integer' unless sample_delay.is_a?(Integer)
+
+          @cache = LRUCache.new(MAX_CACHE_SIZE)
+          @sample_delay_seconds = sample_delay
+        end
+
+        def sample?(request, response)
+          return true if @sample_delay_seconds.zero?
+
+          key = Zlib.crc32("#{request.request_method}#{RouteExtractor.route_pattern(request)}#{response.status}")
+          current_timestamp = Core::Utils::Time.now.to_i
+          cached_timestamp = @cache[key] || 0
+
+          return false if current_timestamp - cached_timestamp <= @sample_delay_seconds
+
+          @cache.store(key, current_timestamp)
+          true
+        end
+      end
+    end
+  end
+end

lib/datadog/appsec/configuration/settings.rb

@@ -311,12 +311,27 @@ def self.add_settings!(base)
               end
 
               settings :api_security do
+                define_method(:enabled?) { get_option(:enabled) }
+
                 option :enabled do |o|
                   o.type :bool
                   o.env 'DD_EXPERIMENTAL_API_SECURITY_ENABLED'
                   o.default false
                 end
 
+                # NOTE: Unfortunately, we have to go with Float due to other libs
+                #       setup, even tho we don't plan to support sub-second delays.
+                #
+                # WARNING: The value will be converted to Integer.
+                option :sample_delay do |o|
+                  o.type :float
+                  o.env 'DD_API_SECURITY_SAMPLE_DELAY'
+                  o.default 30
+                  o.setter do |value|
+                    value.to_i
+                  end
+                end
+
                 option :sample_rate do |o|
                   o.type :float
                   o.env 'DD_API_SECURITY_REQUEST_SAMPLE_RATE'

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

@@ -8,6 +8,7 @@
 require_relative '../../event'
 require_relative '../../response'
 require_relative '../../processor'
+require_relative '../../api_security'
 require_relative '../../security_event'
 require_relative '../../instrumentation/gateway'
 
@@ -100,7 +101,23 @@ def call(env)
               http_response = AppSec::Response.from_interrupt_params(interrupt_params, env['HTTP_ACCEPT']).to_rack
             end
 
-            if AppSec.perform_api_security_check?
+            # NOTE: This is not optimal, but in the current implementation
+            #       `gateway_response` is a container to dispatch response event
+            #       and in case of interruption it suppose to be `nil`.
+            #
+            #       `http_response` is a real response object in both cases, but
+            #       to save us some computations, we will use already pre-computed
+            #       `gateway_response` instead of re-creating it.
+            #
+            # WARNING: This part will be refactored.
+            tmp_response = if interrupt_params
+              Gateway::Response.new(http_response[2], http_response[0], http_response[1], context: ctx)
+            else
+              gateway_response
+            end
+
+            if AppSec::APISecurity.enabled? && AppSec::APISecurity.sample_trace?(ctx.trace) &&
+                AppSec::APISecurity.sample?(gateway_request.request, tmp_response.response)
               ctx.events.push(
                 AppSec::SecurityEvent.new(ctx.extract_schema, trace: ctx.trace, span: ctx.span)
               )

sig/datadog/appsec/api_security.rbs

@@ -1,6 +1,22 @@
 module Datadog
   module AppSec
     module APISecurity
+      interface _Request
+        def env: () -> Hash[String, untyped]
+        def script_name: () -> String?
+        def request_method: () -> String
+        def path: () -> String
+      end
+
+      interface _Response
+        def status: () -> Integer
+      end
+
+      def self.enabled?: () -> bool
+
+      def self.sample?: (_Request request, _Response response) -> bool
+
+      def self.sample_trace?: (Tracing::TraceOperation trace) -> bool
     end
   end
 end

sig/datadog/appsec/api_security/lru_cache.rbs

@@ -12,6 +12,8 @@ module Datadog
 
         def []: (untyped key) -> untyped?
 
+        def store: (untyped key, untyped value) -> untyped
+
         def fetch_or_store: (untyped key) { () -> untyped } -> untyped
 
         def clear: () -> void

sig/datadog/appsec/api_security/route_extractor.rbs

@@ -0,0 +1,21 @@
+module Datadog
+  module AppSec
+    module APISecurity
+      module RouteExtractor
+        SINATRA_ROUTE_KEY: String
+
+        SINATRA_ROUTE_SEPARATOR: String
+
+        GRAPE_ROUTE_KEY: String
+
+        RAILS_ROUTE_KEY: String
+
+        RAILS_ROUTES_KEY: String
+
+        RAILS_FORMAT_SUFFIX: String
+
+        def self.route_pattern: (APISecurity::_Request request) -> String
+      end
+    end
+  end
+end

sig/datadog/appsec/api_security/sampler.rbs

@@ -0,0 +1,27 @@
+module Datadog
+  module AppSec
+    module APISecurity
+      class Sampler
+        THREAD_KEY: Symbol
+
+        MAX_CACHE_SIZE: Integer
+
+        @cache: LRUCache
+
+        @sample_delay_seconds: Integer
+
+        def self.thread_local: () -> Sampler
+
+        def self.reset!: () -> void
+
+        def initialize: (Integer sample_delay) -> void
+
+        def sample?: (_Request request, _Response response) -> bool
+
+        private
+
+        def self.sample_delay: () -> Integer
+      end
+    end
+  end
+end