Add a checksum: option to Zstd.compress and Zstd::StreamingCompress

README.md

@@ -77,6 +77,22 @@ cdict = Zstd::CDict.new(File.read('dictionary_file'), 5)
 compressed_using_dict = Zstd.compress("", dict: cdict)
 ```
 
+#### Compression with Checksum
+
+```ruby
+compressed_data = Zstd.compress(data, checksum: true) # default checksum is false
+```
+
+`checksum: true` adds a checksum of the uncompressed data to the end of each frame, as the `zstd` command line tool does by default. Decompression verifies it and raises an error when the data does not match. `Zstd::StreamingCompress.new` accepts it too:
+
+```ruby
+stream = Zstd::StreamingCompress.new(checksum: true)
+stream << "abc" << "def"
+res = stream.finish
+```
+
+`Zstd::StreamingDecompress` verifies the checksum when it reaches the end of the frame, after it has returned the data of that frame.
+
 #### Streaming Compression
 ```ruby
 stream = Zstd::StreamingCompress.new

ext/zstdruby/common.h

@@ -30,18 +30,30 @@ static int convert_compression_level(ZSTD_CCtx* ctx, VALUE compression_level_val
    dictionary needs no such handling: ZSTD_CCtx_loadDictionary copies it. */
 static VALUE set_compress_params(ZSTD_CCtx* const ctx, VALUE kwargs)
 {
-  ID kwargs_keys[2];
+  ID kwargs_keys[3];
   kwargs_keys[0] = rb_intern("level");
   kwargs_keys[1] = rb_intern("dict");
-  VALUE kwargs_values[2];
-  rb_get_kwargs(kwargs, kwargs_keys, 0, 2, kwargs_values);
+  kwargs_keys[2] = rb_intern("checksum");
+  VALUE kwargs_values[3];
+  rb_get_kwargs(kwargs, kwargs_keys, 0, 3, kwargs_values);
 
   int compression_level = ZSTD_CLEVEL_DEFAULT;
   if (kwargs_values[0] != Qundef && kwargs_values[0] != Qnil) {
     compression_level = convert_compression_level(ctx, kwargs_values[0]);
   }
   ZSTD_CCtx_setParameter(ctx, ZSTD_c_compressionLevel, compression_level);
 
+  /* `checksum: true` writes a checksum of the content at the end of each
+     frame, which decompression verifies. Set before `dict:`, whose CDict
+     branch returns early. */
+  if (kwargs_values[2] != Qundef && kwargs_values[2] != Qnil) {
+    if (kwargs_values[2] != Qtrue && kwargs_values[2] != Qfalse) {
+      ZSTD_freeCCtx(ctx);
+      rb_raise(rb_eArgError, "`checksum:` must be true or false");
+    }
+    ZSTD_CCtx_setParameter(ctx, ZSTD_c_checksumFlag, kwargs_values[2] == Qtrue);
+  }
+
   if (kwargs_values[1] != Qundef && kwargs_values[1] != Qnil) {
     if (CLASS_OF(kwargs_values[1]) == rb_cCDict) {
       ZSTD_CDict* cdict = DATA_PTR(kwargs_values[1]);

spec/zstd-ruby-streaming-compress_spec.rb

@@ -129,6 +129,48 @@
     end
   end
 
+  describe 'checksum' do
+    # Content_Checksum_flag: bit 2 of the frame header descriptor, the byte
+    # after the magic number.
+    def checksum_flag?(frame)
+      frame.getbyte(4) & 0x04 != 0
+    end
+
+    it 'is not written by default' do
+      stream = Zstd::StreamingCompress.new
+      stream << "abc" << "def"
+      res = stream.finish
+      expect(checksum_flag?(res)).to eq(false)
+      expect(Zstd.decompress(res)).to eq('abcdef')
+    end
+
+    it 'covers the content written across flushes' do
+      stream = Zstd::StreamingCompress.new(level: 5, checksum: true)
+      res = stream.compress("abc" * 50_000)
+      res << stream.flush
+      res << stream.compress("def" * 50_000)
+      res << stream.finish
+      expect(checksum_flag?(res)).to eq(true)
+      # Zstd.decompress verifies the checksum
+      expect(Zstd.decompress(res)).to eq(("abc" * 50_000) + ("def" * 50_000))
+    end
+
+    it 'works with a Zstd::CDict' do
+      dictionary = File.read("#{__dir__}/dictionary")
+      user_json = File.read("#{__dir__}/user_springmt.json")
+      stream = Zstd::StreamingCompress.new(dict: Zstd::CDict.new(dictionary, 5), checksum: true)
+      stream << user_json
+      res = stream.finish
+      expect(checksum_flag?(res)).to eq(true)
+      expect(Zstd.decompress(res, dict: dictionary)).to eq(user_json)
+    end
+
+    it 'raises ArgumentError unless true or false' do
+      expect { Zstd::StreamingCompress.new(checksum: 1) }.to raise_error(ArgumentError)
+      expect { Zstd::StreamingCompress.new(checksum: 'true') }.to raise_error(ArgumentError)
+    end
+  end
+
   if Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('3.0.0')
     describe 'Ractor' do
       it 'should be supported' do

spec/zstd-ruby_spec.rb

@@ -55,6 +55,67 @@ def to_str
       decompressed = Zstd.decompress(compressed)
       expect(decompressed).to eq('abc')
     end
+
+    context 'with checksum' do
+      # Content_Checksum_flag: bit 2 of the frame header descriptor, the byte
+      # after the magic number.
+      def checksum_flag?(frame)
+        frame.getbyte(4) & 0x04 != 0
+      end
+
+      # Flips a bit in the middle of the frame.
+      def corrupt(frame)
+        corrupted = frame.dup
+        middle = corrupted.bytesize / 2
+        corrupted.setbyte(middle, corrupted.getbyte(middle) ^ 1)
+        corrupted
+      end
+
+      it 'should not write a checksum by default' do
+        compressed = Zstd.compress(user_json)
+        expect(checksum_flag?(compressed)).to eq(false)
+        expect(Zstd.compress(user_json, checksum: false)).to eq(compressed)
+        expect(Zstd.compress(user_json, checksum: nil)).to eq(compressed)
+      end
+
+      it 'should add a checksum of the content to the end of the frame' do
+        compressed = Zstd.compress('abc', checksum: true)
+        expect(checksum_flag?(compressed)).to eq(true)
+        # The low 4 bytes of XXH64('abc'), 0x44BC2CF5AD770999, little endian,
+        # as the zstd command line tool writes them.
+        expect(compressed.byteslice(-4, 4)).to eq("\x99\x09\x77\xAD".b)
+        expect(Zstd.decompress(compressed)).to eq('abc')
+
+        compressed = Zstd.compress(user_json, level: 5, checksum: true)
+        expect(compressed.bytesize).to eq(Zstd.compress(user_json, level: 5).bytesize + 4)
+        expect(Zstd.decompress(compressed)).to eq(user_json)
+      end
+
+      it 'should add a checksum with a dictionary' do
+        dictionary = File.read("#{__dir__}/dictionary")
+        [dictionary, Zstd::CDict.new(dictionary)].each do |dict|
+          compressed = Zstd.compress(user_json, dict: dict, checksum: true)
+          expect(checksum_flag?(compressed)).to eq(true)
+          expect(Zstd.decompress(compressed, dict: dictionary)).to eq(user_json)
+        end
+      end
+
+      it 'should detect corrupted content on decompress' do
+        # Incompressible data is stored as is, so the flipped bit changes the
+        # content without breaking the frame: only the checksum can detect it.
+        data = Random.new(42).bytes(1000)
+        expect(Zstd.decompress(corrupt(Zstd.compress(data)))).not_to eq(data)
+
+        corrupted = corrupt(Zstd.compress(data, checksum: true))
+        expect { Zstd.decompress(corrupted) }.to raise_error(RuntimeError, /checksum/)
+        expect { Zstd::StreamingDecompress.new.decompress(corrupted) }.to raise_error(RuntimeError, /checksum/)
+      end
+
+      it 'should raise ArgumentError unless true or false' do
+        expect { Zstd.compress(user_json, checksum: 1) }.to raise_error(ArgumentError)
+        expect { Zstd.compress(user_json, checksum: 'true') }.to raise_error(ArgumentError)
+      end
+    end
   end
 
   describe 'decompress' do