feat: Make cosign copy faster (#2901)

* ggcr: Use remote.Pusher to make cosign copy faster Signed-off-by: Jon Johnson <jon.johnson@chainguard.dev> * Address review feedback Signed-off-by: Jon Johnson <jon.johnson@chainguard.dev> --------- Signed-off-by: Jon Johnson <jon.johnson@chainguard.dev>
sigstore · Apr 17, 2023 · 82fcb7a · 82fcb7a
1 parent b2227a4
commit 82fcb7a
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 40 deletions.
diff --git a/cmd/cosign/cli/copy/copy.go b/cmd/cosign/cli/copy/copy.go
@@ -20,6 +20,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"runtime"
 
 	"github.com/google/go-containerregistry/pkg/name"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -29,6 +30,7 @@ import (
 	"github.com/sigstore/cosign/v2/pkg/oci"
 	ociremote "github.com/sigstore/cosign/v2/pkg/oci/remote"
 	"github.com/sigstore/cosign/v2/pkg/oci/walk"
+	"golang.org/x/sync/errgroup"
 )
 
 // CopyCmd implements the logic to copy the supplied container image and signatures.
@@ -48,41 +50,60 @@ func CopyCmd(ctx context.Context, regOpts options.RegistryOptions, srcImg, dstIm
 	dstRepoRef := dstRef.Context()
 
 	remoteOpts := regOpts.GetRegistryClientOpts(ctx)
+
+	pusher, err := remote.NewPusher(remoteOpts...)
+	if err != nil {
+		return err
+	}
+
+	g, gctx := errgroup.WithContext(ctx)
+	g.SetLimit(runtime.GOMAXPROCS(0))
+
 	root, err := ociremote.SignedEntity(srcRef, ociremote.WithRemoteOptions(remoteOpts...))
 	if err != nil {
 		return err
 	}
 
-	if err := walk.SignedEntity(ctx, root, func(ctx context.Context, se oci.SignedEntity) error {
+	if err := walk.SignedEntity(gctx, root, func(ctx context.Context, se oci.SignedEntity) error {
 		// Both of the SignedEntity types implement Digest()
 		h, err := se.(interface{ Digest() (v1.Hash, error) }).Digest()
 		if err != nil {
 			return err
 		}
 		srcDigest := srcRepoRef.Digest(h.String())
 
-		// Copy signatures.
-		if err := copyTagImage(ociremote.SignatureTag, srcDigest, dstRepoRef, force, remoteOpts...); err != nil {
-			return err
-		}
-		if sigOnly {
+		copyTag := func(tm tagMap) error {
+			src, err := tm(srcDigest, ociremote.WithRemoteOptions(remoteOpts...))
+			if err != nil {
+				return err
+			}
+
+			dst := dstRepoRef.Tag(src.Identifier())
+			g.Go(func() error {
+				return remoteCopy(ctx, pusher, src, dst, force, remoteOpts...)
+			})
+
 			return nil
 		}
 
-		// Copy attestations
-		if err := copyTagImage(ociremote.AttestationTag, srcDigest, dstRepoRef, force, remoteOpts...); err != nil {
+		if err := copyTag(ociremote.SignatureTag); err != nil {
 			return err
 		}
+		if sigOnly {
+			return nil
+		}
 
-		// Copy SBOMs
-		if err := copyTagImage(ociremote.SBOMTag, srcDigest, dstRepoRef, force, remoteOpts...); err != nil {
-			return err
+		for _, tm := range []tagMap{ociremote.AttestationTag, ociremote.SBOMTag} {
+			if err := copyTag(tm); err != nil {
+				return err
+			}
 		}
 
 		// Copy the entity itself.
-		if err := copyImage(srcDigest, dstRepoRef.Tag(srcDigest.Identifier()), force, remoteOpts...); err != nil {
-			return err
-		}
+		g.Go(func() error {
+			dst := dstRepoRef.Tag(srcDigest.Identifier())
+			return remoteCopy(ctx, pusher, srcDigest, dst, force, remoteOpts...)
+		})
 
 		return nil
 	}); err != nil {
@@ -92,12 +113,17 @@ func CopyCmd(ctx context.Context, regOpts options.RegistryOptions, srcImg, dstIm
 		return nil
 	}
 
+	// Wait for everything to be copied over.
+	if err := g.Wait(); err != nil {
+		return err
+	}
+
 	// Now that everything has been copied over, update the tag.
 	h, err := root.(interface{ Digest() (v1.Hash, error) }).Digest()
 	if err != nil {
 		return err
 	}
-	return copyImage(srcRepoRef.Digest(h.String()), dstRef, force, remoteOpts...)
+	return remoteCopy(ctx, pusher, srcRepoRef.Digest(h.String()), dstRef, force, remoteOpts...)
 }
 
 func descriptorsEqual(a, b *v1.Descriptor) bool {
@@ -109,15 +135,7 @@ func descriptorsEqual(a, b *v1.Descriptor) bool {
 
 type tagMap func(name.Reference, ...ociremote.Option) (name.Tag, error)
 
-func copyTagImage(tm tagMap, srcDigest name.Digest, dstRepo name.Repository, overwrite bool, opts ...remote.Option) error {
-	src, err := tm(srcDigest, ociremote.WithRemoteOptions(opts...))
-	if err != nil {
-		return err
-	}
-	return copyImage(src, dstRepo.Tag(src.Identifier()), overwrite, opts...)
-}
-
-func copyImage(src, dest name.Reference, overwrite bool, opts ...remote.Option) error {
+func remoteCopy(ctx context.Context, pusher *remote.Pusher, src, dest name.Reference, overwrite bool, opts ...remote.Option) error {
 	got, err := remote.Get(src, opts...)
 	if err != nil {
 		var te *transport.Error
@@ -141,17 +159,5 @@ func copyImage(src, dest name.Reference, overwrite bool, opts ...remote.Option)
 	}
 
 	fmt.Fprintf(os.Stderr, "Copying %s to %s...\n", src, dest)
-	if got.MediaType.IsIndex() {
-		imgIdx, err := got.ImageIndex()
-		if err != nil {
-			return err
-		}
-		return remote.WriteIndex(dest, imgIdx, opts...)
-	}
-
-	img, err := got.Image()
-	if err != nil {
-		return err
-	}
-	return remote.Write(dest, img, opts...)
+	return pusher.Push(ctx, dest, got)
 }
diff --git a/go.mod b/go.mod
@@ -17,7 +17,7 @@ require (
 	github.com/go-piv/piv-go v1.11.0
 	github.com/google/certificate-transparency-go v1.1.4
 	github.com/google/go-cmp v0.5.9
-	github.com/google/go-containerregistry v0.14.0
+	github.com/google/go-containerregistry v0.14.1-0.20230409045903-ed5c185df419
 	github.com/google/go-github/v50 v50.2.0
 	github.com/in-toto/in-toto-golang v0.7.1
 	github.com/kelseyhightower/envconfig v1.4.0

diff --git a/go.sum b/go.sum
@@ -485,8 +485,8 @@ github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8
 github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
-github.com/google/go-containerregistry v0.14.0 h1:z58vMqHxuwvAsVwvKEkmVBz2TlgBgH5k6koEXBtlYkw=
-github.com/google/go-containerregistry v0.14.0/go.mod h1:aiJ2fp/SXvkWgmYHioXnbMdlgB8eXiiYOY55gfN91Wk=
+github.com/google/go-containerregistry v0.14.1-0.20230409045903-ed5c185df419 h1:gMlTWagRJgCJ3EnISyF5+p9phYpFyWEI70Z56T+o2MY=
+github.com/google/go-containerregistry v0.14.1-0.20230409045903-ed5c185df419/go.mod h1:ETSJmRH9iO4Q0WQILIMkDUiKk+CaxItZW+gEDjyw8Ug=
 github.com/google/go-github/v50 v50.2.0 h1:j2FyongEHlO9nxXLc+LP3wuBSVU9mVxfpdYUexMpIfk=
 github.com/google/go-github/v50 v50.2.0/go.mod h1:VBY8FB6yPIjrtKhozXv4FQupxKLS6H4m6xFZlT43q8Q=
 github.com/google/go-querystring v1.1.0 h1:AnCroh3fv4ZBgVIf1Iwtovgjaw/GiKJo8M8yD/fhyJ8=