chunked: add no-compression option for ZstdWriter

a new function NoCompression() is added to provide a way to create
uncompressed zstd:chunked files.

Signed-off-by: Giuseppe Scrivano <[email protected]>

Author: giuseppe

pkg/chunked/compressor/compressor.go

@@ -206,11 +206,11 @@ type tarSplitData struct {
 	packer              storage.Packer
 }
 
-func newTarSplitData(level int) (*tarSplitData, error) {
+func newTarSplitData(createZstdWriter minimal.CreateZstdWriterFunc) (*tarSplitData, error) {
 	compressed := bytes.NewBuffer(nil)
 	digester := digest.Canonical.Digester()
 
-	zstdWriter, err := minimal.ZstdWriterWithLevel(io.MultiWriter(compressed, digester.Hash()), level)
+	zstdWriter, err := createZstdWriter(io.MultiWriter(compressed, digester.Hash()))
 	if err != nil {
 		return nil, err
 	}
@@ -227,11 +227,11 @@ func newTarSplitData(level int) (*tarSplitData, error) {
 	}, nil
 }
 
-func writeZstdChunkedStream(destFile io.Writer, outMetadata map[string]string, reader io.Reader, level int) error {
+func writeZstdChunkedStream(destFile io.Writer, outMetadata map[string]string, reader io.Reader, createZstdWriter minimal.CreateZstdWriterFunc) error {
 	// total written so far.  Used to retrieve partial offsets in the file
 	dest := ioutils.NewWriteCounter(destFile)
 
-	tarSplitData, err := newTarSplitData(level)
+	tarSplitData, err := newTarSplitData(createZstdWriter)
 	if err != nil {
 		return err
 	}
@@ -251,7 +251,7 @@ func writeZstdChunkedStream(destFile io.Writer, outMetadata map[string]string, r
 
 	buf := make([]byte, 4096)
 
-	zstdWriter, err := minimal.ZstdWriterWithLevel(dest, level)
+	zstdWriter, err := createZstdWriter(dest)
 	if err != nil {
 		return err
 	}
@@ -420,7 +420,7 @@ func writeZstdChunkedStream(destFile io.Writer, outMetadata map[string]string, r
 		UncompressedSize: tarSplitData.uncompressedCounter.Count,
 	}
 
-	return minimal.WriteZstdChunkedManifest(dest, outMetadata, uint64(dest.Count), &ts, metadata, level)
+	return minimal.WriteZstdChunkedManifest(dest, outMetadata, uint64(dest.Count), &ts, metadata, createZstdWriter)
 }
 
 type zstdChunkedWriter struct {
@@ -447,7 +447,7 @@ func (w zstdChunkedWriter) Write(p []byte) (int, error) {
 	}
 }
 
-// zstdChunkedWriterWithLevel writes a zstd compressed tarball where each file is
+// makeZstdChunkedWriter writes a zstd compressed tarball where each file is
 // compressed separately so it can be addressed separately.  Idea based on CRFS:
 // https://github.com/google/crfs
 // The difference with CRFS is that the zstd compression is used instead of gzip.
@@ -462,12 +462,12 @@ func (w zstdChunkedWriter) Write(p []byte) (int, error) {
 // [SKIPPABLE FRAME 1]: [ZSTD SKIPPABLE FRAME, SIZE=MANIFEST LENGTH][MANIFEST]
 // [SKIPPABLE FRAME 2]: [ZSTD SKIPPABLE FRAME, SIZE=16][MANIFEST_OFFSET][MANIFEST_LENGTH][MANIFEST_LENGTH_UNCOMPRESSED][MANIFEST_TYPE][CHUNKED_ZSTD_MAGIC_NUMBER]
 // MANIFEST_OFFSET, MANIFEST_LENGTH, MANIFEST_LENGTH_UNCOMPRESSED and CHUNKED_ZSTD_MAGIC_NUMBER are 64 bits unsigned in little endian format.
-func zstdChunkedWriterWithLevel(out io.Writer, metadata map[string]string, level int) (io.WriteCloser, error) {
+func makeZstdChunkedWriter(out io.Writer, metadata map[string]string, createZstdWriter minimal.CreateZstdWriterFunc) (io.WriteCloser, error) {
 	ch := make(chan error, 1)
 	r, w := io.Pipe()
 
 	go func() {
-		ch <- writeZstdChunkedStream(out, metadata, r, level)
+		ch <- writeZstdChunkedStream(out, metadata, r, createZstdWriter)
 		_, _ = io.Copy(io.Discard, r) // Ordinarily writeZstdChunkedStream consumes all of r. If it fails, ensure the write end never blocks and eventually terminates.
 		r.Close()
 		close(ch)
@@ -486,5 +486,40 @@ func ZstdCompressor(r io.Writer, metadata map[string]string, level *int) (io.Wri
 		level = &l
 	}
 
-	return zstdChunkedWriterWithLevel(r, metadata, *level)
+	createZstdWriter := func(dest io.Writer) (minimal.ZstdWriter, error) {
+		return minimal.ZstdWriterWithLevel(dest, *level)
+	}
+
+	return makeZstdChunkedWriter(r, metadata, createZstdWriter)
+}
+
+type noCompression struct {
+	dest io.Writer
+}
+
+func (n *noCompression) Write(p []byte) (int, error) {
+	return n.dest.Write(p)
+}
+
+func (n *noCompression) Close() error {
+	return nil
+}
+
+func (n *noCompression) Flush() error {
+	return nil
+}
+
+func (n *noCompression) Reset(dest io.Writer) {
+	n.dest = dest
+}
+
+// NoCompression writes directly to the output file without any compression
+//
+// Such an output does not follow the zstd:chunked spec and cannot be generally consumed; this function
+// only exists for internal purposes and should not be called from outside c/storage.
+func NoCompression(r io.Writer, metadata map[string]string) (io.WriteCloser, error) {
+	createZstdWriter := func(dest io.Writer) (minimal.ZstdWriter, error) {
+		return &noCompression{dest: dest}, nil
+	}
+	return makeZstdChunkedWriter(r, metadata, createZstdWriter)
 }

pkg/chunked/compressor/compressor_test.go

@@ -3,8 +3,11 @@ package compressor
 import (
 	"bufio"
 	"bytes"
+	"errors"
 	"io"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestHole(t *testing.T) {
@@ -88,3 +91,82 @@ func TestTwoHoles(t *testing.T) {
 		t.Error("didn't receive EOF")
 	}
 }
+
+func TestNoCompressionWrite(t *testing.T) {
+	var buf bytes.Buffer
+	nc := &noCompression{dest: &buf}
+
+	data := []byte("hello world")
+	n, err := nc.Write(data)
+	assert.NoError(t, err)
+	assert.Equal(t, len(data), n)
+	assert.Equal(t, data, buf.Bytes())
+
+	data2 := []byte(" again")
+	n, err = nc.Write(data2)
+	assert.NoError(t, err)
+	assert.Equal(t, len(data2), n)
+	assert.Equal(t, append(data, data2...), buf.Bytes())
+}
+
+func TestNoCompressionClose(t *testing.T) {
+	var buf bytes.Buffer
+	nc := &noCompression{dest: &buf}
+	err := nc.Close()
+	assert.NoError(t, err)
+}
+
+func TestNoCompressionFlush(t *testing.T) {
+	var buf bytes.Buffer
+	nc := &noCompression{dest: &buf}
+	err := nc.Flush()
+	assert.NoError(t, err)
+}
+
+func TestNoCompressionReset(t *testing.T) {
+	var buf1 bytes.Buffer
+	nc := &noCompression{dest: &buf1}
+
+	data1 := []byte("initial data")
+	_, err := nc.Write(data1)
+	assert.NoError(t, err)
+	assert.Equal(t, data1, buf1.Bytes())
+
+	err = nc.Close()
+	assert.NoError(t, err)
+
+	var buf2 bytes.Buffer
+	nc.Reset(&buf2)
+
+	data2 := []byte("new data")
+	_, err = nc.Write(data2)
+	assert.NoError(t, err)
+
+	assert.Equal(t, data1, buf1.Bytes(), "Buffer 1 should remain unchanged")
+	assert.Equal(t, data2, buf2.Bytes(), "Buffer 2 should contain the new data")
+
+	err = nc.Close()
+	assert.NoError(t, err)
+
+	// Test Reset with nil, though Write would panic, Reset itself should work
+	nc.Reset(nil)
+	assert.Nil(t, nc.dest)
+}
+
+// Mock writer that returns an error on Write
+type errorWriter struct{}
+
+func (ew *errorWriter) Write(p []byte) (n int, err error) {
+	return 0, errors.New("mock write error")
+}
+
+func TestNoCompressionWriteError(t *testing.T) {
+	ew := &errorWriter{}
+	nc := &noCompression{dest: ew}
+
+	data := []byte("hello world")
+	n, err := nc.Write(data)
+	assert.Error(t, err)
+	assert.Equal(t, 0, n)
+	assert.Equal(t, "mock write error", err.Error())
+}

pkg/chunked/internal/minimal/compression.go

@@ -20,6 +20,15 @@ import (
 	"github.com/vbatts/tar-split/archive/tar"
 )
 
+// ZstdWriter is an interface that wraps standard io.WriteCloser and Reset() to reuse the compressor with a new writer.
+type ZstdWriter interface {
+	io.WriteCloser
+	Reset(dest io.Writer)
+}
+
+// CreateZstdWriterFunc is a function that creates a ZstdWriter for the provided destination writer.
+type CreateZstdWriterFunc func(dest io.Writer) (ZstdWriter, error)
+
 // TOC is short for Table of Contents and is used by the zstd:chunked
 // file format to effectively add an overall index into the contents
 // of a tarball; it also includes file metadata.
@@ -179,7 +188,7 @@ type TarSplitData struct {
 	UncompressedSize int64
 }
 
-func WriteZstdChunkedManifest(dest io.Writer, outMetadata map[string]string, offset uint64, tarSplitData *TarSplitData, metadata []FileMetadata, level int) error {
+func WriteZstdChunkedManifest(dest io.Writer, outMetadata map[string]string, offset uint64, tarSplitData *TarSplitData, metadata []FileMetadata, createZstdWriter CreateZstdWriterFunc) error {
 	// 8 is the size of the zstd skippable frame header + the frame size
 	const zstdSkippableFrameHeader = 8
 	manifestOffset := offset + zstdSkippableFrameHeader
@@ -198,7 +207,7 @@ func WriteZstdChunkedManifest(dest io.Writer, outMetadata map[string]string, off
 	}
 
 	var compressedBuffer bytes.Buffer
-	zstdWriter, err := ZstdWriterWithLevel(&compressedBuffer, level)
+	zstdWriter, err := createZstdWriter(&compressedBuffer)
 	if err != nil {
 		return err
 	}
@@ -244,7 +253,7 @@ func WriteZstdChunkedManifest(dest io.Writer, outMetadata map[string]string, off
 	return appendZstdSkippableFrame(dest, manifestDataLE)
 }
 
-func ZstdWriterWithLevel(dest io.Writer, level int) (*zstd.Encoder, error) {
+func ZstdWriterWithLevel(dest io.Writer, level int) (ZstdWriter, error) {
 	el := zstd.EncoderLevelFromZstd(level)
 	return zstd.NewWriter(dest, zstd.WithEncoderLevel(el))
 }