From 7b92c7088c4df799043a7ce62a436dea5cc2dd4b Mon Sep 17 00:00:00 2001 From: Reid Morrison Date: Wed, 14 Oct 2026 23:58:04 -0400 Subject: [PATCH] Add StreamingDecompress#frame_complete? to detect truncated input StreamingDecompress#decompress and #decompress_with_pos discard the return value of ZSTD_decompressStream: 0 once a frame is decoded and flushed, and otherwise how much more input the frame needs. Without it, a caller cannot tell input that ends at the end of a frame from input that was cut short. Both return the data decoded so far without an error, and input cut right after a block, or before the content checksum, returns all of the data it holds. Record after each call whether the input consumed so far ends at the end of a frame, and expose it as frame_complete?. Only a call that consumes input or produces output updates it: after a complete frame, libzstd answers a call with no input with the size of the next frame header, which says nothing about the input so far. It is true before any input, and false after a call raises. Co-Authored-By: Claude Opus 5.5 --- README.md | 18 +++ ext/zstdruby/streaming_decompress.c | 30 +++++ spec/zstd-ruby-streaming-decompress_spec.rb | 124 ++++++++++++++++++++ 3 files changed, 172 insertions(+) diff --git a/README.md b/README.md index d20c093..8b137c0 100644 --- a/README.md +++ b/README.md @@ -181,6 +181,24 @@ result, consumed_bytes = stream.decompress_with_pos(cstr[0, 10]) This is particularly useful when processing streaming data where you need to track the exact position in the input stream. +#### Detecting a Truncated Stream + +`frame_complete?` returns `true` when the input decompressed so far ends exactly at the end of a frame, and `false` when it ends within a frame. Once all of the input has been decompressed, `false` means that it was cut short: + +```ruby +stream = Zstd::StreamingDecompress.new +File.open('data.zst', 'rb') do |input| + File.open('data', 'wb') do |output| + while (chunk = input.read(65_536)) + output << stream.decompress(chunk) + end + end +end +raise 'data.zst is truncated' unless stream.frame_complete? +``` + +It is `true` before any input, and a call with an empty String does not change it. With `decompress_with_pos`, only the input it reports as consumed counts. + ### Skippable frame ```ruby diff --git a/ext/zstdruby/streaming_decompress.c b/ext/zstdruby/streaming_decompress.c index 58d623c..8fae7cd 100644 --- a/ext/zstdruby/streaming_decompress.c +++ b/ext/zstdruby/streaming_decompress.c @@ -5,6 +5,7 @@ struct streaming_decompress_t { VALUE buf; size_t buf_size; VALUE dict; /* Zstd::DDict the dctx borrows a pointer into, or Qnil */ + bool frame_complete; /* the input decompressed so far ends at the end of a frame */ }; static void @@ -69,6 +70,7 @@ rb_streaming_decompress_allocate(VALUE klass) RB_OBJ_WRITE(obj, &sd->buf, Qnil); sd->buf_size = 0; RB_OBJ_WRITE(obj, &sd->dict, Qnil); + sd->frame_complete = true; return obj; } @@ -96,6 +98,20 @@ rb_streaming_decompress_initialize(int argc, VALUE *argv, VALUE obj) return obj; } +/* ZSTD_decompressStream returns 0 once a frame is decoded and flushed, and + otherwise how much more input the frame needs. A call that consumes no input + and produces no output says nothing about the input so far: after a complete + frame, it asks for the header of the next one. */ +static void +update_frame_complete(struct streaming_decompress_t* sd, size_t ret, bool progress) +{ + if (ZSTD_isError(ret)) { + sd->frame_complete = false; + } else if (progress) { + sd->frame_complete = ret == 0; + } +} + static VALUE rb_streaming_decompress_decompress(VALUE obj, VALUE src) { @@ -111,7 +127,9 @@ rb_streaming_decompress_decompress(VALUE obj, VALUE src) while (input.pos < input.size) { const char* output_data = RSTRING_PTR(sd->buf); ZSTD_outBuffer output = { (void*)output_data, sd->buf_size, 0 }; + size_t const input_pos = input.pos; size_t const ret = zstd_stream_decompress(sd->dctx, &output, &input, false); + update_frame_complete(sd, ret, input.pos > input_pos || output.pos > 0); if (ZSTD_isError(ret)) { rb_raise(rb_eRuntimeError, "decompress error error code: %s", ZSTD_getErrorName(ret)); @@ -140,6 +158,7 @@ rb_streaming_decompress_decompress_with_pos(VALUE obj, VALUE src) VALUE result = rb_str_new(0, 0); ZSTD_outBuffer output = { (void*)output_data, sd->buf_size, 0 }; size_t const ret = zstd_stream_decompress(sd->dctx, &output, &input, false); + update_frame_complete(sd, ret, input.pos > 0 || output.pos > 0); if (ZSTD_isError(ret)) { rb_raise(rb_eRuntimeError, "decompress error error code: %s", ZSTD_getErrorName(ret)); } @@ -147,6 +166,16 @@ rb_streaming_decompress_decompress_with_pos(VALUE obj, VALUE src) return rb_ary_new_from_args(2, result, ULONG2NUM(input.pos)); } +/* Whether the input decompressed so far ends at the end of a frame: true before + any input, false while it ends within a frame, or after a call raised. */ +static VALUE +rb_streaming_decompress_frame_complete_p(VALUE obj) +{ + struct streaming_decompress_t* sd; + TypedData_Get_Struct(obj, struct streaming_decompress_t, &streaming_decompress_type, sd); + return sd->frame_complete ? Qtrue : Qfalse; +} + extern VALUE rb_mZstd, cStreamingDecompress; void zstd_ruby_streaming_decompress_init(void) @@ -156,4 +185,5 @@ zstd_ruby_streaming_decompress_init(void) rb_define_method(cStreamingDecompress, "initialize", rb_streaming_decompress_initialize, -1); rb_define_method(cStreamingDecompress, "decompress", rb_streaming_decompress_decompress, 1); rb_define_method(cStreamingDecompress, "decompress_with_pos", rb_streaming_decompress_decompress_with_pos, 1); + rb_define_method(cStreamingDecompress, "frame_complete?", rb_streaming_decompress_frame_complete_p, 0); } diff --git a/spec/zstd-ruby-streaming-decompress_spec.rb b/spec/zstd-ruby-streaming-decompress_spec.rb index 053c4ed..09a4f96 100644 --- a/spec/zstd-ruby-streaming-decompress_spec.rb +++ b/spec/zstd-ruby-streaming-decompress_spec.rb @@ -180,6 +180,130 @@ end end + describe 'frame_complete?' do + # Written by the zstd command line tool, which adds a checksum of the + # content to the end of each frame by default: + # ruby -e 'print "foo bar buzz " * 10' | zstd -c + # The magic number (4 bytes), frame header (2), block header (3), + # compressed block (20) and content checksum (4). + let(:frame) { ['28b52ffd0458a5000068666f6f206261722062757a7a20010012d2f905b7745bfe'].pack('H*') } + let(:content) { 'foo bar buzz ' * 10 } + + it 'is true before any input' do + stream = Zstd::StreamingDecompress.new + expect(stream.frame_complete?).to eq(true) + expect(stream.decompress('')).to eq('') + expect(stream.decompress_with_pos('')).to eq(['', 0]) + expect(stream.frame_complete?).to eq(true) + end + + it 'is true when the input ends at the end of a frame' do + stream = Zstd::StreamingDecompress.new + expect(stream.decompress(frame)).to eq(content) + expect(stream.frame_complete?).to eq(true) + end + + { + 'within the magic number' => 2, + 'within the frame header' => 5, + 'within a block header' => 7, + 'within a block' => 15, + 'right before the checksum' => 29, + 'within the checksum' => 31 + }.each do |where, length| + it "is false when the input ends #{where}" do + stream = Zstd::StreamingDecompress.new + result = stream.decompress(frame.byteslice(0, length)) + expect(stream.frame_complete?).to eq(false) + + result << stream.decompress(frame.byteslice(length..-1)) + expect(result).to eq(content) + expect(stream.frame_complete?).to eq(true) + end + end + + it 'is false when the input ends between two blocks' do + compress = Zstd::StreamingCompress.new + # flush ends a block, so that all the data so far can be decompressed + head = compress.compress('abc' * 1000) + compress.flush + tail = compress.compress('def') + compress.finish + + stream = Zstd::StreamingDecompress.new + expect(stream.decompress(head)).to eq('abc' * 1000) + expect(stream.frame_complete?).to eq(false) + expect(stream.decompress(tail)).to eq('def') + expect(stream.frame_complete?).to eq(true) + end + + it 'is true only at the end of each of several frames' do + skippable = [0x184D2A50, 8].pack('VV') + 'metadata' + frames = [Zstd.compress('Hello, '), frame, skippable, Zstd.compress(''), Zstd.compress('World!')] + input = frames.join + ends = frames.each_with_object([]) { |f, sums| sums << (sums.last || 0) + f.bytesize } + + # One byte at a time, so that the input ends at every position. + stream = Zstd::StreamingDecompress.new + stream_with_pos = Zstd::StreamingDecompress.new + result = '' + input.bytesize.times do |pos| + byte = input.byteslice(pos, 1) + result << stream.decompress(byte) + expect(stream_with_pos.decompress_with_pos(byte)[1]).to eq(1) + + expected = ends.include?(pos + 1) + expect(stream.frame_complete?).to eq(expected), "decompress after #{pos + 1} bytes" + expect(stream_with_pos.frame_complete?).to eq(expected), "decompress_with_pos after #{pos + 1} bytes" + end + expect(result).to eq('Hello, ' + content + 'World!') + end + + it 'is false when the input ends within a later frame' do + last = Zstd.compress('World!') + input = Zstd.compress('Hello, ') + frame + last + + stream = Zstd::StreamingDecompress.new + # 5 bytes into the last frame: its magic number and frame header descriptor + expect(stream.decompress(input.byteslice(0, input.bytesize - last.bytesize + 5))).to eq('Hello, ' + content) + expect(stream.frame_complete?).to eq(false) + + stream = Zstd::StreamingDecompress.new + expect(stream.decompress(input)).to eq('Hello, ' + content + 'World!') + expect(stream.frame_complete?).to eq(true) + end + + it 'stays true after calls with empty input' do + stream = Zstd::StreamingDecompress.new + stream.decompress(frame) + 3.times do + expect(stream.decompress('')).to eq('') + expect(stream.decompress_with_pos('')).to eq(['', 0]) + expect(stream.frame_complete?).to eq(true) + end + end + + it 'is false until decompress_with_pos has returned all the output of a frame' do + # decompress_with_pos returns at most ZSTD_DStreamOutSize (128 KiB) per call. + data = 'a' * 1_000_000 + input = Zstd.compress(data) + stream = Zstd::StreamingDecompress.new + result = '' + until input.empty? + output, consumed = stream.decompress_with_pos(input) + result << output + input = input.byteslice(consumed..-1) + expect(stream.frame_complete?).to eq(result.bytesize == data.bytesize) + end + expect(result).to eq(data) + end + + it 'is false after the input fails to decompress' do + stream = Zstd::StreamingDecompress.new + stream.decompress(frame) + expect { stream.decompress('not zstd') }.to raise_error(RuntimeError) + expect(stream.frame_complete?).to eq(false) + end + end + if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('3.0.0') describe 'Ractor' do it 'should be supported' do