Add StreamingDecompress#frame_complete? to detect truncated input

[reidmorrison]

Summary

Adds Zstd::StreamingDecompress#frame_complete?. It returns true when the input decompressed so far ends exactly at the end of a frame, and false when it ends within a frame, so once all of the input has been passed in, false means that it was cut short:

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?

Motivation

IOStreams is a Ruby library that streams files through compression, encryption and storage (local files, S3, SFTP) one block at a time. On MRI it reads .zst files with Zstd::StreamingDecompress#decompress_with_pos, and it cannot detect a truncated file, for example one cut short by an interrupted transfer or a full disk. decompress and decompress_with_pos return the data decoded so far without an error, whether the input ends at the end of a frame or within one. When it ends right after a block (a writer that called flush and then stopped) or just before the content checksum, they even return all of the data the file holds. The zstd command line tool reports these files as truncated ("premature end"), and zstd-jni, which IOStreams uses on JRuby, raises Truncated source.

ZSTD_decompressStream returns exactly this: 0 once a frame is decoded and flushed, and otherwise how much more input the frame needs. Both methods discard it. #88 returned it from decompress2, but it was dropped when that became decompress_with_pos, before 2.0.0.

Change

After each call to ZSTD_decompressStream in decompress and decompress_with_pos, StreamingDecompress records whether the input consumed so far ends at the end of a frame, and frame_complete? returns it.

Returning the raw value would not be enough: after a complete frame, a call with no input returns 9 (the size of the next frame header and block header), although nothing is missing. So only a call that makes progress, by consuming input or producing output (libzstd's own definition of forward progress), updates it:

• It is true for a new stream, before any input.
• It is false while the input ends within a frame: in the magic number, the frame header, a block header, a block, or the checksum, or between two blocks.
• It is true again whenever a frame ends, for any number of frames, including skippable and empty frames.
• A call with an empty String does not change it, so an empty call at the end of the input cannot make a complete stream look truncated.
• It is false after a call raises.

With decompress_with_pos, input that the call did not consume counts once it is passed in again. For example, when the output of a frame needs more than one call, libzstd holds back the last byte of the frame until the rest of its output has been returned, and frame_complete? stays false until then.

Compatibility

This only adds a method. decompress and decompress_with_pos return the same values as before.

Tests

spec/zstd-ruby-streaming-decompress_spec.rb uses a frame written by the zstd command line tool, which has a content checksum, and checks that frame_complete?:

• is true before any input, and after empty input;
• is true when the input ends at the end of a frame;
• is false when the input ends within the magic number, the frame header, a block header, a block, right before the checksum, or within the checksum, and true once the rest arrives;
• is false when the input ends between two blocks, after StreamingCompress#flush, although all of its data has been returned;
• is true only at the end of each of several frames, including a skippable and an empty frame, fed one byte at a time to both decompress and decompress_with_pos;
• is false when the input ends within a later frame;
• stays true over repeated calls with empty input after a complete frame;
• stays false until decompress_with_pos has returned all the output of a 1 MB frame;
• is false after the input fails to decompress.

bundle exec rake passes with 105 examples, 0 failures (91 on main), on Ruby 3.2, 3.4 and 4.0 on macOS (arm64), with no new compiler warnings.

I also checked it against 39 files written by the zstd command line tool (from files and from stdin, with and without checksums, -19, --fast=5, --long=24, -T4, --stream-size) and by this gem (one shot and streaming, levels -5 to 19). They hold empty, small, text, random, all-zero and mixed data, plus a skippable frame on its own and 8 concatenated frames including 2 skippable frames and an empty frame. Each file was cut at many points (every byte for small files, and every byte near each frame boundary for large ones) and read with decompress and decompress_with_pos, in 64 KiB chunks and one byte at a time: 12,816 reads, all correct, apart from 205 reads of the concatenated frames with decompress in 64 KiB chunks, which hit the bug below.

Related

Since v2.0.1, StreamingDecompress#decompress stops when a frame without output (an empty or skippable frame) ends, and drops the rest of its input, so for such input frame_complete? can only describe what decompress decoded. #166 fixes decompress. The two PRs are independent and merge cleanly in either order; with both, all 12,816 reads are correct.

Zstd::StreamReader could use frame_complete? to raise on a truncated stream instead of returning short data. I left that out, since it changes how an experimental class behaves, and can send it separately.

🤖 Generated with Claude Code