Add some specs for thrift data types

lib/rex/proto/thrift.rb

@@ -88,7 +88,7 @@ module Rex::Proto::Thrift
     endian :big
 
     thrift_data_type :data_type
-    int32            :members_size
+    int32            :members_size, initial_value: -> { members.num_bytes }
     choice           :members, onlyif: -> { members_size > 0 }, selection: :data_type do
       array ThriftDataType::T_BOOLEAN, type: :thrift_boolean, read_until: -> { members.num_bytes == members_size }
       array ThriftDataType::T_I16, type: :int16, read_until: -> { members.num_bytes == members_size }

spec/lib/rex/proto/thrift/thrift_array_spec.rb

@@ -0,0 +1,183 @@
+# -*- coding: binary -*-
+require 'spec_helper'
+require 'rex/text'
+
+RSpec.describe Rex::Proto::Thrift::ThriftArray do
+  context 'when the data type is T_BOOLEAN' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_BOOLEAN }
+    let(:value) { {
+      data_type: data_type,
+      members: [ true ]
+    } }
+    let(:binary_s) { [data_type, 1, 1].pack('CNC') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_I16' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_I16 }
+    let(:number) { 0x7fff - rand(0xffff) }
+    let(:value) { {
+      data_type: data_type,
+      members: [ number ]
+    } }
+    let(:binary_s) { [data_type, 2, number].pack('CNs>') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_I32' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_I32 }
+    let(:number) { 0x7fffffff - rand(0xffffffff) }
+    let(:value) { {
+      data_type: data_type,
+      members: [ number ]
+    } }
+    let(:binary_s) { [data_type, 4, number].pack('CNl>') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_I64' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_I64 }
+    let(:number) { 0x7fffffffffffffff - rand(0xffffffffffffffff) }
+    let(:value) { {
+      data_type: data_type,
+      members: [ number ]
+    } }
+    let(:binary_s) { [data_type, 8, number].pack('CNq>') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_UTF7' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_UTF7 }
+    let(:text) { Rex::Text.rand_text_alphanumeric(10) }
+    let(:value) { {
+      data_type: data_type,
+      members: [ text ]
+    } }
+    let(:binary_s) { [data_type, text.length + 4, text.length].pack('CNN') + text }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_STRUCT' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_STRUCT }
+    # use an empty struct
+    let(:object) { [ { data_type: Rex::Proto::Thrift::ThriftDataType::T_STOP } ] }
+    let(:value) { {
+      data_type: data_type,
+      members: [ object ]
+    } }
+    let(:binary_s) { [data_type, 1, 0].pack('CNC') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_SET' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_SET }
+    # use an empty set
+    let(:object) { { data_type: Rex::Proto::Thrift::ThriftDataType::T_I16, members_size: 0 } }
+    let(:value) { {
+      data_type: data_type,
+      members: [ object ]
+    } }
+    let(:binary_s) { [data_type, 5, Rex::Proto::Thrift::ThriftDataType::T_I16, 0].pack('CNCN') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_LIST' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_LIST }
+    # use an empty set
+    let(:object) { { data_type: Rex::Proto::Thrift::ThriftDataType::T_I16, members_size: 0 } }
+    let(:value) { {
+      data_type: data_type,
+      members: [ object ]
+    } }
+    let(:binary_s) { [data_type, 5, Rex::Proto::Thrift::ThriftDataType::T_I16, 0].pack('CNCN') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+end

spec/lib/rex/proto/thrift/thrift_boolean_spec.rb

@@ -0,0 +1,20 @@
+# -*- coding: binary -*-
+require 'spec_helper'
+require 'rex/text'
+
+RSpec.describe Rex::Proto::Thrift::ThriftBoolean do
+  let(:value) { true }
+  let(:binary_s) { "\x01".b }
+
+  describe '#to_binary_s' do
+    it 'should correctly encode' do
+      expect(described_class.new(value).to_binary_s).to eq binary_s
+    end
+  end
+
+  describe '.read' do
+    it 'should correctly decode' do
+      expect(described_class.read(binary_s)).to eq value
+    end
+  end
+end

spec/lib/rex/proto/thrift/thrift_data_spec.rb

@@ -0,0 +1,242 @@
+# -*- coding: binary -*-
+require 'spec_helper'
+require 'rex/text'
+
+RSpec.describe Rex::Proto::Thrift::ThriftData do
+  subject(:instance) { described_class.new }
+
+  it { should respond_to :data_type }
+  it { should respond_to :field_id }
+  it { should respond_to :data_value }
+
+  it 'is big endian' do
+    expect(described_class.fields.instance_variable_get(:@hints)[:endian]).to eq :big
+  end
+
+  it 'tracks the data type in a ThriftDataType field' do
+    expect(instance.data_type).to be_a Rex::Proto::Thrift::ThriftDataType
+  end
+
+  it 'tracks the field ID in a Uint16 field' do
+    expect(instance.field_id).to be_a BinData::Uint16be
+  end
+
+  it 'tracks the data value in a Choice field' do
+    expect(instance.data_value).to be_a BinData::Choice
+  end
+
+  it 'sets the data type correctly by default' do
+    expect(instance.data_type).to eq Rex::Proto::Thrift::ThriftDataType::T_STOP
+  end
+
+  context 'when the data type is T_STOP' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_STOP }
+    let(:value) { {
+      data_type: data_type
+    } }
+    let(:binary_s) { [data_type].pack('C') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_BOOLEAN' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_BOOLEAN }
+    let(:field_id) { rand(0xffff) }
+    let(:value) { {
+      data_type: data_type,
+      field_id: field_id,
+      data_value: true
+    } }
+    let(:binary_s) { [data_type, field_id, 1].pack('CnC') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_I16' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_I16 }
+    let(:field_id) { rand(0xffff) }
+    let(:number) { 0x7fff - rand(0xffff) }
+    let(:value) { {
+      data_type: data_type,
+      field_id: field_id,
+      data_value: number
+    } }
+    let(:binary_s) { [data_type, field_id, number].pack('Cns>') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_I32' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_I32 }
+    let(:field_id) { rand(0xffff) }
+    let(:number) { 0x7fffffff - rand(0xffffffff) }
+    let(:value) { {
+      data_type: data_type,
+      field_id: field_id,
+      data_value: number
+    } }
+    let(:binary_s) { [data_type, field_id, number].pack('Cnl>') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_I64' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_I64 }
+    let(:field_id) { rand(0xffff) }
+    let(:number) { 0x7fffffffffffffff - rand(0xffffffffffffffff) }
+    let(:value) { {
+      data_type: data_type,
+      field_id: field_id,
+      data_value: number
+    } }
+    let(:binary_s) { [data_type, field_id, number].pack('Cnq>') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_UTF7' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_UTF7 }
+    let(:field_id) { rand(0xffff) }
+    let(:text) { Rex::Text.rand_text_alphanumeric(10) }
+    let(:value) { {
+      data_type: data_type,
+      field_id: field_id,
+      data_value: text
+    } }
+    let(:binary_s) { [data_type, field_id, text.length].pack('CnN') + text }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_STRUCT' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_STRUCT }
+    let(:field_id) { rand(0xffff) }
+    let(:object) { [ { data_type: Rex::Proto::Thrift::ThriftDataType::T_STOP } ] }
+    let(:value) { {
+      data_type: data_type,
+      field_id: field_id,
+      data_value: object
+    } }
+    let(:binary_s) { [data_type, field_id, 0].pack('CnC') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_SET' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_SET }
+    let(:field_id) { rand(0xffff) }
+    let(:object) { { data_type: Rex::Proto::Thrift::ThriftDataType::T_I16, members_size: 0 } }
+    let(:value) { {
+      data_type: data_type,
+      field_id: field_id,
+      data_value: object
+    } }
+    let(:binary_s) { [data_type, field_id, Rex::Proto::Thrift::ThriftDataType::T_I16, 0].pack('CnCN') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+
+  context 'when the data type is T_LIST' do
+    let(:data_type) { Rex::Proto::Thrift::ThriftDataType::T_LIST }
+    let(:field_id) { rand(0xffff) }
+    let(:object) { { data_type: Rex::Proto::Thrift::ThriftDataType::T_I16, members_size: 0 } }
+    let(:value) { {
+      data_type: data_type,
+      field_id: field_id,
+      data_value: object
+    } }
+    let(:binary_s) { [data_type, field_id, Rex::Proto::Thrift::ThriftDataType::T_I16, 0].pack('CnCN') }
+
+    describe '#to_binary_s' do
+      it 'should correctly encode' do
+        expect(described_class.new(value).to_binary_s).to eq binary_s
+      end
+    end
+
+    describe '.read' do
+      it 'should correctly decode' do
+        expect(described_class.read(binary_s).snapshot.keep_if { |k,_| value.key? k }).to eq value
+      end
+    end
+  end
+end

spec/lib/rex/proto/thrift/thrift_header_spec.rb

@@ -0,0 +1,46 @@
+RSpec.describe Rex::Proto::Thrift::ThriftHeader do
+  let(:value) { { version: 0x8001, message_type: 1, method_name: '', sequence_id: 1 } }
+  let(:binary_s) { "\x80\x01\x00\x01\x00\x00\x00\x00\x00\x00\x00\x01".b }
+  subject(:instance) { described_class.new }
+
+  it { should respond_to :version }
+  it { should respond_to :message_type }
+  it { should respond_to :method_name }
+  it { should respond_to :sequence_id }
+
+  it 'is big endian' do
+    expect(described_class.fields.instance_variable_get(:@hints)[:endian]).to eq :big
+  end
+
+  it 'tracks the version in a Uint16 field' do
+    expect(instance.version).to be_a BinData::Uint16be
+  end
+
+  it 'tracks the message type in a ThriftMessageType field' do
+    expect(instance.message_type).to be_a Rex::Proto::Thrift::ThriftMessageType
+  end
+
+  it 'tracks the method name in a ThriftString field' do
+    expect(instance.method_name).to be_a Rex::Proto::Thrift::ThriftString
+  end
+
+  it 'tracks the sequence ID in a Uint32 field' do
+    expect(instance.sequence_id).to be_a BinData::Uint32be
+  end
+
+  it 'sets the version correctly by default' do
+    expect(instance.version).to eq 0x8001
+  end
+
+  describe '#to_binary_s' do
+    it 'should correctly encode' do
+      expect(described_class.new(value).to_binary_s).to eq binary_s
+    end
+  end
+
+  describe '.read' do
+    it 'should correctly decode' do
+      expect(described_class.read(binary_s)).to eq value
+    end
+  end
+end

spec/lib/rex/proto/thrift/thrift_string_spec.rb

@@ -0,0 +1,20 @@
+# -*- coding: binary -*-
+require 'spec_helper'
+require 'rex/text'
+
+RSpec.describe Rex::Proto::Thrift::ThriftString do
+  let(:value) { Rex::Text.rand_text_alphanumeric(10) }
+  let(:binary_s) { [value.length].pack('N') + value }
+
+  describe '#to_binary_s' do
+    it 'should correctly encode' do
+      expect(described_class.new(value).to_binary_s).to eq binary_s
+    end
+  end
+
+  describe '.read' do
+    it 'should correctly decode' do
+      expect(described_class.read(binary_s)).to eq value
+    end
+  end
+end

spec/lib/rex/proto/thrift/thrift_struct_spec.rb

@@ -0,0 +1,24 @@
+# -*- coding: binary -*-
+require 'spec_helper'
+require 'rex/text'
+
+RSpec.describe Rex::Proto::Thrift::ThriftStruct do
+  let(:text) { Rex::Text.rand_text_alphanumeric(10) }
+  let(:value) { [
+    { field_id: 1, data_type: Rex::Proto::Thrift::ThriftDataType::T_UTF7, data_value: text },
+    { data_type: Rex::Proto::Thrift::ThriftDataType::T_STOP }
+  ] }
+  let(:binary_s) { [Rex::Proto::Thrift::ThriftDataType::T_UTF7, 1, text.length].pack('CnN') + text + "\x00".b }
+
+  describe '#to_binary_s' do
+    it 'should correctly encode' do
+      expect(described_class.new(value).to_binary_s).to eq binary_s
+    end
+  end
+
+  describe '.read' do
+    it 'should correctly decode' do
+      expect(described_class.read(binary_s)).to eq value
+    end
+  end
+end