WIP: Add meta_struct and stack traces support

lib/datadog/appsec.rb

@@ -5,6 +5,7 @@
 require_relative 'appsec/context'
 require_relative 'appsec/ext'
 require_relative 'appsec/utils'
+require_relative 'appsec/stack_trace'
 
 module Datadog
   # Namespace for Datadog AppSec instrumentation

lib/datadog/appsec/contrib/active_record/instrumentation.rb

@@ -36,6 +36,11 @@ def detect_sql_injection(sql, adapter_name)
                 actions: result.actions
               }
               context.waf_runner.events << event
+
+              # Include the stack trace if waf result includes it
+              if AppSec::StackTrace.include_stack_trace?(result)
+                AppSec::StackTrace.add_stack_trace(context, result)
+              end
             end
           end
 

lib/datadog/appsec/contrib/graphql/gateway/watcher.rb

@@ -39,6 +39,11 @@ def watch_multiplex(gateway = Instrumentation.gateway)
 
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
                       context.waf_runner.events << event
+
+                      # Include the stack trace if waf result includes it
+                      if AppSec::StackTrace.include_stack_trace?(result)
+                        AppSec::StackTrace.add_stack_trace(context, result)
+                      end
                     end
 
                     block = GraphQL::Reactive::Multiplex.publish(engine, gateway_multiplex)

lib/datadog/appsec/contrib/rack/gateway/watcher.rb

@@ -44,6 +44,11 @@ def watch_request(gateway = Instrumentation.gateway)
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
                       context.waf_runner.events << event
+
+                      # Include the stack trace if waf result includes it
+                      if AppSec::StackTrace.include_stack_trace?(result)
+                        AppSec::StackTrace.add_stack_trace(context, result)
+                      end
                     end
                   end
 
@@ -75,6 +80,11 @@ def watch_response(gateway = Instrumentation.gateway)
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
                       context.waf_runner.events << event
+
+                      # Include the stack trace if waf result includes it
+                      if AppSec::StackTrace.include_stack_trace?(result)
+                        AppSec::StackTrace.add_stack_trace(context, result)
+                      end
                     end
                   end
 
@@ -106,6 +116,11 @@ def watch_request_body(gateway = Instrumentation.gateway)
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
                       context.waf_runner.events << event
+
+                      # Include the stack trace if waf result includes it
+                      if AppSec::StackTrace.include_stack_trace?(result)
+                        AppSec::StackTrace.add_stack_trace(context, result)
+                      end
                     end
                   end
 

lib/datadog/appsec/contrib/rails/gateway/watcher.rb

@@ -40,6 +40,11 @@ def watch_request_action(gateway = Instrumentation.gateway)
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
                       context.waf_runner.events << event
+
+                      # Include the stack trace if waf result includes it
+                      if AppSec::StackTrace.include_stack_trace?(result)
+                        AppSec::StackTrace.add_stack_trace(context, result)
+                      end
                     end
                   end
 

lib/datadog/appsec/contrib/sinatra/gateway/watcher.rb

@@ -42,6 +42,11 @@ def watch_request_dispatch(gateway = Instrumentation.gateway)
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
                       context.waf_runner.events << event
+
+                      # Include the stack trace if waf result includes it
+                      if AppSec::StackTrace.include_stack_trace?(result)
+                        AppSec::StackTrace.add_stack_trace(context, result)
+                      end
                     end
                   end
 
@@ -73,6 +78,11 @@ def watch_request_routed(gateway = Instrumentation.gateway)
                       context.trace.keep! if context.trace
                       Datadog::AppSec::Event.tag_and_keep!(context, result)
                       context.waf_runner.events << event
+
+                      # Include the stack trace if waf result includes it
+                      if AppSec::StackTrace.include_stack_trace?(result)
+                        AppSec::StackTrace.add_stack_trace(context, result)
+                      end
                     end
                   end
 

lib/datadog/appsec/ext.rb

@@ -7,10 +7,12 @@ 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'
+      TAG_STACK_TRACE = '_dd.stack'
     end
   end
 end

lib/datadog/appsec/monitor/gateway/watcher.rb

@@ -38,6 +38,11 @@ def watch_user_id(gateway = Instrumentation.gateway)
                     context.trace.keep! if context.trace
                     Datadog::AppSec::Event.tag_and_keep!(context, result)
                     context.waf_runner.events << event
+
+                    # Include the stack trace if waf result includes it
+                    if AppSec::StackTrace.include_stack_trace?(result)
+                      AppSec::StackTrace.add_stack_trace(context, result)
+                    end
                   end
                 end
 

lib/datadog/appsec/stack_trace.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'msgpack'
+
 module Datadog
   module AppSec
     # Stack trace
@@ -33,6 +35,25 @@ def to_h
         }
       end
 
+      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
+
+      def pretty_print(q)
+        q.text to_s
+      end
+
       class << self
         # Create a stack trace from the result of WAF execution
         def from_waf_result(waf_result)
@@ -45,7 +66,29 @@ def from_waf_result(waf_result)
 
         # Whether to include the stack trace in the event
         def include_stack_trace?(waf_result)
-          Datadog.configuration.appsec.stack_trace.enabled  && waf_result.actions.key?('generate_stack')
+          Datadog.configuration.appsec.stack_trace.enabled && waf_result.actions.key?('generate_stack')
+        end
+
+        # Add the stack trace to the trace
+        def add_stack_trace(scope, result)
+          # We use methods defined in Tracing::Metadata::Tagging,
+          # which means we can work on both the trace and the service entry span
+          service_entry_operation = scope.trace || scope.span
+
+          unless service_entry_operation
+            Datadog.logger.debug { "Cannot find trace or service entry span to add stack trace" }
+            return
+          end
+
+          stack_trace = from_waf_result(result)
+
+          service_entry_operation.modify_meta_struct_tag(AppSec::Ext::TAG_STACK_TRACE) do |stack_traces|
+            # _dd.stack is a hash of stack traces (Hash<String, Array<StackTrace>>)
+            stack_traces ||= {}
+            stack_traces[AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY] ||= []
+            stack_traces[AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY] << stack_trace
+            stack_traces
+          end
         end
       end
 
@@ -70,7 +113,24 @@ def to_h
             function: @function
           }
         end
+
+        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/tracing/metadata/tagging.rb

@@ -65,6 +65,10 @@ def set_tags(hash)
           hash.each { |k, v| set_tag(k, v) }
         end
 
+        def set_meta_struct_tags(hash)
+          hash.each { |k, v| set_meta_struct_tag(k, v) }
+        end
+
         # Returns true if the provided `tag` was set to a non-nil value.
         # False otherwise.
         #
@@ -111,29 +115,44 @@ def clear_metric(key)
         end
 
         # Getter for meta_struct
-        def get_struct_tag(key)
+        def get_meta_struct_tag(key)
           meta_struct[key]
         end
 
         # Setter for meta_struct, key must be a string, value can be anything
-        # Keys can be duplicated in meta_struct and meta/metrics
-        def set_struct_tag(key, value)
-          meta_struct[key] = value
+        # Keys can be duplicated in meta_struct and meta/metrics but because each tag have a different structure
+        # we must create a condition for each. Default behaviour erases the previous value.
+        def set_meta_struct_tag(key, value)
+          if key == Datadog::AppSec::Ext::TAG_STACK_TRACE && meta_struct.key?(key)
+            value.each do |event_type, stack_traces|
+              meta_struct[key][event_type] ||= []
+              meta_struct[key][event_type].concat(stack_traces)
+            end
+          else
+            meta_struct[key] = value
+          end
         rescue StandardError => e
           Datadog.logger.debug("Unable to set the tag #{key} in meta_struct, ignoring it. Caused by: #{e}")
         end
 
+        # Modifies the content of a tag in meta_struct
+        def modify_meta_struct_tag(key)
+          meta_struct[key] = yield(meta_struct[key])
+        rescue StandardError => e
+          Datadog.logger.debug("Unable to modify the tag #{key} in meta_struct, ignoring it. Caused by: #{e}")
+        end
+
         # Returns true if the provided `tag` was set to a non-nil value.
         # False otherwise.
         #
         # @param [String] tag the tag or metric to check for presence
         # @return [Boolean] if the tag is present and not nil
-        def has_struct_tag?(tag) # rubocop:disable Naming/PredicateName
-          !get_struct_tag(tag).nil? # nil is considered not present, thus we can't use `Hash#has_key?`
+        def has_meta_struct_tag?(tag) # rubocop:disable Naming/PredicateName
+          !get_meta_struct_tag(tag).nil? # nil is considered not present, thus we can't use `Hash#has_key?`
         end
 
         # This method removes a tag for the given key.
-        def clear_struct_tag(key)
+        def clear_meta_struct_tag(key)
           meta_struct.delete(key)
         end
 

lib/datadog/tracing/span.rb

@@ -142,6 +142,7 @@ def to_hash
           error: @status,
           meta: @meta,
           metrics: @metrics,
+          meta_struct: @meta_struct,
           name: @name,
           parent_id: @parent_id,
           resource: @resource,
@@ -185,12 +186,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

@@ -285,6 +285,7 @@ def to_hash
           id: @id,
           meta: meta,
           metrics: metrics,
+          meta_struct: meta_struct,
           name: @name,
           parent_id: @parent_id,
           resource: @resource,
@@ -324,12 +325,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
 
@@ -452,6 +457,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/trace_operation.rb

@@ -73,6 +73,7 @@ def initialize(
         profiling_enabled: nil,
         tags: nil,
         metrics: nil,
+        meta_struct: nil,
         trace_state: nil,
         trace_state_unknown_fields: nil,
         remote_parent: false,
@@ -105,6 +106,7 @@ def initialize(
         # Generic tags
         set_tags(tags) if tags
         set_tags(metrics) if metrics
+        set_meta_struct_tags(meta_struct) if meta_struct
 
         # State
         @root_span = nil
@@ -369,6 +371,7 @@ def fork_clone
           trace_state_unknown_fields: (@trace_state_unknown_fields && @trace_state_unknown_fields.dup),
           tags: meta.dup,
           metrics: metrics.dup,
+          meta_struct: meta_struct.dup,
           remote_parent: @remote_parent
         )
       end
@@ -508,6 +511,7 @@ def build_trace(spans, partial = false)
           service: service,
           tags: meta,
           metrics: metrics,
+          meta_struct: meta_struct,
           root_span_id: !partial ? root_span && root_span.id : nil,
           profiling_enabled: @profiling_enabled,
         )