Add stack trace collection support

Steepfile

@@ -157,6 +157,7 @@ target :datadog do
   library 'zlib'
   library 'time'
   library 'pp'
+  library 'forwardable'
 
   # Load all dependency signatures from the `vendor/rbs` directory
   repo_path 'vendor/rbs'

lib/datadog/appsec/actions_handler.rb

@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require_relative 'actions_handler/stack_trace_in_metastruct'
+require_relative 'actions_handler/stack_trace_collection'
+
 module Datadog
   module AppSec
     # this module encapsulates functions for handling actions that libddawf returns
@@ -19,7 +22,32 @@ def interrupt_execution(action_params)
         throw(Datadog::AppSec::Ext::INTERRUPT, action_params)
       end
 
-      def generate_stack(_action_params); end
+      def generate_stack(action_params)
+        return unless Datadog.configuration.appsec.stack_trace.enabled
+
+        context = AppSec.active_context
+        if context.nil? || context.trace.nil? && context.span.nil?
+          Datadog.logger.debug { 'Cannot find trace or service entry span to add stack trace' }
+          return
+        end
+
+        # Check that the sum of stack_trace count in trace and entry_span does not exceed configuration
+        span_stack = StackTraceInMetastruct.create(context.span&.metastruct)
+        trace_stack = StackTraceInMetastruct.create(context.trace&.metastruct)
+        config = Datadog.configuration.appsec.stack_trace
+        return if config.max_collect != 0 && span_stack.count + trace_stack.count >= config.max_collect
+
+        # Generate stacktrace
+        stack_id = action_params['stack_id'].encode('UTF-8')
+        stack_frames = StackTraceCollection.collect(
+          max_depth: config.max_depth,
+          top_percent: config.max_depth_top_percent
+        )
+
+        # Add newly created stacktrace to metastruct
+        stack = context.trace.nil? ? span_stack : trace_stack
+        stack.push({ language: 'ruby', id: stack_id, frames: stack_frames })
+      end
 
       def generate_schema(_action_params); end
     end

lib/datadog/appsec/actions_handler/stack_trace_collection.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AppSec
+    module ActionsHandler
+      # Module that collects the stack trace into formatted hash
+      module StackTraceCollection
+        module_function
+
+        def collect(max_depth:, top_percent:)
+          locations = StackTraceCollection.filter_map_datadog_locations(caller_locations || [])
+
+          return [] if locations.empty?
+          return StackTraceCollection.convert(locations) if max_depth.zero? || locations.size <= max_depth
+
+          top_limit = (max_depth * top_percent / 100.0).round
+          bottom_limit = locations.size - (max_depth - top_limit)
+
+          locations.slice!(top_limit...bottom_limit)
+          StackTraceCollection.convert(locations)
+        end
+
+        def filter_map_datadog_locations(locations)
+          locations.each_with_object([]) do |location, result|
+            text = location.to_s
+            next if text.include?('lib/datadog')
+
+            result << {
+              text: text,
+              file: location.absolute_path || location.path,
+              line: location.lineno,
+              function: location.label
+            }
+          end
+        end
+
+        def convert(locations)
+          locations.each_with_index do |location, index|
+            location[:id] = index
+            # Strings can be frozen so we need to copy them
+            location[:text] = location[:text].encode('UTF-8')
+            location[:file] = location[:file]&.encode('UTF-8')
+            location[:function] = location[:function]&.encode('UTF-8')
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/appsec/actions_handler/stack_trace_in_metastruct.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AppSec
+    module ActionsHandler
+      # Object that holds a metastruct, and modify the exploit group stack traces
+      class StackTraceInMetastruct
+        # Implementation with empty metastruct
+        class Noop
+          def count
+            0
+          end
+
+          def push(_)
+            nil
+          end
+        end
+
+        def self.create(metastruct)
+          metastruct.nil? ? Noop.new : new(metastruct)
+        end
+
+        def initialize(metastruct)
+          @metastruct = metastruct
+        end
+
+        def count
+          @metastruct.dig(AppSec::Ext::TAG_STACK_TRACE, AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY)&.size || 0
+        end
+
+        def push(stack_trace)
+          @metastruct[AppSec::Ext::TAG_STACK_TRACE] ||= {}
+          @metastruct[AppSec::Ext::TAG_STACK_TRACE][AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY] ||= []
+          @metastruct[AppSec::Ext::TAG_STACK_TRACE][AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY] << stack_trace
+        end
+      end
+    end
+  end
+end

lib/datadog/appsec/configuration/settings.rb

@@ -267,6 +267,55 @@ def self.add_settings!(base)
                   o.default false
                 end
               end
+
+              settings :stack_trace do
+                option :enabled do |o|
+                  o.type :bool
+                  o.env 'DD_APPSEC_STACK_TRACE_ENABLED'
+                  o.default true
+                end
+
+                # The maximum number of stack frames to collect for each stack trace.
+                # If the number of frames in a stack trace exceeds this value,
+                # max_depth / 4 frames will be collected from the top, and max_depth * 3 / 4 from the bottom.
+                option :max_depth do |o|
+                  o.type :int
+                  o.env 'DD_APPSEC_MAX_STACK_TRACE_DEPTH'
+                  o.default 32
+                  # 0 means no limit
+                  o.setter do |value|
+                    value = 0 if value.negative?
+                    value
+                  end
+                end
+
+                # The percentage that decides the number of top stack frame to collect
+                # for each stack trace if there is more stack frames than max_depth.
+                # number_of_top_frames = max_depth * max_depth_top_percent / 100
+                # Default is 75
+                option :max_depth_top_percent do |o|
+                  o.type :float
+                  o.env 'DD_APPSEC_MAX_STACK_TRACE_DEPTH_TOP_PERCENT'
+                  o.default 75
+                  o.setter do |value|
+                    value = 100 if value > 100
+                    value = 0 if value < 0
+                    value
+                  end
+                end
+
+                # The maximum number of stack traces to collect for each exploit prevention event.
+                option :max_collect do |o|
+                  o.type :int
+                  o.env 'DD_APPSEC_MAX_STACK_TRACES'
+                  o.default 2
+                  # 0 means no limit
+                  o.setter do |value|
+                    value = 0 if value < 0
+                    value
+                  end
+                end
+              end
             end
           end
         end

lib/datadog/appsec/context.rb

@@ -62,8 +62,12 @@ def extract_schema
       def export_metrics
         return if @span.nil?
 
-        Metrics::Exporter.export_waf_metrics(@metrics.waf, @span)
-        Metrics::Exporter.export_rasp_metrics(@metrics.rasp, @span)
+        # This does not caused a steep error previously because
+        # @span was wrongly defined as a SpanOperation that cannot be nil in context.rbs.
+        # Even though we check that @span is not nil, steep consideres that the thread can pause after that check,
+        # and another thread change it to nil. This does not happen in our case, which is why steep:ignore has been added.
+        Metrics::Exporter.export_waf_metrics(@metrics.waf, @span) # steep:ignore ArgumentTypeMismatch
+        Metrics::Exporter.export_rasp_metrics(@metrics.rasp, @span) # steep:ignore ArgumentTypeMismatch
       end
 
       def finalize

lib/datadog/appsec/ext.rb

@@ -10,12 +10,14 @@ module Ext
       INTERRUPT = :datadog_appsec_interrupt
       CONTEXT_KEY = 'datadog.appsec.context'
       ACTIVE_CONTEXT_KEY = :datadog_appsec_active_context
+      EXPLOIT_PREVENTION_EVENT_CATEGORY = 'exploit'
 
       TAG_APPSEC_ENABLED = '_dd.appsec.enabled'
       TAG_APM_ENABLED = '_dd.apm.enabled'
       TAG_DISTRIBUTED_APPSEC_EVENT = '_dd.p.appsec'
 
       TELEMETRY_METRICS_NAMESPACE = 'appsec'
+      TAG_STACK_TRACE = '_dd.stack'
     end
   end
 end

lib/datadog/tracing/metadata.rb

@@ -2,6 +2,7 @@
 
 require_relative 'metadata/analytics'
 require_relative 'metadata/tagging'
+require_relative 'metadata/metastruct'
 require_relative 'metadata/errors'
 
 module Datadog

lib/datadog/tracing/metadata/metastruct.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require 'forwardable'
+
+module Datadog
+  module Tracing
+    module Metadata
+      # Adds complex structures tagging behavior through metastruct
+      class Metastruct
+        extend Forwardable
+
+        MERGER = proc do |_, v1, v2|
+          if v1.is_a?(Hash) && v2.is_a?(Hash)
+            v1.merge(v2, &MERGER)
+          elsif v1.is_a?(Array) && v2.is_a?(Array)
+            v1.concat(v2)
+          elsif v2.nil?
+            v1
+          else
+            v2
+          end
+        end
+
+        def self.empty
+          new({})
+        end
+
+        def initialize(metastruct)
+          @metastruct = metastruct
+        end
+
+        # Deep merge two metastructs
+        # If the types are not both Arrays or Hashes, the second one will overwrite the first one
+        #
+        # Example with same types:
+        # metastruct = { a: { b: [1, 2] } }
+        # second = { a: { b: [3, 4], c: 5 } }
+        # result = { a: { b: [1, 2, 3, 4], c: 5 } }
+        #
+        # Example with different types:
+        # metastruct = { a: { b: 1 } }
+        # second = { a: { b: [2, 3] } }
+        # result = { a: { b: [2, 3] } }
+        def deep_merge!(second)
+          @metastruct.merge!(second.to_h, &MERGER)
+        end
+
+        def_delegators :@metastruct, :[], :[]=, :dig, :to_h
+
+        def pretty_print(q)
+          q.seplist @metastruct.each do |key, value|
+            q.text "#{key} => #{value}\n"
+          end
+        end
+
+        def to_msgpack(packer = nil)
+          packer ||= MessagePack::Packer.new
+
+          packer.write(@metastruct.transform_values(&:to_msgpack))
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/span.rb

@@ -21,6 +21,7 @@ class Span
         :end_time,
         :id,
         :meta,
+        :metastruct,
         :metrics,
         :name,
         :parent_id,
@@ -53,6 +54,7 @@ def initialize(
         end_time: nil,
         id: nil,
         meta: nil,
+        metastruct: {},
         metrics: nil,
         parent_id: 0,
         resource: name,
@@ -75,6 +77,7 @@ def initialize(
         @trace_id = trace_id || Tracing::Utils.next_id
 
         @meta = meta || {}
+        @metastruct = Tracing::Metadata::Metastruct.new(metastruct)
         @metrics = metrics || {}
         @status = status || 0
 
@@ -144,6 +147,7 @@ def to_hash
           error: @status,
           meta: @meta,
           metrics: @metrics,
+          meta_struct: @metastruct.to_h,
           name: @name,
           parent_id: @parent_id,
           resource: @resource,
@@ -185,12 +189,16 @@ def pretty_print(q)
               q.text "#{key} => #{value}"
             end
           end
-          q.group(2, 'Metrics: [', ']') do
+          q.group(2, 'Metrics: [', "]\n") do
             q.breakable
             q.seplist @metrics.each do |key, value|
               q.text "#{key} => #{value}"
             end
           end
+          q.group(2, 'Metastruct: ') do
+            q.breakable
+            q.pp metastruct
+          end
         end
       end