thrift 0.24.0 → 0.25.0

data/lib/thrift/protocol/compact_protocol.rb

@@ -20,15 +20,24 @@
 
 module Thrift
   class CompactProtocol < BaseProtocol
-
-    PROTOCOL_ID = [0x82].pack('c').unpack('c').first
+    PROTOCOL_ID = [0x82].pack("c").unpack1("c")
     VERSION = 1
     VERSION_MASK = 0x1f
     TYPE_MASK = 0xE0
     TYPE_BITS = 0x07
     TYPE_SHIFT_AMOUNT = 5
     MAX_VARINT32_BYTES = 5  # ceil(32/7); matches protobuf wire format
+    VARINT32_PREFIX_BYTES = MAX_VARINT32_BYTES - 1
+    MAX_VARINT32_LAST_BYTE = 0x0f
     MAX_VARINT_BYTES = 10   # ceil(64/7); matches protobuf wire format
+    BYTE_MIN = -(2**7)
+    BYTE_MAX = (2**7) - 1
+    I16_MIN = -(2**15)
+    I16_MAX = (2**15) - 1
+    I32_MIN = -(2**31)
+    I32_MAX = (2**31) - 1
+    I64_MIN = -(2**63)
+    I64_MAX = (2**63) - 1
 
     TSTOP = [nil, Types::STOP, 0]
 
@@ -68,40 +77,53 @@ module Thrift
         SET           => Types::SET,
         MAP           => Types::MAP,
         STRUCT        => Types::STRUCT,
-        UUID          => Types::UUID
+        UUID          => Types::UUID,
       }
 
       TTYPE_TO_COMPACT = {
-        Types::STOP           => Types::STOP,
-        Types::BOOL           => BOOLEAN_TRUE,
-        Types::BYTE           => BYTE,
-        Types::I16            => I16,
-        Types::I32            => I32,
-        Types::I64            => I64,
-        Types::DOUBLE         => DOUBLE,
-        Types::STRING         => BINARY,
-        Types::LIST           => LIST,
-        Types::SET            => SET,
-        Types::MAP            => MAP,
-        Types::STRUCT         => STRUCT,
-        Types::UUID           => UUID
+        Types::STOP   => Types::STOP,
+        Types::BOOL   => BOOLEAN_TRUE,
+        Types::BYTE   => BYTE,
+        Types::I16    => I16,
+        Types::I32    => I32,
+        Types::I64    => I64,
+        Types::DOUBLE => DOUBLE,
+        Types::STRING => BINARY,
+        Types::LIST   => LIST,
+        Types::SET    => SET,
+        Types::MAP    => MAP,
+        Types::STRUCT => STRUCT,
+        Types::UUID   => UUID,
       }
 
       def self.get_ttype(compact_type)
         val = COMPACT_TO_TTYPE[compact_type & 0x0f]
-        raise "don't know what type: #{compact_type & 0x0f}" unless val
-        val
+        return val if val
+
+        raise ProtocolException.new(
+          ProtocolException::INVALID_DATA,
+          "Unknown compact type: #{compact_type & 0x0f}",
+        )
       end
 
       def self.get_compact_type(ttype)
         val = TTYPE_TO_COMPACT[ttype]
-        raise "don't know what type: #{ttype & 0x0f}" unless val
-        val
+        return val if val
+
+        raise ProtocolException.new(
+          ProtocolException::INVALID_DATA,
+          "Unknown compact type: #{ttype}",
+        )
       end
     end
 
-    def initialize(transport)
+    attr_reader :max_string_size
+
+    def initialize(transport, max_string_size: DEFAULT_MAX_STRING_SIZE)
+      BaseProtocol.validate_max_string_size(max_string_size)
       super(transport)
+      @max_string_size = max_string_size
+      @reset_message_size = transport.message_boundaries?
 
       @last_field = [0]
       @boolean_value = nil
@@ -111,8 +133,14 @@ module Thrift
     end
 
     def write_message_begin(name, type, seqid)
-      write_byte(PROTOCOL_ID)
-      write_byte((VERSION & VERSION_MASK) | ((type << TYPE_SHIFT_AMOUNT) & TYPE_MASK))
+      unless seqid.is_a?(Integer)
+        raise "nil argument not allowed!" if seqid.nil?
+        raise ::TypeError, "integer argument expected"
+      end
+      raise RangeError if seqid < I32_MIN || seqid > I32_MAX
+
+      write_byte_direct(PROTOCOL_ID)
+      write_byte_direct((VERSION & VERSION_MASK) | ((type << TYPE_SHIFT_AMOUNT) & TYPE_MASK))
       write_varint32(message_seqid_to_varint32(seqid))
       write_string(name)
       nil
@@ -129,6 +157,12 @@ module Thrift
     end
 
     def write_field_begin(name, type, id)
+      unless id.is_a?(Integer)
+        raise "nil argument not allowed!" if id.nil?
+        raise ::TypeError, "integer argument expected"
+      end
+      raise RangeError if id < I16_MIN || id > I16_MAX
+
       if type == Types::BOOL
         # we want to possibly include the value, so we'll wait.
         @boolean_field = [type, id]
@@ -147,16 +181,16 @@ module Thrift
       last_id = @last_field.pop
 
       # if there's a type override, use that.
-      typeToWrite = type_override || CompactTypes.get_compact_type(type)
+      type_to_write = type_override || CompactTypes.get_compact_type(type)
 
       # check if we can use delta encoding for the field id
       if id > last_id && id - last_id <= 15
         # write them together
-        write_byte((id - last_id) << 4 | typeToWrite)
+        write_byte_direct(((id - last_id) << 4) | type_to_write)
       else
         # write them separate
-        write_byte(typeToWrite)
-        write_i16(id)
+        write_byte_direct(type_to_write)
+        write_varint32(int_to_zig_zag(id))
       end
 
       @last_field.push(id)
@@ -164,15 +198,16 @@ module Thrift
     end
 
     def write_field_stop
-      write_byte(Types::STOP)
+      write_byte_direct(Types::STOP)
     end
 
     def write_map_begin(ktype, vtype, size)
+      size = validate_size(size)
       if (size == 0)
-        write_byte(0)
+        write_byte_direct(0)
       else
         write_varint32(size)
-        write_byte(CompactTypes.get_compact_type(ktype) << 4 | CompactTypes.get_compact_type(vtype))
+        write_byte_direct((CompactTypes.get_compact_type(ktype) << 4) | CompactTypes.get_compact_type(vtype))
       end
     end
 
@@ -185,6 +220,7 @@ module Thrift
     end
 
     def write_bool(bool)
+      # Preserve Ruby truthiness: only false and nil encode as false.
       type = bool ? CompactTypes::BOOLEAN_TRUE : CompactTypes::BOOLEAN_FALSE
       unless @boolean_field.nil?
         # we haven't written the field header yet
@@ -192,23 +228,47 @@ module Thrift
         @boolean_field = nil
       else
         # we're not part of a field, so just write the value.
-        write_byte(type)
+        write_byte_direct(type)
       end
     end
 
     def write_byte(byte)
-      @trans.write([byte].pack('c'))
+      unless byte.is_a?(Integer)
+        raise "nil argument not allowed!" if byte.nil?
+        raise ::TypeError, "integer argument expected"
+      end
+      raise RangeError if byte < BYTE_MIN || byte > BYTE_MAX
+
+      write_byte_direct(byte)
     end
 
     def write_i16(i16)
+      unless i16.is_a?(Integer)
+        raise "nil argument not allowed!" if i16.nil?
+        raise ::TypeError, "integer argument expected"
+      end
+      raise RangeError if i16 < I16_MIN || i16 > I16_MAX
+
       write_varint32(int_to_zig_zag(i16))
     end
 
     def write_i32(i32)
+      unless i32.is_a?(Integer)
+        raise "nil argument not allowed!" if i32.nil?
+        raise ::TypeError, "integer argument expected"
+      end
+      raise RangeError if i32 < I32_MIN || i32 > I32_MAX
+
       write_varint32(int_to_zig_zag(i32))
     end
 
     def write_i64(i64)
+      unless i64.is_a?(Integer)
+        raise "nil argument not allowed!" if i64.nil?
+        raise ::TypeError, "integer argument expected"
+      end
+      raise RangeError if i64 < I64_MIN || i64 > I64_MAX
+
       write_varint64(long_to_zig_zag(i64))
     end
 
@@ -222,7 +282,8 @@ module Thrift
     end
 
     def write_binary(buf)
-      write_varint32(buf.bytesize)
+      size = validate_size(buf.bytesize)
+      write_varint32(size)
       @trans.write(buf)
     end
 
@@ -232,20 +293,28 @@ module Thrift
     end
 
     def read_message_begin
-      protocol_id = read_byte()
+      trans.reset_message_size if @reset_message_size
+
+      protocol_id = read_byte
       if protocol_id != PROTOCOL_ID
-        raise ProtocolException.new("Expected protocol id #{PROTOCOL_ID} but got #{protocol_id}")
+        raise ProtocolException.new(
+          ProtocolException::BAD_VERSION,
+          "Expected protocol id #{PROTOCOL_ID} but got #{protocol_id}",
+        )
       end
 
-      version_and_type = read_byte()
+      version_and_type = read_byte
       version = version_and_type & VERSION_MASK
       if (version != VERSION)
-        raise ProtocolException.new("Expected version #{VERSION} but got #{version}");
+        raise ProtocolException.new(
+          ProtocolException::BAD_VERSION,
+          "Expected version #{VERSION} but got #{version}",
+        )
       end
 
       type = (version_and_type >> TYPE_SHIFT_AMOUNT) & TYPE_BITS
-      seqid = message_seqid_from_varint32(read_varint32())
-      messageName = read_string()
+      seqid = message_seqid_from_varint32(read_varint32)
+      messageName = read_string
       [messageName, type, seqid]
     end
 
@@ -255,12 +324,12 @@ module Thrift
     end
 
     def read_struct_end
-      @last_field.pop()
+      @last_field.pop
       nil
     end
 
     def read_field_begin
-      type = read_byte()
+      type = read_byte
 
       # if it's a stop, then we can return immediately, as the struct is over.
       if (type & 0x0f) == Types::STOP
@@ -273,7 +342,7 @@ module Thrift
         if modifier == 0
           # not a delta. look ahead for the zigzag varint field id.
           @last_field.pop
-          field_id = read_i16()
+          field_id = read_i16
         else
           # has a delta. add the delta to the last read field id.
           field_id = @last_field.pop + modifier
@@ -292,17 +361,19 @@ module Thrift
     end
 
     def read_map_begin
-      size = read_varint32()
-      key_and_value_type = size == 0 ? 0 : read_byte()
+      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
 
     def read_list_begin
-      size_and_type = read_byte()
+      size_and_type = read_byte
       size = (size_and_type >> 4) & 0x0f
       if size == 15
-        size = read_varint32()
+        size = read_varint32
       end
+      validate_container_size(size)
       type = CompactTypes.get_ttype(size_and_type)
       [type, size]
     end
@@ -317,7 +388,7 @@ module Thrift
         @bool_value = nil
         bv
       else
-        read_byte() == CompactTypes::BOOLEAN_TRUE
+        read_byte == CompactTypes::BOOLEAN_TRUE
       end
     end
 
@@ -330,21 +401,21 @@ module Thrift
     end
 
     def read_i16
-      zig_zag_to_int(read_varint32())
+      value = zig_zag_to_int(read_varint32) & 0xffff
+      value > I16_MAX ? value - 0x10000 : value
     end
 
     def read_i32
-      zig_zag_to_int(read_varint32())
+      zig_zag_to_int(read_varint32)
     end
 
     def read_i64
-      zig_zag_to_long(read_varint64())
+      zig_zag_to_long(read_varint64 & 0xffffffffffffffff)
     end
 
     def read_double
       trans.read_into_buffer(@rbuf, 8)
-      val = @rbuf.reverse.unpack('G').first
-      val
+      @rbuf.reverse.unpack1("G")
     end
 
     def read_string
@@ -353,16 +424,24 @@ module Thrift
     end
 
     def read_binary
-      size = read_varint32()
+      size = read_varint32
+      if size > I32_MAX
+        raise ProtocolException.new(ProtocolException::SIZE_LIMIT, "Binary size limit exceeded")
+      end
+      check_string_size(size)
       trans.read_all(size)
     end
 
+    def skip_string
+      read_binary
+    end
+
     def read_uuid
       UUID.uuid_from_bytes(trans.read_all(16))
     end
 
     def to_s
-      "compact(#{super.to_s})"
+      "compact(#{super})"
     end
 
     private
@@ -372,68 +451,84 @@ module Thrift
     # the wire differ only by the type indicator.
     #
     def write_collection_begin(elem_type, size)
+      size = validate_size(size)
+      compact_type = CompactTypes.get_compact_type(elem_type)
       if size <= 14
-        write_byte(size << 4 | CompactTypes.get_compact_type(elem_type))
+        write_byte_direct((size << 4) | compact_type)
       else
-        write_byte(0xf0 | CompactTypes.get_compact_type(elem_type))
+        write_byte_direct(0xf0 | compact_type)
         write_varint32(size)
       end
     end
 
+    def write_byte_direct(byte)
+      @trans.write([byte].pack("C"))
+    end
+
     def write_varint32(n)
-      # int idx = 0;
-      while true
-        if (n & ~0x7F) == 0
-          # i32buf[idx++] = (byte)n;
-          write_byte(n)
-          break
-          # return;
-        else
-          # i32buf[idx++] = (byte)((n & 0x7F) | 0x80);
-          write_byte((n & 0x7F) | 0x80)
-          n = n >> 7
-        end
+      if (n & ~0x7F) == 0
+        write_byte_direct(n)
+        return
       end
-      # trans_.write(i32buf, 0, idx);
+
+      buffer = String.new(capacity: MAX_VARINT32_BYTES, encoding: Encoding::BINARY)
+      while (n & ~0x7F) != 0
+        buffer << ((n & 0x7F) | 0x80)
+        n >>= 7
+      end
+      buffer << n
+      @trans.write(buffer)
     end
 
     SEVEN_BIT_MASK = 0x7F
     EVERYTHING_ELSE_MASK = ~SEVEN_BIT_MASK
 
     def write_varint64(n)
-      while true
-        if (n & EVERYTHING_ELSE_MASK) == 0 # TODO need to find a way to make this into a long...
-          write_byte(n)
-          break
-        else
-          write_byte((n & SEVEN_BIT_MASK) | 0x80)
-          n >>= 7
-        end
+      if (n & EVERYTHING_ELSE_MASK) == 0
+        write_byte_direct(n)
+        return
+      end
+
+      buffer = String.new(capacity: MAX_VARINT_BYTES, encoding: Encoding::BINARY)
+      while (n & EVERYTHING_ELSE_MASK) != 0
+        buffer << ((n & SEVEN_BIT_MASK) | 0x80)
+        n >>= 7
       end
+      buffer << n
+      @trans.write(buffer)
     end
 
-    def read_varint32()
+    def read_varint32
       shift = 0
       result = 0
-      MAX_VARINT32_BYTES.times do
-        b = read_byte()
+      VARINT32_PREFIX_BYTES.times do
+        b = read_byte
         result |= (b & 0x7f) << shift
         return result if (b & 0x80) != 0x80
         shift += 7
       end
-      raise ProtocolException.new(ProtocolException::INVALID_DATA, 'Variable-length int over 5 bytes.')
+
+      b = read_byte
+      if (b & 0x80) != 0
+        raise ProtocolException.new(ProtocolException::INVALID_DATA, "Variable-length int over 5 bytes.")
+      end
+      if (b & ~MAX_VARINT32_LAST_BYTE) != 0
+        raise ProtocolException.new(ProtocolException::INVALID_DATA, "Variable-length int overflows uint32.")
+      end
+
+      result | (b << shift)
     end
 
-    def read_varint64()
+    def read_varint64
       shift = 0
       result = 0
       MAX_VARINT_BYTES.times do
-        b = read_byte()
+        b = read_byte
         result |= (b & 0x7f) << shift
         return result if (b & 0x80) != 0x80
         shift += 7
       end
-      raise ProtocolException.new(ProtocolException::INVALID_DATA, 'Variable-length int over 10 bytes.')
+      raise ProtocolException.new(ProtocolException::INVALID_DATA, "Variable-length int over 10 bytes.")
     end
 
     def int_to_zig_zag(n)
@@ -454,11 +549,17 @@ module Thrift
     end
 
     def message_seqid_to_varint32(seqid)
-      if seqid < -(2**31) || seqid > (2**31) - 1
-        raise RangeError, "seqid must be a signed int32"
+      seqid < 0 ? seqid + (2**32) : seqid
+    end
+
+    def validate_size(size)
+      unless size.is_a?(Integer)
+        raise "nil argument not allowed!" if size.nil?
+        raise ::TypeError, "integer argument expected"
       end
+      raise RangeError if size < 0 || size > I32_MAX
 
-      seqid < 0 ? seqid + (2**32) : seqid
+      size
     end
 
     def message_seqid_from_varint32(seqid)
@@ -467,8 +568,12 @@ module Thrift
   end
 
   class CompactProtocolFactory < BaseProtocolFactory
+    def initialize(max_string_size: BaseProtocol::DEFAULT_MAX_STRING_SIZE)
+      @max_string_size = max_string_size
+    end
+
     def get_protocol(trans)
-      CompactProtocol.new(trans)
+      CompactProtocol.new(trans, max_string_size: @max_string_size)
     end
 
     def to_s

data/lib/thrift/protocol/header_protocol.rb

@@ -178,8 +178,9 @@ module Thrift
         @header_transport.reset_protocol
         reset_protocol_if_needed
       rescue ProtocolException => ex
+        use_default_protocol
         app_ex = ApplicationException.new(ApplicationException::INVALID_PROTOCOL, ex.message)
-        write_message_begin("", MessageTypes::EXCEPTION, 0)
+        write_message_begin("", MessageTypes::EXCEPTION, @header_transport.sequence_id)
         app_ex.write(self)
         write_message_end
         @header_transport.flush
@@ -264,12 +265,16 @@ module Thrift
       @protocol.read_binary
     end
 
+    def skip_string
+      @protocol.skip_string
+    end
+
     def read_uuid
       @protocol.read_uuid
     end
 
     def to_s
-      "header(#{@protocol.to_s})"
+      "header(#{@protocol})"
     end
 
     private
@@ -283,6 +288,12 @@ module Thrift
       end
     end
 
+    def use_default_protocol
+      @protocol = create_protocol(@default_protocol)
+      @current_protocol_id = @default_protocol
+      @header_transport.protocol_id = @default_protocol
+    end
+
     # Creates a protocol instance based on protocol ID
     def create_protocol(protocol_id)
       case protocol_id
@@ -293,7 +304,7 @@ module Thrift
       else
         raise ProtocolException.new(
           ProtocolException::INVALID_DATA,
-          "Unknown protocol ID: #{protocol_id}"
+          "Unknown protocol ID: #{protocol_id}",
         )
       end
     end