Merge pull request #1572 from ktock/layers

mikebrow · web-flow · commit ffe79709e4de · 2020-09-15T11:50:52.000-05:00
Limit size of additional annotation for avoiding unpack failure
diff --git a/pkg/server/image_pull.go b/pkg/server/image_pull.go
@@ -31,6 +31,7 @@ import (
 	"github.com/containerd/containerd"
 	"github.com/containerd/containerd/errdefs"
 	containerdimages "github.com/containerd/containerd/images"
+	"github.com/containerd/containerd/labels"
 	"github.com/containerd/containerd/log"
 	distribution "github.com/containerd/containerd/reference/docker"
 	"github.com/containerd/containerd/remotes/docker"
@@ -460,7 +461,7 @@ const (
 	targetDigestLabel = "containerd.io/snapshot/cri.layer-digest"
 	// targetImageLayersLabel is a label which contains layer digests contained in
 	// the target image and will be passed to snapshotters for preparing layers in
-	// parallel.
+	// parallel. Skipping some layers is allowed and only affects performance.
 	targetImageLayersLabel = "containerd.io/snapshot/cri.image-layers"
 )
 
@@ -478,15 +479,6 @@ func appendInfoHandlerWrapper(ref string) func(f containerdimages.Handler) conta
 			}
 			switch desc.MediaType {
 			case imagespec.MediaTypeImageManifest, containerdimages.MediaTypeDockerSchema2Manifest:
-				var layers string
-				for _, c := range children {
-					if containerdimages.IsLayerType(c.MediaType) {
-						layers += fmt.Sprintf("%s,", c.Digest.String())
-					}
-				}
-				if len(layers) >= 1 {
-					layers = layers[:len(layers)-1]
-				}
 				for i := range children {
 					c := &children[i]
 					if containerdimages.IsLayerType(c.MediaType) {
@@ -495,11 +487,33 @@ func appendInfoHandlerWrapper(ref string) func(f containerdimages.Handler) conta
 						}
 						c.Annotations[targetRefLabel] = ref
 						c.Annotations[targetDigestLabel] = c.Digest.String()
-						c.Annotations[targetImageLayersLabel] = layers
+						c.Annotations[targetImageLayersLabel] = getLayers(ctx, targetImageLayersLabel, children[i:], labels.Validate)
 					}
 				}
 			}
 			return children, nil
 		})
 	}
 }
+
+// getLayers returns comma-separated digests based on the passed list of
+// descriptors. The returned list contains as many digests as possible as well
+// as meets the label validation.
+func getLayers(ctx context.Context, key string, descs []imagespec.Descriptor, validate func(k, v string) error) (layers string) {
+	var item string
+	for _, l := range descs {
+		if containerdimages.IsLayerType(l.MediaType) {
+			item = l.Digest.String()
+			if layers != "" {
+				item = "," + item
+			}
+			// This avoids the label hits the size limitation.
+			if err := validate(key, layers+item); err != nil {
+				log.G(ctx).WithError(err).WithField("label", key).Debugf("%q is omitted in the layers list", l.Digest.String())
+				break
+			}
+			layers += item
+		}
+	}
+	return
+}
diff --git a/pkg/server/image_pull_test.go b/pkg/server/image_pull_test.go
@@ -17,9 +17,14 @@
 package server
 
 import (
+	"context"
 	"encoding/base64"
+	"fmt"
+	"strings"
 	"testing"
 
+	digest "github.com/opencontainers/go-digest"
+	imagespec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/stretchr/testify/assert"
 	runtime "k8s.io/cri-api/pkg/apis/runtime/v1alpha2"
 
@@ -327,3 +332,48 @@ func TestEncryptedImagePullOpts(t *testing.T) {
 		assert.Equal(t, test.expectedOpts, got)
 	}
 }
+
+func TestImageLayersLabel(t *testing.T) {
+	sampleKey := "sampleKey"
+	sampleDigest, err := digest.Parse("sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
+	assert.NoError(t, err)
+	sampleMaxSize := 300
+	sampleValidate := func(k, v string) error {
+		if (len(k) + len(v)) > sampleMaxSize {
+			return fmt.Errorf("invalid: %q: %q", k, v)
+		}
+		return nil
+	}
+
+	tests := []struct {
+		name      string
+		layersNum int
+		wantNum   int
+	}{
+		{
+			name:      "valid number of layers",
+			layersNum: 2,
+			wantNum:   2,
+		},
+		{
+			name:      "many layers",
+			layersNum: 5, // hits sampleMaxSize (300 chars).
+			wantNum:   4, // layers should be ommitted for avoiding invalid label.
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var sampleLayers []imagespec.Descriptor
+			for i := 0; i < tt.layersNum; i++ {
+				sampleLayers = append(sampleLayers, imagespec.Descriptor{
+					MediaType: imagespec.MediaTypeImageLayerGzip,
+					Digest:    sampleDigest,
+				})
+			}
+			gotS := getLayers(context.Background(), sampleKey, sampleLayers, sampleValidate)
+			got := len(strings.Split(gotS, ","))
+			assert.Equal(t, tt.wantNum, got)
+		})
+	}
+}
