Reject SafeMarshal collection lengths longer than the remaining input

lib/rubygems/safe_marshal/reader.rb

@@ -26,6 +26,9 @@ class DataTooShortError < Error
       class NegativeLengthError < Error
       end
 
+      class LengthTooLongError < Error
+      end
+
       def initialize(io)
         @io = io
         @object_links = {}
@@ -93,6 +96,18 @@ def read_integer
         end
       end
 
+      # Reads an element count and validates it against the number of bytes
+      # remaining in the input, since each element to be read consumes at
+      # least one byte. This prevents allocating huge backing stores for
+      # maliciously crafted lengths that could never be satisfied.
+      def read_count
+        count = read_integer
+        raise NegativeLengthError if count < 0
+        remaining = @io.size - @io.pos
+        raise LengthTooLongError, "expected #{count} elements, but only #{remaining} bytes remain" if count > remaining
+        count
+      end
+
       def read_element
         type = read_byte
         case type
@@ -172,9 +187,8 @@ def read_user_defined
       private_constant :EMPTY_ARRAY
 
       def read_array
-        length = read_integer
+        length = read_count
         return EMPTY_ARRAY if length == 0
-        raise NegativeLengthError if length < 0
         elements = Array.new(length) do
           read_element
         end
@@ -183,8 +197,7 @@ def read_array
 
       def read_object_with_ivars
         object = read_element
-        length = read_integer
-        raise NegativeLengthError if length < 0
+        length = read_count
         ivars = Array.new(length) do
           [read_element, read_element]
         end
@@ -211,7 +224,7 @@ def read_object_link
       private_constant :EMPTY_HASH
 
       def read_hash
-        length = read_integer
+        length = read_count
         return EMPTY_HASH if length == 0
         pairs = Array.new(length) do
           [read_element, read_element]
@@ -220,8 +233,7 @@ def read_hash
       end
 
       def read_hash_with_default_value
-        length = read_integer
-        raise NegativeLengthError if length < 0
+        length = read_count
         pairs = Array.new(length) do
           [read_element, read_element]
         end
@@ -232,8 +244,7 @@ def read_hash_with_default_value
       def read_object
         name = read_element
         object = Elements::Object.new(name)
-        length = read_integer
-        raise NegativeLengthError if length < 0
+        length = read_count
         ivars = Array.new(length) do
           [read_element, read_element]
         end

test/rubygems/test_gem_safe_marshal.rb

@@ -423,10 +423,10 @@ def test_unexpected_eof
     end
     assert_equal e.message, "Unexpected EOF"
 
-    e = assert_raise(Gem::SafeMarshal::Reader::EOFError) do
+    e = assert_raise(Gem::SafeMarshal::Reader::LengthTooLongError) do
       Gem::SafeMarshal.safe_load("\x04\x08[\x06")
     end
-    assert_equal e.message, "Unexpected EOF"
+    assert_equal e.message, "expected 1 elements, but only 0 bytes remain"
 
     e = assert_raise(Gem::SafeMarshal::Reader::EOFError) do
       Gem::SafeMarshal.safe_load("\004\010:\012")
@@ -459,6 +459,33 @@ def test_negative_length
     assert_raise(Gem::SafeMarshal::Reader::EOFError) do
       Gem::SafeMarshal.safe_load("\004\010@\377")
     end
+    assert_raise(Gem::SafeMarshal::Reader::NegativeLengthError) do
+      Gem::SafeMarshal.safe_load("\004\010{\325")
+    end
+  end
+
+  def test_length_too_long
+    huge_length = "\x04#{[2_000_000_000].pack("V")}".b
+
+    assert_raise(Gem::SafeMarshal::Reader::LengthTooLongError) do
+      Gem::SafeMarshal.safe_load("\x04\x08[#{huge_length}")
+    end
+    assert_raise(Gem::SafeMarshal::Reader::LengthTooLongError) do
+      Gem::SafeMarshal.safe_load("\x04\x08{#{huge_length}")
+    end
+    assert_raise(Gem::SafeMarshal::Reader::LengthTooLongError) do
+      Gem::SafeMarshal.safe_load("\x04\x08}#{huge_length}")
+    end
+    assert_raise(Gem::SafeMarshal::Reader::LengthTooLongError) do
+      Gem::SafeMarshal.safe_load("\x04\x08I\"\x00#{huge_length}")
+    end
+    assert_raise(Gem::SafeMarshal::Reader::LengthTooLongError) do
+      Gem::SafeMarshal.safe_load("\x04\x08o:\x06C#{huge_length}")
+    end
+
+    # lengths that fit within the remaining input still parse
+    assert_equal [1, 2, 3], Gem::SafeMarshal.safe_load("\x04\x08[\x08i\x06i\ai\x08")
+    assert_equal({ 1 => 2 }, Gem::SafeMarshal.safe_load("\x04\x08{\x06i\x06i\a"))
   end
 
   def assert_safe_load_marshal(dumped, additional_methods: [], permitted_ivars: nil, equality: true, marshal_dump_equality: true,