Add StreamingDecompress#frame_complete? to detect truncated input

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

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,13 +158,24 @@ 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));
   }
   rb_str_cat(result, output.dst, output.pos);
   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);
 }

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