THRIFT-6129: Stream Ruby HeaderTransport ZLIB output

Client: rb

Co-Authored-By: OpenAI Codex (GPT-5.6) <codex@openai.com>

Author: kpumuk

lib/rb/lib/thrift/transport/header_transport.rb

@@ -77,9 +77,6 @@ class HeaderTransport < BaseTransport
     # Default decompressed-size cap for ZLIB transform (~15.6 MB, matches other Thrift bindings)
     DEFAULT_MAX_DECOMPRESSED_SIZE = 16_384_000
 
-    # Chunk size for streaming inflate loop
-    ZLIB_INFLATE_CHUNK_SIZE = 16_384
-
     # Binary protocol version mask and version 1
     BINARY_VERSION_MASK = 0xffff0000
     BINARY_VERSION_1 = 0x80010000
@@ -408,25 +405,19 @@ def parse_header_format(buf)
     def bounded_inflate(compressed)
       inflater = Zlib::Inflate.new
       buffer = Bytes.empty_byte_buffer
-      offset = 0
-      begin
-        while offset < compressed.bytesize
-          buffer << inflater.inflate(compressed.byteslice(offset, ZLIB_INFLATE_CHUNK_SIZE))
-          if buffer.bytesize > @max_decompressed_size
-            raise TransportException.new(
-              TransportException::SIZE_LIMIT,
-              "Decompressed size exceeds limit of #{@max_decompressed_size}"
-            )
-          end
-          offset += ZLIB_INFLATE_CHUNK_SIZE
-        end
-        buffer << inflater.finish
-        if buffer.bytesize > @max_decompressed_size
+      append_chunk = lambda do |chunk|
+        if buffer.bytesize + chunk.bytesize > @max_decompressed_size
           raise TransportException.new(
             TransportException::SIZE_LIMIT,
             "Decompressed size exceeds limit of #{@max_decompressed_size}"
           )
         end
+
+        buffer << chunk
+      end
+      begin
+        inflater.inflate(compressed, &append_chunk)
+        inflater.finish(&append_chunk)
         buffer
       ensure
         inflater.close rescue nil

lib/rb/spec/header_transport_spec.rb

@@ -201,6 +201,31 @@
           expect(e.message).to match(/limit/)
         end
       end
+
+      it "should stream oversized ZLIB output before enforcing the limit" do
+        write_buf = Thrift::MemoryBufferTransport.new
+        writer = Thrift::HeaderTransport.new(write_buf)
+        writer.add_transform(Thrift::HeaderTransformID::ZLIB)
+        writer.write("A" * 8_000_000)
+        writer.flush
+
+        written_data = write_buf.read(write_buf.available)
+        read_trans = Thrift::HeaderTransport.new(Thrift::MemoryBufferTransport.new(written_data))
+        read_trans.set_max_decompressed_size(100)
+
+        expect(Zlib::Inflate).to receive(:new).and_wrap_original do |new, *args|
+          inflater = new.call(*args)
+          expect(inflater).to receive(:inflate).and_wrap_original do |inflate, compressed, &block|
+            expect(block).not_to be_nil
+            inflate.call(compressed, &block)
+          end
+          inflater
+        end
+
+        expect { read_trans.read(1) }.to raise_error(Thrift::TransportException) do |e|
+          expect(e.type).to eq(Thrift::TransportException::SIZE_LIMIT)
+        end
+      end
     end
 
     describe "read and frame detection" do