Add stack trace collection to meta_struct and actions_handler

lib/datadog/appsec/actions_handler.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative 'actions_handler/stack_trace'
+
 module Datadog
   module AppSec
     # this module encapsulates functions for handling actions that libddawf returns
@@ -19,7 +21,26 @@ def interrupt_execution(action_params)
         throw(Datadog::AppSec::Ext::INTERRUPT, action_params)
       end
 
-      def generate_stack(_action_params); end
+      def generate_stack(action_params)
+        if Datadog.configuration.appsec.stack_trace.enabled
+          context = AppSec::Context.active
+          return if context.nil? ||
+            ActionsHandler::StackTrace.skip_stack_trace?(context, group: AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY)
+
+          collected_stack_frames = ActionsHandler::StackTrace.collect_stack_frames
+          utf8_stack_id = action_params['stack_id'].encode('UTF-8') if action_params['stack_id']
+          stack_trace = ActionsHandler::StackTrace::Representor.new(
+            id: utf8_stack_id,
+            frames: collected_stack_frames
+          )
+
+          ActionsHandler::StackTrace.add_stack_trace_to_context(
+            stack_trace,
+            context,
+            group: AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY
+          )
+        end
+      end
 
       def generate_schema(_action_params); end
     end

lib/datadog/appsec/actions_handler/stack_trace.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+require_relative 'stack_trace/representor'
+require_relative 'stack_trace/collector'
+
+require_relative '../../tracing/metadata/metastruct'
+
+module Datadog
+  module AppSec
+    module ActionsHandler
+      # Adds stack traces to meta_struct
+      module StackTrace
+        module_function
+
+        def skip_stack_trace?(context, group:)
+          if context.trace.nil? && context.span.nil?
+            Datadog.logger.debug { 'Cannot find trace or service entry span to add stack trace' }
+            return true
+          end
+
+          max_collect = Datadog.configuration.appsec.stack_trace.max_collect
+          return false if max_collect == 0
+
+          stack_traces_count = 0
+
+          unless context.trace.nil?
+            trace_dd_stack = context.trace.metastruct[AppSec::Ext::TAG_STACK_TRACE]
+            stack_traces_count += trace_dd_stack[group].size unless trace_dd_stack.nil? || trace_dd_stack[group].nil?
+          end
+
+          unless context.span.nil?
+            span_dd_stack = context.span.metastruct[AppSec::Ext::TAG_STACK_TRACE]
+            stack_traces_count += span_dd_stack[group].size unless span_dd_stack.nil? || span_dd_stack[group].nil?
+          end
+
+          stack_traces_count >= max_collect
+        end
+
+        def collect_stack_frames
+          # caller_locations without params always returns an array but steep still thinks it can be nil
+          # So we add || [] but it will never run the second part anyway (either this or steep:ignore)
+          stack_frames = caller_locations || []
+          # Steep thinks that path can still be nil and that include? is not a method of nil
+          # We must add a variable assignment to avoid this
+          stack_frames.reject! do |loc|
+            path = loc.path
+            next true if path.nil?
+
+            path.include?('lib/datadog')
+          end
+
+          StackTrace::Collector.collect(stack_frames)
+        end
+
+        def add_stack_trace_to_context(stack_trace, context, group:)
+          # We use methods defined in Tracing::Metadata::Tagging,
+          # which means we can use both the trace and the service entry span
+          service_entry_op = (context.trace || context.span)
+
+          dd_stack = service_entry_op.metastruct[AppSec::Ext::TAG_STACK_TRACE]
+          if dd_stack.nil?
+            service_entry_op.metastruct[AppSec::Ext::TAG_STACK_TRACE] = {}
+            dd_stack = service_entry_op.metastruct[AppSec::Ext::TAG_STACK_TRACE]
+          end
+
+          dd_stack[AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY] ||= []
+          stack_group = dd_stack[AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY]
+
+          stack_group << stack_trace
+        rescue StandardError => e
+          Datadog.logger.debug("Unable to add stack_trace #{stack_trace.id} in metastruct, ignoring it. Caused by: #{e}")
+        end
+      end
+    end
+  end
+end

lib/datadog/appsec/actions_handler/stack_trace/collector.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require_relative 'frame'
+
+module Datadog
+  module AppSec
+    module ActionsHandler
+      module StackTrace
+        # Represent a stack trace with its id and message in message pack
+        module Collector
+          class << self
+            def collect(locations)
+              return [] if locations.nil? || locations.empty?
+
+              skip_frames = skip_frames(locations.size)
+              frames = []
+
+              locations.each_with_index do |location, index|
+                next if skip_frames.include?(index)
+
+                frames << StackTrace::Frame.new(
+                  id: index,
+                  text: location.to_s.encode('UTF-8'),
+                  file: file_path(location),
+                  line: location.lineno,
+                  function: function_label(location)
+                )
+              end
+              frames
+            end
+
+            private
+
+            def skip_frames(locations_size)
+              max_depth = Datadog.configuration.appsec.stack_trace.max_depth
+              return [] if max_depth == 0 || locations_size <= max_depth
+
+              top_frames_limit = (max_depth * Datadog.configuration.appsec.stack_trace.max_depth_top_percent / 100.0).round
+              bottom_frames_limit = locations_size - (max_depth - top_frames_limit)
+              (top_frames_limit...bottom_frames_limit)
+            end
+
+            def file_path(location)
+              path = location.absolute_path || location.path
+              return if path.nil?
+
+              path.encode('UTF-8')
+            end
+
+            def function_label(location)
+              label = location.label
+              return if label.nil?
+
+              label.encode('UTF-8')
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/appsec/actions_handler/stack_trace/frame.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AppSec
+    module ActionsHandler
+      module StackTrace
+        # Formatted stack frame.
+        # This class extends a Struct as it's required by Steep to be able to add a method to it.
+        class Frame < Struct.new(:id, :text, :file, :line, :function, keyword_init: true) # rubocop:disable Style/StructInheritance
+          def to_msgpack(packer = nil)
+            packer ||= MessagePack::Packer.new
+
+            packer.write_map_header(5)
+            packer.write('id')
+            packer.write(id)
+            packer.write('text')
+            packer.write(text)
+            packer.write('file')
+            packer.write(file)
+            packer.write('line')
+            packer.write(line)
+            packer.write('function')
+            packer.write(function)
+            packer
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/appsec/actions_handler/stack_trace/representor.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Datadog
+  module AppSec
+    module ActionsHandler
+      module StackTrace
+        # Represent a stack trace with its id and message in message pack
+        class Representor < Struct.new(:id, :message, :frames, keyword_init: true) # rubocop:disable Style/StructInheritance
+          def to_msgpack(packer = nil)
+            packer ||= MessagePack::Packer.new
+
+            packer.write_map_header(4)
+            packer.write('language')
+            packer.write('ruby')
+            packer.write('id')
+            packer.write(id)
+            packer.write('message')
+            packer.write(message)
+            packer.write('frames')
+            packer.write(frames)
+            packer
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/appsec/configuration/settings.rb

@@ -214,6 +214,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

@@ -60,10 +60,13 @@ def extract_schema
       end
 
       def export_metrics
-        return if @span.nil?
+        # Required to satisfy steep, as @span is defined as nilable.
+        # According to soutaro, this is because instance variable can be changed by other threads.
+        span = @span
+        return if span.nil?
 
-        Metrics::Exporter.export_waf_metrics(@metrics.waf, @span)
-        Metrics::Exporter.export_rasp_metrics(@metrics.rasp, @span)
+        Metrics::Exporter.export_waf_metrics(@metrics.waf, span)
+        Metrics::Exporter.export_rasp_metrics(@metrics.rasp, span)
       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

sig/datadog/appsec/actions_handler/stack_trace.rbs

@@ -0,0 +1,13 @@
+module Datadog
+  module AppSec
+    module ActionsHandler
+      module StackTrace
+        def self.skip_stack_trace?: (Datadog::AppSec::Context context, group: String) -> bool
+
+        def self.collect_stack_frames: () -> Array[StackTrace::Frame]?
+
+        def self.add_stack_trace_to_context: (Datadog::AppSec::ActionsHandler::StackTrace::Representor stack_trace, Datadog::AppSec::Context context, group: String) -> void
+      end
+    end
+  end
+end

sig/datadog/appsec/actions_handler/stack_trace/collector.rbs

@@ -0,0 +1,19 @@
+module Datadog
+  module AppSec
+    module ActionsHandler
+      module StackTrace
+        module Collector
+          def self.collect: (Array[Thread::Backtrace::Location] locations) -> Array[Datadog::AppSec::ActionsHandler::StackTrace::Frame]
+
+          private
+
+          def self.skip_frames: (Integer locations_size) -> (Range[Integer] | Array[untyped])
+
+          def self.file_path: (Thread::Backtrace::Location location) -> String?
+
+          def self.function_label: (Thread::Backtrace::Location location) -> String?
+        end
+      end
+    end
+  end
+end

sig/datadog/appsec/actions_handler/stack_trace/frame.rbs

@@ -0,0 +1,19 @@
+module Datadog
+  module AppSec
+    module ActionsHandler
+      module StackTrace
+        class Frame
+          attr_reader id: Integer
+          attr_reader text: String?
+          attr_reader file: String?
+          attr_reader line: Integer?
+          attr_reader function: String?
+
+          def initialize: (?id: Integer, ?text: String?, ?file: String?, ?line: Integer?, ?function: String?) -> void
+
+          def to_msgpack: ((::MessagePack::Packer | nil) packer) -> ::MessagePack::Packer
+        end
+      end
+    end
+  end
+end

sig/datadog/appsec/actions_handler/stack_trace/representor.rbs

@@ -0,0 +1,17 @@
+module Datadog
+  module AppSec
+    module ActionsHandler
+      module StackTrace
+        class Representor
+          attr_reader id: String?
+          attr_reader message: String?
+          attr_reader frames: Array[StackTrace::Frame]?
+
+          def initialize: (?id: String?, ?message: String?, ?frames: Array[StackTrace::Frame]?) -> void
+
+          def to_msgpack: ((::MessagePack::Packer | nil) packer) -> ::MessagePack::Packer
+        end
+      end
+    end
+  end
+end