diff --git a/lib/rb/lib/thrift/processor.rb b/lib/rb/lib/thrift/processor.rb index 15d5ba88f6..23c216c7f2 100644 --- a/lib/rb/lib/thrift/processor.rb +++ b/lib/rb/lib/thrift/processor.rb @@ -34,6 +34,17 @@ def initialize(handler, logger = nil) def process(iprot, oprot) name, type, seqid = iprot.read_message_begin + unless type == MessageTypes::CALL || type == MessageTypes::ONEWAY + iprot.skip(Types::STRUCT) + iprot.read_message_end + x = ApplicationException.new( + ApplicationException::INVALID_MESSAGE_TYPE, + "Invalid message type #{type} for function #{name}" + ) + write_error(x, oprot, name, seqid) + return false + end + if respond_to?("process_#{name}") begin send("process_#{name}", seqid, iprot, oprot) diff --git a/lib/rb/spec/processor_spec.rb b/lib/rb/spec/processor_spec.rb index 444174be91..68a9afcda1 100644 --- a/lib/rb/spec/processor_spec.rb +++ b/lib/rb/spec/processor_spec.rb @@ -23,6 +23,14 @@ describe 'Processor' do class ProcessorSpec include Thrift::Processor + + attr_reader :processed + + def process_work(seqid, iprot, _oprot) + iprot.skip(Thrift::Types::STRUCT) + iprot.read_message_end + @processed = seqid + end end describe Thrift::Processor do @@ -39,12 +47,93 @@ def mock_trans(obj) end end + def input_protocol(name, type, seqid, args = nil) + transport = Thrift::MemoryBufferTransport.new + protocol = Thrift::BinaryProtocol.new(transport) + protocol.write_message_begin(name, type, seqid) + if args + args.write(protocol) + else + protocol.write_struct_begin("args") + protocol.write_field_stop + protocol.write_struct_end + end + protocol.write_message_end + Thrift::BinaryProtocol.new(transport) + end + + def output_protocol + transport = Thrift::MemoryBufferTransport.new + [transport, Thrift::BinaryProtocol.new(transport)] + end + it "should call process_ when it receives that message" do expect(@prot).to receive(:read_message_begin).ordered.and_return ['testMessage', Thrift::MessageTypes::CALL, 17] expect(@processor).to receive(:process_testMessage).with(17, @prot, @prot).ordered expect(@processor.process(@prot, @prot)).to eq(true) end + [Thrift::MessageTypes::REPLY, Thrift::MessageTypes::EXCEPTION].each do |message_type| + it "rejects message type #{message_type} before dispatching" do + input = input_protocol("work", message_type, 11) + output_transport, output = output_protocol + + expect(@processor.process(input, output)).to be false + expect(@processor.processed).to be_nil + + response = Thrift::BinaryProtocol.new(output_transport) + name, type, seqid = response.read_message_begin + exception = Thrift::ApplicationException.new + exception.read(response) + response.read_message_end + + expect(name).to eq("work") + expect(type).to eq(Thrift::MessageTypes::EXCEPTION) + expect(seqid).to eq(11) + expect(exception.type).to eq(Thrift::ApplicationException::INVALID_MESSAGE_TYPE) + expect(exception.message).to eq("Invalid message type #{message_type} for function work") + end + end + + [Thrift::MessageTypes::CALL, Thrift::MessageTypes::ONEWAY].each do |message_type| + it "dispatches valid message type #{message_type}" do + input = input_protocol("work", message_type, 12) + output_transport, output = output_protocol + + expect(@processor.process(input, output)).to be true + expect(@processor.processed).to eq(12) + expect(output_transport.available).to eq(0) + end + end + + it "keeps generated oneway behavior when its envelope is CALL" do + handler = double("Handler") + expect(handler).to receive(:unblock).with(9) + processor = SpecNamespace::NonblockingService::Processor.new(handler) + args = SpecNamespace::NonblockingService::Unblock_args.new(:n => 9) + input = input_protocol("unblock", Thrift::MessageTypes::CALL, 13, args) + output_transport, output = output_protocol + + expect(processor.process(input, output)).to be true + expect(output_transport.available).to eq(0) + end + + it "keeps generated reply behavior when a normal method envelope is ONEWAY" do + handler = double("Handler") + expect(handler).to receive(:sleep).with(3.0) + processor = SpecNamespace::NonblockingService::Processor.new(handler) + args = SpecNamespace::NonblockingService::Sleep_args.new(:seconds => 3.0) + input = input_protocol("sleep", Thrift::MessageTypes::ONEWAY, 14, args) + output_transport, output = output_protocol + + expect(processor.process(input, output)).to be true + + response = Thrift::BinaryProtocol.new(output_transport) + expect(response.read_message_begin).to eq(["sleep", Thrift::MessageTypes::REPLY, 14]) + response.skip(Thrift::Types::STRUCT) + response.read_message_end + end + it "should raise an ApplicationException when the received message cannot be processed" do expect(@prot).to receive(:read_message_begin).ordered.and_return ['testMessage', Thrift::MessageTypes::CALL, 4] expect(@prot).to receive(:skip).with(Thrift::Types::STRUCT).ordered