Add stack trace collection support

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

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

lib/datadog/tracing/metadata.rb

@@ -2,6 +2,7 @@
 
 require_relative 'metadata/analytics'
 require_relative 'metadata/tagging'
+require_relative 'metadata/meta_struct'
 require_relative 'metadata/errors'
 
 module Datadog
@@ -11,6 +12,7 @@ module Metadata
       def self.included(base)
         base.include(Metadata::Tagging)
         base.include(Metadata::Errors)
+        base.include(Metadata::MetaStruct)
 
         # Additional extensions
         base.prepend(Metadata::Analytics)

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

lib/datadog/tracing/metadata/meta_struct.rb

@@ -0,0 +1,25 @@
+# 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
+
+        protected
+
+        def meta_struct
+          @meta_struct ||= {}
+        end
+      end
+    end
+  end
+end

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

lib/datadog/tracing/span.rb

@@ -21,6 +21,7 @@ class Span
         :end_time,
         :id,
         :meta,
+        :meta_struct,
         :metrics,
         :name,
         :parent_id,
@@ -53,6 +54,7 @@ def initialize(
         end_time: nil,
         id: nil,
         meta: nil,
+        meta_struct: nil,
         metrics: nil,
         parent_id: 0,
         resource: name,
@@ -75,6 +77,7 @@ def initialize(
         @trace_id = trace_id || Tracing::Utils.next_id
 
         @meta = meta || {}
+        @meta_struct = meta_struct || {}
         @metrics = metrics || {}
         @status = status || 0
 
@@ -144,6 +147,7 @@ def to_hash
           error: @status,
           meta: @meta,
           metrics: @metrics,
+          meta_struct: @meta_struct,
           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, 'Meta-Struct: [', ']') do
+            q.breakable
+            q.pp meta_struct
+          end
         end
       end
 

lib/datadog/tracing/span_operation.rb

@@ -289,6 +289,7 @@ def to_hash
           id: @id,
           meta: meta,
           metrics: metrics,
+          meta_struct: meta_struct,
           name: @name,
           parent_id: @parent_id,
           resource: @resource,
@@ -328,12 +329,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, 'Meta-Struct: [', ']') do
+            q.breakable
+            q.pp meta_struct
+          end
         end
       end
 
@@ -456,6 +461,7 @@ def build_span
           id: @id,
           meta: Core::Utils::SafeDup.frozen_or_dup(meta),
           metrics: Core::Utils::SafeDup.frozen_or_dup(metrics),
+          meta_struct: Core::Utils::SafeDup.frozen_or_dup(meta_struct),
           parent_id: @parent_id,
           resource: @resource,
           service: @service,

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

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