Add option to write checksums

writer.go

@@ -143,6 +143,9 @@ type WriterParams struct {
 
 	// Dict is optional dictionary used for compression.
 	Dict *CDict
+
+	// Checksum. Enable writes of uncompressed data checksum at end of the frame.
+	Checksum bool
 }
 
 // NewWriterParams returns new zstd writer writing compressed data to w
@@ -230,6 +233,16 @@ func initCStream(cs *C.ZSTD_CStream, params WriterParams) {
 		C.ZSTD_cParameter(C.ZSTD_c_windowLog),
 		C.int(params.WindowLog))
 	ensureNoError("ZSTD_CCtx_setParameter", result)
+
+	checksumFlag := 0
+	if params.Checksum {
+		checksumFlag = 1
+	}
+	result = C.ZSTD_CCtx_setParameter_wrapper(
+		C.uintptr_t(uintptr(unsafe.Pointer(cs))),
+		C.ZSTD_cParameter(C.ZSTD_c_checksumFlag),
+		C.int(checksumFlag))
+	ensureNoError("ZSTD_CCtx_setParameter", result)
 }
 
 func freeCStream(v interface{}) {

writer_test.go

@@ -496,3 +496,28 @@ func TestWriterBig(t *testing.T) {
 		t.Fatalf("unequal writtenBB and readBB\nwrittenBB=\n%X\nreadBB=\n%X", writtenBB.Bytes(), readBB.Bytes())
 	}
 }
+
+func TestCorruptData(t *testing.T) {
+	var bb bytes.Buffer
+	params := &WriterParams{
+		Checksum: true,
+	}
+	src := []byte(newTestString(512, 3))
+	zw := NewWriterParams(&bb, params)
+	if _, err := zw.Write(src); err != nil {
+		t.Fatalf("error while writing with Checksum param on: %s", err)
+	}
+	if err := zw.Close(); err != nil {
+		t.Fatalf("unexpected error when closing zw: %s", err)
+	}
+
+	cd := bb.Bytes()
+	const checksumFieldLen = 4
+	cd[len(cd)-checksumFieldLen-1] ^= 0xff // flip data bits
+
+	r := NewReader(bytes.NewReader(cd))
+	_, err := io.ReadAll(r)
+	if err == nil {
+		t.Fatalf("expected an error on corrupted data")
+	}
+}