From 4a00e0f8080b05fad57c275b937e4ae67d3a1f93 Mon Sep 17 00:00:00 2001
From: CrazyMax <1951866+crazy-max@users.noreply.github.com>
Date: Fri, 23 Feb 2024 00:48:24 +0100
Subject: [PATCH] build: set record provenance in response

Signed-off-by: CrazyMax <1951866+crazy-max@users.noreply.github.com>
---
 build/build.go      |   5 +++
 build/provenance.go | 104 ++++++++++++++++++++++++++++++++++++++++++++
 go.mod              |   2 +-
 tests/bake.go       |  16 +++++--
 tests/build.go      |  16 +++++--
 5 files changed, 136 insertions(+), 7 deletions(-)
 create mode 100644 build/provenance.go

diff --git a/build/build.go b/build/build.go
index e10efa93e546..4809bd170a1c 100644
--- a/build/build.go
+++ b/build/build.go
@@ -816,6 +816,11 @@ func BuildWithResultHandler(ctx context.Context, nodes []builder.Node, opt map[s
 						rr.ExporterResponse[k] = string(v)
 					}
 					rr.ExporterResponse["buildx.build.ref"] = buildRef
+					if node.Driver.HistoryAPISupported(ctx) {
+						if err := setRecordProvenance(ctx, c, rr, so.Ref, pw); err != nil {
+							return err
+						}
+					}
 
 					node := dp.Node().Driver
 					if node.IsMobyDriver() {
diff --git a/build/provenance.go b/build/provenance.go
new file mode 100644
index 000000000000..0ee399d49f1f
--- /dev/null
+++ b/build/provenance.go
@@ -0,0 +1,104 @@
+package build
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"io"
+	"strings"
+	"sync"
+
+	"github.com/containerd/containerd/content"
+	"github.com/containerd/containerd/content/proxy"
+	"github.com/docker/buildx/util/progress"
+	controlapi "github.com/moby/buildkit/api/services/control"
+	"github.com/moby/buildkit/client"
+	"github.com/opencontainers/go-digest"
+	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
+	"github.com/pkg/errors"
+	"golang.org/x/sync/errgroup"
+)
+
+func setRecordProvenance(ctx context.Context, c *client.Client, sr *client.SolveResponse, ref string, pw progress.Writer) error {
+	var mu sync.Mutex
+	pw = progress.ResetTime(pw)
+	return progress.Wrap("fetching build record provenance", pw.Write, func(l progress.SubLogger) error {
+		cl, err := c.ControlClient().ListenBuildHistory(ctx, &controlapi.BuildHistoryRequest{
+			Ref:       ref,
+			EarlyExit: true,
+		})
+		if err != nil {
+			return err
+		}
+
+		eg, ctx := errgroup.WithContext(ctx)
+		for {
+			ev, err := cl.Recv()
+			if errors.Is(err, io.EOF) {
+				break
+			} else if err != nil {
+				return err
+			}
+			if ev.Record == nil {
+				continue
+			}
+			if ev.Record.Result != nil {
+				provenanceDgst := provenanceDigest(ev.Record.Result)
+				if provenanceDgst == nil {
+					continue
+				}
+				eg.Go(func() error {
+					buf := &bytes.Buffer{}
+					if err := writeBlob(ctx, c, *provenanceDgst, buf); err != nil {
+						return errors.Wrapf(err, "failed to load provenance from build record")
+					}
+					mu.Lock()
+					sr.ExporterResponse["buildx.build.provenance"] = base64.StdEncoding.EncodeToString(buf.Bytes())
+					mu.Unlock()
+					return nil
+				})
+			} else if ev.Record.Results != nil {
+				for platform, res := range ev.Record.Results {
+					platform := platform
+					provenanceDgst := provenanceDigest(res)
+					if provenanceDgst == nil {
+						continue
+					}
+					eg.Go(func() error {
+						buf := &bytes.Buffer{}
+						if err := writeBlob(ctx, c, *provenanceDgst, buf); err != nil {
+							return errors.Wrapf(err, "failed to load provenance from build record")
+						}
+						mu.Lock()
+						sr.ExporterResponse["buildx.build.provenance/"+platform] = base64.StdEncoding.EncodeToString(buf.Bytes())
+						mu.Unlock()
+						return nil
+					})
+				}
+			}
+		}
+		return eg.Wait()
+	})
+}
+
+func provenanceDigest(res *controlapi.BuildResultInfo) *digest.Digest {
+	for _, a := range res.Attestations {
+		if a.MediaType == "application/vnd.in-toto+json" && strings.HasPrefix(a.Annotations["in-toto.io/predicate-type"], "https://slsa.dev/provenance/") {
+			return &a.Digest
+		}
+	}
+	return nil
+}
+
+func writeBlob(ctx context.Context, c *client.Client, dgst digest.Digest, w io.Writer) error {
+	store := proxy.NewContentStore(c.ContentClient())
+	ra, err := store.ReaderAt(ctx, ocispecs.Descriptor{
+		Digest: dgst,
+	})
+	if err != nil {
+		return err
+	}
+	defer ra.Close()
+	_, err = io.Copy(w, content.NewReader(ra))
+	return err
+}
diff --git a/go.mod b/go.mod
index b71467ce8ef7..07dbfa9ea74c 100644
--- a/go.mod
+++ b/go.mod
@@ -24,6 +24,7 @@ require (
 	github.com/google/uuid v1.5.0
 	github.com/hashicorp/go-cty-funcs v0.0.0-20230405223818-a090f58aa992
 	github.com/hashicorp/hcl/v2 v2.19.1
+	github.com/in-toto/in-toto-golang v0.5.0
 	github.com/moby/buildkit v0.13.0-rc1.0.20240222164755-8e3fe35738c2 // master (v0.13.0-dev)
 	github.com/moby/sys/mountinfo v0.7.1
 	github.com/moby/sys/signal v0.7.0
@@ -105,7 +106,6 @@ require (
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.16.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
 	github.com/imdario/mergo v0.3.16 // indirect
-	github.com/in-toto/in-toto-golang v0.5.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
diff --git a/tests/bake.go b/tests/bake.go
index 1cc20ef75f4f..9b69061d4935 100644
--- a/tests/bake.go
+++ b/tests/bake.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/containerd/continuity/fs/fstest"
 	"github.com/docker/buildx/util/gitutil"
+	slsa02 "github.com/in-toto/in-toto-golang/in_toto/slsa_provenance/v0.2"
 	"github.com/moby/buildkit/util/testutil/integration"
 	"github.com/stretchr/testify/require"
 )
@@ -35,7 +36,7 @@ var bakeTests = []func(t *testing.T, sb integration.Sandbox){
 	testBakeEmpty,
 	testBakeShmSize,
 	testBakeUlimits,
-	testBakeRefs,
+	testBakeMetadata,
 }
 
 func testBakeLocal(t *testing.T, sb integration.Sandbox) {
@@ -589,7 +590,7 @@ target "default" {
 	require.Contains(t, string(dt), `1024`)
 }
 
-func testBakeRefs(t *testing.T, sb integration.Sandbox) {
+func testBakeMetadata(t *testing.T, sb integration.Sandbox) {
 	dockerfile := []byte(`
 FROM scratch
 COPY foo /foo
@@ -622,7 +623,8 @@ target "default" {
 
 	type mdT struct {
 		Default struct {
-			BuildRef string `json:"buildx.build.ref"`
+			BuildRef        string                 `json:"buildx.build.ref"`
+			BuildProvenance map[string]interface{} `json:"buildx.build.provenance"`
 		} `json:"default"`
 	}
 	var md mdT
@@ -630,4 +632,12 @@ target "default" {
 	require.NoError(t, err)
 
 	require.NotEmpty(t, md.Default.BuildRef)
+	require.NotEmpty(t, md.Default.BuildProvenance)
+
+	dtprv, err := json.Marshal(md.Default.BuildProvenance)
+	require.NoError(t, err)
+
+	var prv slsa02.ProvenancePredicate
+	require.NoError(t, json.Unmarshal(dtprv, &prv))
+	require.Equal(t, "https://mobyproject.org/buildkit@v1", prv.BuildType)
 }
diff --git a/tests/build.go b/tests/build.go
index 4500ceabc377..57bf7643afc2 100644
--- a/tests/build.go
+++ b/tests/build.go
@@ -16,6 +16,7 @@ import (
 	"github.com/containerd/containerd/platforms"
 	"github.com/containerd/continuity/fs/fstest"
 	"github.com/creack/pty"
+	slsa02 "github.com/in-toto/in-toto-golang/in_toto/slsa_provenance/v0.2"
 	"github.com/moby/buildkit/util/appdefaults"
 	"github.com/moby/buildkit/util/contentutil"
 	"github.com/moby/buildkit/util/testutil"
@@ -53,7 +54,7 @@ var buildTests = []func(t *testing.T, sb integration.Sandbox){
 	testBuildNetworkModeBridge,
 	testBuildShmSize,
 	testBuildUlimit,
-	testBuildRef,
+	testBuildMetadata,
 }
 
 func testBuild(t *testing.T, sb integration.Sandbox) {
@@ -515,7 +516,7 @@ COPY --from=build /ulimit /
 	require.Contains(t, string(dt), `1024`)
 }
 
-func testBuildRef(t *testing.T, sb integration.Sandbox) {
+func testBuildMetadata(t *testing.T, sb integration.Sandbox) {
 	dir := createTestProject(t)
 	dirDest := t.TempDir()
 
@@ -533,13 +534,22 @@ func testBuildRef(t *testing.T, sb integration.Sandbox) {
 	require.NoError(t, err)
 
 	type mdT struct {
-		BuildRef string `json:"buildx.build.ref"`
+		BuildRef        string                 `json:"buildx.build.ref"`
+		BuildProvenance map[string]interface{} `json:"buildx.build.provenance"`
 	}
 	var md mdT
 	err = json.Unmarshal(dt, &md)
 	require.NoError(t, err)
 
 	require.NotEmpty(t, md.BuildRef)
+	require.NotEmpty(t, md.BuildProvenance)
+
+	dtprv, err := json.Marshal(md.BuildProvenance)
+	require.NoError(t, err)
+
+	var prv slsa02.ProvenancePredicate
+	require.NoError(t, json.Unmarshal(dtprv, &prv))
+	require.Equal(t, "https://mobyproject.org/buildkit@v1", prv.BuildType)
 }
 
 func createTestProject(t *testing.T) string {