AVRO-4297: [ruby] Bound allocation when decoding length-prefixed values and collections

lang/ruby/lib/avro/io.rb

@@ -37,6 +37,13 @@ def initialize(writers_schema, readers_schema)
     class BinaryDecoder
       # Read leaf values
 
+      # Reads with a declared length above this many bytes are validated
+      # against the number of bytes actually remaining (when the reader can
+      # report its size) before allocating, to guard against an out-of-memory
+      # attack from a malicious or truncated input. Smaller reads skip the
+      # check to avoid per-value overhead.
+      MAX_UNCHECKED_READ = 1024 * 1024
+
       # reader is an object on which we can call read, seek and tell.
       attr_reader :reader
       def initialize(reader)
@@ -103,10 +110,43 @@ def read_string
       end
 
       def read(len)
-        # Read n bytes
+        # Read n bytes. Reject a declared length that exceeds the bytes
+        # actually remaining before allocating for it, to guard against an
+        # out-of-memory attack from a malicious or truncated input. The check
+        # is only applied to larger reads; smaller reads and stream readers that
+        # cannot report their size fall back to reading directly.
+        if len < 0
+          # A negative length would make IO#read return the rest of the stream,
+          # which bypasses the size check and can allocate without bound.
+          raise AvroError, "Cannot read a negative number of bytes: #{len}"
+        end
+        if len > MAX_UNCHECKED_READ
+          remaining = bytes_remaining
+          if remaining && len > remaining
+            raise AvroError, "Cannot read #{len} bytes, only #{remaining} remaining"
+          end
+        end
         @reader.read(len)
       end
 
+      # Number of bytes still available to read, or nil when the reader cannot
+      # report its size. Used to reject a declared length or collection block
+      # count that exceeds the data actually available before allocating for it.
+      def bytes_remaining
+        return nil unless @reader.respond_to?(:size) && @reader.respond_to?(:tell)
+        size = @reader.size
+        pos = @reader.tell
+        return nil unless size.is_a?(Integer) && pos.is_a?(Integer)
+        # Clamp negative (e.g. the file was truncated below the current position)
+        # to 0 so callers see "no bytes available" rather than a confusing
+        # negative count.
+        [size - pos, 0].max
+      rescue IOError, SystemCallError, NotImplementedError
+        # The reader responds to #size/#tell but cannot actually report them;
+        # treat the remaining size as unknown and skip the check.
+        nil
+      end
+
       def skip_null
         nil
       end
@@ -312,12 +352,14 @@ def read_enum(writers_schema, readers_schema, decoder)
 
       def read_array(writers_schema, readers_schema, decoder)
         read_items = []
+        min_bytes = min_bytes_per_element(writers_schema.items)
         block_count = decoder.read_long
         while block_count != 0
           if block_count < 0
             block_count = -block_count
             _block_size = decoder.read_long
           end
+          ensure_collection_available(decoder, block_count, min_bytes)
           block_count.times do
             read_items << read_data(writers_schema.items,
                                     readers_schema.items,
@@ -331,12 +373,15 @@ def read_array(writers_schema, readers_schema, decoder)
 
       def read_map(writers_schema, readers_schema, decoder)
         read_items = {}
+        # Map keys are strings (>= 1 byte length prefix) plus the value.
+        min_bytes = 1 + min_bytes_per_element(writers_schema.values)
         block_count = decoder.read_long
         while block_count != 0
           if block_count < 0
             block_count = -block_count
             _block_size = decoder.read_long
           end
+          ensure_collection_available(decoder, block_count, min_bytes)
           block_count.times do
             key = decoder.read_string
             read_items[key] = read_data(writers_schema.values,
@@ -496,6 +541,45 @@ def skip_record(writers_schema, decoder)
       end
 
       private
+      # Minimum number of bytes a single value of the given schema can occupy on
+      # the wire. Used to reject an array/map block count that could not be
+      # backed by the bytes remaining. A type that can encode to zero bytes
+      # (null) returns 0, which disables the check for it (so an array of nulls
+      # is not falsely rejected).
+      def min_bytes_per_element(schema, visited = nil)
+        visited ||= {}.compare_by_identity
+        case schema.type_sym
+        when :null then 0
+        when :float then 4
+        when :double then 8
+        when :fixed then schema.size
+        when :record, :error, :request
+          return 0 if visited[schema]
+          visited[schema] = true
+          total = schema.fields.sum { |field| min_bytes_per_element(field.type, visited) }
+          visited.delete(schema)
+          total
+        else
+          # boolean, int, long, bytes, string, enum, union, array, map: >= 1 byte
+          # (a union encodes at least a 1-byte branch index).
+          1
+        end
+      end
+
+      # Reject a collection (array or map) block whose declared element count
+      # could not be backed by the bytes actually remaining, before iterating.
+      # Skipped when the per-element minimum is zero or when the reader cannot
+      # report how many bytes remain.
+      def ensure_collection_available(decoder, count, min_bytes_per_element)
+        return if count <= 0 || min_bytes_per_element <= 0
+        remaining = decoder.bytes_remaining
+        # Compare via integer division rather than multiplying, so a huge count
+        # does not create a large intermediate product.
+        if remaining && count > remaining / min_bytes_per_element
+          raise AvroError, "Collection claims #{count} elements with at least #{min_bytes_per_element} bytes each, but only #{remaining} bytes are available"
+        end
+      end
+
       def skip_blocks(decoder, &blk)
         block_count = decoder.read_long
         while block_count != 0

lang/ruby/test/test_io.rb

@@ -42,6 +42,70 @@ def test_bytes
     check_default('"bytes"', '"foo"', "foo")
   end
 
+  # A bytes/string value declares a length prefix; a malicious or truncated
+  # input can declare far more bytes than actually exist. On a reader that can
+  # report its size, that is rejected before allocating for it.
+  def test_read_bytes_rejects_length_beyond_stream
+    writer = StringIO.new
+    Avro::IO::BinaryEncoder.new(writer).write_long(100 * 1024 * 1024)
+    reader = StringIO.new(writer.string)
+    decoder = Avro::IO::BinaryDecoder.new(reader)
+    assert_raise(Avro::AvroError) { decoder.read_bytes }
+  end
+
+  def test_read_bytes_within_stream_still_reads
+    payload = 'x' * (2 * 1024 * 1024)
+    writer = StringIO.new
+    Avro::IO::BinaryEncoder.new(writer).write_bytes(payload)
+    reader = StringIO.new(writer.string)
+    decoder = Avro::IO::BinaryDecoder.new(reader)
+    assert_equal(payload, decoder.read_bytes)
+  end
+
+  # An array/map block declares an element count; a malicious or truncated input
+  # can declare far more elements than the remaining bytes could hold. The count
+  # is validated against the bytes remaining before iterating, using the minimum
+  # on-wire size of the element schema (so 0-byte elements like null are not
+  # falsely rejected).
+  def decode(schema_json, encoded)
+    schema = Avro::Schema.parse(schema_json)
+    reader = Avro::IO::DatumReader.new(schema)
+    reader.read(Avro::IO::BinaryDecoder.new(StringIO.new(encoded)))
+  end
+
+  def test_read_array_rejects_count_beyond_stream
+    writer = StringIO.new
+    Avro::IO::BinaryEncoder.new(writer).write_long(1_000_000)
+    assert_raise(Avro::AvroError) do
+      decode('{"type":"array","items":"long"}', writer.string)
+    end
+  end
+
+  def test_read_map_rejects_count_beyond_stream
+    writer = StringIO.new
+    Avro::IO::BinaryEncoder.new(writer).write_long(1_000_000)
+    assert_raise(Avro::AvroError) do
+      decode('{"type":"map","values":"long"}', writer.string)
+    end
+  end
+
+  def test_read_array_of_null_not_falsely_rejected
+    count = 100_000
+    writer = StringIO.new
+    encoder = Avro::IO::BinaryEncoder.new(writer)
+    encoder.write_long(count) # one block of `count` nulls (zero bytes each)
+    encoder.write_long(0)     # end-of-array marker
+    result = decode('{"type":"array","items":"null"}', writer.string)
+    assert_equal([nil] * count, result)
+  end
+
+  def test_read_array_within_stream_still_reads
+    schema = Avro::Schema.parse('{"type":"array","items":"long"}')
+    writer = StringIO.new
+    Avro::IO::DatumWriter.new(schema).write([1, 2, 3], Avro::IO::BinaryEncoder.new(writer))
+    assert_equal([1, 2, 3], decode('{"type":"array","items":"long"}', writer.string))
+  end
+
   def test_int
     check('"int"')
     check_default('"int"', "5", 5)