Cap nesting depth in Gem::SafeMarshal::Reader

lib/rubygems/safe_marshal/reader.rb

@@ -29,10 +29,14 @@ class NegativeLengthError < Error
       class LengthTooLongError < Error
       end
 
+      class TooDeeplyNestedError < Error
+      end
+
       def initialize(io)
         @io = io
         @object_links = {}
         @symbol_links = {}
+        @depth = 0
       end
 
       def read!
@@ -47,6 +51,9 @@ def read!
       MARSHAL_VERSION = [Marshal::MAJOR_VERSION, Marshal::MINOR_VERSION].map(&:chr).join.freeze
       private_constant :MARSHAL_VERSION
 
+      MAX_NESTING_DEPTH = 1_000
+      private_constant :MAX_NESTING_DEPTH
+
       def read_header
         v = @io.read(2)
         raise UnsupportedVersionError, "Unsupported marshal version #{v.bytes.map(&:ord).join(".")}, expected #{Marshal::MAJOR_VERSION}.#{Marshal::MINOR_VERSION}" unless v == MARSHAL_VERSION
@@ -109,6 +116,9 @@ def read_count
       end
 
       def read_element
+        @depth += 1
+        raise TooDeeplyNestedError, "exceeded maximum nesting depth (#{MAX_NESTING_DEPTH})" if @depth > MAX_NESTING_DEPTH
+
         type = read_byte
         case type
         when 34 then read_string # ?"
@@ -139,6 +149,8 @@ def read_element
         else
           raise Error, "Unknown marshal type discriminator #{type.chr.inspect} (#{type})"
         end
+      ensure
+        @depth -= 1
       end
 
       STRING_E_SYMBOL = Elements::Symbol.new("E").freeze

test/rubygems/test_gem_safe_marshal.rb

@@ -478,6 +478,35 @@ def test_unexpected_eof
     assert_equal e.message, "expected 1 bytes, got EOF"
   end
 
+  def test_nesting_depth_is_capped
+    # 2 bytes per level on the wire, so this is a small payload
+    payload = "\x04\x08".b + ("[\x06".b * 2_400) + "0".b
+
+    e = assert_raise(Gem::SafeMarshal::Reader::TooDeeplyNestedError) do
+      Gem::SafeMarshal.safe_load(payload)
+    end
+    assert_equal "exceeded maximum nesting depth (1000)", e.message
+
+    # the cap raises a StandardError, so callers that already rescue
+    # StandardError around safe_load keep working
+    assert_kind_of StandardError, e
+  end
+
+  def test_nesting_below_the_cap_still_parses
+    payload = "\x04\x08".b + ("[\x06".b * 998) + "0".b
+    # deliberately not asserting on the value itself: inspecting a 998-deep
+    # array is what blows the stack, not parsing it
+    assert_equal ::Array, Gem::SafeMarshal.safe_load(payload).class
+  end
+
+  def test_repeated_siblings_do_not_count_as_depth
+    # 5_000 elements, all at depth 2 -- must not trip the cap
+    payload = "\x04\x08[".b + "\x02\x88\x13".b + ("0".b * 5_000)
+    parsed = Gem::SafeMarshal.safe_load(payload)
+    assert_equal ::Array, parsed.class
+    assert_equal 5_000, parsed.size
+  end
+
   def test_negative_length
     assert_raise(Gem::SafeMarshal::Reader::NegativeLengthError) do
       Gem::SafeMarshal.safe_load("\004\010}\325")