Add zstd compression for image files and tarballs (#71)

README.md

@@ -127,7 +127,7 @@ Features of Buf's include:
 - **Check anything from anywhere**. Buf allows you to not only check a Protobuf schema stored
   locally as `.proto` files, but allows you to check many different [Inputs](https://buf.build/docs/inputs):
 
-  - Tarballs containing `.proto` files, both local and remote.
+  - Tar or zip archives containing `.proto` files, both local and remote.
   - Git repository branches or tags containing `.proto` files, both local and remote.
   - Pre-built [Images](https://buf.build/docs/build-images) or FileDescriptorSets from `protoc`, from both local and remote
     (http/https) locations.

internal/buf/buffetch/ref_parser.go

@@ -160,7 +160,7 @@ func (a *refParser) checkDeprecated(parsedRef fetch.ParsedRef) {
 	format := parsedRef.Format()
 	if replacementFormat, ok := deprecatedCompressionFormatToReplacementFormat[format]; ok {
 		a.logger.Sugar().Warnf(
-			`Format %q is deprecated. Use "format=%s,compression=gz" instead. This will continue to work forever, but updating is recommended.`,
+			`Format %q is deprecated. Use "format=%s,compression=gzip" instead. This will continue to work forever, but updating is recommended.`,
 			format,
 			replacementFormat,
 		)
@@ -195,6 +195,18 @@ func processRawRef(rawRef *fetch.RawRef) error {
 			default:
 				return fmt.Errorf("path %q had .gz extension with unknown format", rawRef.Path)
 			}
+		case ".zst":
+			compressionType = fetch.CompressionTypeZstd
+			switch filepath.Ext(strings.TrimSuffix(rawRef.Path, filepath.Ext(rawRef.Path))) {
+			case ".bin":
+				format = formatBin
+			case ".json":
+				format = formatJSON
+			case ".tar":
+				format = formatTar
+			default:
+				return fmt.Errorf("path %q had .zst extension with unknown format", rawRef.Path)
+			}
 		case ".tgz":
 			format = formatTar
 			compressionType = fetch.CompressionTypeGzip

internal/pkg/fetch/fetch.go

@@ -57,6 +57,8 @@ const (
 	CompressionTypeNone CompressionType = iota + 1
 	// CompressionTypeGzip is gzip compression.
 	CompressionTypeGzip
+	// CompressionTypeZstd is zstd compression.
+	CompressionTypeZstd
 )
 
 // FileScheme is a file scheme.

internal/pkg/fetch/fetch_test.go

@@ -0,0 +1,116 @@
+// Copyright 2020 Buf Technologies Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package fetch
+
+import (
+	"context"
+	"io/ioutil"
+	"path/filepath"
+	"testing"
+
+	"github.com/bufbuild/buf/internal/pkg/app"
+	"github.com/bufbuild/buf/internal/pkg/tmp"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+)
+
+func TestRoundTripBin(t *testing.T) {
+	testRoundTripLocalFile(
+		t,
+		"file.bin",
+		[]byte("one"),
+		testFormatBin,
+		CompressionTypeNone,
+	)
+}
+
+func TestRoundTripBinGz(t *testing.T) {
+	testRoundTripLocalFile(
+		t,
+		"file.bin.gz",
+		[]byte("one"),
+		testFormatBin,
+		CompressionTypeGzip,
+	)
+}
+
+func TestRoundTripBinZst(t *testing.T) {
+	testRoundTripLocalFile(
+		t,
+		"file.bin.zst",
+		[]byte("one"),
+		testFormatBin,
+		CompressionTypeZstd,
+	)
+}
+
+func testRoundTripLocalFile(
+	t *testing.T,
+	filename string,
+	expectedData []byte,
+	expectedFormat string,
+	expectedCompressionType CompressionType,
+) {
+	t.Parallel()
+
+	logger := zap.NewNop()
+	refParser := testNewRefParser(logger)
+	reader := testNewReader(logger)
+	writer := testNewWriter(logger)
+
+	ctx := context.Background()
+	container := app.NewContainer(nil, nil, nil, nil)
+
+	tmpDir, err := tmp.NewDir()
+	require.NoError(t, err)
+	filePath := filepath.Join(tmpDir.AbsPath(), filename)
+
+	parsedRef, err := refParser.GetParsedRef(ctx, filePath)
+	require.NoError(t, err)
+	require.Equal(t, expectedFormat, parsedRef.Format())
+	fileRef, ok := parsedRef.(FileRef)
+	require.True(t, ok)
+	require.Equal(t, expectedCompressionType, fileRef.CompressionType())
+
+	writeCloser, err := writer.PutFile(ctx, container, fileRef)
+	require.NoError(t, err)
+	_, err = writeCloser.Write(expectedData)
+	require.NoError(t, err)
+	require.NoError(t, writeCloser.Close())
+
+	readCloser, err := reader.GetFile(ctx, container, fileRef)
+	require.NoError(t, err)
+	actualData, err := ioutil.ReadAll(readCloser)
+	require.NoError(t, err)
+	require.NoError(t, readCloser.Close())
+
+	require.Equal(t, string(expectedData), string(actualData))
+
+	require.NoError(t, tmpDir.Close())
+}
+
+func testNewReader(logger *zap.Logger) Reader {
+	return NewReader(
+		logger,
+		WithReaderLocal(),
+	)
+}
+
+func testNewWriter(logger *zap.Logger) Writer {
+	return NewWriter(
+		logger,
+		WithWriterLocal(),
+	)
+}

internal/pkg/fetch/reader.go

@@ -34,6 +34,7 @@ import (
 	"github.com/bufbuild/buf/internal/pkg/storage/storagearchive"
 	"github.com/bufbuild/buf/internal/pkg/storage/storagemem"
 	"github.com/bufbuild/buf/internal/pkg/storage/storageos"
+	"github.com/klauspost/compress/zstd"
 	"github.com/klauspost/pgzip"
 	"go.uber.org/multierr"
 	"go.uber.org/zap"
@@ -288,11 +289,30 @@ func (r *reader) getFileReadCloserAndSize(
 	case CompressionTypeNone:
 		return readCloser, size, nil
 	case CompressionTypeGzip:
-		gzipReader, err := pgzip.NewReader(readCloser)
+		gzipReadCloser, err := pgzip.NewReader(readCloser)
 		if err != nil {
 			return nil, -1, err
 		}
-		return ioutilextended.CompositeReadCloser(gzipReader, readCloser), -1, nil
+		return ioutilextended.CompositeReadCloser(
+			gzipReadCloser,
+			ioutilextended.ChainCloser(
+				gzipReadCloser,
+				readCloser,
+			),
+		), -1, nil
+	case CompressionTypeZstd:
+		zstdDecoder, err := zstd.NewReader(readCloser)
+		if err != nil {
+			return nil, -1, err
+		}
+		zstdReadCloser := zstdDecoder.IOReadCloser()
+		return ioutilextended.CompositeReadCloser(
+			zstdReadCloser,
+			ioutilextended.ChainCloser(
+				zstdReadCloser,
+				readCloser,
+			),
+		), -1, nil
 	default:
 		return nil, -1, fmt.Errorf("unknown CompressionType: %v", compressionType)
 	}

internal/pkg/fetch/ref_parser.go

@@ -28,6 +28,7 @@ var (
 	knownCompressionTypeStrings = []string{
 		"none",
 		"gzip",
+		"zstd",
 	}
 )
 
@@ -129,8 +130,10 @@ func (a *refParser) getRawRef(value string) (*RawRef, error) {
 			switch value {
 			case "none":
 				rawRef.CompressionType = CompressionTypeNone
-			case "gz":
+			case "gzip":
 				rawRef.CompressionType = CompressionTypeGzip
+			case "zstd":
+				rawRef.CompressionType = CompressionTypeZstd
 			default:
 				return nil, newCompressionUnknownError(value, knownCompressionTypeStrings...)
 			}

internal/pkg/fetch/ref_parser_test.go

@@ -372,7 +372,7 @@ func TestGetParsedRefSuccess(t *testing.T) {
 			FileSchemeLocal,
 			CompressionTypeGzip,
 		),
-		"path/to/file.json.gz#compression=gz",
+		"path/to/file.json.gz#compression=gzip",
 	)
 	testGetParsedRefSuccess(
 		t,
@@ -432,7 +432,7 @@ func TestGetParsedRefSuccess(t *testing.T) {
 			FileSchemeLocal,
 			CompressionTypeGzip,
 		),
-		"path/to/dir#format=bin,compression=gz",
+		"path/to/dir#format=bin,compression=gzip",
 	)
 	testGetParsedRefSuccess(
 		t,
@@ -546,7 +546,7 @@ func TestGetParsedRefSuccess(t *testing.T) {
 			CompressionTypeGzip,
 			1,
 		),
-		"path/to/file#format=tar,strip_components=1,compression=gz",
+		"path/to/file#format=tar,strip_components=1,compression=gzip",
 	)
 	testGetParsedRefSuccess(
 		t,
@@ -560,6 +560,62 @@ func TestGetParsedRefSuccess(t *testing.T) {
 		),
 		"path/to/file#format=zip,strip_components=1",
 	)
+	testGetParsedRefSuccess(
+		t,
+		buildArchiveRef(
+			testFormatTar,
+			"path/to/file.tar.zst",
+			FileSchemeLocal,
+			ArchiveTypeTar,
+			CompressionTypeZstd,
+			0,
+		),
+		"path/to/file.tar.zst",
+	)
+	testGetParsedRefSuccess(
+		t,
+		buildArchiveRef(
+			testFormatTar,
+			"path/to/file.tar.zst",
+			FileSchemeLocal,
+			ArchiveTypeTar,
+			CompressionTypeZstd,
+			1,
+		),
+		"path/to/file.tar.zst#strip_components=1",
+	)
+	testGetParsedRefSuccess(
+		t,
+		buildArchiveRef(
+			testFormatTar,
+			"path/to/file",
+			FileSchemeLocal,
+			ArchiveTypeTar,
+			CompressionTypeZstd,
+			1,
+		),
+		"path/to/file#format=tar,strip_components=1,compression=zstd",
+	)
+	testGetParsedRefSuccess(
+		t,
+		buildSingleRef(
+			testFormatBin,
+			"path/to/file",
+			FileSchemeLocal,
+			CompressionTypeZstd,
+		),
+		"path/to/file#format=bin,compression=zstd",
+	)
+	testGetParsedRefSuccess(
+		t,
+		buildSingleRef(
+			testFormatBin,
+			"path/to/file.bin.zst",
+			FileSchemeLocal,
+			CompressionTypeZstd,
+		),
+		"path/to/file.bin.zst",
+	)
 }
 
 func TestGetParsedRefError(t *testing.T) {
@@ -696,7 +752,7 @@ func TestGetParsedRefError(t *testing.T) {
 	testGetParsedRefError(
 		t,
 		newCannotSpecifyCompressionForZipError(),
-		"path/to/foo.zip#compression=gz",
+		"path/to/foo.zip#compression=gzip",
 	)
 	testGetParsedRefError(
 		t,
@@ -706,7 +762,7 @@ func TestGetParsedRefError(t *testing.T) {
 	testGetParsedRefError(
 		t,
 		newCannotSpecifyCompressionForZipError(),
-		"path/to/foo#format=zip,compression=gz",
+		"path/to/foo#format=zip,compression=gzip",
 	)
 }
 
@@ -818,6 +874,18 @@ func testRawRefProcessor(rawRef *RawRef) error {
 			default:
 				return fmt.Errorf("path %q had .gz extension with unknown format", rawRef.Path)
 			}
+		case ".zst":
+			compressionType = CompressionTypeZstd
+			switch filepath.Ext(strings.TrimSuffix(rawRef.Path, filepath.Ext(rawRef.Path))) {
+			case ".bin":
+				format = testFormatBin
+			case ".json":
+				format = testFormatJSON
+			case ".tar":
+				format = testFormatTar
+			default:
+				return fmt.Errorf("path %q had .zst extension with unknown format", rawRef.Path)
+			}
 		case ".tgz":
 			format = testFormatTar
 			compressionType = CompressionTypeGzip

internal/pkg/fetch/writer.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/bufbuild/buf/internal/pkg/app"
 	"github.com/bufbuild/buf/internal/pkg/ioutilextended"
+	"github.com/klauspost/compress/zstd"
 	"go.uber.org/multierr"
 	"go.uber.org/zap"
 )
@@ -127,6 +128,18 @@ func (w *writer) putFileWriteCloser(
 				writeCloser,
 			),
 		), nil
+	case CompressionTypeZstd:
+		zstdWriteCloser, err := zstd.NewWriter(writeCloser)
+		if err != nil {
+			return nil, err
+		}
+		return ioutilextended.CompositeWriteCloser(
+			zstdWriteCloser,
+			ioutilextended.ChainCloser(
+				zstdWriteCloser,
+				writeCloser,
+			),
+		), nil
 	default:
 		return nil, fmt.Errorf("unknown CompressionType: %v", compressionType)
 	}