From 0e5e082b88988406a6258f50687b836acbba63c5 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Thu, 2 Jan 2025 11:54:39 +0100
Subject: [PATCH] Add StackTrace class methods from_waf_result and
 include_stack_trace?

---
 lib/datadog/appsec/stack_trace.rb  | 36 +++++++++++++++++++++++-------
 sig/datadog/appsec/stack_trace.rbs | 10 ++++++++-
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/lib/datadog/appsec/stack_trace.rb b/lib/datadog/appsec/stack_trace.rb
index 1211ef9380b..7543b91dbd5 100644
--- a/lib/datadog/appsec/stack_trace.rb
+++ b/lib/datadog/appsec/stack_trace.rb
@@ -9,14 +9,18 @@ class StackTrace
       def initialize(id, stack_trace, message: nil)
         @id = id
         @message = message
-        @frames = stack_trace.each_with_index.map do |frame, index|
-          Frame.new(
-            index,
-            text: frame.to_s,
-            file: frame.absolute_path, # path will show 'main' instead of the actual file
-            line: frame.lineno,
-            function: frame.base_label # label will show 'block in function' instead of just 'function'
-          )
+        if stack_trace
+          @frames = stack_trace.each_with_index.map do |frame, index|
+            Frame.new(
+              index,
+              text: frame.to_s,
+              file: frame.absolute_path, # path will show 'main' instead of the actual file
+              line: frame.lineno,
+              function: frame.base_label # label will show 'block in function' instead of just 'function'
+            )
+          end
+        else
+          @frames = []
         end
       end
 
@@ -29,6 +33,22 @@ def to_h
         }
       end
 
+      class << self
+        # Create a stack trace from the result of WAF execution
+        def from_waf_result(waf_result)
+          stack_frames = caller_locations || []
+          # caller_locations without params always returns an array but rbs still thinks it can be nil
+          stack_frames.reject! { |loc| loc.path && loc.path.include?('lib/datadog') }
+          stack_trace_id = waf_result.actions['generate_stack']['stack_id']
+          stack_trace = AppSec::StackTrace.new(stack_trace_id, stack_frames)
+        end
+
+        # 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')
+        end
+      end
+
       # Stack frame
       class Frame
         attr_reader :id, :text, :file, :line, :function
diff --git a/sig/datadog/appsec/stack_trace.rbs b/sig/datadog/appsec/stack_trace.rbs
index c940b9d776f..5f62de55ac1 100644
--- a/sig/datadog/appsec/stack_trace.rbs
+++ b/sig/datadog/appsec/stack_trace.rbs
@@ -7,7 +7,15 @@ module Datadog
 
       @to_h: ::Hash[::Symbol, untyped]
 
-      def initialize: (::String id, ::Array[::Thread::Backtrace::Location] stack_trace, ?message: ::String?) -> void
+      def initialize: (::String id, ::Array[::Thread::Backtrace::Location]? stack_trace, ?message: ::String?) -> void
+                    
+      def to_h: () -> ::Hash[::Symbol, untyped]
+              
+      private
+
+      def self.from_waf_result: (Datadog::AppSec::WAF::Result waf_result) -> Datadog::AppSec::StackTrace
+
+      def self.include_stack_trace?: (Datadog::AppSec::WAF::Result waf_result) -> bool
 
       class Frame
         attr_reader id: ::Integer