From 82f8165c0fe40ecfa327608706b212a734072b33 Mon Sep 17 00:00:00 2001
From: Bala Kumar <sbalakumar@gitlab.com>
Date: Sun, 9 Aug 2026 18:51:53 +0530
Subject: [PATCH] Align StreamReader with Ruby's IO conventions

StreamReader#read passed its length argument straight to the underlying
IO, so it read that many *compressed* bytes and returned however many
decompressed bytes fell out. The size of the return value tracked the
compression ratio rather than the caller's request: a small frame could
return far more than asked for, while a larger one returned an empty
String because zstd was still filling an internal block. That makes the
reader unusable for consumers that need a specific number of bytes, such
as Gem::Package::TarReader.

Buffer decompressed output so length means decompressed bytes, and serve
reads from that buffer. Refills use decompress_with_pos rather than
decompress: it writes at most ZSTD_DStreamOutSize bytes per call and
reports how much input it consumed, so a high compression ratio cannot
balloon the buffer.

Also return nil at EOF instead of raising StandardError, matching
IO#read. Rescuing StandardError to detect EOF would otherwise swallow
genuine decompression failures, which the extension raises as
RuntimeError.

Add eof?, support read with no length and an outbuf argument, and fix
close, which called finish on StreamingDecompress (a method it does not
define) and wrote to a read-only IO.

Correct the spelling of the experimental marker on both StreamReader and
StreamWriter. The marker stays in place: this does not promote either
class to a stable API.
---
 lib/zstd-ruby/stream_reader.rb       |  92 +++++++++++--
 lib/zstd-ruby/stream_writer.rb       |   2 +-
 spec/zstd-ruby-stream_reader_spec.rb | 195 ++++++++++++++++++++++++++-
 3 files changed, 273 insertions(+), 16 deletions(-)

diff --git a/lib/zstd-ruby/stream_reader.rb b/lib/zstd-ruby/stream_reader.rb
index 0efef9c..f3b4c19 100644
--- a/lib/zstd-ruby/stream_reader.rb
+++ b/lib/zstd-ruby/stream_reader.rb
@@ -1,22 +1,98 @@
 module Zstd
-  # @todo Exprimental
+  # @todo Experimental
+  #
+  # Wraps a readable IO containing zstd-compressed data and exposes the
+  # decompressed bytes through a standard read interface.
+  #
+  #   File.open('data.zst') do |file|
+  #     reader = Zstd::StreamReader.new(file)
+  #     while (chunk = reader.read(16 * 1024))
+  #       # ...
+  #     end
+  #   end
+  #
+  # +length+ is a number of *decompressed* bytes, and +read+ returns +nil+ once
+  # the stream is exhausted, so the reader can drive any consumer that expects
+  # an IO-like object.
   class StreamReader
-    def initialize(io)
+    # Number of compressed bytes pulled from the underlying IO per refill.
+    DEFAULT_CHUNK_SIZE = 64 * 1024
+
+    def initialize(io, chunk_size: DEFAULT_CHUNK_SIZE)
+      raise ArgumentError, "chunk_size must be positive" unless chunk_size.to_int > 0
+
       @io = io
+      @chunk_size = chunk_size.to_int
       @stream = Zstd::StreamingDecompress.new
+      @buffer = +''.b
+      @pending = +''.b
+      @source_eof = false
     end
 
-    def read(length)
-      if @io.eof?
-        raise StandardError, "EOF"
+    # Reads and returns up to +length+ decompressed bytes, or every remaining
+    # byte when +length+ is nil.
+    #
+    # Returns nil at the end of the stream, or an empty String when +length+ is
+    # zero, mirroring IO#read.
+    def read(length = nil, outbuf = nil)
+      if length.nil?
+        fill_until(Float::INFINITY)
+        data = @buffer
+        @buffer = +''.b
+        return outbuf ? outbuf.replace(data) : data
+      end
+
+      length = length.to_int
+      raise ArgumentError, "negative length #{length} given" if length < 0
+      return outbuf ? outbuf.replace(+''.b) : +''.b if length == 0
+
+      fill_until(length)
+      if @buffer.empty?
+        outbuf&.replace(+''.b)
+        return nil
       end
-      data = @io.read(length)
-      @stream.decompress(data)
+
+      data = @buffer.byteslice(0, length)
+      @buffer = @buffer.byteslice(data.bytesize..) || +''.b
+      outbuf ? outbuf.replace(data) : data
+    end
+
+    def eof?
+      fill_until(1)
+      @buffer.empty?
     end
+    alias_method :eof, :eof?
 
     def close
-      @io.write(@stream.finish)
       @io.close
+      nil
+    end
+
+    private
+
+    def fill_until(n)
+      while @buffer.bytesize < n
+        if @pending.empty?
+          break if @source_eof
+
+          chunk = @io.read(@chunk_size)
+          if chunk.nil? || chunk.empty?
+            @source_eof = true
+            break
+          end
+          @pending = chunk
+        end
+
+        # decompress_with_pos writes at most ZSTD_DStreamOutSize bytes per call
+        # and reports how much input it consumed, so a high compression ratio
+        # cannot balloon the buffer the way decompress would.
+        decompressed, consumed = @stream.decompress_with_pos(@pending)
+        @pending = @pending.byteslice(consumed..) || +''.b
+        @buffer << decompressed
+
+        # No forward progress: the frame is truncated or complete.
+        break if consumed == 0 && decompressed.empty?
+      end
     end
   end
 end
diff --git a/lib/zstd-ruby/stream_writer.rb b/lib/zstd-ruby/stream_writer.rb
index 78ac4db..e7d7d92 100644
--- a/lib/zstd-ruby/stream_writer.rb
+++ b/lib/zstd-ruby/stream_writer.rb
@@ -1,5 +1,5 @@
 module Zstd
-  # @todo Exprimental
+  # @todo Experimental
   class StreamWriter
     def initialize(io, level: nil)
       @io = io
diff --git a/spec/zstd-ruby-stream_reader_spec.rb b/spec/zstd-ruby-stream_reader_spec.rb
index 4d25709..a6b29cc 100644
--- a/spec/zstd-ruby-stream_reader_spec.rb
+++ b/spec/zstd-ruby-stream_reader_spec.rb
@@ -1,10 +1,27 @@
 require "spec_helper"
 require 'zstd-ruby'
-require 'pry'
+require 'stringio'
 
 RSpec.describe Zstd::StreamReader do
-  describe 'read' do
-    it 'shoud work' do
+  # Compressible payload large enough to span several internal zstd blocks.
+  let(:compressible) { (1..20_000).map { |n| %({"name":"pkg-#{n}","version":"1.0.#{n}"}) }.join("\n") << "\n" }
+  # Incompressible payload: compressed size ends up close to the original size.
+  let(:incompressible) { Random.new(42).bytes(512 * 1024) }
+
+  def reader_for(data, **opts)
+    described_class.new(StringIO.new(Zstd.compress(data)), **opts)
+  end
+
+  def read_all(reader, length)
+    +''.b.tap do |out|
+      while (chunk = reader.read(length))
+        out << chunk
+      end
+    end
+  end
+
+  describe '#read' do
+    it 'reads the data written by StreamWriter' do
       io = StringIO.new
       writer = Zstd::StreamWriter.new(io)
       writer.write("abc")
@@ -12,10 +29,174 @@
       writer.finish
       io.rewind
 
-      reader = Zstd::StreamReader.new(io)
-      expect(reader.read(10)).to eq('a')
-      expect(reader.read(10)).to eq('bcdef')
-      expect(reader.read(10)).to eq('')
+      reader = described_class.new(io)
+      expect(reader.read(10)).to eq('abcdef')
+      expect(reader.read(10)).to be_nil
+    end
+
+    it 'returns exactly the requested number of decompressed bytes' do
+      reader = reader_for(compressible)
+
+      5.times { expect(reader.read(512).bytesize).to eq(512) }
+    end
+
+    it 'returns the requested length regardless of the compression ratio' do
+      # A tiny compressed frame expands to far more than the requested length;
+      # the reader must not hand back the whole frame at once.
+      reader = reader_for("hello world\n" * 20_000)
+
+      expect(reader.read(512).bytesize).to eq(512)
+    end
+
+    it 'round-trips compressible data' do
+      expect(read_all(reader_for(compressible), 512)).to eq(compressible)
+    end
+
+    it 'round-trips incompressible data' do
+      expect(read_all(reader_for(incompressible), 4096)).to eq(incompressible)
+    end
+
+    it 'round-trips data smaller than a single read' do
+      expect(read_all(reader_for('sample data'), 4096)).to eq('sample data')
+    end
+
+    it 'returns a short final read rather than padding' do
+      reader = reader_for('0123456789')
+
+      expect(reader.read(4)).to eq('0123')
+      expect(reader.read(100)).to eq('456789')
+      expect(reader.read(100)).to be_nil
+    end
+
+    it 'returns nil once the stream is exhausted' do
+      reader = reader_for('sample data')
+      reader.read(4096)
+
+      expect(reader.read(4096)).to be_nil
+      expect(reader.read(4096)).to be_nil
+    end
+
+    it 'reads the whole stream when no length is given' do
+      expect(reader_for(compressible).read).to eq(compressible)
+    end
+
+    it 'returns an empty String at EOF when no length is given' do
+      reader = reader_for('sample data')
+      reader.read
+
+      expect(reader.read).to eq('')
+    end
+
+    it 'returns an empty String for a zero length' do
+      expect(reader_for('sample data').read(0)).to eq('')
+    end
+
+    it 'raises ArgumentError for a negative length' do
+      expect { reader_for('sample data').read(-1) }.to raise_error(ArgumentError)
+    end
+
+    it 'writes into outbuf when given' do
+      reader = reader_for(compressible)
+      outbuf = +'previous contents'
+
+      result = reader.read(512, outbuf)
+
+      expect(result).to equal(outbuf)
+      expect(outbuf.bytesize).to eq(512)
+    end
+
+    it 'clears outbuf and returns nil at EOF' do
+      reader = reader_for('sample data')
+      reader.read(4096)
+      outbuf = +'previous contents'
+
+      expect(reader.read(4096, outbuf)).to be_nil
+      expect(outbuf).to be_empty
+    end
+
+    it 'honours a custom chunk_size' do
+      reader = reader_for(compressible, chunk_size: 1024)
+
+      expect(read_all(reader, 512)).to eq(compressible)
+    end
+
+    it 'rejects a non-positive chunk_size' do
+      expect { reader_for('sample data', chunk_size: 0) }.to raise_error(ArgumentError)
+    end
+  end
+
+  describe '#eof?' do
+    it 'is false while data remains and true once drained' do
+      reader = reader_for('sample data')
+
+      expect(reader.eof?).to eq(false)
+      expect(reader.read(4096)).to eq('sample data')
+      expect(reader.eof?).to eq(true)
+    end
+
+    it 'does not consume data' do
+      reader = reader_for('sample data')
+
+      expect(reader.eof?).to eq(false)
+      expect(reader.read(4096)).to eq('sample data')
+    end
+  end
+
+  describe '#close' do
+    it 'closes the underlying IO' do
+      io = StringIO.new(Zstd.compress('sample data'))
+      reader = described_class.new(io)
+
+      expect { reader.close }.not_to raise_error
+      expect(io).to be_closed
+    end
+  end
+
+  describe 'driving an IO consumer' do
+    it 'can be wrapped to feed Gem::Package::TarReader' do
+      require 'rubygems/package'
+
+      tar = StringIO.new
+      Gem::Package::TarWriter.new(tar) do |writer|
+        3.times { |i| writer.add_file("0#{i}.txt", 0644) { |f| f.write(compressible) } }
+      end
+
+      io_like = Class.new do
+        def initialize(reader)
+          @reader = reader
+          @pos = 0
+        end
+
+        attr_reader :pos
+
+        def read(length = nil, outbuf = nil)
+          data = @reader.read(length, outbuf)
+          @pos += data.bytesize if data
+          data
+        end
+
+        def eof? = @reader.eof?
+
+        def seek(amount, whence = IO::SEEK_SET)
+          raise Errno::EINVAL unless whence == IO::SEEK_CUR
+
+          while amount > 0
+            chunk = read(amount < 65_536 ? amount : 65_536)
+            break unless chunk
+
+            amount -= chunk.bytesize
+          end
+          0
+        end
+      end
+
+      entries = []
+      source = StringIO.new(Zstd.compress(tar.string))
+      Gem::Package::TarReader.new(io_like.new(described_class.new(source))) do |reader|
+        reader.each { |entry| entries << [entry.full_name, entry.read.bytesize] }
+      end
+
+      expect(entries).to eq(3.times.map { |i| ["0#{i}.txt", compressible.bytesize] })
     end
   end
 end