From 9f65f8c1fa462da2dc1444d01df924df637800c9 Mon Sep 17 00:00:00 2001
From: "Jonathan A. Sternberg" <jonathan.sternberg@docker.com>
Date: Thu, 14 Nov 2024 13:38:00 -0600
Subject: [PATCH] gateway: ensure llb digests are deterministic when sent by
 frontends

This ensures different valid protobuf serializations that are sent by
frontends will be rewritten into digests that are normalized for the
buildkit solver.

The most recent example of this is that older frontends would generate
protobuf with gogo and the newer buildkit is using the google protobuf
library. These produce different serializations and cause the solver to
think that identical operations are actually different.

Signed-off-by: Jonathan A. Sternberg <jonathan.sternberg@docker.com>
---
 solver/llbsolver/testdata/gogoproto.data | Bin 0 -> 284 bytes
 solver/llbsolver/vertex.go               |  43 ++++------------
 solver/llbsolver/vertex_test.go          |  61 +++++++++++++++++++++++
 solver/pb/ops.go                         |   4 +-
 4 files changed, 75 insertions(+), 33 deletions(-)
 create mode 100644 solver/llbsolver/testdata/gogoproto.data

diff --git a/solver/llbsolver/testdata/gogoproto.data b/solver/llbsolver/testdata/gogoproto.data
new file mode 100644
index 0000000000000000000000000000000000000000..3c93d3626c76f3c25d4b338299f5209e7676b277
GIT binary patch
literal 284
zcmb`>Jr2S!3_xM#jEvj?8prvmdjW(vg=5E10jku31h<EQfdeqU<<oPx0oUmymo{(5
z8=u<HeRDTX-F2)p&#$gNmRFsg!<FxCxx1q~xczeGkn&JslbW=U7eo$#mA$v#B5Dp4
ziyCOu#wA#40*DEjDn+!I$-J$P|D#qvn5hzZLhv5AQ81taxu98q;A6AQNkn=9j;fz8
P?LRkE&2nLsuwK0ZrL0h+

literal 0
HcmV?d00001

diff --git a/solver/llbsolver/vertex.go b/solver/llbsolver/vertex.go
index 9052a727a6fd..c2f659c5788e 100644
--- a/solver/llbsolver/vertex.go
+++ b/solver/llbsolver/vertex.go
@@ -14,7 +14,6 @@ import (
 	digest "github.com/opencontainers/go-digest"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
-	"google.golang.org/protobuf/proto"
 )
 
 type vertex struct {
@@ -208,7 +207,6 @@ func recomputeDigests(ctx context.Context, all map[digest.Digest]*pb.Op, visited
 		return "", errors.Errorf("invalid missing input digest %s", dgst)
 	}
 
-	var mutated bool
 	for _, input := range op.Inputs {
 		select {
 		case <-ctx.Done():
@@ -220,25 +218,20 @@ func recomputeDigests(ctx context.Context, all map[digest.Digest]*pb.Op, visited
 		if err != nil {
 			return "", err
 		}
-		if digest.Digest(input.Digest) != iDgst {
-			mutated = true
-			input.Digest = string(iDgst)
-		}
-	}
-
-	if !mutated {
-		visited[dgst] = dgst
-		return dgst, nil
+		input.Digest = string(iDgst)
 	}
 
-	dt, err := deterministicMarshal(op)
+	dt, err := op.Marshal()
 	if err != nil {
 		return "", err
 	}
+
 	newDgst := digest.FromBytes(dt)
+	if newDgst != dgst {
+		all[newDgst] = op
+		delete(all, dgst)
+	}
 	visited[dgst] = newDgst
-	all[newDgst] = op
-	delete(all, dgst)
 	return newDgst, nil
 }
 
@@ -250,7 +243,6 @@ func loadLLB(ctx context.Context, def *pb.Definition, polEngine SourcePolicyEval
 	}
 
 	allOps := make(map[digest.Digest]*pb.Op)
-	mutatedDigests := make(map[digest.Digest]digest.Digest) // key: old, val: new
 
 	var lastDgst digest.Digest
 
@@ -261,27 +253,18 @@ func loadLLB(ctx context.Context, def *pb.Definition, polEngine SourcePolicyEval
 		}
 		dgst := digest.FromBytes(dt)
 		if polEngine != nil {
-			mutated, err := polEngine.Evaluate(ctx, op.GetSource())
-			if err != nil {
+			if _, err := polEngine.Evaluate(ctx, op.GetSource()); err != nil {
 				return solver.Edge{}, errors.Wrap(err, "error evaluating the source policy")
 			}
-			if mutated {
-				dtMutated, err := deterministicMarshal(&op)
-				if err != nil {
-					return solver.Edge{}, err
-				}
-				dgstMutated := digest.FromBytes(dtMutated)
-				mutatedDigests[dgst] = dgstMutated
-				dgst = dgstMutated
-			}
 		}
+
 		allOps[dgst] = &op
 		lastDgst = dgst
 	}
 
+	mutatedDigests := make(map[digest.Digest]digest.Digest) // key: old, val: new
 	for dgst := range allOps {
-		_, err := recomputeDigests(ctx, allOps, mutatedDigests, dgst)
-		if err != nil {
+		if _, err := recomputeDigests(ctx, allOps, mutatedDigests, dgst); err != nil {
 			return solver.Edge{}, err
 		}
 	}
@@ -400,7 +383,3 @@ func fileOpName(actions []*pb.FileAction) string {
 
 	return strings.Join(names, ", ")
 }
-
-func deterministicMarshal[Message proto.Message](m Message) ([]byte, error) {
-	return proto.MarshalOptions{Deterministic: true}.Marshal(m)
-}
diff --git a/solver/llbsolver/vertex_test.go b/solver/llbsolver/vertex_test.go
index 5ecdd62edeb5..3123207221f2 100644
--- a/solver/llbsolver/vertex_test.go
+++ b/solver/llbsolver/vertex_test.go
@@ -2,6 +2,8 @@ package llbsolver
 
 import (
 	"context"
+	_ "embed"
+	"fmt"
 	"testing"
 
 	"github.com/moby/buildkit/solver/pb"
@@ -53,3 +55,62 @@ func TestRecomputeDigests(t *testing.T) {
 	require.Equal(t, newDigest, digest.Digest(op2.Inputs[0].Digest))
 	assert.NotEqual(t, op2Digest, updated)
 }
+
+//go:embed testdata/gogoproto.data
+var gogoprotoData []byte
+
+func TestIngestDigest(t *testing.T) {
+	op1 := &pb.Op{
+		Op: &pb.Op_Source{
+			Source: &pb.SourceOp{
+				Identifier: "docker-image://docker.io/library/busybox:latest",
+			},
+		},
+	}
+	op1Data, err := op1.Marshal()
+	require.NoError(t, err)
+	op1Digest := digest.FromBytes(op1Data)
+
+	op2 := &pb.Op{
+		Inputs: []*pb.Input{
+			{Digest: string(op1Digest)}, // Input is the old digest, this should be updated after recomputeDigests
+		},
+	}
+	op2Data, err := op2.Marshal()
+	require.NoError(t, err)
+	op2Digest := digest.FromBytes(op2Data)
+
+	var def pb.Definition
+	err = def.Unmarshal(gogoprotoData)
+	require.NoError(t, err)
+	require.Len(t, def.Def, 2)
+
+	// Read the definition from the test data and ensure it uses the
+	// canonical digests after recompute.
+	var lastDgst digest.Digest
+	all := map[digest.Digest]*pb.Op{}
+	for _, in := range def.Def {
+		op := new(pb.Op)
+		err := op.Unmarshal(in)
+		require.NoError(t, err)
+
+		lastDgst = digest.FromBytes(in)
+		all[lastDgst] = op
+	}
+	fmt.Println(all, lastDgst)
+
+	visited := map[digest.Digest]digest.Digest{}
+	newDgst, err := recomputeDigests(context.Background(), all, visited, lastDgst)
+	require.NoError(t, err)
+	require.Len(t, visited, 2)
+	require.Equal(t, op2Digest, newDgst)
+	require.Equal(t, op2Digest, visited[newDgst])
+	delete(visited, newDgst)
+
+	// Last element should correspond to op1.
+	// The old digest doesn't really matter.
+	require.Len(t, visited, 1)
+	for _, newDgst := range visited {
+		require.Equal(t, op1Digest, newDgst)
+	}
+}
diff --git a/solver/pb/ops.go b/solver/pb/ops.go
index ef9de664b049..c6aeaead0b16 100644
--- a/solver/pb/ops.go
+++ b/solver/pb/ops.go
@@ -1,5 +1,7 @@
 package pb
 
+import proto "google.golang.org/protobuf/proto"
+
 func (m *Definition) IsNil() bool {
 	return m == nil || m.Metadata == nil
 }
@@ -13,7 +15,7 @@ func (m *Definition) Unmarshal(dAtA []byte) error {
 }
 
 func (m *Op) Marshal() ([]byte, error) {
-	return m.MarshalVT()
+	return proto.MarshalOptions{Deterministic: true}.Marshal(m)
 }
 
 func (m *Op) Unmarshal(dAtA []byte) error {