THRIFT-6136: Handle Ruby container size conversion errors

lib/rb/ext/compact_protocol.c

@@ -425,6 +425,15 @@ static int8_t get_ttype(int8_t ctype) {
   }
 }
 
+static inline void validate_container_size(uint32_t size) {
+  if (RB_UNLIKELY(size > INT32_MAX)) {
+    rb_exc_raise(get_protocol_exception(
+      INT2FIX(PROTOERR_SIZE_LIMIT),
+      rb_str_new2("Container size limit exceeded")
+    ));
+  }
+}
+
 static char read_byte_direct(VALUE self) {
   VALUE byte = rb_funcall(GET_TRANSPORT(self), read_byte_method_id, 0);
   return (char)(FIX2INT(byte));
@@ -561,6 +570,7 @@ VALUE rb_thrift_compact_proto_read_field_begin(VALUE self) {
 
 VALUE rb_thrift_compact_proto_read_map_begin(VALUE self) {
   uint32_t size = read_varint32(self);
+  validate_container_size(size);
   uint8_t key_and_value_type = size == 0 ? 0 : read_byte_direct(self);
   return rb_ary_new3(3, INT2FIX(get_ttype(key_and_value_type >> 4)), INT2FIX(get_ttype(key_and_value_type & 0xf)), UINT2NUM(size));
 }
@@ -571,6 +581,7 @@ VALUE rb_thrift_compact_proto_read_list_begin(VALUE self) {
   if (size == 15) {
     size = read_varint32(self);
   }
+  validate_container_size(size);
   uint8_t type = get_ttype(size_and_type & 0x0f);
   return rb_ary_new3(2, INT2FIX(type), UINT2NUM(size));
 }

lib/rb/ext/struct.c

@@ -17,6 +17,8 @@
  * under the License.
  */
 
+#include <stdint.h>
+
 #include "struct.h"
 #include "constants.h"
 #include "macros.h"
@@ -31,24 +33,39 @@ ID setvalue_id;
 ID to_s_method_id;
 ID name_to_id_method_id;
 static ID sorted_field_ids_method_id;
+static ID validate_container_size_method_id;
 static VALUE default_sym;
 
 #define IS_CONTAINER(ttype) ((ttype) == TTYPE_MAP || (ttype) == TTYPE_LIST || (ttype) == TTYPE_SET)
 #define STRUCT_FIELDS(obj) rb_const_get(CLASS_OF(obj), fields_const_id)
 
-static void validate_container_size(int size) {
-  if (RB_UNLIKELY(size < 0)) {
-    rb_exc_raise(
-      get_protocol_exception(
-        INT2FIX(PROTOERR_NEGATIVE_SIZE),
-        rb_str_new2("Negative container size")
-      )
-    );
+static int container_size(VALUE protocol, VALUE size) {
+  if (RB_LIKELY(FIXNUM_P(size))) {
+    long value = FIX2LONG(size);
+    if (RB_UNLIKELY(value < 0)) {
+      rb_exc_raise(
+        get_protocol_exception(
+          INT2FIX(PROTOERR_NEGATIVE_SIZE),
+          rb_str_new2("Negative size")
+        )
+      );
+    }
+    if (RB_UNLIKELY((unsigned long)value > INT32_MAX)) {
+      rb_exc_raise(
+        get_protocol_exception(
+          INT2FIX(PROTOERR_SIZE_LIMIT),
+          rb_str_new2("Container size limit exceeded")
+        )
+      );
+    }
+    return (int)value;
   }
+
+  rb_funcall(protocol, validate_container_size_method_id, 1, size);
+  return NUM2INT(size);
 }
 
 static VALUE new_container_array(int size) {
-  validate_container_size(size);
   return rb_ary_new2(size > 1024 ? 1024 : size);
 }
 
@@ -527,11 +544,7 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
     VALUE map_header = default_read_map_begin(protocol);
     int key_ttype = FIX2INT(rb_ary_entry(map_header, 0));
     int value_ttype = FIX2INT(rb_ary_entry(map_header, 1));
-    int num_entries = FIX2INT(rb_ary_entry(map_header, 2));
-
-    if (num_entries < 0) {
-      rb_exc_raise(get_protocol_exception(INT2FIX(PROTOERR_NEGATIVE_SIZE), rb_str_new2("Negative container size")));
-    }
+    int num_entries = container_size(protocol, rb_ary_entry(map_header, 2));
 
     // Check the declared key and value types against the expected ones and skip the map contents
     // if the types don't match.
@@ -563,7 +576,7 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
   } else if (ttype == TTYPE_LIST) {
     VALUE list_header = default_read_list_begin(protocol);
     int element_ttype = FIX2INT(rb_ary_entry(list_header, 0));
-    int num_elements = FIX2INT(rb_ary_entry(list_header, 1));
+    int num_elements = container_size(protocol, rb_ary_entry(list_header, 1));
 
     // Check the declared element type against the expected one and skip the list contents
     // if the types don't match.
@@ -577,11 +590,9 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
           rb_ary_push(result, read_anything(protocol, element_ttype, rb_hash_aref(field_info, element_sym)));
         }
       } else {
-        validate_container_size(num_elements);
         skip_list_or_set_contents(protocol, INT2FIX(element_ttype), num_elements);
       }
     } else {
-      validate_container_size(num_elements);
       skip_list_or_set_contents(protocol, INT2FIX(element_ttype), num_elements);
     }
 
@@ -591,7 +602,7 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
 
     VALUE set_header = default_read_set_begin(protocol);
     int element_ttype = FIX2INT(rb_ary_entry(set_header, 0));
-    int num_elements = FIX2INT(rb_ary_entry(set_header, 1));
+    int num_elements = container_size(protocol, rb_ary_entry(set_header, 1));
 
     // Check the declared element type against the expected one and skip the set contents
     // if the types don't match.
@@ -607,11 +618,9 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
 
         result = rb_class_new_instance(1, &items, rb_cSet);
       } else {
-        validate_container_size(num_elements);
         skip_list_or_set_contents(protocol, INT2FIX(element_ttype), num_elements);
       }
     } else {
-      validate_container_size(num_elements);
       skip_list_or_set_contents(protocol, INT2FIX(element_ttype), num_elements);
     }
 
@@ -798,6 +807,8 @@ void Init_struct(void) {
   sorted_field_ids_method_id = rb_intern("sorted_field_ids");
   rb_global_variable(&sorted_field_ids_method_id);
 
+  validate_container_size_method_id = rb_intern("validate_container_size");
+
   default_sym = ID2SYM(rb_intern("default"));
   rb_global_variable(&default_sym);
 }

lib/rb/lib/thrift/protocol/base_protocol.rb

@@ -42,6 +42,8 @@ def initialize(type = UNKNOWN, message = nil)
 
   class BaseProtocol
 
+    MAX_CONTAINER_SIZE = (1 << 31) - 1
+
     attr_reader :trans
 
     def initialize(trans)
@@ -398,6 +400,7 @@ def skip(type, max_depth = 64)
 
     def validate_container_size(size)
       raise ProtocolException.new(ProtocolException::NEGATIVE_SIZE, 'Negative size') unless size >= 0
+      raise ProtocolException.new(ProtocolException::SIZE_LIMIT, 'Container size limit exceeded') if size > MAX_CONTAINER_SIZE
     end
 
     def to_s

lib/rb/lib/thrift/protocol/compact_protocol.rb

@@ -308,6 +308,7 @@ def read_field_begin
 
     def read_map_begin
       size = read_varint32()
+      validate_container_size(size)
       key_and_value_type = size == 0 ? 0 : read_byte()
       [CompactTypes.get_ttype(key_and_value_type >> 4), CompactTypes.get_ttype(key_and_value_type & 0xf), size]
     end
@@ -318,6 +319,7 @@ def read_list_begin
       if size == 15
         size = read_varint32()
       end
+      validate_container_size(size)
       type = CompactTypes.get_ttype(size_and_type)
       [type, size]
     end

lib/rb/lib/thrift/struct_union.rb

@@ -56,7 +56,7 @@ def read_field(iprot, field = {})
         value.read(iprot)
       when Types::MAP
         key_type, val_type, size = iprot.read_map_begin
-        raise ProtocolException.new(ProtocolException::NEGATIVE_SIZE, 'Negative size') unless size >= 0
+        iprot.validate_container_size(size)
         # Skip the map contents if the declared key or value types don't match the expected ones.
         if (size != 0 && (key_type != field[:key][:type] || val_type != field[:value][:type]))
           size.times do
@@ -75,7 +75,7 @@ def read_field(iprot, field = {})
         iprot.read_map_end
       when Types::LIST
         e_type, size = iprot.read_list_begin
-        raise ProtocolException.new(ProtocolException::NEGATIVE_SIZE, 'Negative size') unless size >= 0
+        iprot.validate_container_size(size)
         # Skip the list contents if the declared element type doesn't match the expected one.
         if (e_type != field[:element][:type])
           size.times do
@@ -91,7 +91,7 @@ def read_field(iprot, field = {})
         iprot.read_list_end
       when Types::SET
         e_type, size = iprot.read_set_begin
-        raise ProtocolException.new(ProtocolException::NEGATIVE_SIZE, 'Negative size') unless size >= 0
+        iprot.validate_container_size(size)
         # Skip the set contents if the declared element type doesn't match the expected one.
         if (e_type != field[:element][:type])
           size.times do

lib/rb/spec/base_protocol_spec.rb

@@ -241,6 +241,14 @@
         expect(e.type).to eq(Thrift::ProtocolException::NEGATIVE_SIZE)
       end
     end
+
+    it "should reject container sizes above the signed 32-bit range while skipping" do
+      expect(@prot).to receive(:read_list_begin).and_return([Thrift::Types::I32, 1 << 31])
+
+      expect { @prot.skip(Thrift::Types::LIST) }.to raise_error(Thrift::ProtocolException, "Container size limit exceeded") do |e|
+        expect(e.type).to eq(Thrift::ProtocolException::SIZE_LIMIT)
+      end
+    end
   end
 
   describe Thrift::BaseProtocolFactory do

lib/rb/spec/compact_protocol_spec.rb

@@ -32,6 +32,12 @@
     :i64 => [0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01]
   }
 
+  CONTAINER_SIZE_ENCODINGS = {
+    (2**31) - 1 => [0xff, 0xff, 0xff, 0xff, 0x07],
+    2**31 => [0x80, 0x80, 0x80, 0x80, 0x08],
+    (2**32) - 1 => [0xff, 0xff, 0xff, 0xff, 0x0f]
+  }
+
   TESTS = {
     :byte => (-127..127).to_a,
     :i16 => (0..14).map { |shift| [1 << shift, -(1 << shift)] }.flatten.sort,
@@ -150,6 +156,36 @@
     end
   end
 
+  it "should accept container sizes within the signed 32-bit range" do
+    bytes = [0xf5, *CONTAINER_SIZE_ENCODINGS.fetch((2**31) - 1)]
+    trans = Thrift::MemoryBufferTransport.new(bytes.pack("C*"))
+    proto = Thrift::CompactProtocol.new(trans)
+
+    expect(proto.read_list_begin).to eq([Thrift::Types::I32, (2**31) - 1])
+  end
+
+  it "should reject container sizes above the signed 32-bit range" do
+    {
+      :read_map_begin => [0x55],
+      :read_list_begin => [0xf5],
+      :read_set_begin => [0xf5]
+    }.each do |reader_method, container_header|
+      [2**31, (2**32) - 1].each do |size|
+        bytes = if reader_method == :read_map_begin
+                  [*CONTAINER_SIZE_ENCODINGS.fetch(size), *container_header]
+                else
+                  [*container_header, *CONTAINER_SIZE_ENCODINGS.fetch(size)]
+                end
+        trans = Thrift::MemoryBufferTransport.new(bytes.pack("C*"))
+        proto = Thrift::CompactProtocol.new(trans)
+
+        expect { proto.public_send(reader_method) }.to raise_error(Thrift::ProtocolException, "Container size limit exceeded") do |error|
+          expect(error.type).to eq(Thrift::ProtocolException::SIZE_LIMIT)
+        end
+      end
+    end
+  end
+
   it "should report the original unknown type when writing fields and containers" do
     {
       :write_field_begin => [nil, 99, 1],

lib/rb/spec/struct_spec.rb

@@ -212,6 +212,26 @@ def validate_default_arguments(object)
       }.to raise_error(sentinel)
     end
 
+    [
+      ['map', Thrift::Types::MAP, 5, :read_map_begin, [Thrift::Types::I32, Thrift::Types::MAP]],
+      ['list', Thrift::Types::LIST, 4, :read_list_begin, [Thrift::Types::I32]],
+      ['set', Thrift::Types::SET, 6, :read_set_begin, [Thrift::Types::I16]]
+    ].each do |name, field_type, field_id, begin_method, header|
+      it "rejects #{name} sizes above the signed 32-bit range" do
+        struct = SpecNamespace::Foo.new
+        prot = Thrift::BaseProtocol.new(double("transport"))
+        declared_size = 1 << 31
+
+        expect(prot).to receive(:read_struct_begin)
+        expect(prot).to receive(:read_field_begin).and_return([name, field_type, field_id])
+        expect(prot).to receive(begin_method).and_return([*header, declared_size])
+
+        expect { struct.read(prot) }.to raise_error(Thrift::ProtocolException, "Container size limit exceeded") do |error|
+          expect(error.type).to eq(Thrift::ProtocolException::SIZE_LIMIT)
+        end
+      end
+    end
+
     it "should serialize false boolean fields correctly" do
       b = SpecNamespace::BoolStruct.new(:yesno => false)
       prot = Thrift::BinaryProtocol.new(Thrift::MemoryBufferTransport.new)

lib/rb/test/fuzz/fuzz_common.rb

@@ -30,13 +30,7 @@
 # Ruzzy.enable_branch_coverage_hooks
 
 def ignorable_fuzz_exception?(error)
-  return true if error.is_a?(Thrift::ProtocolException) ||
-    error.is_a?(EOFError)
-
-  [
-    /too big to convert to '(?:int|long)'/,
-    /bignum too big to convert into 'long'/
-  ].any? { |pattern| error.message =~ pattern }
+  error.is_a?(Thrift::ProtocolException) || error.is_a?(EOFError)
 end
 
 def read_fuzz_test(protocol, read_message_begin)