Add image handler to skip non-distributable blobs.

This makes it easy for share functionality across tools to prevent
pushing layers that are not supposed to be re-distributed.

Signed-off-by: Brian Goff <cpuguy83@gmail.com>

remotes/handlers.go

@@ -253,6 +253,43 @@ func PushContent(ctx context.Context, pusher Pusher, desc ocispec.Descriptor, st
 	return nil
 }
 
+// SkipNonDistributableBlobs returns a handler that skips blobs that have a media type that is "non-distributeable".
+// An example of this kind of content would be a Windows base layer, which is not supposed to be redistributed.
+//
+// This is based on the media type of the content:
+// 	- application/vnd.oci.image.layer.nondistributable
+// 	- application/vnd.docker.image.rootfs.foreign
+func SkipNonDistributableBlobs(f images.HandlerFunc) images.HandlerFunc {
+	return func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		if images.IsNonDistributable(desc.MediaType) {
+			log.G(ctx).WithField("digest", desc.Digest).WithField("mediatype", desc.MediaType).Debug("Skipping non-distributable blob")
+			return nil, images.ErrSkipDesc
+		}
+
+		if images.IsLayerType(desc.MediaType) {
+			return nil, nil
+		}
+
+		children, err := f(ctx, desc)
+		if err != nil {
+			return nil, err
+		}
+		if len(children) == 0 {
+			return nil, nil
+		}
+
+		out := make([]ocispec.Descriptor, 0, len(children))
+		for _, child := range children {
+			if !images.IsNonDistributable(child.MediaType) {
+				out = append(out, child)
+			} else {
+				log.G(ctx).WithField("digest", child.Digest).WithField("mediatype", child.MediaType).Debug("Skipping non-distributable blob")
+			}
+		}
+		return out, nil
+	}
+}
+
 // FilterManifestByPlatformHandler allows Handler to handle non-target
 // platform's manifest and configuration data.
 func FilterManifestByPlatformHandler(f images.HandlerFunc, m platforms.Matcher) images.HandlerFunc {

remotes/handlers_test.go

@@ -18,9 +18,15 @@ package remotes
 
 import (
 	"context"
+	_ "crypto/sha256"
+	"encoding/json"
+	"sync"
 	"testing"
 
+	"github.com/containerd/containerd/content"
+	"github.com/containerd/containerd/content/local"
 	"github.com/containerd/containerd/images"
+	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
@@ -71,3 +77,141 @@ func TestContextCustomKeyPrefix(t *testing.T) {
 		}
 	})
 }
+
+func TestSkipNonDistributableBlobs(t *testing.T) {
+	ctx := context.Background()
+
+	out, err := SkipNonDistributableBlobs(images.HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		return []ocispec.Descriptor{
+			{MediaType: images.MediaTypeDockerSchema2Layer, Digest: "test:1"},
+			{MediaType: images.MediaTypeDockerSchema2LayerForeign, Digest: "test:2"},
+			{MediaType: images.MediaTypeDockerSchema2LayerForeignGzip, Digest: "test:3"},
+			{MediaType: ocispec.MediaTypeImageLayerNonDistributable, Digest: "test:4"},
+			{MediaType: ocispec.MediaTypeImageLayerNonDistributableGzip, Digest: "test:5"},
+			{MediaType: ocispec.MediaTypeImageLayerNonDistributableZstd, Digest: "test:6"},
+		}, nil
+	}))(ctx, ocispec.Descriptor{MediaType: images.MediaTypeDockerSchema2Manifest})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(out) != 1 {
+		t.Fatalf("unexpected number of descriptors returned: %d", len(out))
+	}
+	if out[0].Digest != "test:1" {
+		t.Fatalf("unexpected digest returned: %s", out[0].Digest)
+	}
+
+	dir := t.TempDir()
+	cs, err := local.NewLabeledStore(dir, newMemoryLabelStore())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	write := func(i interface{}, ref string) digest.Digest {
+		t.Helper()
+
+		data, err := json.Marshal(i)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		w, err := cs.Writer(ctx, content.WithRef(ref))
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer w.Close()
+
+		dgst := digest.SHA256.FromBytes(data)
+
+		n, err := w.Write(data)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if err := w.Commit(ctx, int64(n), dgst); err != nil {
+			t.Fatal(err)
+		}
+
+		return dgst
+	}
+
+	configDigest := write(ocispec.ImageConfig{}, "config")
+
+	manifest := ocispec.Manifest{
+		Config:    ocispec.Descriptor{Digest: configDigest, MediaType: ocispec.MediaTypeImageConfig},
+		MediaType: ocispec.MediaTypeImageManifest,
+		Layers: []ocispec.Descriptor{
+			{MediaType: images.MediaTypeDockerSchema2Layer, Digest: "test:1"},
+			{MediaType: images.MediaTypeDockerSchema2LayerForeign, Digest: "test:2"},
+			{MediaType: images.MediaTypeDockerSchema2LayerForeignGzip, Digest: "test:3"},
+			{MediaType: ocispec.MediaTypeImageLayerNonDistributable, Digest: "test:4"},
+			{MediaType: ocispec.MediaTypeImageLayerNonDistributableGzip, Digest: "test:5"},
+			{MediaType: ocispec.MediaTypeImageLayerNonDistributableZstd, Digest: "test:6"},
+		},
+	}
+
+	manifestDigest := write(manifest, "manifest")
+
+	out, err = SkipNonDistributableBlobs(images.ChildrenHandler(cs))(ctx, ocispec.Descriptor{MediaType: manifest.MediaType, Digest: manifestDigest})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(out) != 2 {
+		t.Fatalf("unexpected number of descriptors returned: %v", out)
+	}
+
+	if out[0].Digest != configDigest {
+		t.Fatalf("unexpected digest returned: %v", out[0])
+	}
+	if out[1].Digest != manifest.Layers[0].Digest {
+		t.Fatalf("unexpected digest returned: %v", out[1])
+	}
+}
+
+type memoryLabelStore struct {
+	l      sync.Mutex
+	labels map[digest.Digest]map[string]string
+}
+
+func newMemoryLabelStore() local.LabelStore {
+	return &memoryLabelStore{
+		labels: map[digest.Digest]map[string]string{},
+	}
+}
+
+func (mls *memoryLabelStore) Get(d digest.Digest) (map[string]string, error) {
+	mls.l.Lock()
+	labels := mls.labels[d]
+	mls.l.Unlock()
+
+	return labels, nil
+}
+
+func (mls *memoryLabelStore) Set(d digest.Digest, labels map[string]string) error {
+	mls.l.Lock()
+	mls.labels[d] = labels
+	mls.l.Unlock()
+
+	return nil
+}
+
+func (mls *memoryLabelStore) Update(d digest.Digest, update map[string]string) (map[string]string, error) {
+	mls.l.Lock()
+	labels, ok := mls.labels[d]
+	if !ok {
+		labels = map[string]string{}
+	}
+	for k, v := range update {
+		if v == "" {
+			delete(labels, k)
+		} else {
+			labels[k] = v
+		}
+	}
+	mls.labels[d] = labels
+	mls.l.Unlock()
+
+	return labels, nil
+}