From 7443fbdce9c1c691048958aea143cb416ebceefd Mon Sep 17 00:00:00 2001
From: Eli Treuherz <et@arenko.group>
Date: Fri, 21 Jun 2024 11:56:54 +0100
Subject: [PATCH] Pass in index annotations from builds on multiple nodes

Fixes #2540

Signed-off-by: Eli Treuherz <et@arenko.group>
---
 build/build.go                | 11 ++++++++++-
 commands/imagetools/create.go |  8 +++++++-
 controller/build/build.go     |  5 +++++
 util/imagetools/create.go     |  9 ++-------
 4 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/build/build.go b/build/build.go
index 1d06824ac831..938421f7fc36 100644
--- a/build/build.go
+++ b/build/build.go
@@ -61,6 +61,7 @@ type Options struct {
 
 	Ref           string
 	Allow         []entitlements.Entitlement
+	Annotations   map[exptypes.AnnotationKey]string
 	Attests       map[string]*string
 	BuildArgs     map[string]string
 	CacheFrom     []client.CacheOptionsEntry
@@ -607,7 +608,15 @@ func BuildWithResultHandler(ctx context.Context, nodes []builder.Node, opt map[s
 								}
 							}
 
-							dt, desc, err := itpull.Combine(ctx, srcs, nil, false)
+							filteredAnnotations := map[exptypes.AnnotationKey]string{}
+							for k, v := range opt.Annotations {
+								switch k.Type {
+								case exptypes.AnnotationIndex, exptypes.AnnotationManifestDescriptor:
+									filteredAnnotations[k] = v
+								}
+							}
+
+							dt, desc, err := itpull.Combine(ctx, srcs, filteredAnnotations, false)
 							if err != nil {
 								return err
 							}
diff --git a/commands/imagetools/create.go b/commands/imagetools/create.go
index 934602ee6ef1..05f460d777df 100644
--- a/commands/imagetools/create.go
+++ b/commands/imagetools/create.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/distribution/reference"
 	"github.com/docker/buildx/builder"
+	"github.com/docker/buildx/util/buildflags"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/buildx/util/imagetools"
 	"github.com/docker/buildx/util/progress"
@@ -154,7 +155,12 @@ func runCreate(ctx context.Context, dockerCli command.Cli, in createOptions, arg
 		}
 	}
 
-	dt, desc, err := r.Combine(ctx, srcs, in.annotations, in.preferIndex)
+	annotations, err := buildflags.ParseAnnotations(in.annotations)
+	if err != nil {
+		return errors.Wrapf(err, "failed to parse annotations")
+	}
+
+	dt, desc, err := r.Combine(ctx, srcs, annotations, in.preferIndex)
 	if err != nil {
 		return err
 	}
diff --git a/controller/build/build.go b/controller/build/build.go
index 1cae9b6c70d9..8732e0dccb7e 100644
--- a/controller/build/build.go
+++ b/controller/build/build.go
@@ -153,6 +153,11 @@ func RunBuild(ctx context.Context, dockerCli command.Cli, in controllerapi.Build
 
 	opts.SourcePolicy = in.SourcePolicy
 
+	opts.Annotations, err = buildflags.ParseAnnotations(in.Annotations)
+	if err != nil {
+		return nil, nil, errors.Wrapf(err, "failed to parse annotations")
+	}
+
 	allow, err := buildflags.ParseEntitlements(in.Allow)
 	if err != nil {
 		return nil, nil, err
diff --git a/util/imagetools/create.go b/util/imagetools/create.go
index 08148ad28c66..54cbf344b294 100644
--- a/util/imagetools/create.go
+++ b/util/imagetools/create.go
@@ -14,7 +14,6 @@ import (
 	"github.com/containerd/containerd/remotes"
 	"github.com/containerd/errdefs"
 	"github.com/distribution/reference"
-	"github.com/docker/buildx/util/buildflags"
 	"github.com/moby/buildkit/exporter/containerimage/exptypes"
 	"github.com/moby/buildkit/util/contentutil"
 	"github.com/opencontainers/go-digest"
@@ -29,7 +28,7 @@ type Source struct {
 	Ref  reference.Named
 }
 
-func (r *Resolver) Combine(ctx context.Context, srcs []*Source, ann []string, preferIndex bool) ([]byte, ocispec.Descriptor, error) {
+func (r *Resolver) Combine(ctx context.Context, srcs []*Source, ann map[exptypes.AnnotationKey]string, preferIndex bool) ([]byte, ocispec.Descriptor, error) {
 	eg, ctx := errgroup.WithContext(ctx)
 
 	dts := make([][]byte, len(srcs))
@@ -152,11 +151,7 @@ func (r *Resolver) Combine(ctx context.Context, srcs []*Source, ann []string, pr
 	// annotations are only allowed on OCI indexes
 	indexAnnotation := make(map[string]string)
 	if mt == ocispec.MediaTypeImageIndex {
-		annotations, err := buildflags.ParseAnnotations(ann)
-		if err != nil {
-			return nil, ocispec.Descriptor{}, err
-		}
-		for k, v := range annotations {
+		for k, v := range ann {
 			switch k.Type {
 			case exptypes.AnnotationIndex:
 				indexAnnotation[k.Key] = v