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.

This is done by rewriting the incoming definition sent by the llb bridge
forwarder when a gateway calls solve with a protobuf definition.

Signed-off-by: Jonathan A. Sternberg <jonathan.sternberg@docker.com>

client/llb/diff.go

@@ -67,7 +67,7 @@ func (m *DiffOp) Marshal(ctx context.Context, constraints *Constraints) (digest.
 
 	proto.Op = &pb.Op_Diff{Diff: op}
 
-	dt, err := deterministicMarshal(proto)
+	dt, err := proto.MarshalDeterministic()
 	if err != nil {
 		return "", nil, nil, nil, err
 	}

client/llb/exec.go

@@ -442,7 +442,7 @@ func (e *ExecOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 		peo.Mounts = append(peo.Mounts, pm)
 	}
 
-	dt, err := deterministicMarshal(pop)
+	dt, err := pop.MarshalDeterministic()
 	if err != nil {
 		return "", nil, nil, nil, err
 	}

client/llb/fileop.go

@@ -811,7 +811,7 @@ func (f *FileOp) Marshal(ctx context.Context, c *Constraints) (digest.Digest, []
 		})
 	}
 
-	dt, err := deterministicMarshal(pop)
+	dt, err := pop.MarshalDeterministic()
 	if err != nil {
 		return "", nil, nil, nil, err
 	}

client/llb/marshal.go

@@ -8,7 +8,6 @@ import (
 	"github.com/containerd/platforms"
 	"github.com/moby/buildkit/solver/pb"
 	digest "github.com/opencontainers/go-digest"
-	"google.golang.org/protobuf/proto"
 )
 
 // Definition is the LLB definition structure with per-vertex metadata entries
@@ -147,7 +146,3 @@ type marshalCacheResult struct {
 	md     *pb.OpMetadata
 	srcs   []*SourceLocation
 }
-
-func deterministicMarshal[Message proto.Message](m Message) ([]byte, error) {
-	return proto.MarshalOptions{Deterministic: true}.Marshal(m)
-}

client/llb/merge.go

@@ -54,7 +54,7 @@ func (m *MergeOp) Marshal(ctx context.Context, constraints *Constraints) (digest
 	}
 	pop.Op = &pb.Op_Merge{Merge: op}
 
-	dt, err := deterministicMarshal(pop)
+	dt, err := pop.MarshalDeterministic()
 	if err != nil {
 		return "", nil, nil, nil, err
 	}

client/llb/source.go

@@ -77,7 +77,7 @@ func (s *SourceOp) Marshal(ctx context.Context, constraints *Constraints) (diges
 		proto.Platform = nil
 	}
 
-	dt, err := deterministicMarshal(proto)
+	dt, err := proto.MarshalDeterministic()
 	if err != nil {
 		return "", nil, nil, nil, err
 	}

frontend/gateway/gateway.go

@@ -760,6 +760,17 @@ func (lbf *llbBridgeForwarder) Solve(ctx context.Context, req *pb.SolveRequest)
 		}
 	}
 
+	if req.Definition != nil {
+		// Rewrite digests in the definition. This ensures the digests are validated
+		// and standardized that were sent by a plugin so digests can be compared
+		// even when the frontend is running a different version of protobuf.
+		def, err := rewriteDigests(ctx, req.Definition)
+		if err != nil {
+			return nil, err
+		}
+		req.Definition = def
+	}
+
 	ctx = tracing.ContextWithSpanFromContext(ctx, lbf.callCtx)
 	res, err := lbf.llbBridge.Solve(ctx, frontend.SolveRequest{
 		Evaluate:       req.Evaluate,
@@ -1663,3 +1674,91 @@ func addCapsForKnownFrontends(caps map[string]struct{}, dgst digest.Digest) {
 		caps["moby.buildkit.frontend.inputs"] = struct{}{}
 	}
 }
+
+func rewriteDigests(ctx context.Context, def *opspb.Definition) (*opspb.Definition, error) {
+	clone := def.CloneVT()
+
+	var rewriter digestRewriter
+	if err := rewriter.Rewrite(ctx, clone); err != nil {
+		return nil, err
+	}
+	return clone, nil
+}
+
+type digestRewriter struct {
+	def      *opspb.Definition
+	byIndex  map[digest.Digest]int
+	byDigest map[digest.Digest]*opspb.Op
+	visited  map[digest.Digest]digest.Digest
+}
+
+func (r *digestRewriter) Rewrite(ctx context.Context, def *opspb.Definition) error {
+	byIndex := make(map[digest.Digest]int)
+	byDigest := make(map[digest.Digest]*opspb.Op)
+	for i, in := range def.Def {
+		dgst := digest.FromBytes(in)
+
+		var op opspb.Op
+		if err := op.UnmarshalVT(in); err != nil {
+			return err
+		}
+		byDigest[dgst] = &op
+		byIndex[dgst] = i
+	}
+
+	r.def = def
+	r.byDigest = byDigest
+	r.byIndex = byIndex
+	r.visited = make(map[digest.Digest]digest.Digest)
+
+	for dgst := range r.byIndex {
+		if _, err := r.rewrite(ctx, dgst); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (r *digestRewriter) rewrite(ctx context.Context, dgst digest.Digest) (digest.Digest, error) {
+	if dgst, ok := r.visited[dgst]; ok {
+		return dgst, nil
+	}
+
+	op, ok := r.byDigest[dgst]
+	if !ok {
+		return "", errors.Errorf("invalid missing input digest %s", dgst)
+	}
+
+	// Recompute input digests.
+	for _, input := range op.Inputs {
+		select {
+		case <-ctx.Done():
+			return "", context.Cause(ctx)
+		default:
+		}
+
+		iDgst, err := r.rewrite(ctx, digest.Digest(input.Digest))
+		if err != nil {
+			return "", err
+		}
+		if digest.Digest(input.Digest) != iDgst {
+			input.Digest = string(iDgst)
+		}
+	}
+
+	// Must use deterministic marshal here so the digest is consistent.
+	data, err := op.MarshalDeterministic()
+	if err != nil {
+		return "", err
+	}
+	newDgst := digest.FromBytes(data)
+
+	r.visited[dgst] = newDgst
+	if dgst != newDgst {
+		r.visited[newDgst] = newDgst
+	}
+
+	index := r.byIndex[dgst]
+	r.def.Def[index] = data
+	return newDgst, nil
+}

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 {
@@ -231,7 +230,7 @@ func recomputeDigests(ctx context.Context, all map[digest.Digest]*pb.Op, visited
 		return dgst, nil
 	}
 
-	dt, err := deterministicMarshal(op)
+	dt, err := op.MarshalDeterministic()
 	if err != nil {
 		return "", err
 	}
@@ -266,7 +265,7 @@ func loadLLB(ctx context.Context, def *pb.Definition, polEngine SourcePolicyEval
 				return solver.Edge{}, errors.Wrap(err, "error evaluating the source policy")
 			}
 			if mutated {
-				dtMutated, err := deterministicMarshal(&op)
+				dtMutated, err := op.MarshalDeterministic()
 				if err != nil {
 					return solver.Edge{}, err
 				}
@@ -400,7 +399,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)
-}

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
 }
@@ -12,10 +14,23 @@ func (m *Definition) Unmarshal(dAtA []byte) error {
 	return m.UnmarshalVT(dAtA)
 }
 
-func (m *Op) Marshal() ([]byte, error) {
-	return m.MarshalVT()
+// MarshalDeterministic will marshal this Op in
+// a deterministic manner.
+//
+// This will cause a message to be serialized in the same
+// way within the same binary. This can be used when the bytes
+// need to be the same for multiple marshal operations.
+//
+// There is no guarantee that the bytes will be the same within
+// different binaries (aka different versions of protobuf).
+func (m *Op) MarshalDeterministic() ([]byte, error) {
+	return deterministicMarshal(m)
 }
 
 func (m *Op) Unmarshal(dAtA []byte) error {
 	return m.UnmarshalVT(dAtA)
 }
+
+func deterministicMarshal[Message proto.Message](m Message) ([]byte, error) {
+	return proto.MarshalOptions{Deterministic: true}.Marshal(m)
+}