Fix crashes caused by unchecked memory allocation in Ruby library

Author: kpumuk

lib/rb/ext/struct.c

@@ -35,6 +35,19 @@ static ID sorted_field_ids_method_id;
 #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 VALUE new_container_array(int size) {
+  if (size < 0) {
+    rb_exc_raise(
+      get_protocol_exception(
+        INT2FIX(PROTOERR_NEGATIVE_SIZE),
+        rb_str_new2("Negative container size")
+      )
+    );
+  }
+
+  return rb_ary_new2(size > 1024 ? 1024 : size);
+}
+
 //-------------------------------------------
 // Writing section
 //-------------------------------------------
@@ -483,6 +496,10 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
     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")));
+    }
+
     // Check the declared key and value types against the expected ones and skip the map contents
     // if the types don't match.
     VALUE key_info = rb_hash_aref(field_info, key_sym);
@@ -523,7 +540,7 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
     if (!NIL_P(element_info)) {
       int specified_element_type = FIX2INT(rb_hash_aref(element_info, type_sym));
       if (specified_element_type == element_ttype) {
-        result = rb_ary_new2(num_elements);
+        result = new_container_array(num_elements);
 
         for (i = 0; i < num_elements; ++i) {
           rb_ary_push(result, read_anything(protocol, element_ttype, rb_hash_aref(field_info, element_sym)));
@@ -550,7 +567,7 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
     if (!NIL_P(element_info)) {
       int specified_element_type = FIX2INT(rb_hash_aref(element_info, type_sym));
       if (specified_element_type == element_ttype) {
-        items = rb_ary_new2(num_elements);
+        items = new_container_array(num_elements);
 
         for (i = 0; i < num_elements; ++i) {
           rb_ary_push(items, read_anything(protocol, element_ttype, rb_hash_aref(field_info, element_sym)));

lib/rb/lib/thrift/struct_union.rb

@@ -20,6 +20,12 @@
 
 module Thrift
   module Struct_Union
+    def validate_container_size(size)
+      return size unless size < 0
+
+      raise ProtocolException.new(ProtocolException::NEGATIVE_SIZE, 'Negative size')
+    end
+
     def name_to_id(name)
       names_to_ids = self.class.instance_variable_get(:@names_to_ids)
       unless names_to_ids
@@ -55,6 +61,7 @@ def read_field(iprot, field = {})
         value.read(iprot)
       when Types::MAP
         key_type, val_type, size = iprot.read_map_begin
+        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
@@ -73,20 +80,23 @@ def read_field(iprot, field = {})
         iprot.read_map_end
       when Types::LIST
         e_type, size = iprot.read_list_begin
+        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
             iprot.skip(e_type)
           end
           value = nil
         else
-          value = Array.new(size) do |n|
-            read_field(iprot, field_info(field[:element]))
+          value = []
+          size.times do
+            value << read_field(iprot, field_info(field[:element]))
           end
         end
         iprot.read_list_end
       when Types::SET
         e_type, size = iprot.read_set_begin
+        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/struct_spec.rb

@@ -142,6 +142,34 @@ def validate_default_arguments(object)
       expect(struct.shorts).to eq(Set.new([3, 2]))
     end
 
+    it "rejects negative container sizes while reading" do
+      struct = SpecNamespace::Foo.new
+      prot = Thrift::BaseProtocol.new(double("transport"))
+
+      expect(prot).to receive(:read_list_begin).and_return([Thrift::Types::I32, -1])
+
+      expect {
+        struct.send(:read_field, prot, SpecNamespace::Foo::FIELDS[4])
+      }.to raise_error(Thrift::ProtocolException, "Negative size") { |error|
+        expect(error.type).to eq(Thrift::ProtocolException::NEGATIVE_SIZE)
+      }
+    end
+
+    it "does not preallocate arrays from declared list sizes" do
+      struct = SpecNamespace::Foo.new
+      prot = Thrift::BaseProtocol.new(double("transport"))
+      declared_size = 1 << 30
+      sentinel = RuntimeError.new("stop after first element")
+
+      expect(prot).to receive(:read_list_begin).and_return([Thrift::Types::I32, declared_size])
+      expect(prot).to receive(:read_i32).and_raise(sentinel)
+      expect(Array).not_to receive(:new).with(declared_size)
+
+      expect {
+        struct.send(:read_field, prot, SpecNamespace::Foo::FIELDS[4])
+      }.to raise_error(sentinel)
+    end
+
     it "should serialize false boolean fields correctly" do
       b = SpecNamespace::BoolStruct.new(:yesno => false)
       prot = Thrift::BinaryProtocol.new(Thrift::MemoryBufferTransport.new)