Skip to content

Commit f847fac

Browse files
vpellanvpellan
committed
Rework stacktrace collection architecture
1 parent a6ab36f commit f847fac

File tree

2 files changed

+127
-107
lines changed

2 files changed

+127
-107
lines changed

‎lib/datadog/appsec/actions_handler.rb

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,44 @@ def interrupt_execution(action_params)
1919
throw(Datadog::AppSec::Ext::INTERRUPT, action_params)
2020
end
2121

22+
# We should change the name of this method as we also add the stack trace to the trace/span
2223
def generate_stack(action_params)
23-
AppSec::StackTrace.add_stack_trace(action_params) if Datadog.configuration.appsec.stack_trace.enabled
24+
if Datadog.configuration.appsec.stack_trace.enabled
25+
context = AppSec::Context.active
26+
# We use methods defined in Tracing::Metadata::Tagging,
27+
# which means we can work on both the trace and the service entry span
28+
service_entry_operation = context && (context.trace || context.span)
29+
30+
unless service_entry_operation
31+
Datadog.logger.debug { 'Cannot find trace or service entry span to add stack trace' }
32+
return
33+
end
34+
35+
stack_frames = caller_locations || []
36+
# caller_locations without params always returns an array but rbs still thinks it can be nil
37+
stack_frames.reject! { |loc| loc.path && loc.path.include?('lib/datadog') }
38+
39+
collected_stack_frames = AppSec::StackTrace::Collection.new(stack_frames)
40+
41+
# ASCII-8BIT is encoded as byte array and backend cannot decode it properly so we enforce UTF-8 on stack_id.
42+
# We must copy it as it can be frozen.
43+
utf8_stack_id = action_params['stack_id'] && action_params['stack_id'].encode('UTF-8')
44+
stack_trace = AppSec::StackTrace::Representation.new(utf8_stack_id, collected_stack_frames)
45+
46+
# Add stack trace to trace or service entry span by modifying meta_struct['_dd.stack']
47+
service_entry_operation.modify_meta_struct_tag(AppSec::Ext::TAG_STACK_TRACE) do |stack_trace_categories|
48+
# _dd.stack is a hash of stack traces (Hash<String, Array<StackTrace>>)
49+
stack_trace_categories ||= {}
50+
exploit_stack_traces = (stack_trace_categories[AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY] ||= [])
51+
# We only keep the first 'max_collect' stack traces
52+
# In practice we should also check the vulnerability category size, but we don't have IAST in Ruby
53+
if Datadog.configuration.appsec.stack_trace.max_collect == 0 ||
54+
exploit_stack_traces.size < Datadog.configuration.appsec.stack_trace.max_collect
55+
exploit_stack_traces << stack_trace
56+
end
57+
stack_trace_categories
58+
end
59+
end
2460
end
2561

2662
def generate_schema(_action_params); end

‎lib/datadog/appsec/stack_trace.rb

Lines changed: 90 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -4,117 +4,62 @@
44

55
module Datadog
66
module AppSec
7-
# Stack trace
8-
class StackTrace
9-
attr_reader :id, :message, :frames
10-
11-
def initialize(id, stack_trace, message: nil)
12-
# ASCII-8BIT is encoded as byte array and backend cannot decode it properly
13-
# so we enforce UTF-8 on strings
14-
# We must copy them as they can be frozen
15-
@id = id && id.encode('UTF-8')
16-
@message = message && message.encode('UTF-8')
17-
@frames = if stack_trace
18-
top_frames = if Datadog.configuration.appsec.stack_trace.max_depth == 0
19-
stack_trace.size
20-
else
21-
(Datadog.configuration.appsec.stack_trace.max_depth *
22-
Datadog.configuration.appsec.stack_trace.max_depth_top_percent / 100.0).round
23-
end
24-
bottom_frames = Datadog.configuration.appsec.stack_trace.max_depth - top_frames
25-
cropped_stack_trace = []
26-
stack_trace.each_with_index do |frame, index|
27-
next if index >= top_frames && index < stack_trace.size - bottom_frames
28-
29-
cropped_stack_trace << Frame.new(
30-
# We must keep the original stack trace index, to show that it was cropped
31-
index,
32-
text: frame.to_s,
33-
# path will show 'main' instead of the actual file
34-
# but absolute_path does not work for e.g. <internal:kernel>
35-
file: frame.absolute_path || frame.path,
36-
line: frame.lineno,
37-
# label will show 'block in function' instead of just 'function'
38-
function: frame.label
39-
)
40-
end
41-
cropped_stack_trace
42-
else
43-
[]
44-
end
45-
end
46-
47-
def to_h
48-
@to_h ||= {
49-
language: 'ruby',
50-
id: @id,
51-
message: @message,
52-
frames: @frames.map(&:to_h)
53-
}
54-
end
55-
56-
def to_msgpack(packer = nil)
57-
packer ||= MessagePack::Packer.new
58-
59-
packer.write_map_header(4)
60-
packer.write('language')
61-
packer.write('ruby')
62-
packer.write('id')
63-
packer.write(@id)
64-
packer.write('message')
65-
packer.write(@message)
66-
packer.write('frames')
67-
packer.write(@frames)
68-
packer
69-
end
70-
71-
def pretty_print(q)
72-
q.text to_s
73-
end
74-
75-
class << self
76-
# Add the stack trace to the trace
77-
def add_stack_trace(action_params)
78-
# As ActionsHandler is only called when there is a context, we shouldn't need to check if it exists
79-
context = AppSec::Context.active
80-
# We use methods defined in Tracing::Metadata::Tagging,
81-
# which means we can work on both the trace and the service entry span
82-
service_entry_operation = context.trace || context.span
83-
84-
unless service_entry_operation
85-
Datadog.logger.debug { 'Cannot find trace or service entry span to add stack trace' }
86-
return
7+
module StackTrace
8+
# Stack trace collection class, including Enumerable (iterates on cropped frames array)
9+
class Collection
10+
include Enumerable
11+
12+
def initialize(frames)
13+
max_depth = Datadog.configuration.appsec.stack_trace.max_depth
14+
# 0 means no limit, setting the top_frames limit to the size of the stack trace
15+
# will ensure that we keep all the frames.
16+
# If not set to 0, "max_depth_top_percent" % of the max_depth value
17+
# corresponds to the number of top frames we want to keep.
18+
top_frames = if max_depth == 0
19+
frames.size
20+
else
21+
(max_depth * Datadog.configuration.appsec.stack_trace.max_depth_top_percent / 100.0).round
22+
end
23+
bottom_frames = max_depth - top_frames
24+
cropped_frames = []
25+
frames.each_with_index do |frame, index|
26+
next if index >= top_frames && index < frames.size - bottom_frames
27+
28+
# ASCII-8BIT is encoded as byte array and backend cannot decode it properly
29+
# so we enforce UTF-8 on strings.
30+
# We must copy them as they can be frozen
31+
cropped_frames << Frame.new(
32+
# We must keep the original stack trace index, to show that it was cropped
33+
index,
34+
text: frame.to_s && frame.to_s.encode('UTF-8'),
35+
# path will show 'main' instead of the actual file
36+
# but absolute_path does not work for e.g. <internal:kernel>
37+
file: (frame.absolute_path || frame.path) && (frame.absolute_path || frame.path).encode('UTF-8'),
38+
line: frame.lineno,
39+
# label will show 'block in function' instead of just 'function'
40+
function: frame.label && frame.label.encode('UTF-8')
41+
)
8742
end
43+
@frames = cropped_frames
44+
end
8845

89-
stack_trace = from_waf_result(action_params)
90-
91-
service_entry_operation.modify_meta_struct_tag(AppSec::Ext::TAG_STACK_TRACE) do |stack_traces|
92-
# _dd.stack is a hash of stack traces (Hash<String, Array<StackTrace>>)
93-
stack_traces ||= {}
94-
exploit_stack_traces = (stack_traces[AppSec::Ext::EXPLOIT_PREVENTION_EVENT_CATEGORY] ||= [])
95-
# We only keep the first 'max_collect' stack traces
96-
# In practice we should also check the vulnerability category size, but we don't have IAST in Ruby
97-
if Datadog.configuration.appsec.stack_trace.max_collect == 0 ||
98-
exploit_stack_traces.size < Datadog.configuration.appsec.stack_trace.max_collect
99-
exploit_stack_traces << stack_trace
100-
end
101-
stack_traces
102-
end
46+
def each(&block)
47+
@frames.each(&block)
10348
end
10449

105-
private
50+
def to_msgpack(packer = nil)
51+
packer ||= MessagePack::Packer.new
10652

107-
# Create a stack trace from the result of WAF execution
108-
def from_waf_result(action_params)
109-
stack_frames = caller_locations || []
110-
# caller_locations without params always returns an array but rbs still thinks it can be nil
111-
stack_frames.reject! { |loc| loc.path && loc.path.include?('lib/datadog') }
112-
stack_trace_id = action_params['stack_id']
113-
AppSec::StackTrace.new(stack_trace_id, stack_frames)
53+
packer.write_array_header(@frames.size)
54+
@frames.each do |frame|
55+
packer.write(frame)
56+
end
57+
packer
11458
end
11559
end
11660

117-
# Stack frame
61+
# Stack frame.
62+
# We cannot use Struct as we need to encode it to MessagePack, and Struct does not provide a 'to_msgpack' method.
11863
class Frame
11964
attr_reader :id, :text, :file, :line, :function
12065

@@ -123,10 +68,10 @@ class Frame
12368
# We must copy them as they can be frozen
12469
def initialize(id, text: nil, file: nil, line: nil, function: nil)
12570
@id = id
126-
@text = text && text.encode('UTF-8')
127-
@file = file && file.encode('UTF-8')
71+
@text = text
72+
@file = file
12873
@line = line
129-
@function = function && function.encode('UTF-8')
74+
@function = function
13075
end
13176

13277
def to_h
@@ -156,6 +101,45 @@ def to_msgpack(packer = nil)
156101
packer
157102
end
158103
end
104+
105+
# Stack trace representation.
106+
class Representation
107+
attr_reader :id, :message, :stack_trace
108+
109+
def initialize(id, stack_trace, message: nil)
110+
@id = id
111+
@message = message
112+
@stack_trace = stack_trace
113+
end
114+
115+
def to_h
116+
@to_h ||= {
117+
language: 'ruby',
118+
id: @id,
119+
message: @message,
120+
frames: @stack_trace.map(&:to_h)
121+
}
122+
end
123+
124+
def to_msgpack(packer = nil)
125+
packer ||= MessagePack::Packer.new
126+
127+
packer.write_map_header(4)
128+
packer.write('language')
129+
packer.write('ruby')
130+
packer.write('id')
131+
packer.write(@id)
132+
packer.write('message')
133+
packer.write(@message)
134+
packer.write('frames')
135+
packer.write(@stack_trace)
136+
packer
137+
end
138+
139+
def pretty_print(q)
140+
q.text to_s
141+
end
142+
end
159143
end
160144
end
161145
end

0 commit comments

Comments
 (0)