imageutil: validate source on resolving existing blob by digest

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
(cherry picked from commit 04fc13480d8ef253e0c3cc4c1f889455aefbdcf7)

client/client_test.go

@@ -161,6 +161,7 @@ func TestIntegration(t *testing.T) {
 		testUncompressedLocalCacheImportExport,
 		testUncompressedRegistryCacheImportExport,
 		testStargzLazyRegistryCacheImportExport,
+		testValidateDigestOrigin,
 	)
 	tests = append(tests, diffOpTestCases()...)
 	integration.Run(t, tests, mirrors)
@@ -5728,6 +5729,72 @@ func testBuildInfoNoExport(t *testing.T, sb integration.Sandbox) {
 	require.Equal(t, exbi.Sources[0].Ref, "docker.io/library/busybox:latest")
 }
 
+func testValidateDigestOrigin(t *testing.T, sb integration.Sandbox) {
+	integration.SkipIfDockerd(t, sb, "direct push")
+	requiresLinux(t)
+	c, err := New(sb.Context(), sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	st := llb.Image("busybox:latest").Run(llb.Shlex("touch foo"), llb.Dir("/wd")).AddMount("/wd", llb.Scratch())
+
+	def, err := st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	registry, err := sb.NewRegistry()
+	if errors.Is(err, integration.ErrRequirements) {
+		t.Skip(err.Error())
+	}
+	require.NoError(t, err)
+
+	target := registry + "/buildkit/testdigest:latest"
+
+	resp, err := c.Solve(sb.Context(), def, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type: ExporterImage,
+				Attrs: map[string]string{
+					"name": target,
+					"push": "true",
+				},
+			},
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	dgst, ok := resp.ExporterResponse[exptypes.ExporterImageDigestKey]
+	require.True(t, ok)
+
+	err = c.Prune(sb.Context(), nil, PruneAll)
+	require.NoError(t, err)
+
+	st = llb.Image(target + "@" + dgst)
+
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.NoError(t, err)
+
+	// accessing the digest from invalid names should fail
+	st = llb.Image("example.invalid/nosuchrepo@" + dgst)
+
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.Error(t, err)
+
+	// also check repo that does exists but not digest
+	st = llb.Image("docker.io/library/ubuntu@" + dgst)
+
+	def, err = st.Marshal(sb.Context())
+	require.NoError(t, err)
+
+	_, err = c.Solve(sb.Context(), def, SolveOpt{}, nil)
+	require.Error(t, err)
+}
+
 func tmpdir(appliers ...fstest.Applier) (string, error) {
 	tmpdir, err := ioutil.TempDir("", "buildkit-client")
 	if err != nil {

util/contentutil/buffer.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"io/ioutil"
+	"strings"
 	"sync"
 	"time"
 
@@ -18,22 +19,74 @@ import (
 type Buffer interface {
 	content.Provider
 	content.Ingester
+	content.Manager
 }
 
 // NewBuffer returns a new buffer
 func NewBuffer() Buffer {
 	return &buffer{
 		buffers: map[digest.Digest][]byte{},
+		infos:   map[digest.Digest]content.Info{},
 		refs:    map[string]struct{}{},
 	}
 }
 
 type buffer struct {
 	mu      sync.Mutex
 	buffers map[digest.Digest][]byte
+	infos   map[digest.Digest]content.Info
 	refs    map[string]struct{}
 }
 
+func (b *buffer) Info(ctx context.Context, dgst digest.Digest) (content.Info, error) {
+	b.mu.Lock()
+	v, ok := b.infos[dgst]
+	b.mu.Unlock()
+	if !ok {
+		return content.Info{}, errdefs.ErrNotFound
+	}
+	return v, nil
+}
+
+func (b *buffer) Update(ctx context.Context, new content.Info, fieldpaths ...string) (content.Info, error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	updated, ok := b.infos[new.Digest]
+	if !ok {
+		return content.Info{}, errdefs.ErrNotFound
+	}
+
+	if len(fieldpaths) == 0 {
+		fieldpaths = []string{"labels"}
+	}
+
+	for _, path := range fieldpaths {
+		if strings.HasPrefix(path, "labels.") {
+			if updated.Labels == nil {
+				updated.Labels = map[string]string{}
+			}
+			key := strings.TrimPrefix(path, "labels.")
+			updated.Labels[key] = new.Labels[key]
+			continue
+		}
+		if path == "labels" {
+			updated.Labels = new.Labels
+		}
+	}
+
+	b.infos[new.Digest] = updated
+	return updated, nil
+}
+
+func (b *buffer) Walk(ctx context.Context, fn content.WalkFunc, filters ...string) error {
+	return nil // not implemented
+}
+
+func (b *buffer) Delete(ctx context.Context, dgst digest.Digest) error {
+	return nil // not implemented
+}
+
 func (b *buffer) Writer(ctx context.Context, opts ...content.WriterOpt) (content.Writer, error) {
 	var wOpts content.WriterOpts
 	for _, opt := range opts {
@@ -82,6 +135,7 @@ func (b *buffer) addValue(k digest.Digest, dt []byte) {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 	b.buffers[k] = dt
+	b.infos[k] = content.Info{Digest: k, Size: int64(len(dt))}
 }
 
 type bufferedWriter struct {

util/contentutil/buffer_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/errdefs"
+	"github.com/containerd/containerd/remotes/docker"
 	digest "github.com/opencontainers/go-digest"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/pkg/errors"
@@ -71,3 +72,49 @@ func TestReaderAt(t *testing.T) {
 	require.Equal(t, err, io.EOF)
 	require.Equal(t, "bar", string(buf[:n]))
 }
+
+func TestLabels(t *testing.T) {
+	t.Parallel()
+	ctx := context.TODO()
+
+	b := NewBuffer()
+
+	err := content.WriteBlob(ctx, b, "foo", bytes.NewBuffer([]byte("foobar")), ocispecs.Descriptor{Size: -1})
+	require.NoError(t, err)
+
+	_, err = b.Info(ctx, digest.FromBytes([]byte("abc")))
+	require.Error(t, err)
+
+	info, err := b.Info(ctx, digest.FromBytes([]byte("foobar")))
+	require.NoError(t, err)
+
+	require.Equal(t, info.Digest, digest.FromBytes([]byte("foobar")))
+
+	hf, err := docker.AppendDistributionSourceLabel(b, "docker.io/library/busybox:latest")
+	require.NoError(t, err)
+	_, err = hf.Handle(ctx, ocispecs.Descriptor{Digest: digest.FromBytes([]byte("foobar"))})
+	require.NoError(t, err)
+
+	info, err = b.Info(ctx, digest.FromBytes([]byte("foobar")))
+	require.NoError(t, err)
+	require.Equal(t, info.Digest, digest.FromBytes([]byte("foobar")))
+
+	require.Equal(t, "library/busybox", info.Labels["containerd.io/distribution.source.docker.io"])
+
+	hf, err = docker.AppendDistributionSourceLabel(b, "docker.io/library/alpine:3.15")
+	require.NoError(t, err)
+	_, err = hf.Handle(ctx, ocispecs.Descriptor{Digest: digest.FromBytes([]byte("foobar"))})
+	require.NoError(t, err)
+
+	hf, err = docker.AppendDistributionSourceLabel(b, "ghcr.io/repos/alpine:3.11")
+	require.NoError(t, err)
+	_, err = hf.Handle(ctx, ocispecs.Descriptor{Digest: digest.FromBytes([]byte("foobar"))})
+	require.NoError(t, err)
+
+	info, err = b.Info(ctx, digest.FromBytes([]byte("foobar")))
+	require.NoError(t, err)
+	require.Equal(t, info.Digest, digest.FromBytes([]byte("foobar")))
+
+	require.Equal(t, "library/alpine,library/busybox", info.Labels["containerd.io/distribution.source.docker.io"])
+	require.Equal(t, "repos/alpine", info.Labels["containerd.io/distribution.source.ghcr.io"])
+}

util/contentutil/source.go

@@ -0,0 +1,34 @@
+package contentutil
+
+import (
+	"net/url"
+	"strings"
+
+	"github.com/containerd/containerd/content"
+	"github.com/containerd/containerd/reference"
+)
+
+func HasSource(info content.Info, refspec reference.Spec) (bool, error) {
+	u, err := url.Parse("dummy://" + refspec.Locator)
+	if err != nil {
+		return false, err
+	}
+
+	if info.Labels == nil {
+		return false, nil
+	}
+
+	source, target := u.Hostname(), strings.TrimPrefix(u.Path, "/")
+	repoLabel, ok := info.Labels["containerd.io/distribution.source."+source]
+	if !ok || repoLabel == "" {
+		return false, nil
+	}
+
+	for _, repo := range strings.Split(repoLabel, ",") {
+		// the target repo is not a candidate
+		if repo == target {
+			return true, nil
+		}
+	}
+	return false, nil
+}

util/contentutil/source_test.go

@@ -0,0 +1,57 @@
+package contentutil
+
+import (
+	"testing"
+
+	"github.com/containerd/containerd/content"
+	"github.com/containerd/containerd/reference"
+	"github.com/stretchr/testify/require"
+)
+
+func TestHasSource(t *testing.T) {
+	info := content.Info{
+		Labels: map[string]string{
+			"containerd.io/distribution.source.docker.io": "library/alpine",
+		},
+	}
+	ref, err := reference.Parse("docker.io/library/alpine:latest")
+	require.NoError(t, err)
+	b, err := HasSource(info, ref)
+	require.NoError(t, err)
+	require.True(t, b)
+
+	info = content.Info{
+		Labels: map[string]string{
+			"containerd.io/distribution.source.docker.io": "library/alpine,library/ubuntu",
+		},
+	}
+	b, err = HasSource(info, ref)
+	require.NoError(t, err)
+	require.True(t, b)
+
+	info = content.Info{}
+	b, err = HasSource(info, ref)
+	require.NoError(t, err)
+	require.False(t, b)
+
+	info = content.Info{Labels: map[string]string{}}
+	b, err = HasSource(info, ref)
+	require.NoError(t, err)
+	require.False(t, b)
+
+	info = content.Info{
+		Labels: map[string]string{
+			"containerd.io/distribution.source.docker.io": "library/ubuntu",
+		},
+	}
+	b, err = HasSource(info, ref)
+	require.NoError(t, err)
+	require.False(t, b)
+
+	info = content.Info{Labels: map[string]string{
+		"containerd.io/distribution.source.ghcr.io": "library/alpine",
+	}}
+	b, err = HasSource(info, ref)
+	require.NoError(t, err)
+	require.False(t, b)
+}

util/imageutil/config.go

@@ -13,6 +13,7 @@ import (
 	"github.com/containerd/containerd/reference"
 	"github.com/containerd/containerd/remotes"
 	"github.com/containerd/containerd/remotes/docker"
+	"github.com/moby/buildkit/util/contentutil"
 	"github.com/moby/buildkit/util/leaseutil"
 	"github.com/moby/buildkit/util/resolver/limited"
 	"github.com/moby/buildkit/util/resolver/retryhandler"
@@ -24,6 +25,7 @@ import (
 type ContentCache interface {
 	content.Ingester
 	content.Provider
+	content.Manager
 }
 
 var leasesMu sync.Mutex
@@ -75,10 +77,15 @@ func Config(ctx context.Context, str string, resolver remotes.Resolver, cache Co
 	if desc.Digest != "" {
 		ra, err := cache.ReaderAt(ctx, desc)
 		if err == nil {
-			desc.Size = ra.Size()
-			mt, err := DetectManifestMediaType(ra)
+			info, err := cache.Info(ctx, desc.Digest)
 			if err == nil {
-				desc.MediaType = mt
+				if ok, err := contentutil.HasSource(info, ref); err == nil && ok {
+					desc.Size = ra.Size()
+					mt, err := DetectManifestMediaType(ra)
+					if err == nil {
+						desc.MediaType = mt
+					}
+				}
 			}
 		}
 	}
@@ -101,8 +108,14 @@ func Config(ctx context.Context, str string, resolver remotes.Resolver, cache Co
 
 	children := childrenConfigHandler(cache, platform)
 
+	dslHandler, err := docker.AppendDistributionSourceLabel(cache, ref.String())
+	if err != nil {
+		return "", nil, err
+	}
+
 	handlers := []images.Handler{
 		retryhandler.New(limited.FetchHandler(cache, fetcher, str), func(_ []byte) {}),
+		dslHandler,
 		children,
 	}
 	if err := images.Dispatch(ctx, images.Handlers(handlers...), nil, desc); err != nil {

util/pull/pull.go

@@ -90,6 +90,11 @@ func (p *Puller) tryLocalResolve(ctx context.Context) error {
 	if err != nil {
 		return err
 	}
+
+	if ok, err := contentutil.HasSource(info, p.Src); err != nil || !ok {
+		return errors.Errorf("no matching source")
+	}
+
 	desc.Size = info.Size
 	p.ref = p.Src.String()
 	ra, err := p.ContentStore.ReaderAt(ctx, desc)