Refactor stack_trace in multiple classes, and move it in tracing

Author: vpellan

Diff for: lib/datadog/appsec.rb

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

Diff for: 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,37 +22,34 @@ def interrupt_execution(action_params)
         throw(Datadog::AppSec::Ext::INTERRUPT, action_params)
       end
 
-      # We should change the name of this method as we also add the stack trace to the trace/span
       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
 
-          unless service_entry_operation
+          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 that path can still be nil and that include? is not a method of nil
-          # A way to fix this is to assign loc.path to a variable and use it instead, but it adds an operation
-          # Which is why we are ignoring this line
+          # 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|
-            loc.path &&
-              loc.path.include?('lib/datadog') # steep:ignore NoMethod
+            path = loc.path
+            next true if path.nil?
+
+            path.include?('lib/datadog')
           end
 
-          collected_stack_frames = AppSec::StackTrace::Collection.new(stack_frames)
+          collected_stack_frames = Datadog::Tracing::StackTrace::Collector.collect(stack_frames)
 
-          # ASCII-8BIT is encoded as byte array and backend cannot decode it properly so we enforce UTF-8 on stack_id.
-          # We must copy it as it can be frozen.
           utf8_stack_id = action_params['stack_id'].encode('UTF-8') if action_params['stack_id']
-          # We should always have a stack_id, but we should still send the stack trace if not (visible with feature flags)
-          stack_trace = AppSec::StackTrace::Representation.new(utf8_stack_id, collected_stack_frames)
+          stack_trace = Datadog::Tracing::StackTrace::Representor.new(id: utf8_stack_id, message: nil, frames: collected_stack_frames)
 
           # Add stack trace to trace or service entry span by modifying meta_struct['_dd.stack']
           service_entry_operation.modify_meta_struct_tag(AppSec::Ext::TAG_STACK_TRACE) do |stack_trace_categories|

Diff for: lib/datadog/appsec/stack_trace.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 for: lib/datadog/tracing/stack_trace/collector.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)
+        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 for: lib/datadog/tracing/stack_trace/frame.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)
+        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