builder-next: apply source policies on config resolve

SourcePolicy was accounted for in https://github.com/moby/buildkit/commit/330cf7ae7d9ac983432ab13f8bc9e7d4aff75695

TODO: replace applySourcePolicies with BuildKit's implementation, which
is currently unexported.

Co-authored-by: Tonis Tiigi <tonistiigi@gmail.com>
Signed-off-by: Bjorn Neergaard <bjorn.neergaard@docker.com>

builder/builder-next/adapters/containerimage/pull.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"path"
+	"strings"
 	"sync"
 	"time"
 
@@ -16,6 +17,7 @@ import (
 	"github.com/containerd/containerd/leases"
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/platforms"
+	cdreference "github.com/containerd/containerd/reference"
 	ctdreference "github.com/containerd/containerd/reference"
 	"github.com/containerd/containerd/remotes"
 	"github.com/containerd/containerd/remotes/docker"
@@ -32,8 +34,12 @@ import (
 	"github.com/moby/buildkit/client/llb"
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/solver"
+	"github.com/moby/buildkit/solver/pb"
 	"github.com/moby/buildkit/source"
 	srctypes "github.com/moby/buildkit/source/types"
+	"github.com/moby/buildkit/sourcepolicy"
+	policy "github.com/moby/buildkit/sourcepolicy/pb"
+	spb "github.com/moby/buildkit/sourcepolicy/pb"
 	"github.com/moby/buildkit/util/flightcontrol"
 	"github.com/moby/buildkit/util/imageutil"
 	"github.com/moby/buildkit/util/leaseutil"
@@ -92,8 +98,9 @@ func (is *Source) resolveLocal(refStr string) (*image.Image, error) {
 	return img, nil
 }
 
-func (is *Source) resolveRemote(ctx context.Context, ref string, platform *ocispec.Platform, sm *session.Manager, g session.Group) (digest.Digest, []byte, error) {
+func (is *Source) resolveRemote(ctx context.Context, ref string, platform *ocispec.Platform, sm *session.Manager, g session.Group) (string, digest.Digest, []byte, error) {
 	type t struct {
+		ref  string
 		dgst digest.Digest
 		dt   []byte
 	}
@@ -105,32 +112,36 @@ func (is *Source) resolveRemote(ctx context.Context, ref string, platform *ocisp
 	key := "getconfig::" + ref + "::" + platforms.Format(p)
 	res, err := is.g.Do(ctx, key, func(ctx context.Context) (interface{}, error) {
 		res := resolver.DefaultPool.GetResolver(is.RegistryHosts, ref, "pull", sm, g)
-		dgst, dt, err := imageutil.Config(ctx, ref, res, is.ContentStore, is.LeaseManager, platform)
+		ref, dgst, dt, err := imageutil.Config(ctx, ref, res, is.ContentStore, is.LeaseManager, platform, []*policy.Policy{})
 		if err != nil {
 			return nil, err
 		}
-		return &t{dgst: dgst, dt: dt}, nil
+		return &t{ref: ref, dgst: dgst, dt: dt}, nil
 	})
 	var typed *t
 	if err != nil {
-		return "", nil, err
+		return ref, "", nil, err
 	}
 	typed = res.(*t)
-	return typed.dgst, typed.dt, nil
+	return typed.ref, typed.dgst, typed.dt, nil
 }
 
 // ResolveImageConfig returns image config for an image
-func (is *Source) ResolveImageConfig(ctx context.Context, ref string, opt llb.ResolveImageConfigOpt, sm *session.Manager, g session.Group) (digest.Digest, []byte, error) {
+func (is *Source) ResolveImageConfig(ctx context.Context, ref string, opt llb.ResolveImageConfigOpt, sm *session.Manager, g session.Group) (string, digest.Digest, []byte, error) {
+	ref, err := applySourcePolicies(ctx, ref, opt.SourcePolicies)
+	if err != nil {
+		return "", "", nil, err
+	}
 	resolveMode, err := source.ParseImageResolveMode(opt.ResolveMode)
 	if err != nil {
-		return "", nil, err
+		return ref, "", nil, err
 	}
 	switch resolveMode {
 	case source.ResolveModeForcePull:
-		dgst, dt, err := is.resolveRemote(ctx, ref, opt.Platform, sm, g)
+		ref, dgst, dt, err := is.resolveRemote(ctx, ref, opt.Platform, sm, g)
 		// TODO: pull should fallback to local in case of failure to allow offline behavior
 		// the fallback doesn't work currently
-		return dgst, dt, err
+		return ref, dgst, dt, err
 		/*
 			if err == nil {
 				return dgst, dt, err
@@ -152,14 +163,14 @@ func (is *Source) ResolveImageConfig(ctx context.Context, ref string, opt llb.Re
 					path.Join(img.OS, img.Architecture, img.Variant),
 				)
 			} else {
-				return "", img.RawJSON(), err
+				return ref, "", img.RawJSON(), err
 			}
 		}
 		// fallback to remote
 		return is.resolveRemote(ctx, ref, opt.Platform, sm, g)
 	}
 	// should never happen
-	return "", nil, fmt.Errorf("builder cannot resolve image %s: invalid mode %q", ref, opt.ResolveMode)
+	return ref, "", nil, fmt.Errorf("builder cannot resolve image %s: invalid mode %q", ref, opt.ResolveMode)
 }
 
 // Resolve returns access to pulling for an identifier
@@ -289,11 +300,12 @@ func (p *puller) resolve(ctx context.Context, g session.Group) error {
 			if err != nil {
 				return struct{}{}, err
 			}
-			_, dt, err := p.is.ResolveImageConfig(ctx, ref.String(), llb.ResolveImageConfigOpt{Platform: &p.platform, ResolveMode: resolveModeToString(p.src.ResolveMode)}, p.sm, g)
+			newRef, _, dt, err := p.is.ResolveImageConfig(ctx, ref.String(), llb.ResolveImageConfigOpt{Platform: &p.platform, ResolveMode: resolveModeToString(p.src.ResolveMode)}, p.sm, g)
 			if err != nil {
 				return struct{}{}, err
 			}
 
+			p.ref = newRef
 			p.config = dt
 		}
 		return struct{}{}, nil
@@ -860,3 +872,41 @@ func platformMatches(img *image.Image, p *ocispec.Platform) bool {
 		Variant:      img.Variant,
 	})
 }
+
+func applySourcePolicies(ctx context.Context, str string, spls []*spb.Policy) (string, error) {
+	ref, err := cdreference.Parse(str)
+	if err != nil {
+		return "", errors.WithStack(err)
+	}
+	op := &pb.Op{
+		Op: &pb.Op_Source{
+			Source: &pb.SourceOp{
+				Identifier: srctypes.DockerImageScheme + "://" + ref.String(),
+			},
+		},
+	}
+
+	mut, err := sourcepolicy.NewEngine(spls).Evaluate(ctx, op)
+	if err != nil {
+		return "", errors.Wrap(err, "could not resolve image due to policy")
+	}
+
+	if mut {
+		var (
+			t  string
+			ok bool
+		)
+		t, newRef, ok := strings.Cut(op.GetSource().GetIdentifier(), "://")
+		if !ok {
+			return "", errors.Errorf("could not parse ref: %s", op.GetSource().GetIdentifier())
+		}
+		if ok && t != srctypes.DockerImageScheme {
+			return "", &imageutil.ResolveToNonImageError{Ref: str, Updated: newRef}
+		}
+		ref, err = cdreference.Parse(newRef)
+		if err != nil {
+			return "", errors.WithStack(err)
+		}
+	}
+	return ref.String(), nil
+}

builder/builder-next/worker/worker.go

@@ -236,7 +236,7 @@ func (w *Worker) ResolveOp(v solver.Vertex, s frontend.FrontendLLBBridge, sm *se
 }
 
 // ResolveImageConfig returns image config for an image
-func (w *Worker) ResolveImageConfig(ctx context.Context, ref string, opt llb.ResolveImageConfigOpt, sm *session.Manager, g session.Group) (digest.Digest, []byte, error) {
+func (w *Worker) ResolveImageConfig(ctx context.Context, ref string, opt llb.ResolveImageConfigOpt, sm *session.Manager, g session.Group) (string, digest.Digest, []byte, error) {
 	return w.ImageSource.ResolveImageConfig(ctx, ref, opt, sm, g)
 }