Add support for zstd compression (#1487)

* Add support for zstd compression

Squashed commits:
[6563fc8] Run ./hack/update-codegen.sh
[c86a013] Fix linting issues
[353dd69] Remove redundant break statements
[a06b4cc] Fix some comments
[fb0678f] Expose compression algorithms enum
[3321d0c] Update mediaType according to compression for tarball layers

(cherry picked from commit 56d4df5)
[5124d2c] Fix unit tests

(cherry picked from commit 7ce08d1)
[2958f30] Optimize compression detection

(cherry picked from commit 0a3437f)
[c211a61] Add support for zstd-compressedd layers

(cherry picked from commit d8fb0f3)
[0b9a996] Add zstd compression utilities

(cherry picked from commit c645201)
[7dd71ab] Add layer types with zstd compression

(cherry picked from commit 328a14b)
Signed-off-by: Lavrenti Frobeen <lavrenti@northflank.com>

* Remove OCIRestrictedLayerZStd from MediaType enum

* Remove unused checkCompression method

* Add missing copyright headers

* Fix goimport lint errors

* Run "go mod tidy" in cmd/krane

* Rename gw to zw in zstd.go

* Fix typos and clean up comments

* Create public compression package for Compression enum; disable "none" compression for tarball layer

* Fix WithCompression method

* Print warning if mediaType does not match selected compression

* Add unit tests for new zstd compression implementation

* Run go mod tidy in pkg/authn/k8schain and pkg/authn/kubernetes

Author: LFrobeen

cmd/krane/go.sum

@@ -8,6 +8,7 @@ require (
 	github.com/docker/distribution v2.8.1+incompatible
 	github.com/docker/docker v20.10.20+incompatible
 	github.com/google/go-cmp v0.5.9
+	github.com/klauspost/compress v1.15.11
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.1.0-rc2
@@ -28,7 +29,6 @@ require (
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
 	github.com/inconshreveable/mousetrap v1.0.1 // indirect
-	github.com/klauspost/compress v1.15.11 // indirect
 	github.com/moby/term v0.0.0-20210610120745-9d4ed1856297 // indirect
 	github.com/morikuni/aec v1.0.0 // indirect
 	github.com/pkg/errors v0.9.1 // indirect

go.mod

@@ -0,0 +1,95 @@
+// Copyright 2022 Google LLC All Rights Reserved.
+//
+// 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 compression
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+
+	"github.com/google/go-containerregistry/internal/gzip"
+	"github.com/google/go-containerregistry/internal/zstd"
+	"github.com/google/go-containerregistry/pkg/compression"
+)
+
+type Opener = func() (io.ReadCloser, error)
+
+// GetCompression detects whether an Opener is compressed and which algorithm is used.
+func GetCompression(opener Opener) (compression.Compression, error) {
+	rc, err := opener()
+	if err != nil {
+		return compression.None, err
+	}
+	defer rc.Close()
+
+	cp, _, err := PeekCompression(rc)
+	if err != nil {
+		return compression.None, err
+	}
+
+	return cp, nil
+}
+
+// PeekCompression detects whether the input stream is compressed and which algorithm is used.
+//
+// If r implements Peek, we will use that directly, otherwise a small number
+// of bytes are buffered to Peek at the gzip/zstd header, and the returned
+// PeekReader can be used as a replacement for the consumed input io.Reader.
+func PeekCompression(r io.Reader) (compression.Compression, PeekReader, error) {
+	pr := intoPeekReader(r)
+
+	if isGZip, _, err := checkHeader(pr, gzip.MagicHeader); err != nil {
+		return compression.None, pr, err
+	} else if isGZip {
+		return compression.GZip, pr, nil
+	}
+
+	if isZStd, _, err := checkHeader(pr, zstd.MagicHeader); err != nil {
+		return compression.None, pr, err
+	} else if isZStd {
+		return compression.ZStd, pr, nil
+	}
+
+	return compression.None, pr, nil
+}
+
+// PeekReader is an io.Reader that also implements Peek a la bufio.Reader.
+type PeekReader interface {
+	io.Reader
+	Peek(n int) ([]byte, error)
+}
+
+// IntoPeekReader creates a PeekReader from an io.Reader.
+// If the reader already has a Peek method, it will just return the passed reader.
+func intoPeekReader(r io.Reader) PeekReader {
+	if p, ok := r.(PeekReader); ok {
+		return p
+	} else {
+		return bufio.NewReader(r)
+	}
+}
+
+// CheckHeader checks whether the first bytes from a PeekReader match an expected header
+func checkHeader(pr PeekReader, expectedHeader []byte) (bool, PeekReader, error) {
+	header, err := pr.Peek(len(expectedHeader))
+	if err != nil {
+		// https://github.com/google/go-containerregistry/issues/367
+		if err == io.EOF {
+			return false, pr, nil
+		}
+		return false, pr, err
+	}
+	return bytes.Equal(header, expectedHeader), pr, nil
+}

internal/compression/compression.go

@@ -0,0 +1,76 @@
+// Copyright 2020 Google LLC All Rights Reserved.
+//
+// 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 compression
+
+import (
+	"bytes"
+	"io"
+	"testing"
+
+	"github.com/google/go-containerregistry/internal/and"
+	"github.com/google/go-containerregistry/internal/gzip"
+	"github.com/google/go-containerregistry/internal/zstd"
+)
+
+type Compressor = func(rc io.ReadCloser) io.ReadCloser
+type Decompressor = func(rc io.ReadCloser) (io.ReadCloser, error)
+
+func testPeekCompression(t *testing.T,
+	compressionExpected string,
+	compress Compressor,
+	decompress Decompressor,
+) {
+	content := "This is the input string."
+	contentBuf := bytes.NewBufferString(content)
+
+	compressed := compress(io.NopCloser(contentBuf))
+	compressionDetected, pr, err := PeekCompression(compressed)
+
+	if err != nil {
+		t.Error("PeekCompression() =", err)
+	}
+
+	if got := string(compressionDetected); got != compressionExpected {
+		t.Errorf("PeekCompression(); got %q, content %q", got, compressionExpected)
+	}
+
+	decompressed, err := decompress(withCloser(pr, compressed))
+
+	b, err := io.ReadAll(decompressed)
+	if err != nil {
+		t.Error("ReadAll() =", err)
+	}
+
+	if got := string(b); got != content {
+		t.Errorf("ReadAll(); got %q, content %q", got, content)
+	}
+}
+
+func TestPeekCompression(t *testing.T) {
+	testPeekCompression(t, "gzip", gzip.ReadCloser, gzip.UnzipReadCloser)
+	testPeekCompression(t, "zstd", zstd.ReadCloser, zstd.UnzipReadCloser)
+
+	nopCompress := func(rc io.ReadCloser) io.ReadCloser { return rc }
+	nopDecompress := func(rc io.ReadCloser) (io.ReadCloser, error) { return rc, nil }
+
+	testPeekCompression(t, "none", nopCompress, nopDecompress)
+}
+
+func withCloser(pr PeekReader, rc io.ReadCloser) io.ReadCloser {
+	return &and.ReadCloser{
+		Reader:    pr,
+		CloseFunc: rc.Close,
+	}
+}

internal/compression/compression_test.go

@@ -24,7 +24,7 @@ import (
 	"github.com/google/go-containerregistry/internal/and"
 )
 
-var gzipMagicHeader = []byte{'\x1f', '\x8b'}
+var MagicHeader = []byte{'\x1f', '\x8b'}
 
 // ReadCloser reads uncompressed input data from the io.ReadCloser and
 // returns an io.ReadCloser from which compressed data may be read.
@@ -84,7 +84,7 @@ func ReadCloserLevel(r io.ReadCloser, level int) io.ReadCloser {
 }
 
 // UnzipReadCloser reads compressed input data from the io.ReadCloser and
-// returns an io.ReadCloser from which uncompessed data may be read.
+// returns an io.ReadCloser from which uncompressed data may be read.
 func UnzipReadCloser(r io.ReadCloser) (io.ReadCloser, error) {
 	gr, err := gzip.NewReader(r)
 	if err != nil {
@@ -113,34 +113,5 @@ func Is(r io.Reader) (bool, error) {
 	if err != nil {
 		return false, err
 	}
-	return bytes.Equal(magicHeader, gzipMagicHeader), nil
-}
-
-// PeekReader is an io.Reader that also implements Peek a la bufio.Reader.
-type PeekReader interface {
-	io.Reader
-	Peek(n int) ([]byte, error)
-}
-
-// Peek detects whether the input stream is gzip compressed.
-//
-// If r implements Peek, we will use that directly, otherwise a small number
-// of bytes are buffered to Peek at the gzip header, and the returned
-// PeekReader can be used as a replacement for the consumed input io.Reader.
-func Peek(r io.Reader) (bool, PeekReader, error) {
-	var pr PeekReader
-	if p, ok := r.(PeekReader); ok {
-		pr = p
-	} else {
-		pr = bufio.NewReader(r)
-	}
-	header, err := pr.Peek(2)
-	if err != nil {
-		// https://github.com/google/go-containerregistry/issues/367
-		if err == io.EOF {
-			return false, pr, nil
-		}
-		return false, pr, err
-	}
-	return bytes.Equal(header, gzipMagicHeader), pr, nil
+	return bytes.Equal(magicHeader, MagicHeader), nil
 }

internal/gzip/zip.go

@@ -0,0 +1,114 @@
+// Copyright 2022 Google LLC All Rights Reserved.
+//
+// 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 zstd
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+
+	"github.com/google/go-containerregistry/internal/and"
+	"github.com/klauspost/compress/zstd"
+)
+
+var MagicHeader = []byte{'\x28', '\xb5', '\x2f', '\xfd'}
+
+// ReadCloser reads uncompressed input data from the io.ReadCloser and
+// returns an io.ReadCloser from which compressed data may be read.
+// This uses zstd level 1 for the compression.
+func ReadCloser(r io.ReadCloser) io.ReadCloser {
+	return ReadCloserLevel(r, 1)
+}
+
+// ReadCloserLevel reads uncompressed input data from the io.ReadCloser and
+// returns an io.ReadCloser from which compressed data may be read.
+func ReadCloserLevel(r io.ReadCloser, level int) io.ReadCloser {
+	pr, pw := io.Pipe()
+
+	// For highly compressible layers, zstd.Writer will output a very small
+	// number of bytes per Write(). This is normally fine, but when pushing
+	// to a registry, we want to ensure that we're taking full advantage of
+	// the available bandwidth instead of sending tons of tiny writes over
+	// the wire.
+	// 64K ought to be small enough for anybody.
+	bw := bufio.NewWriterSize(pw, 2<<16)
+
+	// Returns err so we can pw.CloseWithError(err)
+	go func() error {
+		// TODO(go1.14): Just defer {pw,zw,r}.Close like you'd expect.
+		// Context: https://golang.org/issue/24283
+		zw, err := zstd.NewWriter(bw, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(level)))
+		if err != nil {
+			return pw.CloseWithError(err)
+		}
+
+		if _, err := io.Copy(zw, r); err != nil {
+			defer r.Close()
+			defer zw.Close()
+			return pw.CloseWithError(err)
+		}
+
+		// Close zstd writer to Flush it and write zstd trailers.
+		if err := zw.Close(); err != nil {
+			return pw.CloseWithError(err)
+		}
+
+		// Flush bufio writer to ensure we write out everything.
+		if err := bw.Flush(); err != nil {
+			return pw.CloseWithError(err)
+		}
+
+		// We don't really care if these fail.
+		defer pw.Close()
+		defer r.Close()
+
+		return nil
+	}()
+
+	return pr
+}
+
+// UnzipReadCloser reads compressed input data from the io.ReadCloser and
+// returns an io.ReadCloser from which uncompressed data may be read.
+func UnzipReadCloser(r io.ReadCloser) (io.ReadCloser, error) {
+	gr, err := zstd.NewReader(r)
+	if err != nil {
+		return nil, err
+	}
+	return &and.ReadCloser{
+		Reader: gr,
+		CloseFunc: func() error {
+			// If the unzip fails, then this seems to return the same
+			// error as the read.  We don't want this to interfere with
+			// us closing the main ReadCloser, since this could leave
+			// an open file descriptor (fails on Windows).
+			gr.Close()
+			return r.Close()
+		},
+	}, nil
+}
+
+// Is detects whether the input stream is compressed.
+func Is(r io.Reader) (bool, error) {
+	magicHeader := make([]byte, 4)
+	n, err := r.Read(magicHeader)
+	if n == 0 && err == io.EOF {
+		return false, nil
+	}
+	if err != nil {
+		return false, err
+	}
+	return bytes.Equal(magicHeader, MagicHeader), nil
+}