Pull optimization: use klauspost compress/gzip instead of stdlib compress/gzip

internal/imagecache/imagecache_test.go

@@ -59,12 +59,12 @@ func newWrappedRegistry(t *testing.T, wrap func(http.Handler) http.Handler) (*ht
 	return srv, u.Host
 }
 
-func layerFromEntries(t *testing.T, entries []tarEntry) v1.Layer {
+func layerFromEntries(t *testing.T, entries []tarEntry, opts ...tarball.LayerOption) v1.Layer {
 	t.Helper()
 	b := buildTar(t, entries)
 	l, err := tarball.LayerFromOpener(func() (io.ReadCloser, error) {
 		return io.NopCloser(bytes.NewReader(b)), nil
-	})
+	}, opts...)
 	if err != nil {
 		t.Fatalf("tarball.LayerFromOpener: %v", err)
 	}

internal/imagecache/verify.go

@@ -15,18 +15,28 @@
 package imagecache
 
 import (
+	"bufio"
+	"bytes"
 	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/klauspost/compress/gzip"
+	"github.com/klauspost/compress/zstd"
 )
 
 // errDiffIDMismatch marks layer content that does not hash to its claimed
 // diffID. It is a property of the image, so retrying cannot fix it.
 var errDiffIDMismatch = errors.New("layer content does not match its claimed diffID")
 
+// Magic numbers at the start of gzip (RFC 1952) and zstd (RFC 8878) streams.
+var (
+	gzipMagic = []byte{0x1f, 0x8b}
+	zstdMagic = []byte{0x28, 0xb5, 0x2f, 0xfd}
+)
+
 // readVerifiedLayer passes the layer's uncompressed stream to consume, then
 // reads whatever consume left unread and checks that the whole stream
 // hashes to diffID. It returns only the number of bytes consume read. The
@@ -39,17 +49,42 @@ var errDiffIDMismatch = errors.New("layer content does not match its claimed dif
 // go-containerregistry check the compressed blob against the manifest
 // digest, which it does only at EOF.
 func readVerifiedLayer(layer v1.Layer, diffID v1.Hash, consume func(io.Reader) error) (int64, error) {
-	rc, err := layer.Uncompressed()
+	// Decompress here because layer.Uncompressed() uses the slower stdlib gzip.
+	rc, err := layer.Compressed()
 	if err != nil {
 		return 0, fmt.Errorf("while opening layer stream: %w", err)
 	}
 	defer rc.Close()
 
+	buf := bufio.NewReader(rc)
+	// A stream shorter than zstdMagic returns io.EOF with the bytes it has.
+	magic, err := buf.Peek(len(zstdMagic))
+	if err != nil && !errors.Is(err, io.EOF) {
+		return 0, fmt.Errorf("while peeking layer stream: %w", err)
+	}
+	var r io.Reader = buf
+	switch {
+	case bytes.HasPrefix(magic, gzipMagic):
+		gzr, err := gzip.NewReader(buf)
+		if err != nil {
+			return 0, fmt.Errorf("while opening gzip layer stream: %w", err)
+		}
+		defer gzr.Close()
+		r = gzr
+	case bytes.HasPrefix(magic, zstdMagic):
+		zr, err := zstd.NewReader(buf)
+		if err != nil {
+			return 0, fmt.Errorf("while opening zstd layer stream: %w", err)
+		}
+		defer zr.Close()
+		r = zr
+	}
+
 	hasher, err := v1.Hasher(diffID.Algorithm)
 	if err != nil {
 		return 0, fmt.Errorf("while creating hasher for %s: %w", diffID, err)
 	}
-	tee := io.TeeReader(rc, hasher)
+	tee := io.TeeReader(r, hasher)
 	counter := &countingReader{r: tee}
 	if err := consume(counter); err != nil {
 		return 0, err

internal/imagecache/verify_test.go

@@ -31,6 +31,7 @@ import (
 	"testing/synctest"
 
 	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/compression"
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/empty"
@@ -235,7 +236,7 @@ func TestDiffIDMismatchNotRetriedPerCredential(t *testing.T) {
 	}
 }
 
-// gatedLayer holds a layer flight open: its uncompressed stream closes
+// gatedLayer holds a layer flight open: its compressed stream closes
 // started, blocks until release is closed, then fails with err if set.
 type gatedLayer struct {
 	v1.Layer
@@ -247,24 +248,24 @@ func newGatedLayer(layer v1.Layer, err error) gatedLayer {
 	return gatedLayer{Layer: layer, started: make(chan struct{}), release: make(chan struct{}), err: err}
 }
 
-func (g gatedLayer) Uncompressed() (io.ReadCloser, error) {
+func (g gatedLayer) Compressed() (io.ReadCloser, error) {
 	close(g.started)
 	<-g.release
 	if g.err != nil {
 		return nil, g.err
 	}
-	return g.Layer.Uncompressed()
+	return g.Layer.Compressed()
 }
 
-// openCountingLayer counts opens of its uncompressed stream.
+// openCountingLayer counts opens of its compressed stream.
 type openCountingLayer struct {
 	v1.Layer
 	opens *atomic.Int32
 }
 
-func (c openCountingLayer) Uncompressed() (io.ReadCloser, error) {
+func (c openCountingLayer) Compressed() (io.ReadCloser, error) {
 	c.opens.Add(1)
-	return c.Layer.Uncompressed()
+	return c.Layer.Compressed()
 }
 
 // joinFlight runs ensureLayer for leader, then for joiner once the leader's
@@ -351,6 +352,10 @@ func TestSameBlobPullsShareOneFlight(t *testing.T) {
 	})
 }
 
+type uncompressedLayer struct{ v1.Layer }
+
+func (u uncompressedLayer) Compressed() (io.ReadCloser, error) { return u.Uncompressed() }
+
 // The hash covers the whole uncompressed stream, while the returned length
 // covers only what the consumer read.
 func TestReadVerifiedLayerDrainsStream(t *testing.T) {
@@ -359,19 +364,33 @@ func TestReadVerifiedLayerDrainsStream(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	zstdLayer := func(body string) v1.Layer {
+		return layerFromEntries(t, []tarEntry{{name: "lib", typeflag: tar.TypeReg, mode: 0o644, body: body}}, tarball.WithCompression(compression.ZStd))
+	}
 	readTen := func(r io.Reader) error {
 		_, err := io.CopyN(io.Discard, r, 10)
 		return err
 	}
-	n, err := readVerifiedLayer(good, diffID, readTen)
-	if err != nil {
-		t.Fatalf("readVerifiedLayer(good): %v", err)
-	}
-	if n != 10 {
-		t.Errorf("length = %d, want 10", n)
-	}
-	if _, err := readVerifiedLayer(evil, diffID, readTen); !errors.Is(err, errDiffIDMismatch) {
-		t.Errorf("readVerifiedLayer(evil) error = %v, want errDiffIDMismatch", err)
+	for _, tc := range []struct {
+		name       string
+		good, evil v1.Layer
+	}{
+		{name: "gzip", good: good, evil: evil},
+		{name: "zstd", good: zstdLayer("good"), evil: zstdLayer("EVIL")},
+		{name: "uncompressed", good: uncompressedLayer{good}, evil: uncompressedLayer{evil}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			n, err := readVerifiedLayer(tc.good, diffID, readTen)
+			if err != nil {
+				t.Fatalf("readVerifiedLayer(good): %v", err)
+			}
+			if n != 10 {
+				t.Errorf("length = %d, want 10", n)
+			}
+			if _, err := readVerifiedLayer(tc.evil, diffID, readTen); !errors.Is(err, errDiffIDMismatch) {
+				t.Errorf("readVerifiedLayer(evil) error = %v, want errDiffIDMismatch", err)
+			}
+		})
 	}
 }