diff --git a/lib/rb/lib/thrift/transport/buffered_transport.rb b/lib/rb/lib/thrift/transport/buffered_transport.rb index 11abc8d8221..32bcec05465 100644 --- a/lib/rb/lib/thrift/transport/buffered_transport.rb +++ b/lib/rb/lib/thrift/transport/buffered_transport.rb @@ -54,19 +54,30 @@ def read(sz) end ret + ensure + clear_consumed_read_buffer if @rbuf && @index >= @rbuf.size end def read_byte - # If the read buffer is exhausted, try to read up to DEFAULT_BUFFER more bytes into it. - if @index >= @rbuf.size - @rbuf = @transport.read(DEFAULT_BUFFER) - @index = 0 + if @index < @rbuf.size - 1 + @index += 1 + return Bytes.get_string_byte(@rbuf, @index - 1) end - # The read buffer has some data now, read a single byte. Using get_string_byte() avoids - # allocating a temp string of size 1 unnecessarily. - @index += 1 - Bytes.get_string_byte(@rbuf, @index - 1) + begin + # If the read buffer is exhausted, try to read up to DEFAULT_BUFFER more bytes into it. + if @index >= @rbuf.size + @rbuf = @transport.read(DEFAULT_BUFFER) + @index = 0 + end + + # The read buffer has some data now, read a single byte. Using get_string_byte() avoids + # allocating a temp string of size 1 unnecessarily. + @index += 1 + Bytes.get_string_byte(@rbuf, @index - 1) + ensure + clear_consumed_read_buffer if @rbuf && @index >= @rbuf.size + end end # Reads a number of bytes from the transport into the buffer passed. @@ -91,6 +102,8 @@ def read_into_buffer(buffer, size) i += 1 end i + ensure + clear_consumed_read_buffer if @rbuf && @index >= @rbuf.size end def write(buf) @@ -109,6 +122,16 @@ def flush def to_s "buffered(#{@transport})" end + + private + + # Replace exhausted storage; clearing it in place could change returned slices. + def clear_consumed_read_buffer + unless @rbuf.empty? + @rbuf = Bytes.empty_byte_buffer + @index = 0 + end + end end class BufferedTransportFactory < BaseTransportFactory diff --git a/lib/rb/lib/thrift/transport/framed_transport.rb b/lib/rb/lib/thrift/transport/framed_transport.rb index 574705c5424..ec2da521f6f 100644 --- a/lib/rb/lib/thrift/transport/framed_transport.rb +++ b/lib/rb/lib/thrift/transport/framed_transport.rb @@ -66,17 +66,28 @@ def read(sz) @index += sz @rbuf.slice(@index - sz, sz) || Bytes.empty_byte_buffer + ensure + clear_consumed_read_buffer if @rbuf && @index >= @rbuf.size end def read_byte return @transport.read_byte unless @read - read_frame if @index >= @rbuf.length + if @index < @rbuf.size - 1 + @index += 1 + return Bytes.get_string_byte(@rbuf, @index - 1) + end - # The read buffer has some data now, read a single byte. Using get_string_byte() avoids - # allocating a temp string of size 1 unnecessarily. - @index += 1 - Bytes.get_string_byte(@rbuf, @index - 1) + begin + read_frame if @index >= @rbuf.length + + # The read buffer has some data now, read a single byte. Using get_string_byte() avoids + # allocating a temp string of size 1 unnecessarily. + @index += 1 + Bytes.get_string_byte(@rbuf, @index - 1) + ensure + clear_consumed_read_buffer if @rbuf && @index >= @rbuf.size + end end def read_into_buffer(buffer, size) @@ -91,6 +102,8 @@ def read_into_buffer(buffer, size) i += 1 end i + ensure + clear_consumed_read_buffer if @rbuf && @index >= @rbuf.size end def write(buf, sz = nil) @@ -121,7 +134,16 @@ def to_s private + # Replace exhausted storage; clearing it in place could change returned slices. + def clear_consumed_read_buffer + unless @rbuf.empty? + @rbuf = Bytes.empty_byte_buffer + @index = 0 + end + end + def read_frame + clear_consumed_read_buffer if @index >= @rbuf.size sz = @transport.read_all(4).unpack1("N") raise TransportException.new(TransportException::END_OF_FILE, "Cannot read from a zero-length frame") if sz == 0 @@ -134,8 +156,8 @@ def read_frame ) end - @index = 0 @rbuf = @transport.read_all(sz) + @index = 0 end end diff --git a/lib/rb/lib/thrift/transport/header_transport.rb b/lib/rb/lib/thrift/transport/header_transport.rb index d707c9dcb38..d4e260b28e7 100644 --- a/lib/rb/lib/thrift/transport/header_transport.rb +++ b/lib/rb/lib/thrift/transport/header_transport.rb @@ -224,6 +224,19 @@ def read(sz) read_frame(bytes_left) additional = @read_buffer.read(bytes_left) data + (additional || Bytes.empty_byte_buffer) + ensure + # Keep frame metadata for replies, but release the exhausted payload. + if @read_buffer.eof? && !@read_buffer.string.empty? + @read_buffer.string = Bytes.empty_byte_buffer + end + end + + def read_byte + return super unless @read_buffer && !@read_buffer.eof? + + byte = @read_buffer.getbyte + @read_buffer.string = Bytes.empty_byte_buffer if @read_buffer.eof? + byte end def write(buf) diff --git a/lib/rb/lib/thrift/transport/http_client_transport.rb b/lib/rb/lib/thrift/transport/http_client_transport.rb index a88c989da7c..471f7940798 100644 --- a/lib/rb/lib/thrift/transport/http_client_transport.rb +++ b/lib/rb/lib/thrift/transport/http_client_transport.rb @@ -42,6 +42,18 @@ def read(sz) return data unless data.nil? raise TransportException.new(TransportException::END_OF_FILE, "#{self.class.name} reached EOF reading response from #{self}, HTTP status code #{@response_code}") + ensure + if @inbuf && @inbuf.eof? && !@inbuf.string.empty? + @inbuf.string = Bytes.empty_byte_buffer + end + end + + def read_byte + return super unless @inbuf && !@inbuf.eof? + + byte = @inbuf.getbyte + @inbuf.string = Bytes.empty_byte_buffer if @inbuf.eof? + byte end def write(buf); @outbuf << Bytes.force_binary_encoding(buf) end diff --git a/lib/rb/spec/base_transport_spec.rb b/lib/rb/spec/base_transport_spec.rb index 80bdf787392..f42338ae7f2 100644 --- a/lib/rb/spec/base_transport_spec.rb +++ b/lib/rb/spec/base_transport_spec.rb @@ -19,6 +19,8 @@ # require "spec_helper" +require "weakref" +require "timeout" describe "BaseTransport" do describe Thrift::TransportException do @@ -655,4 +657,196 @@ @trans.close end end + + context "input lifetime" do + def input_transport(kind, payload) + case kind + when :buffered + Thrift::BufferedTransport.new(Thrift::MemoryBufferTransport.new(payload)) + when :framed + Thrift::FramedTransport.new(Thrift::MemoryBufferTransport.new([payload.bytesize].pack("N") + payload)) + when :header + wire = Thrift::MemoryBufferTransport.new + writer = Thrift::HeaderTransport.new(wire) + writer.set_header("name", "value") + writer.write(payload) + writer.flush + Thrift::HeaderTransport.new(wire) + when :http + transport = Thrift::HTTPClientTransport.new("http://localhost/") + transport.instance_variable_set(:@inbuf, StringIO.new(payload)) + transport.instance_variable_set(:@response_code, 200) + transport + end + end + + def input_storage(transport) + if transport.is_a?(Thrift::HeaderTransport) + transport.instance_variable_get(:@read_buffer).string + elsif transport.is_a?(Thrift::HTTPClientTransport) + transport.instance_variable_get(:@inbuf).string + else + transport.instance_variable_get(:@rbuf) + end + end + + [:buffered, :framed, :header, :http].each do |kind| + size = kind == :buffered ? 4096 : 32_768 + eof_error = [:header, :http].include?(kind) ? Thrift::TransportException : EOFError + [:read, :read_all, :read_byte, :read_into_buffer].each do |operation| + it "releases exhausted #{kind} storage after #{operation} while idle" do + transport, reference = Thread.new do + transport = input_transport(kind, "x".b * size) + # Leave one byte for each public read entry point to consume. + expect(transport.read_all(size - 1)).to eq("x" * (size - 1)) + reference = WeakRef.new(input_storage(transport)) + case operation + when :read, :read_all + expect(transport.public_send(operation, 1)).to eq("x") + when :read_byte + expect(transport.read_byte).to eq(120) + when :read_into_buffer + output = "\0".b + expect(transport.read_into_buffer(output, 1)).to eq(1) + expect(output).to eq("x") + end + [transport, reference] + end.value + + GC.start + expect(input_storage(transport)).to be_empty + expect(reference.weakref_alive?).to be_falsey + transport.close + GC.start + expect(reference.weakref_alive?).to be_falsey + end + end + + it "keeps returned #{kind} data valid after exhaustion and close" do + transport = input_transport(kind, "x".b * size) + output = transport.read_all(size) + transport.close + GC.start + expect(output).to eq("x" * size) + expect(output.encoding).to eq(Encoding::BINARY) + end + + it "preserves unread #{kind} input across close and a failed destination write" do + transport = input_transport(kind, "x".b * size) + expect(transport.read_all(size - 1)).to eq("x" * (size - 1)) + expect { transport.read_into_buffer("".b, 1) }.to raise_error(IndexError) + transport.close + # Buffered/Framed copy before advancing the cursor; the base implementation + # used by Header/HTTP consumes before copying into the destination. + if [:buffered, :framed].include?(kind) + expect(transport.read(1)).to eq("x") + end + expect(input_storage(transport)).to be_empty + end + + it "releases consumed #{kind} input when a subsequent read fails" do + transport, reference = Thread.new do + transport = input_transport(kind, "x".b * size) + transport.read_all(1) + reference = WeakRef.new(input_storage(transport)) + transport.read_all(size - 1) + expect { transport.read_all(1) }.to raise_error(eof_error) + [transport, reference] + end.value + GC.start + expect(input_storage(transport)).to be_empty + expect(reference.weakref_alive?).to be_falsey + end + + unless kind == :http + it "preserves unread #{kind} input when the underlying close fails" do + transport = input_transport(kind, "x".b * size) + source = transport.instance_variable_get(:@transport) + allow(source).to receive(:close).and_raise(IOError, "close failed") + expect(transport.read(1)).to eq("x") + expect { transport.close }.to raise_error(IOError, "close failed") + expect(transport.read_all(size - 1)).to eq("x" * (size - 1)) + expect(input_storage(transport)).to be_empty + end + end + + it "preserves unsigned #{kind} bytes across mixed reads and releases the last byte" do + transport = input_transport(kind, "\x00\x80\xff".b + ("x" * (size - 4)) + "z") + + expect(transport.read_byte).to eq(0) + expect(transport.read_byte).to eq(128) + expect(transport.read(1)).to eq("\xff".b) + expect(transport.read_all(size - 4)).to eq("x" * (size - 4)) + expect(transport.read_byte).to eq(122) + expect(input_storage(transport)).to be_empty + expect { transport.read_byte }.to raise_error(eof_error) + end + end + + [Thrift::FramedTransport, Thrift::HeaderTransport].each do |transport_class| + it "round trips Binary payloads over TCP using #{transport_class}" do + listener = TCPServer.new("127.0.0.1", 0) + server = Thread.new do + peer = listener.accept + begin + transport = transport_class.new(Thrift::IOStreamTransport.new(peer, peer)) + transport.protocol_id = Thrift::HeaderSubprotocolID::BINARY if transport.is_a?(Thrift::HeaderTransport) + protocol = Thrift::BinaryProtocol.new(transport) + value = protocol.read_binary + expect(input_storage(transport)).to be_empty + protocol.write_binary(value.reverse) + transport.flush + ensure + peer.close + end + end + socket = nil + Timeout.timeout(10) do + socket = Thrift::Socket.new("127.0.0.1", listener.addr[1]) + socket.open + client = transport_class.new(socket) + client.protocol_id = Thrift::HeaderSubprotocolID::BINARY if client.is_a?(Thrift::HeaderTransport) + protocol = Thrift::BinaryProtocol.new(client) + protocol.write_binary("xyz" * 32_768) + client.flush + expect(protocol.read_binary).to eq("zyx" * 32_768) + expect(input_storage(client)).to be_empty + server.value + end + ensure + socket.close if socket + listener.close if listener + server.kill.join if server && server.alive? + end + end + + [:buffered, :framed].each do |kind| + it "preserves #{kind} unread bytes when read_into_buffer fails after crossing a refill" do + source = Thrift::MemoryBufferTransport.new(("x".b * 4096) + ("y".b * 4096)) + if kind == :framed + source = Thrift::MemoryBufferTransport.new([4096].pack("N") + ("x".b * 4096) + [2].pack("N") + "yz") + end + transport = (kind == :buffered ? Thrift::BufferedTransport : Thrift::FramedTransport).new(source) + expect { transport.read_into_buffer("\0".b * 4096, 4097) }.to raise_error(IndexError) + expect(transport.read_all(kind == :buffered ? 4096 : 2)).to eq(kind == :buffered ? "y" * 4096 : "yz") + expect(input_storage(transport)).to be_empty + end + end + + [["frame header", "\0"], ["frame payload", [3].pack("N") + "c"]].each do |stage, truncated| + it "releases the consumed frame when the next #{stage} read fails" do + source = Thrift::MemoryBufferTransport.new([2].pack("N") + "ab" + truncated) + transport = Thrift::FramedTransport.new(source) + output = "\0".b * 3 + + expect { transport.read_into_buffer(output, 3) }.to raise_error(EOFError) + expect(output).to eq("ab\0") + expect(input_storage(transport)).to be_empty + + source.reset_buffer([2].pack("N") + "de") + expect(transport.read_all(2)).to eq("de") + expect(input_storage(transport)).to be_empty + end + end + end end diff --git a/lib/rb/spec/header_transport_spec.rb b/lib/rb/spec/header_transport_spec.rb index 51612b6f83f..dc8a5ffd792 100644 --- a/lib/rb/spec/header_transport_spec.rb +++ b/lib/rb/spec/header_transport_spec.rb @@ -882,4 +882,55 @@ def read(size) expect(trans.send(:read_varint32, io)).to eq(0xFFFFFFFF) end end + + context "input lifetime" do + it "releases compressed and plain payloads without consuming the next frame or reply metadata" do + wire = Thrift::MemoryBufferTransport.new + writer = Thrift::HeaderTransport.new(wire) + writer.sequence_id = 42 + writer.set_header("name", "value") + writer.write("x" * 32_768) + writer.flush + writer.add_transform(Thrift::HeaderTransformID::ZLIB) + writer.sequence_id = 43 + writer.set_header("second", "frame") + writer.write("y" * 32_768) + writer.flush + + reader = Thrift::HeaderTransport.new(wire) + expect(reader.read_all(32_768)).to eq("x" * 32_768) + expect(reader.instance_variable_get(:@read_buffer).string).to be_empty + expect(reader.get_headers).to eq("name" => "value") + expect(reader.sequence_id).to eq(42) + expect(wire.available).to be > 0 + reader.reset_protocol + expect(reader.sequence_id).to eq(43) + expect(reader.read_all(32_768)).to eq("y" * 32_768) + expect(reader.instance_variable_get(:@read_buffer).string).to be_empty + expect(reader.get_headers).to eq("second" => "frame") + end + + it "reads unsigned bytes across frame boundaries while preserving reply metadata" do + wire = Thrift::MemoryBufferTransport.new + writer = Thrift::HeaderTransport.new(wire) + writer.sequence_id = 42 + writer.set_header("name", "first") + writer.write("\x00\x80\xff".b) + writer.flush + writer.sequence_id = 43 + writer.set_header("name", "second") + writer.write("\x81".b) + writer.flush + + reader = Thrift::HeaderTransport.new(wire) + expect(Array.new(3) { reader.read_byte }).to eq([0, 128, 255]) + expect(reader.instance_variable_get(:@read_buffer).string).to be_empty + expect(reader.get_headers).to eq("name" => "first") + expect(reader.sequence_id).to eq(42) + expect(reader.read_byte).to eq(129) + expect(reader.instance_variable_get(:@read_buffer).string).to be_empty + expect(reader.get_headers).to eq("name" => "second") + expect(reader.sequence_id).to eq(43) + end + end end diff --git a/lib/rb/spec/http_client_spec.rb b/lib/rb/spec/http_client_spec.rb index ba867619976..dfaf17bf0fd 100644 --- a/lib/rb/spec/http_client_spec.rb +++ b/lib/rb/spec/http_client_spec.rb @@ -19,6 +19,8 @@ # require "spec_helper" +require "timeout" +require "weakref" describe "Thrift::HTTPClientTransport" do describe Thrift::HTTPClientTransport do @@ -271,4 +273,79 @@ expect(client.read(4)).to eq("data") end end + + context "input lifetime" do + it "releases real responses while idle and preserves unread data after a failed request" do + listener = TCPServer.new("127.0.0.1", 0) + server = Thread.new do + [200, 503, 200].each_with_index do |status, index| + socket = listener.accept + begin + socket.gets + headers = {} + while (line = socket.gets) != "\r\n" + key, value = line.split(":", 2) + headers[key.downcase] = value.strip + end + socket.read(headers.fetch("content-length", "0").to_i) + size = index == 0 ? 32_768 : 4 + socket.write("HTTP/1.1 #{status} Result\r\nContent-Length: #{size}\r\nConnection: close\r\n\r\n") + if index == 0 + 8.times { socket.write("x" * 4096) } + else + socket.write("next") + end + ensure + socket.close + end + end + end + + client_thread = nil + Timeout.timeout(10) do + # Discard request/response stack roots before checking transport ownership. + client_thread = Thread.new do + transport = Thrift::HTTPClientTransport.new("http://127.0.0.1:#{listener.addr[1]}/") + transport.write("request") + transport.flush + expect(transport.read(1)).to eq("x") + reference = WeakRef.new(transport.instance_variable_get(:@inbuf).string) + transport.close + transport.write("request") + expect { transport.flush }.to raise_error(Thrift::TransportException, /HTTP status code 503/) + 7.times { expect(transport.read(4096)).to eq("x" * 4096) } + expect(transport.read(4095)).to eq("x" * 4095) + [transport, reference] + end + transport, reference = client_thread.value + GC.start + expect(transport.instance_variable_get(:@inbuf).string).to be_empty + expect(reference.weakref_alive?).to be_falsey + expect(transport.read(0)).to eq("") + expect { transport.read(1) }.to raise_error(Thrift::TransportException, /HTTP status code 200/) + transport.write("request") + transport.flush + expect(transport.read_all(4)).to eq("next") + transport.close + expect(transport.instance_variable_get(:@inbuf).string).to be_empty + server.value + end + ensure + client_thread.kill.join if client_thread && client_thread.alive? + server.kill.join if server && server.alive? + listener.close if listener + end + + it "preserves response status when a byte read reaches EOF" do + transport = Thrift::HTTPClientTransport.new("http://localhost/") + transport.instance_variable_set(:@inbuf, StringIO.new("\xff".b)) + transport.instance_variable_set(:@response_code, 206) + + expect(transport.read_byte).to eq(255) + expect(transport.instance_variable_get(:@inbuf).string).to be_empty + expect { transport.read_byte }.to raise_error(Thrift::TransportException, /HTTP status code 206/) do |error| + expect(error.type).to eq(Thrift::TransportException::END_OF_FILE) + end + end + end end