Keep the input in the result of Zstd.write_skippable_frame

ext/zstdruby/skippable_frame.c

@@ -21,15 +21,20 @@ static VALUE rb_write_skippable_frame(int argc, VALUE *argv, VALUE self)
   char* skip_data = RSTRING_PTR(skip_value);
   size_t skip_size = RSTRING_LEN(skip_value);
 
-  size_t dst_size = input_size + ZSTD_SKIPPABLEHEADERSIZE + skip_size;
-  VALUE output = rb_str_new(NULL, dst_size);
+  /* The skippable frame goes first, followed by the input: read_skippable_frame
+     reads a frame at the start, and decompress skips it to reach the input. */
+  size_t frame_capacity = ZSTD_SKIPPABLEHEADERSIZE + skip_size;
+  VALUE output = rb_str_new(NULL, frame_capacity + input_size);
   char* output_data = RSTRING_PTR(output);
-  size_t output_size = ZSTD_writeSkippableFrame((void*)output_data, dst_size, (const void*)skip_data, skip_size, magic_variant);
-  if (ZSTD_isError(output_size)) {
-    rb_raise(rb_eRuntimeError, "%s: %s", "write skippable frame failed", ZSTD_getErrorName(output_size));
+  size_t frame_size = ZSTD_writeSkippableFrame((void*)output_data, frame_capacity, (const void*)skip_data, skip_size, magic_variant);
+  if (ZSTD_isError(frame_size)) {
+    rb_raise(rb_eRuntimeError, "%s: %s", "write skippable frame failed", ZSTD_getErrorName(frame_size));
   }
+  memcpy(output_data + frame_size, RSTRING_PTR(input_value), input_size);
 
-  rb_str_resize(output, output_size);
+  rb_str_resize(output, frame_size + input_size);
+  RB_GC_GUARD(input_value);
+  RB_GC_GUARD(skip_value);
   return output;
 }
 

spec/zstd-skippable_frame_spec.rb

@@ -57,4 +57,36 @@
       end
     end
   end
+
+  describe 'write_skippable_frame' do
+    let(:payload) { SecureRandom.hex(150) }
+    let(:compressed_data) { Zstd.compress(payload) }
+
+    it 'returns the skippable frame followed by the input' do
+      result = Zstd.write_skippable_frame(compressed_data, "sample data")
+
+      expect(result.bytesize).to eq(8 + "sample data".bytesize + compressed_data.bytesize)
+      expect(result).to end_with(compressed_data)
+    end
+
+    it 'keeps the compressed data decodable by Zstd.decompress' do
+      result = Zstd.write_skippable_frame(compressed_data, "sample data")
+
+      expect(Zstd.decompress(result)).to eq payload
+    end
+
+    it 'keeps the compressed data with a magic_variant' do
+      result = Zstd.write_skippable_frame(compressed_data, "sample data", magic_variant: 15)
+
+      expect(Zstd.read_skippable_frame(result)).to eq "sample data"
+      expect(Zstd.decompress(result)).to eq payload
+    end
+
+    it 'accepts an empty input' do
+      result = Zstd.write_skippable_frame("", "sample data")
+
+      expect(result.bytesize).to eq(8 + "sample data".bytesize)
+      expect(Zstd.read_skippable_frame(result)).to eq "sample data"
+    end
+  end
 end