attestations: merge attestation refs into result refs map

This patch reworks the attestations code to place its refs into the
correponding Result.Refs map, to allow for easier conversions between
gateway boundaries. This ensures that almost all conversions (notably,
except the tranformation to protobuf), can be performed in a single
line, similar to how Metadata is copied. Additionally, this prevents
multiple different Attestation types, and allows a single one defined in
a common package.

To acheive this, instead of including a Ref type directly in each
Attestation, we include a key to a ref, which can be looked up in the
Result's Ref map. This key must be uniquely generated when attestations
are added - note that we cannot use the same Ref ID as seen in the
protobuf structure (which would be nice) as the gateway forwarder does
not convert to protobuf and simply applies the forwarding in-memory, so
we need a uniquely generated id. To distinguish between these IDs and
past/present/future contents of the ref map, we use the "attestation:"
prefix, which would allow for future extension in a similar way.

(Note: as written, the gateway protobuf for the in-toto attestation is
modified to contain a ref-key instead of a ref directly. This isn't a
requirement, and we could easily preserve compatibility at the protobuf
level by performing a transformation in the helper functions in
frontend/gateway/pb/attestation.go.)

Aside from the change in types and conversion logic, the only major
changes are in the creation of attestations and the export of
attestations. For the creation of attestations, we need to change the
AddAttestation function into an InTotoAttestationAttestation, so as to
allow passing a Ref alongside, which allows for encapsulation of the
prefix logic, and would allow us to possibly change the transport
mechanism for attestations in the future. For the export of
attestations, we simply introduce a layer of indirection to lookup
the attestation in the refs map.

A caveat of this approach is that the behavior of the Refs map will
change. This will require frontends to detect the presence of specific
caps in the buildkit backend (not added yet in this patch), to ensure
that attestations are only attached if they are supported, otherwise the
local/tar exporters will behave unexpectedly.

Signed-off-by: Justin Chadwell <[email protected]>

client/client_test.go

@@ -6329,16 +6329,14 @@ func testExportAttestations(t *testing.T, sb integration.Sandbox) {
 			if err != nil {
 				return nil, err
 			}
-			res.AddAttestation(pk, &gateway.InTotoAttestation{
-				PredicateRef:  refAttest,
+			res.AddInTotoAttestation(pk, &attestation.InTotoAttestation{
 				PredicatePath: "/attestation.json",
 				PredicateType: "https://example.com/attestations/v1.0",
 				Subjects: []attestation.InTotoSubject{
 					&attestation.InTotoSubjectSelf{},
 				},
-			})
-			res.AddAttestation(pk, &gateway.InTotoAttestation{
-				PredicateRef:  refAttest,
+			}, refAttest)
+			res.AddInTotoAttestation(pk, &attestation.InTotoAttestation{
 				PredicatePath: "/attestation2.json",
 				PredicateType: "https://example.com/attestations2/v1.0",
 				Subjects: []attestation.InTotoSubject{
@@ -6347,7 +6345,7 @@ func testExportAttestations(t *testing.T, sb integration.Sandbox) {
 						Digest: []digest.Digest{successDigest},
 					},
 				},
-			})
+			}, refAttest)
 		}
 
 		dt, err := json.Marshal(expPlatforms)

exporter/containerimage/writer.go

@@ -98,8 +98,12 @@ func (ic *ImageWriter) Commit(ctx context.Context, inp exporter.Source, sessionI
 		return mfstDesc, nil
 	}
 
-	if len(p.Platforms) != len(inp.Refs) {
-		return nil, errors.Errorf("number of platforms does not match references %d %d", len(p.Platforms), len(inp.Refs))
+	refCount := len(p.Platforms)
+	for _, attests := range inp.Attestations {
+		refCount += len(attests)
+	}
+	if refCount != len(inp.Refs) {
+		return nil, errors.Errorf("number of required refs does not match references %d %d", refCount, len(inp.Refs))
 	}
 
 	refs := make([]cache.ImmutableRef, 0, len(inp.Refs))
@@ -166,7 +170,7 @@ func (ic *ImageWriter) Commit(ctx context.Context, inp exporter.Source, sessionI
 		labels[fmt.Sprintf("containerd.io/gc.ref.content.%d", i)] = desc.Digest.String()
 
 		if attestations, ok := inp.Attestations[p.ID]; ok {
-			inTotos, err := ic.extractAttestations(ctx, session.NewGroup(sessionID), desc, attestations...)
+			inTotos, err := ic.extractAttestations(ctx, session.NewGroup(sessionID), desc, inp.Refs, attestations)
 			if err != nil {
 				return nil, err
 			}
@@ -244,16 +248,25 @@ func (ic *ImageWriter) exportLayers(ctx context.Context, refCfg cacheconfig.RefC
 	return out, err
 }
 
-func (ic *ImageWriter) extractAttestations(ctx context.Context, s session.Group, desc *ocispecs.Descriptor, attestations ...exporter.Attestation) ([]intoto.Statement, error) {
+func (ic *ImageWriter) extractAttestations(ctx context.Context, s session.Group, desc *ocispecs.Descriptor, refs map[string]cache.ImmutableRef, attestations []attestation.Attestation) ([]intoto.Statement, error) {
 	eg, ctx := errgroup.WithContext(ctx)
 	statements := make([]intoto.Statement, len(attestations))
 
+	if len(attestations) > 0 && refs == nil {
+		return nil, errors.Errorf("no refs map provided to lookup attestation keys")
+	}
+
 	for i, att := range attestations {
 		i, att := i, att
 		eg.Go(func() error {
 			switch att := att.(type) {
-			case *exporter.InTotoAttestation:
-				mount, err := att.PredicateRef.Mount(ctx, true, s)
+			case *attestation.InTotoAttestation:
+				ref, ok := refs[att.PredicateRefKey]
+				if !ok {
+					return errors.Errorf("key %s not found in refs map", att.PredicateRefKey)
+				}
+
+				mount, err := ref.Mount(ctx, true, s)
 				if err != nil {
 					return err
 				}
@@ -264,7 +277,6 @@ func (ic *ImageWriter) extractAttestations(ctx context.Context, s session.Group,
 					return err
 				}
 				defer lm.Unmount()
-
 				predicate, err := os.ReadFile(path.Join(src, att.PredicatePath))
 				if err != nil {
 					return err

exporter/exporter.go

@@ -26,18 +26,5 @@ type Source struct {
 	Ref          cache.ImmutableRef
 	Refs         map[string]cache.ImmutableRef
 	Metadata     map[string][]byte
-	Attestations map[string][]Attestation
+	Attestations map[string][]attestation.Attestation
 }
-
-type Attestation interface {
-	isExporterAttestation()
-}
-
-type InTotoAttestation struct {
-	PredicateType string
-	PredicateRef  cache.ImmutableRef
-	PredicatePath string
-	Subjects      []attestation.InTotoSubject
-}
-
-func (a *InTotoAttestation) isExporterAttestation() {}

exporter/local/export.go

@@ -2,17 +2,20 @@ package local
 
 import (
 	"context"
+	"encoding/json"
 	"os"
 	"strings"
 	"time"
 
 	"github.com/docker/docker/pkg/idtools"
 	"github.com/moby/buildkit/cache"
 	"github.com/moby/buildkit/exporter"
+	"github.com/moby/buildkit/exporter/containerimage/exptypes"
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/session/filesync"
 	"github.com/moby/buildkit/snapshot"
 	"github.com/moby/buildkit/util/progress"
+	"github.com/pkg/errors"
 	"github.com/tonistiigi/fsutil"
 	fstypes "github.com/tonistiigi/fsutil/types"
 	"golang.org/x/sync/errgroup"
@@ -60,6 +63,18 @@ func (e *localExporterInstance) Export(ctx context.Context, inp exporter.Source,
 
 	isMap := len(inp.Refs) > 0
 
+	platformsBytes, ok := inp.Metadata[exptypes.ExporterPlatformsKey]
+	if isMap && !ok {
+		return nil, errors.Errorf("unable to export multiple refs, missing platforms mapping")
+	}
+
+	var p exptypes.Platforms
+	if ok && len(platformsBytes) > 0 {
+		if err := json.Unmarshal(platformsBytes, &p); err != nil {
+			return nil, errors.Wrapf(err, "failed to parse platforms passed to exporter")
+		}
+	}
+
 	export := func(ctx context.Context, k string, ref cache.ImmutableRef) func() error {
 		return func() error {
 			var src string
@@ -130,8 +145,12 @@ func (e *localExporterInstance) Export(ctx context.Context, inp exporter.Source,
 	eg, ctx := errgroup.WithContext(ctx)
 
 	if isMap {
-		for k, ref := range inp.Refs {
-			eg.Go(export(ctx, k, ref))
+		for _, p := range p.Platforms {
+			r, ok := inp.Refs[p.ID]
+			if !ok {
+				return nil, errors.Errorf("failed to find ref for ID %s", p.ID)
+			}
+			eg.Go(export(ctx, p.ID, r))
 		}
 	} else {
 		eg.Go(export(ctx, "", inp.Ref))

exporter/tar/export.go

@@ -2,6 +2,7 @@ package local
 
 import (
 	"context"
+	"encoding/json"
 	"os"
 	"strconv"
 	"strings"
@@ -10,6 +11,7 @@ import (
 	"github.com/docker/docker/pkg/idtools"
 	"github.com/moby/buildkit/cache"
 	"github.com/moby/buildkit/exporter"
+	"github.com/moby/buildkit/exporter/containerimage/exptypes"
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/session/filesync"
 	"github.com/moby/buildkit/snapshot"
@@ -131,12 +133,28 @@ func (e *localExporterInstance) Export(ctx context.Context, inp exporter.Source,
 		}, nil
 	}
 
+	platformsBytes, ok := inp.Metadata[exptypes.ExporterPlatformsKey]
+	if len(inp.Refs) > 0 && !ok {
+		return nil, errors.Errorf("unable to export multiple refs, missing platforms mapping")
+	}
+
+	var p exptypes.Platforms
+	if ok && len(platformsBytes) > 0 {
+		if err := json.Unmarshal(platformsBytes, &p); err != nil {
+			return nil, errors.Wrapf(err, "failed to parse platforms passed to exporter")
+		}
+	}
+
 	var fs fsutil.FS
 
 	if len(inp.Refs) > 0 {
-		dirs := make([]fsutil.Dir, 0, len(inp.Refs))
-		for k, ref := range inp.Refs {
-			d, err := getDir(ctx, k, ref)
+		dirs := make([]fsutil.Dir, 0, len(p.Platforms))
+		for _, p := range p.Platforms {
+			r, ok := inp.Refs[p.ID]
+			if !ok {
+				return nil, errors.Errorf("failed to find ref for ID %s", p.ID)
+			}
+			d, err := getDir(ctx, p.ID, r)
 			if err != nil {
 				return nil, err
 			}

frontend/gateway/client/result.go

@@ -2,33 +2,22 @@ package client
 
 import (
 	"context"
+	"fmt"
 	"sync"
 
+	"github.com/moby/buildkit/identity"
 	"github.com/moby/buildkit/util/attestation"
 	"github.com/pkg/errors"
 )
 
 type BuildFunc func(context.Context, Client) (*Result, error)
 
-type Attestation interface {
-	isClientAttestation()
-}
-
-type InTotoAttestation struct {
-	PredicateType string
-	PredicateRef  Reference
-	PredicatePath string
-	Subjects      []attestation.InTotoSubject
-}
-
-func (a *InTotoAttestation) isClientAttestation() {}
-
 type Result struct {
 	mu           sync.Mutex
 	Ref          Reference
 	Refs         map[string]Reference
 	Metadata     map[string][]byte
-	Attestations map[string][]Attestation
+	Attestations map[string][]attestation.Attestation
 }
 
 func NewResult() *Result {
@@ -53,11 +42,13 @@ func (r *Result) AddRef(k string, ref Reference) {
 	r.mu.Unlock()
 }
 
-func (r *Result) AddAttestation(k string, v Attestation) {
+func (r *Result) AddInTotoAttestation(k string, v *attestation.InTotoAttestation, predicateRef Reference) {
 	r.mu.Lock()
 	if r.Attestations == nil {
-		r.Attestations = map[string][]Attestation{}
+		r.Attestations = map[string][]attestation.Attestation{}
 	}
+	v.PredicateRefKey = fmt.Sprintf("attestation:%s", identity.NewID())
+	r.Refs[v.PredicateRefKey] = predicateRef
 	r.Attestations[k] = append(r.Attestations[k], v)
 	r.mu.Unlock()
 }

frontend/gateway/forwarder/forward.go

@@ -18,7 +18,6 @@ import (
 	llberrdefs "github.com/moby/buildkit/solver/llbsolver/errdefs"
 	opspb "github.com/moby/buildkit/solver/pb"
 	"github.com/moby/buildkit/util/apicaps"
-	"github.com/moby/buildkit/util/attestation"
 	"github.com/moby/buildkit/worker"
 	digest "github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
@@ -84,18 +83,9 @@ func (c *bridgeClient) Solve(ctx context.Context, req client.SolveRequest) (*cli
 		c.refs = append(c.refs, rr)
 		cRes.SetRef(rr)
 	}
-	for k, as := range res.Attestations {
-		for _, a := range as {
-			att, rrs, err := c.newAttestation(a, session.NewGroup(c.sid))
-			if err != nil {
-				return nil, err
-			}
-			c.refs = append(c.refs, rrs...)
-			cRes.AddAttestation(k, att)
-		}
-	}
 	c.mu.Unlock()
 	cRes.Metadata = res.Metadata
+	cRes.Attestations = res.Attestations
 
 	return cRes, nil
 }
@@ -214,43 +204,7 @@ func (c *bridgeClient) toFrontendResult(r *client.Result) (*frontend.Result, err
 		res.Ref = rr.acquireResultProxy()
 	}
 	res.Metadata = r.Metadata
-	if r.Attestations != nil {
-		res.Attestations = make(map[string][]frontend.Attestation)
-		for k, as := range r.Attestations {
-			for _, a := range as {
-				switch a := a.(type) {
-				case *client.InTotoAttestation:
-					rr, ok := a.PredicateRef.(*ref)
-					if !ok {
-						return nil, errors.Errorf("invalid reference type for forward %T", r)
-					}
-
-					subjects := make([]attestation.InTotoSubject, len(a.Subjects))
-					for i, s := range a.Subjects {
-						switch s := s.(type) {
-						case *attestation.InTotoSubjectSelf:
-							subjects[i] = &attestation.InTotoSubjectSelf{}
-						case *attestation.InTotoSubjectRaw:
-							subjects[i] = &attestation.InTotoSubjectRaw{
-								Name:   s.Name,
-								Digest: s.Digest,
-							}
-						default:
-							return nil, errors.Errorf("unknown attestation subject type %T", s)
-						}
-					}
-					res.Attestations[k] = append(res.Attestations[k], &frontend.InTotoAttestation{
-						PredicateRef:  rr.acquireResultProxy(),
-						PredicatePath: a.PredicatePath,
-						PredicateType: a.PredicateType,
-						Subjects:      subjects,
-					})
-				default:
-					return nil, errors.Errorf("unknown attestation type %T", a)
-				}
-			}
-		}
-	}
+	res.Attestations = r.Attestations
 
 	return res, nil
 }
@@ -354,39 +308,6 @@ func (c *bridgeClient) newRef(r solver.ResultProxy, s session.Group) (*ref, erro
 	return &ref{resultProxy: r, session: s, c: c}, nil
 }
 
-func (c *bridgeClient) newAttestation(a frontend.Attestation, s session.Group) (client.Attestation, []*ref, error) {
-	switch a := a.(type) {
-	case *frontend.InTotoAttestation:
-		rr, err := c.newRef(a.PredicateRef, session.NewGroup(c.sid))
-		if err != nil {
-			return nil, nil, err
-		}
-
-		subjects := make([]attestation.InTotoSubject, len(a.Subjects))
-		for i, subject := range a.Subjects {
-			switch subject := subject.(type) {
-			case *attestation.InTotoSubjectSelf:
-				subjects[i] = &attestation.InTotoSubjectSelf{}
-			case *attestation.InTotoSubjectRaw:
-				subjects[i] = &attestation.InTotoSubjectRaw{
-					Name:   subject.Name,
-					Digest: subject.Digest,
-				}
-			default:
-				return nil, nil, errors.Errorf("unknown attestation subject type %T", s)
-			}
-		}
-		return &client.InTotoAttestation{
-			PredicateType: a.PredicateType,
-			PredicateRef:  rr,
-			PredicatePath: a.PredicatePath,
-			Subjects:      subjects,
-		}, []*ref{rr}, nil
-	default:
-		return nil, nil, errors.Errorf("unknown attestation type %T", a)
-	}
-}
-
 type ref struct {
 	resultProxy       solver.ResultProxy
 	resultProxyClones []solver.ResultProxy