diff --git a/internal/imagecache/imagecache_test.go b/internal/imagecache/imagecache_test.go index 07666b0725..37b70a0131 100644 --- a/internal/imagecache/imagecache_test.go +++ b/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) } diff --git a/internal/imagecache/verify.go b/internal/imagecache/verify.go index f2b5151eb8..b65198ca8d 100644 --- a/internal/imagecache/verify.go +++ b/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 diff --git a/internal/imagecache/verify_test.go b/internal/imagecache/verify_test.go index 5fd9687523..8028c8aa68 100644 --- a/internal/imagecache/verify_test.go +++ b/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) + } + }) } }