Add extract payloads max depth

lib/aikido/zen/config.rb

@@ -56,6 +56,10 @@ class Config
     #   each initial heartbeat event.
     attr_accessor :initial_heartbeat_delays
 
+    # @return [Integer] the maximum number of levels deep we want to traverse
+    #   when extracting payloads from a nested request property. Defaults to 64.
+    attr_accessor :extract_payloads_max_depth
+
     # @return [Symbol] the agent mode for forked worker processes. Can be set
     #   through the AIKIDO_AGENT_MODE environment variable.
     attr_reader :agent_mode
@@ -249,6 +253,7 @@ def initialize
       self.api_timeouts = 10
       self.polling_interval = 60 # 1 min
       self.initial_heartbeat_delays = [30, 60 * 2] # 30 sec, 2 min
+      self.extract_payloads_max_depth = 64
       self.agent_mode = ENV.fetch("AIKIDO_AGENT_MODE", "shared")
       self.worker_process_polling_interval = 10
       self.worker_process_polling_jitter = 10

lib/aikido/zen/context.rb

@@ -45,11 +45,14 @@ def self.from_rack_env(env, config = Aikido::Zen.config)
     # @yieldparam request [Rack::Request] the given request object.
     # @yieldreturn [Hash<Symbol, #flat_map>] map of payload source types
     #   to the actual data from the request to populate them.
-    def initialize(request, settings: Aikido::Zen.runtime_settings, &sources)
+    def initialize(request, zen: Aikido::Zen, config: zen.config, settings: zen.runtime_settings, &sources)
       @request = request
+      @config = config
       @settings = settings
       @payload_sources = sources
 
+      @max_depth = @config.extract_payloads_max_depth
+
       @metadata = {}
       @scanning = false
       @request_bypassed = nil
@@ -102,31 +105,46 @@ def payload_sources
     private
 
     # @!visibility private
-    def extract_payloads_from(data, source_type, prefix = nil)
+    def extract_payloads_from(data, source_type, prefix = nil, depth = 0)
+      return [] if depth > @max_depth
+
       if data.is_a?(String)
         [Payload.new(data, source_type, prefix.to_s)]
       elsif data.respond_to?(:to_hash)
         data.to_hash.flat_map do |key, value|
-          extract_payloads_from(value, source_type, [prefix, key].compact.join("."))
+          extract_payloads_from(value, source_type, [prefix, key].compact.join("."), depth + 1)
         end
       elsif data.respond_to?(:to_ary)
         array = data.to_ary
         return array if array.empty?
 
         payloads = array.flat_map.with_index do |value, index|
-          extract_payloads_from(value, source_type, [prefix, index].compact.join("."))
+          extract_payloads_from(value, source_type, [prefix, index].compact.join("."), depth + 1)
         end
 
         unless Aikido::Zen.config.harden?
           # Special case for File.join given a possibly nested array of strings,
           # as might occur when a query parameter is an array.
-          begin
-            string = File.join__internal_for_aikido_zen(*array)
-            if unsafe_path?(string)
-              payloads << Payload.new(string, source_type, [prefix, "__File.join__"].compact.join("."))
+
+          # File.join recursively joins nested string arrays, and can overflow
+          # the stack given deeply nested arrays.
+          #
+          # Flatten array to max depth and check that the flattened array is an
+          # array of strings before calling File.join__internal_for_aikido_zen,
+          # only if the array was fully flattened, to prevent a stack overflow.
+          #
+          # Checking that all values are Strings handily prevents a TypeError
+          # from being raised.
+          flattened_array = array.flatten(@max_depth)
+          if flattened_array.all? { |string| string.is_a?(String) }
+            begin
+              string = File.join__internal_for_aikido_zen(*flattened_array)
+              if unsafe_path?(string)
+                payloads << Payload.new(string, source_type, [prefix, "__File.join__"].compact.join("."))
+              end
+            rescue
+              # Could not create special payload for File.join.
             end
-          rescue
-            # Could not create special payload for File.join.
           end
         end
 

test/aikido/zen/config_test.rb

@@ -20,6 +20,7 @@ class Aikido::Zen::ConfigTest < ActiveSupport::TestCase
     assert_equal 10, @config.api_timeouts[:write_timeout]
     assert_equal 60, @config.polling_interval
     assert_equal [30, 120], @config.initial_heartbeat_delays
+    assert_equal 64, @config.extract_payloads_max_depth
     assert_equal :shared, @config.agent_mode
     assert_equal 10, @config.worker_process_polling_interval
     assert_equal 10, @config.worker_process_polling_jitter

test/aikido/zen/context_test.rb

@@ -14,6 +14,39 @@ class Aikido::Zen::ContextTest < ActiveSupport::TestCase
     assert_equal framework_request, context.request
   end
 
+  test "payload extraction does not recurse past the configured maximum depth" do
+    Aikido::Zen.config.extract_payloads_max_depth = 2
+
+    request = DummyRequest.new({})
+    data = {"a" => {"b" => "kept at the depth limit", "c" => {"d" => "dropped past the depth limit"}}}
+    context = Aikido::Zen::Context.new(request) { {body: data} }
+
+    assert_includes context.payloads, Aikido::Zen::Payload.new("kept at the depth limit", :body, "a.b")
+    refute_includes context.payloads, Aikido::Zen::Payload.new("dropped past the depth limit", :body, "a.c.d")
+  end
+
+  test "payload extraction caps recursion depth to prevent a stack overflow" do
+    data = "leaf"
+    10_000.times { data = {"k" => data} }
+
+    request = DummyRequest.new({})
+    context = Aikido::Zen::Context.new(request) { {body: data} }
+
+    assert_nothing_raised { context.payloads }
+  end
+
+  test "the File.join special case for arrays does not stack overflow on deeply nested arrays" do
+    Aikido::Zen.config.harden = false
+
+    data = "leaf"
+    200_000.times { data = [data] }
+
+    request = DummyRequest.new({})
+    context = Aikido::Zen::Context.new(request) { {query: data} }
+
+    assert_nothing_raised { context.payloads }
+  end
+
   module GenericTests
     extend ActiveSupport::Testing::Declarative