Add stack trace collection to meta_struct and actions_handler

vpellan · vpellan · commit 7aa25ec02dc0 · 2025-01-31T13:10:22.000+01:00
diff --git a/lib/datadog/appsec/actions_handler.rb b/lib/datadog/appsec/actions_handler.rb
@@ -1,5 +1,8 @@
 # frozen_string_literal: true
 
+require_relative '../tracing/stack_trace/collector'
+require_relative '../tracing/stack_trace/representor'
+
 module Datadog
   module AppSec
     # this module encapsulates functions for handling actions that libddawf returns
@@ -19,7 +22,41 @@ 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
+          # We use methods defined in Tracing::Metadata::Tagging,
+          # which means we can use both the trace and the service entry span
+          service_entry_operation = (context.trace || context.span) if context
+
+          if service_entry_operation.nil?
+            Datadog.logger.debug { 'Cannot find trace or service entry span to add stack trace' }
+            return
+          end
+
+          # 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
+
+          collected_stack_frames = Datadog::Tracing::StackTrace::Collector.collect(stack_frames)
+          utf8_stack_id = action_params['stack_id'].encode('UTF-8') if action_params['stack_id']
+          stack_trace = Datadog::Tracing::StackTrace::Representor.new(
+            id: utf8_stack_id,
+            message: nil,
+            frames: collected_stack_frames
+          )
+
+          service_entry_operation.set_stack_trace(stack_trace, group: AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY)
+        end
+      end
 
       def generate_schema(_action_params); end
     end
diff --git a/lib/datadog/appsec/configuration/settings.rb b/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
diff --git a/lib/datadog/appsec/ext.rb b/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
diff --git a/lib/datadog/tracing/metadata/ext.rb b/lib/datadog/tracing/metadata/ext.rb
@@ -193,6 +193,10 @@ module SpanKind
           TAG_INTERNAL = 'internal'
         end
 
+        module MetaStruct
+          TAG_STACK_TRACE = '_dd.stack'
+        end
+
         # @public_api
       end
     end
diff --git a/lib/datadog/tracing/metadata/meta_struct.rb b/lib/datadog/tracing/metadata/meta_struct.rb
@@ -1,11 +1,15 @@
 # frozen_string_literal: true
 
+require_relative 'stack_trace'
+
 module Datadog
   module Tracing
     module Metadata
       # Adds complex structures tagging behavior through meta_struct
       # @public_api
       module MetaStruct
+        include StackTrace
+
         def set_meta_struct(meta_struct)
           self.meta_struct.merge!(meta_struct)
         end
diff --git a/lib/datadog/tracing/metadata/stack_trace.rb b/lib/datadog/tracing/metadata/stack_trace.rb
@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Metadata
+      # Adds stack traces to meta_struct
+      # @public_api
+      module StackTrace
+        def set_stack_trace(stack_trace, group:)
+          meta_struct[Metadata::Ext::MetaStruct::TAG_STACK_TRACE] ||= {}
+          meta_struct[Metadata::Ext::MetaStruct::TAG_STACK_TRACE][group] ||= []
+
+          stack_trace_group = meta_struct[Metadata::Ext::MetaStruct::TAG_STACK_TRACE][group]
+          max_collect = Datadog.configuration.appsec.stack_trace.max_collect
+          return if max_collect > 0 && stack_trace_group.size >= max_collect
+
+          stack_trace_group << stack_trace
+        rescue StandardError => e
+          Datadog.logger.debug("Unable to add stack_trace #{stack_trace.id} in meta_struct, ignoring it. Caused by: #{e}")
+        end
+      end
+    end
+  end
+end
diff --git a/lib/datadog/tracing/stack_trace/collector.rb b/lib/datadog/tracing/stack_trace/collector.rb
@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require_relative 'frame'
+
+module Datadog
+  module Tracing
+    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
diff --git a/lib/datadog/tracing/stack_trace/frame.rb b/lib/datadog/tracing/stack_trace/frame.rb
@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    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
diff --git a/lib/datadog/tracing/stack_trace/representor.rb b/lib/datadog/tracing/stack_trace/representor.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    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
diff --git a/sig/datadog/appsec/ext.rbs b/sig/datadog/appsec/ext.rbs
@@ -14,6 +14,7 @@ module Datadog
       CONTEXT_KEY: ::String
 
       ACTIVE_CONTEXT_KEY: ::Symbol
+      EXPLOIT_PREVENTION_EVENT_CATEGORY: ::String
 
       TAG_APPSEC_ENABLED: ::String
 
@@ -22,6 +23,7 @@ module Datadog
       TAG_DISTRIBUTED_APPSEC_EVENT: ::String
 
       TELEMETRY_METRICS_NAMESPACE: ::String
+      TAG_STACK_TRACE: ::String
     end
   end
 end
diff --git a/sig/datadog/tracing/metadata/ext.rbs b/sig/datadog/tracing/metadata/ext.rbs
@@ -99,6 +99,10 @@ module Datadog
           TAG_CONSUMER: ::String
           TAG_INTERNAL: ::String
         end
+
+        module MetaStruct
+          TAG_STACK_TRACE: ::String
+        end
       end
     end
   end
diff --git a/sig/datadog/tracing/metadata/meta_struct.rbs b/sig/datadog/tracing/metadata/meta_struct.rbs
@@ -2,6 +2,8 @@ module Datadog
   module Tracing
     module Metadata
       module MetaStruct
+        include StackTrace
+
         def set_meta_struct: (Hash[String, untyped] meta_struct) -> void
 
         def meta_struct: () -> Hash[String, untyped]
diff --git a/sig/datadog/tracing/metadata/stack_trace.rbs b/sig/datadog/tracing/metadata/stack_trace.rbs
@@ -0,0 +1,14 @@
+module Datadog
+  module Tracing
+    module Metadata
+      module StackTrace
+        def set_stack_trace: (Datadog::Tracing::StackTrace::Representor stack_trace, group: ::String) -> void
+
+        private
+
+        # Class/Module that includes this module should implement this method
+        def meta_struct: () -> Hash[String, untyped]
+      end
+    end
+  end
+end
diff --git a/sig/datadog/tracing/stack_trace/collector.rbs b/sig/datadog/tracing/stack_trace/collector.rbs
@@ -0,0 +1,17 @@
+module Datadog
+  module Tracing
+    module StackTrace
+      module Collector
+        def self.collect: (Array[Thread::Backtrace::Location] locations) -> Array[Datadog::Tracing::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
diff --git a/sig/datadog/tracing/stack_trace/frame.rbs b/sig/datadog/tracing/stack_trace/frame.rbs
@@ -0,0 +1,17 @@
+module Datadog
+  module Tracing
+    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
diff --git a/sig/datadog/tracing/stack_trace/representor.rbs b/sig/datadog/tracing/stack_trace/representor.rbs
@@ -0,0 +1,15 @@
+module Datadog
+  module Tracing
+    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
diff --git a/spec/datadog/appsec/configuration/settings_spec.rb b/spec/datadog/appsec/configuration/settings_spec.rb
diff --git a/spec/datadog/tracing/stack_trace/collector_spec.rb b/spec/datadog/tracing/stack_trace/collector_spec.rb
diff --git a/spec/datadog/tracing/stack_trace/frame_spec.rb b/spec/datadog/tracing/stack_trace/frame_spec.rb
diff --git a/spec/datadog/tracing/stack_trace/representor_spec.rb b/spec/datadog/tracing/stack_trace/representor_spec.rb
diff --git a/spec/support/thread_backtrace_helpers.rb b/spec/support/thread_backtrace_helpers.rb
