THRIFT-6124: Reset reused Ruby deserialization targets

lib/rb/ext/struct.c

@@ -31,6 +31,7 @@ ID setvalue_id;
 ID to_s_method_id;
 ID name_to_id_method_id;
 static ID sorted_field_ids_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)
@@ -239,15 +240,17 @@ static VALUE rb_thrift_union_write (VALUE self, VALUE protocol);
 static VALUE rb_thrift_struct_write(VALUE self, VALUE protocol);
 static void write_anything(int ttype, VALUE value, VALUE protocol, VALUE field_info);
 
-VALUE get_field_value(VALUE obj, VALUE field_name) {
+static inline ID field_ivar_id(VALUE field_name) {
   char name_buf[RSTRING_LEN(field_name) + 2];
 
   name_buf[0] = '@';
   strlcpy(&name_buf[1], RSTRING_PTR(field_name), RSTRING_LEN(field_name) + 1);
 
-  VALUE value = rb_ivar_get(obj, rb_intern(name_buf));
+  return rb_intern(name_buf);
+}
 
-  return value;
+VALUE get_field_value(VALUE obj, VALUE field_name) {
+  return rb_ivar_get(obj, field_ivar_id(field_name));
 }
 
 static void write_container(int ttype, VALUE field_info, VALUE value, VALUE protocol) {
@@ -430,13 +433,44 @@ static VALUE rb_thrift_struct_read(VALUE self, VALUE protocol);
 static void skip_map_contents(VALUE protocol, VALUE key_type_value, VALUE value_type_value, int size);
 static void skip_list_or_set_contents(VALUE protocol, VALUE element_type_value, int size);
 
-static void set_field_value(VALUE obj, VALUE field_name, VALUE value) {
-  char name_buf[RSTRING_LEN(field_name) + 2];
+static inline void set_field_value(VALUE obj, VALUE field_name, VALUE value) {
+  rb_ivar_set(obj, field_ivar_id(field_name), value);
+}
 
-  name_buf[0] = '@';
-  strlcpy(&name_buf[1], RSTRING_PTR(field_name), RSTRING_LEN(field_name)+1);
+/* Ruby's rescue and iteration APIs require addressable callback functions. */
+static VALUE duplicate_default_value(VALUE value) {
+  return rb_obj_dup(value);
+}
 
-  rb_ivar_set(obj, rb_intern(name_buf), value);
+static VALUE retain_default_value(VALUE value, VALUE exception) {
+  (void)exception;
+  return value;
+}
+
+static int reset_struct_field(VALUE field_id, VALUE field_info, VALUE self) {
+  (void)field_id;
+
+  VALUE field_name = rb_hash_aref(field_info, name_sym);
+  ID ivar_id = field_ivar_id(field_name);
+  VALUE default_value = rb_hash_aref(field_info, default_sym);
+
+  if (NIL_P(default_value)) {
+    if (!NIL_P(rb_ivar_get(self, ivar_id))) {
+      rb_ivar_set(self, ivar_id, Qnil);
+    }
+  } else {
+    VALUE value = rb_rescue2(
+      duplicate_default_value,
+      default_value,
+      retain_default_value,
+      default_value,
+      rb_eStandardError,
+      (VALUE)0
+    );
+    rb_ivar_set(self, ivar_id, value);
+  }
+
+  return ST_CONTINUE;
 }
 
 // Helper method to skip the contents of a map (assumes the map header has been read).
@@ -590,11 +624,16 @@ static VALUE read_anything(VALUE protocol, int ttype, VALUE field_info) {
 }
 
 static VALUE rb_thrift_struct_read(VALUE self, VALUE protocol) {
+  VALUE struct_fields = STRUCT_FIELDS(self);
+
+  if (RHASH_SIZE(struct_fields) > 0 && rb_ivar_count(self) > 0) {
+    rb_check_frozen(self);
+    rb_hash_foreach(struct_fields, reset_struct_field, self);
+  }
+
   // read struct begin
   default_read_struct_begin(protocol);
 
-  VALUE struct_fields = STRUCT_FIELDS(self);
-
   // read each field
   while (true) {
     VALUE field_header = default_read_field_begin(protocol);
@@ -640,6 +679,15 @@ static VALUE rb_thrift_struct_read(VALUE self, VALUE protocol) {
 // --------------------------------
 
 static VALUE rb_thrift_union_read(VALUE self, VALUE protocol) {
+  rb_check_frozen(self);
+
+  if (!NIL_P(rb_ivar_get(self, setfield_id))) {
+    rb_ivar_set(self, setfield_id, Qnil);
+  }
+  if (!NIL_P(rb_ivar_get(self, setvalue_id))) {
+    rb_ivar_set(self, setvalue_id, Qnil);
+  }
+
   // read struct begin
   default_read_struct_begin(protocol);
 
@@ -657,8 +705,9 @@ static VALUE rb_thrift_union_read(VALUE self, VALUE protocol) {
     if (field_type == specified_type) {
       // read the value
       VALUE name = rb_hash_aref(field_info, name_sym);
+      VALUE value = read_anything(protocol, field_type, field_info);
       rb_iv_set(self, "@setfield", rb_str_intern(name));
-      rb_iv_set(self, "@value", read_anything(protocol, field_type, field_info));
+      rb_iv_set(self, "@value", value);
     } else {
       rb_funcall(protocol, skip_method_id, 1, field_type_value);
     }
@@ -748,4 +797,7 @@ void Init_struct(void) {
 
   sorted_field_ids_method_id = rb_intern("sorted_field_ids");
   rb_global_variable(&sorted_field_ids_method_id);
+
+  default_sym = ID2SYM(rb_intern("default"));
+  rb_global_variable(&default_sym);
 }

lib/rb/lib/thrift/struct.rb

@@ -82,6 +82,16 @@ def inspect(skip_optional_nulls = true)
     end
 
     def read(iprot)
+      unless instance_variables.empty?
+        defaults = fields_with_default_values
+        struct_fields.each_value do |field_info|
+          instance_variable_set("@#{field_info[:name]}", nil)
+        end
+        defaults.each do |name, default_value|
+          instance_variable_set("@#{name}", (default_value.dup rescue default_value))
+        end
+      end
+
       iprot.read_struct_begin
       loop do
         fname, ftype, fid = iprot.read_field_begin

lib/rb/lib/thrift/union.rb

@@ -55,6 +55,9 @@ def inspect
     end
 
     def read(iprot)
+      @setfield = nil
+      @value = nil
+
       iprot.read_struct_begin
       fname, ftype, fid = iprot.read_field_begin
       handle_message(iprot, fid, ftype)

lib/rb/spec/serializer_spec.rb

@@ -56,6 +56,61 @@ def write(protocol)
   end
 end
 
+module DeserializerResetFixtures
+  class ResetFieldsStruct
+    include Thrift::Struct, Thrift::Struct_Union
+
+    FIELDS = {
+      1 => {type: Thrift::Types::STRING, name: 'reset_fields', optional: true}
+    }.freeze
+
+    def struct_fields
+      FIELDS
+    end
+
+    def validate
+    end
+
+    Thrift::Struct.generate_accessors(self)
+  end
+
+  class RequiredStruct
+    include Thrift::Struct, Thrift::Struct_Union
+
+    FIELDS = {
+      1 => {type: Thrift::Types::STRING, name: 'required_value'}
+    }.freeze
+
+    def struct_fields
+      FIELDS
+    end
+
+    def validate
+      raise Thrift::ProtocolException.new(Thrift::ProtocolException::INVALID_DATA, 'Required field required_value is unset!') unless @required_value
+    end
+
+    Thrift::Struct.generate_accessors(self)
+  end
+
+  class ResetFieldsUnion < Thrift::Union
+    include Thrift::Struct_Union
+
+    FIELDS = {
+      1 => {type: Thrift::Types::STRING, name: 'reset_fields', optional: true}
+    }.freeze
+
+    def struct_fields
+      FIELDS
+    end
+
+    def validate
+      raise Thrift::ProtocolException.new(Thrift::ProtocolException::INVALID_DATA, 'Union fields are not set.') if get_set_field.nil? || get_value.nil?
+    end
+
+    Thrift::Union.generate_accessors(self)
+  end
+end
+
 describe 'Serializer' do
   describe Thrift::Serializer do
     it "should serialize structs to binary by default" do
@@ -139,5 +194,149 @@ def write(protocol)
       deserializer = Thrift::Deserializer.new(protocol_factory)
       expect(deserializer.deserialize(SpecNamespace::Hello.new, "")).to eq(SpecNamespace::Hello.new(:greeting => "Good day"))
     end
+
+    it "resets absent struct fields and restores defaults before reuse" do
+      target = SpecNamespace::Foo.new
+      first_payload = binary_payload do |protocol|
+        protocol.write_field_begin("simple", Thrift::Types::I32, 1)
+        protocol.write_i32(99)
+        protocol.write_field_end
+        protocol.write_field_begin("opt_string", Thrift::Types::STRING, 7)
+        protocol.write_string("old")
+        protocol.write_field_end
+      end
+      deserializer = Thrift::Deserializer.new
+
+      deserializer.deserialize(target, first_payload)
+      expect(target.simple).to eq(99)
+      expect(target.opt_string).to eq("old")
+      target.ints << 99
+      target.complex = {1 => {"old" => 1.0}}
+
+      mismatched_payload = binary_payload do |protocol|
+        protocol.write_field_begin("simple", Thrift::Types::STRING, 1)
+        protocol.write_string("ignored")
+        protocol.write_field_end
+      end
+      deserializer.deserialize(target, mismatched_payload)
+
+      expect(target.simple).to eq(53)
+      expect(target.ints).to eq([1, 2, 2, 3])
+      expect(target.complex).to be_nil
+      expect(target.opt_string).to be_nil
+    end
+
+    it "does not retain previous struct state when reading fails" do
+      target = SpecNamespace::Foo.new(:simple => 99, :opt_string => "old")
+      payload = binary_payload(finish: false) do |protocol, transport|
+        protocol.write_field_begin("opt_string", Thrift::Types::STRING, 7)
+        transport.write([5].pack('N'))
+      end
+
+      expect {
+        Thrift::Deserializer.new.deserialize(target, payload)
+      }.to raise_error(EOFError)
+      expect(target.simple).to eq(53)
+      expect(target.opt_string).to be_nil
+    end
+
+    it "resets fields even when an accessor has the same name as the reset operation" do
+      target = DeserializerResetFixtures::ResetFieldsStruct.new(:reset_fields => "old")
+
+      Thrift::Deserializer.new.deserialize(target, binary_payload)
+
+      expect(target.reset_fields).to be_nil
+    end
+
+    it "does not deserialize into a frozen struct" do
+      target = DeserializerResetFixtures::ResetFieldsStruct.new.freeze
+      payload = binary_payload do |protocol|
+        protocol.write_field_begin("reset_fields", Thrift::Types::STRING, 1)
+        protocol.write_string("new")
+        protocol.write_field_end
+      end
+
+      expect {
+        Thrift::Deserializer.new.deserialize(target, payload)
+      }.to raise_error(FrozenError)
+    end
+
+    it "does not satisfy required-field validation with a value from the previous read" do
+      target = DeserializerResetFixtures::RequiredStruct.new(:required_value => "old")
+
+      expect {
+        Thrift::Deserializer.new.deserialize(target, binary_payload)
+      }.to raise_error(Thrift::ProtocolException, "Required field required_value is unset!")
+      expect(target.required_value).to be_nil
+    end
+
+    it "does not retain a union variant when the next value has an unknown field" do
+      target = SpecNamespace::TestUnion.new(:string_field, "old")
+      payload = binary_payload do |protocol|
+        protocol.write_field_begin("future", Thrift::Types::I32, 99)
+        protocol.write_i32(1)
+        protocol.write_field_end
+      end
+
+      expect {
+        Thrift::Deserializer.new.deserialize(target, payload)
+      }.to raise_error(Thrift::ProtocolException, "Union fields are not set.")
+      expect(target.get_set_field).to be_nil
+      expect(target.get_value).to be_nil
+    end
+
+    it "resets a union when an accessor has the same name as the reset operation" do
+      target = DeserializerResetFixtures::ResetFieldsUnion.new(:reset_fields, "old")
+      payload = binary_payload do |protocol|
+        protocol.write_field_begin("future", Thrift::Types::I32, 99)
+        protocol.write_i32(1)
+        protocol.write_field_end
+      end
+
+      expect {
+        Thrift::Deserializer.new.deserialize(target, payload)
+      }.to raise_error(Thrift::ProtocolException, "Union fields are not set.")
+      expect(target.get_set_field).to be_nil
+      expect(target.get_value).to be_nil
+    end
+
+    it "does not deserialize into a frozen union" do
+      target = DeserializerResetFixtures::ResetFieldsUnion.new.freeze
+      payload = binary_payload do |protocol|
+        protocol.write_field_begin("future", Thrift::Types::I32, 99)
+        protocol.write_i32(1)
+        protocol.write_field_end
+      end
+
+      expect {
+        Thrift::Deserializer.new.deserialize(target, payload)
+      }.to raise_error(FrozenError)
+    end
+
+    it "does not retain the previous union variant when reading fails" do
+      target = SpecNamespace::TestUnion.new(:string_field, "old")
+      payload = binary_payload(finish: false) do |protocol, transport|
+        protocol.write_field_begin("string_field", Thrift::Types::STRING, 1)
+        transport.write([5].pack('N'))
+      end
+
+      expect {
+        Thrift::Deserializer.new.deserialize(target, payload)
+      }.to raise_error(EOFError)
+      expect(target.get_set_field).to be_nil
+      expect(target.get_value).to be_nil
+    end
+
+    def binary_payload(finish: true)
+      transport = Thrift::MemoryBufferTransport.new
+      protocol = Thrift::BinaryProtocol.new(transport)
+      protocol.write_struct_begin("value")
+      yield protocol, transport if block_given?
+      if finish
+        protocol.write_field_stop
+        protocol.write_struct_end
+      end
+      transport.read(transport.available)
+    end
   end
 end