From 9f527c30f7813917eac08cc41f9d14cfeab7cbec Mon Sep 17 00:00:00 2001 From: Dmytro Shteflyuk Date: Sun, 18 Oct 2026 14:11:21 -0400 Subject: [PATCH] THRIFT-6434: Release exhausted input buffers in Ruby transports Client: rb Release consumed input storage in BufferedTransport, FramedTransport, HeaderTransport and HTTPClientTransport without waiting for another read or response. Preserve unread data, returned strings and Header metadata. Add regression coverage for input ownership, EOF and error handling, pipelined frames and transport round trips. Co-Authored-By: OpenAI Codex (gpt-6.1-sol) --- .../thrift/transport/buffered_transport.rb | 39 +++- .../lib/thrift/transport/framed_transport.rb | 34 ++- .../lib/thrift/transport/header_transport.rb | 13 ++ .../thrift/transport/http_client_transport.rb | 12 ++ lib/rb/spec/base_transport_spec.rb | 194 ++++++++++++++++++ lib/rb/spec/header_transport_spec.rb | 51 +++++ lib/rb/spec/http_client_spec.rb | 77 +++++++ 7 files changed, 406 insertions(+), 14 deletions(-) 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