imagetools: fix concurrent map write from containerd issue

Workaround for containerd issue that can cause concurrent map
write when WithMediaTypeKeyPrefix is called in parallel
goroutines for context that originated from same base context.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>

Author: tonistiigi

commands/imagetools/create.go

@@ -168,11 +168,6 @@ func runCreate(ctx context.Context, dockerCli command.Cli, in createOptions, arg
 		return errors.Wrapf(err, "failed to parse annotations")
 	}
 
-	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.oci.empty.v1+json", "empty")
-	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.dev.cosign.artifact.sig.v1+json", "cosign")
-	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.dev.cosign.simplesigning.v1+json", "simplesigning")
-	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.dev.sigstore.bundle.v0.3+json", "sigstore-bundle")
-
 	dt, desc, manifests, err := r.Combine(ctx, srcs, annotations, in.preferIndex, platforms)
 	if err != nil {
 		return err
@@ -208,16 +203,19 @@ func runCreate(ctx context.Context, dockerCli command.Cli, in createOptions, arg
 	for _, t := range tags {
 		eg.Go(func() error {
 			return progress.Wrap(fmt.Sprintf("pushing %s", t.String()), pw.Write, func(sub progress.SubLogger) error {
+				baseCtx := ctx
 				eg2, _ := errgroup.WithContext(ctx)
 				for _, desc := range manifests {
 					eg2.Go(func() error {
+						ctx = withMediaTypeKeyPrefix(baseCtx)
 						sub.Log(1, fmt.Appendf(nil, "copying %s from %s to %s\n", desc.Digest.String(), desc.Source.Ref.String(), t.String()))
 						return r.Copy(ctx, desc.Source, t)
 					})
 				}
 				if err := eg2.Wait(); err != nil {
 					return err
 				}
+				ctx = withMediaTypeKeyPrefix(ctx) // because of containerd bug this needs to be called separately for each ctx/goroutine pair to avoid concurrent map write
 				sub.Log(1, fmt.Appendf(nil, "pushing %s to %s\n", desc.Digest.String(), t.String()))
 				return r.Push(ctx, t, desc, dt)
 			})
@@ -265,6 +263,14 @@ func parsePlatforms(in []string) ([]ocispecs.Platform, error) {
 	return out, nil
 }
 
+func withMediaTypeKeyPrefix(ctx context.Context) context.Context {
+	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.oci.empty.v1+json", "empty")
+	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.dev.cosign.artifact.sig.v1+json", "cosign")
+	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.dev.cosign.simplesigning.v1+json", "simplesigning")
+	ctx = remotes.WithMediaTypeKeyPrefix(ctx, "application/vnd.dev.sigstore.bundle.v0.3+json", "sigstore-bundle")
+	return ctx
+}
+
 func parseRefs(in []string) ([]reference.Named, error) {
 	refs := make([]reference.Named, len(in))
 	for i, in := range in {