llb: add checksum option to llb.Image

client/client_test.go

@@ -158,6 +158,7 @@ var allTests = []func(t *testing.T, sb integration.Sandbox){
 	testHostnameLookup,
 	testHostnameSpecifying,
 	testPushByDigest,
+	testPullWithDigestCheck,
 	testBasicInlineCacheImportExport,
 	testExportBusyboxLocal,
 	testBridgeNetworking,
@@ -1126,6 +1127,119 @@ func testPushByDigest(t *testing.T, sb integration.Sandbox) {
 	require.Greater(t, desc.Size, int64(0))
 }
 
+func testPullWithDigestCheck(t *testing.T, sb integration.Sandbox) {
+	workers.CheckFeatureCompat(t, sb, workers.FeatureDirectPush)
+	requiresLinux(t)
+	c, err := New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	registry, err := sb.NewRegistry()
+	if errors.Is(err, integration.ErrRequirements) {
+		t.Skip(err.Error())
+	}
+	require.NoError(t, err)
+
+	st := llb.Scratch().File(llb.Mkfile("foo", 0600, []byte("data1")))
+
+	def, err := st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	name1 := registry + "/foo/bar:v1.0.0"
+
+	resp, err := c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: "image",
+				Attrs: map[string]string{
+					"name": name1,
+					"push": "true",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dgst1Str := resp.ExporterResponse[exptypes.ExporterImageDigestKey]
+	dgst1, err := digest.Parse(dgst1Str)
+	require.NoError(t, err)
+
+	st = llb.Scratch().File(llb.Mkfile("foo", 0600, []byte("data2")))
+
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	name2 := registry + "/foo/bar:v2.0.0"
+
+	resp, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: "image",
+				Attrs: map[string]string{
+					"name": name2,
+					"push": "true",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dgst2str := resp.ExporterResponse[exptypes.ExporterImageDigestKey]
+	dgst2, err := digest.Parse(dgst2str)
+	require.NoError(t, err)
+
+	require.NotEqual(t, dgst1, dgst2)
+
+	// if digest is set in ref then pull happens only by the digest
+	st = llb.Image(name2 + "@" + dgst1.String())
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	destDir := t.TempDir()
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type:      ExporterLocal,
+				OutputDir: destDir,
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dt, err := os.ReadFile(filepath.Join(destDir, "foo"))
+	require.NoError(t, err)
+	require.Equal(t, "data1", string(dt))
+
+	// if digest is set by checksum then pull happens by tag
+	st = llb.Image(name2, llb.WithImageChecksum(dgst2))
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+	destDir = t.TempDir()
+	_, err = c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type:      ExporterLocal,
+				OutputDir: destDir,
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dt, err = os.ReadFile(filepath.Join(destDir, "foo"))
+	require.NoError(t, err)
+	require.Equal(t, "data2", string(dt))
+
+	// if checksum doesn't match then pull fails
+	st = llb.Image(name2, llb.WithImageChecksum(dgst1))
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), fmt.Sprintf("image digest %s for %s does not match expected checksum %s", dgst2, name2, dgst1))
+}
+
 func testSecurityMode(t *testing.T, sb integration.Sandbox) {
 	integration.SkipOnPlatform(t, "windows")
 	workers.CheckFeatureCompat(t, sb, workers.FeatureSecurityMode)

client/llb/resolver.go

@@ -2,6 +2,7 @@ package llb
 
 import (
 	"github.com/moby/buildkit/client/llb/sourceresolver"
+	digest "github.com/opencontainers/go-digest"
 )
 
 // WithMetaResolver adds a metadata resolver to an image
@@ -26,5 +27,11 @@ func WithLayerLimit(l int) ImageOption {
 	})
 }
 
+func WithImageChecksum(dgst digest.Digest) ImageOption {
+	return imageOptionFunc(func(ii *ImageInfo) {
+		ii.checksum = dgst
+	})
+}
+
 // ImageMetaResolver can resolve image config metadata from a reference
 type ImageMetaResolver = sourceresolver.ImageMetaResolver

client/llb/source.go

@@ -130,6 +130,11 @@ func Image(ref string, opts ...ImageOption) State {
 		addCap(&info.Constraints, pb.CapSourceImageLayerLimit)
 	}
 
+	if info.checksum != "" {
+		attrs[pb.AttrImageChecksum] = info.checksum.String()
+		addCap(&info.Constraints, pb.CapSourceImageChecksum)
+	}
+
 	src := NewSource("docker-image://"+ref, attrs, info.Constraints) // controversial
 	if err != nil {
 		src.err = err
@@ -227,6 +232,7 @@ type ImageInfo struct {
 	resolveDigest bool
 	resolveMode   ResolveMode
 	layerLimit    *int
+	checksum      digest.Digest
 	RecordType    string
 }
 
@@ -623,11 +629,18 @@ func OCILayerLimit(limit int) OCILayoutOption {
 	})
 }
 
+func OCIChecksum(dgst digest.Digest) OCILayoutOption {
+	return ociLayoutOptionFunc(func(oi *OCILayoutInfo) {
+		oi.checksum = dgst
+	})
+}
+
 type OCILayoutInfo struct {
 	constraintsWrapper
 	sessionID  string
 	storeID    string
 	layerLimit *int
+	checksum   digest.Digest
 }
 
 type DiffType string

solver/pb/attr.go

@@ -34,6 +34,7 @@ const AttrImageResolveModeForcePull = "pull"
 const AttrImageResolveModePreferLocal = "local"
 const AttrImageRecordType = "image.recordtype"
 const AttrImageLayerLimit = "image.layerlimit"
+const AttrImageChecksum = "image.checksum"
 
 const AttrOCILayoutSessionID = "oci.session"
 const AttrOCILayoutStoreID = "oci.store"

solver/pb/caps.go

@@ -12,6 +12,7 @@ const (
 	CapSourceImage            apicaps.CapID = "source.image"
 	CapSourceImageResolveMode apicaps.CapID = "source.image.resolvemode"
 	CapSourceImageLayerLimit  apicaps.CapID = "source.image.layerlimit"
+	CapSourceImageChecksum    apicaps.CapID = "source.image.checksum"
 
 	CapSourceLocal                apicaps.CapID = "source.local"
 	CapSourceLocalUnique          apicaps.CapID = "source.local.unique"
@@ -128,6 +129,12 @@ func init() {
 		Status:  apicaps.CapStatusExperimental,
 	})
 
+	Caps.Init(apicaps.Cap{
+		ID:      CapSourceImageChecksum,
+		Enabled: true,
+		Status:  apicaps.CapStatusExperimental,
+	})
+
 	Caps.Init(apicaps.Cap{
 		ID:      CapSourceLocal,
 		Enabled: true,

source/containerimage/identifier.go

@@ -19,6 +19,7 @@ type ImageIdentifier struct {
 	ResolveMode resolver.ResolveMode
 	RecordType  client.UsageRecordType
 	LayerLimit  *int
+	Checksum    digest.Digest
 }
 
 func NewImageIdentifier(str string) (*ImageIdentifier, error) {

source/containerimage/pull.go

@@ -46,6 +46,7 @@ type puller struct {
 	Ref            string
 	SessionManager *session.Manager
 	layerLimit     *int
+	checksum       digest.Digest
 	vtx            solver.Vertex
 	ResolverType
 	store sourceresolver.ResolveImageConfigOptStore
@@ -130,6 +131,10 @@ func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (cach
 			return struct{}{}, err
 		}
 
+		if p.checksum != "" && p.manifest.MainManifestDesc.Digest != p.checksum {
+			return struct{}{}, errors.Errorf("image digest %s for %s does not match expected checksum %s", p.manifest.MainManifestDesc.Digest, p.Ref, p.checksum)
+		}
+
 		if ll := p.layerLimit; ll != nil {
 			if *ll > len(p.manifest.Descriptors) {
 				return struct{}{}, errors.Errorf("layer limit %d is greater than the number of layers in the image %d", *ll, len(p.manifest.Descriptors))

source/containerimage/source.go

@@ -93,6 +93,7 @@ func (is *Source) Resolve(ctx context.Context, id source.Identifier, sm *session
 		ref        reference.Spec
 		store      sourceresolver.ResolveImageConfigOptStore
 		layerLimit *int
+		checksum   digest.Digest
 	)
 	switch is.ResolverType {
 	case ResolverTypeRegistry:
@@ -108,6 +109,7 @@ func (is *Source) Resolve(ctx context.Context, id source.Identifier, sm *session
 		recordType = imageIdentifier.RecordType
 		ref = imageIdentifier.Reference
 		layerLimit = imageIdentifier.LayerLimit
+		checksum = imageIdentifier.Checksum
 	case ResolverTypeOCILayout:
 		ociIdentifier, ok := id.(*OCIIdentifier)
 		if !ok {
@@ -146,6 +148,7 @@ func (is *Source) Resolve(ctx context.Context, id source.Identifier, sm *session
 		vtx:            vtx,
 		store:          store,
 		layerLimit:     layerLimit,
+		checksum:       checksum,
 	}
 	return p, nil
 }
@@ -245,6 +248,12 @@ func (is *Source) registryIdentifier(ref string, attrs map[string]string, platfo
 				return nil, errors.Errorf("invalid layer limit %s", v)
 			}
 			id.LayerLimit = &l
+		case pb.AttrImageChecksum:
+			dgst, err := digest.Parse(v)
+			if err != nil {
+				return nil, errors.Wrapf(err, "invalid image checksum %s", v)
+			}
+			id.Checksum = dgst
 		}
 	}