Pinned image support

Signed-off-by: Aditi Sharma <adi.sky17@gmail.com>

pkg/cri/server/helpers.go

@@ -80,10 +80,6 @@ const (
 	containerKindSandbox = "sandbox"
 	// containerKindContainer is a label value indicating container is application container
 	containerKindContainer = "container"
-	// imageLabelKey is the label key indicating the image is managed by cri plugin.
-	imageLabelKey = criContainerdPrefix + ".image"
-	// imageLabelValue is the label value indicating the image is managed by cri plugin.
-	imageLabelValue = "managed"
 	// sandboxMetadataExtension is an extension name that identify metadata of sandbox in CreateContainerRequest
 	sandboxMetadataExtension = criContainerdPrefix + ".sandbox.metadata"
 	// containerMetadataExtension is an extension name that identify metadata of container in CreateContainerRequest

pkg/cri/server/image_pull.go

@@ -43,6 +43,7 @@ import (
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/pkg/cri/annotations"
 	criconfig "github.com/containerd/containerd/pkg/cri/config"
+	crilabels "github.com/containerd/containerd/pkg/cri/labels"
 	snpkg "github.com/containerd/containerd/pkg/snapshotters"
 	distribution "github.com/containerd/containerd/reference/docker"
 	"github.com/containerd/containerd/remotes/docker"
@@ -152,12 +153,15 @@ func (c *criService) PullImage(ctx context.Context, r *runtime.PullImageRequest)
 		tracing.Attribute("image.ref", ref),
 		tracing.Attribute("snapshotter.name", snapshotter),
 	)
+
+	labels := c.getLabels(ctx, ref)
+
 	pullOpts := []containerd.RemoteOpt{
 		containerd.WithSchema1Conversion, //nolint:staticcheck // Ignore SA1019. Need to keep deprecated package for compatibility.
 		containerd.WithResolver(resolver),
 		containerd.WithPullSnapshotter(snapshotter),
 		containerd.WithPullUnpack,
-		containerd.WithPullLabel(imageLabelKey, imageLabelValue),
+		containerd.WithPullLabels(labels),
 		containerd.WithMaxConcurrentDownloads(c.config.MaxConcurrentDownloads),
 		containerd.WithImageHandler(imageHandler),
 		containerd.WithUnpackOpts([]containerd.UnpackOpt{
@@ -196,7 +200,7 @@ func (c *criService) PullImage(ctx context.Context, r *runtime.PullImageRequest)
 		if r == "" {
 			continue
 		}
-		if err := c.createImageReference(ctx, r, image.Target()); err != nil {
+		if err := c.createImageReference(ctx, r, image.Target(), labels); err != nil {
 			return nil, fmt.Errorf("failed to create image reference %q: %w", r, err)
 		}
 		// Update image store to reflect the newest state in containerd.
@@ -265,12 +269,12 @@ func ParseAuth(auth *runtime.AuthConfig, host string) (string, string, error) {
 // Note that because create and update are not finished in one transaction, there could be race. E.g.
 // the image reference is deleted by someone else after create returns already exists, but before update
 // happens.
-func (c *criService) createImageReference(ctx context.Context, name string, desc imagespec.Descriptor) error {
+func (c *criService) createImageReference(ctx context.Context, name string, desc imagespec.Descriptor, labels map[string]string) error {
 	img := containerdimages.Image{
 		Name:   name,
 		Target: desc,
 		// Add a label to indicate that the image is managed by the cri plugin.
-		Labels: map[string]string{imageLabelKey: imageLabelValue},
+		Labels: labels,
 	}
 	// TODO(random-liu): Figure out which is the more performant sequence create then update or
 	// update then create.
@@ -278,13 +282,31 @@ func (c *criService) createImageReference(ctx context.Context, name string, desc
 	if err == nil || !errdefs.IsAlreadyExists(err) {
 		return err
 	}
-	if oldImg.Target.Digest == img.Target.Digest && oldImg.Labels[imageLabelKey] == imageLabelValue {
+	if oldImg.Target.Digest == img.Target.Digest && oldImg.Labels[crilabels.ImageLabelKey] == labels[crilabels.ImageLabelKey] {
 		return nil
 	}
-	_, err = c.client.ImageService().Update(ctx, img, "target", "labels."+imageLabelKey)
+	_, err = c.client.ImageService().Update(ctx, img, "target", "labels."+crilabels.ImageLabelKey)
 	return err
 }
 
+// getLabels get image labels to be added on CRI image
+func (c *criService) getLabels(ctx context.Context, name string) map[string]string {
+	labels := map[string]string{crilabels.ImageLabelKey: crilabels.ImageLabelValue}
+	configSandboxImage := c.config.SandboxImage
+	// parse sandbox image
+	sandboxNamedRef, err := distribution.ParseDockerRef(configSandboxImage)
+	if err != nil {
+		log.G(ctx).Errorf("failed to parse sandbox image from config %s", sandboxNamedRef)
+		return nil
+	}
+	sandboxRef := sandboxNamedRef.String()
+	// Adding pinned image label to sandbox image
+	if sandboxRef == name {
+		labels[crilabels.PinnedImageLabelKey] = crilabels.PinnedImageLabelValue
+	}
+	return labels
+}
+
 // updateImage updates image store to reflect the newest state of an image reference
 // in containerd. If the reference is not managed by the cri plugin, the function also
 // generates necessary metadata for the image and make it managed.
@@ -293,7 +315,7 @@ func (c *criService) updateImage(ctx context.Context, r string) error {
 	if err != nil && !errdefs.IsNotFound(err) {
 		return fmt.Errorf("get image by reference: %w", err)
 	}
-	if err == nil && img.Labels()[imageLabelKey] != imageLabelValue {
+	if err == nil && img.Labels()[crilabels.ImageLabelKey] != crilabels.ImageLabelValue {
 		// Make sure the image has the image id as its unique
 		// identifier that references the image in its lifetime.
 		configDesc, err := img.Config(ctx)
@@ -301,14 +323,15 @@ func (c *criService) updateImage(ctx context.Context, r string) error {
 			return fmt.Errorf("get image id: %w", err)
 		}
 		id := configDesc.Digest.String()
-		if err := c.createImageReference(ctx, id, img.Target()); err != nil {
+		labels := c.getLabels(ctx, id)
+		if err := c.createImageReference(ctx, id, img.Target(), labels); err != nil {
 			return fmt.Errorf("create image id reference %q: %w", id, err)
 		}
 		if err := c.imageStore.Update(ctx, id); err != nil {
 			return fmt.Errorf("update image store for %q: %w", id, err)
 		}
 		// The image id is ready, add the label to mark the image as managed.
-		if err := c.createImageReference(ctx, r, img.Target()); err != nil {
+		if err := c.createImageReference(ctx, r, img.Target(), labels); err != nil {
 			return fmt.Errorf("create managed label: %w", err)
 		}
 	}

pkg/cri/server/image_pull_test.go

@@ -22,11 +22,12 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
-
 	"github.com/containerd/containerd/pkg/cri/annotations"
 	criconfig "github.com/containerd/containerd/pkg/cri/config"
+	"github.com/containerd/containerd/pkg/cri/labels"
+
+	"github.com/stretchr/testify/assert"
+	runtime "k8s.io/cri-api/pkg/apis/runtime/v1"
 )
 
 func TestParseAuth(t *testing.T) {
@@ -434,3 +435,56 @@ func TestSnapshotterFromPodSandboxConfig(t *testing.T) {
 		})
 	}
 }
+func TestImageGetLabels(t *testing.T) {
+
+	criService := newTestCRIService()
+
+	tests := []struct {
+		name               string
+		expectedLabel      map[string]string
+		configSandboxImage string
+		pullImageName      string
+	}{
+		{
+			name:               "pinned image labels should get added on sandbox image",
+			expectedLabel:      map[string]string{labels.ImageLabelKey: labels.ImageLabelValue, labels.PinnedImageLabelKey: labels.PinnedImageLabelValue},
+			configSandboxImage: "k8s.gcr.io/pause:3.9",
+			pullImageName:      "k8s.gcr.io/pause:3.9",
+		},
+		{
+			name:               "pinned image labels should get added on sandbox image without tag",
+			expectedLabel:      map[string]string{labels.ImageLabelKey: labels.ImageLabelValue, labels.PinnedImageLabelKey: labels.PinnedImageLabelValue},
+			configSandboxImage: "k8s.gcr.io/pause",
+			pullImageName:      "k8s.gcr.io/pause:latest",
+		},
+		{
+			name:               "pinned image labels should get added on sandbox image specified with tag and digest both",
+			expectedLabel:      map[string]string{labels.ImageLabelKey: labels.ImageLabelValue, labels.PinnedImageLabelKey: labels.PinnedImageLabelValue},
+			configSandboxImage: "k8s.gcr.io/pause:3.9@sha256:45b23dee08af5e43a7fea6c4cf9c25ccf269ee113168c19722f87876677c5cb2",
+			pullImageName:      "k8s.gcr.io/pause@sha256:45b23dee08af5e43a7fea6c4cf9c25ccf269ee113168c19722f87876677c5cb2",
+		},
+
+		{
+			name:               "pinned image labels should get added on sandbox image specified with digest",
+			expectedLabel:      map[string]string{labels.ImageLabelKey: labels.ImageLabelValue, labels.PinnedImageLabelKey: labels.PinnedImageLabelValue},
+			configSandboxImage: "k8s.gcr.io/pause@sha256:45b23dee08af5e43a7fea6c4cf9c25ccf269ee113168c19722f87876677c5cb2",
+			pullImageName:      "k8s.gcr.io/pause@sha256:45b23dee08af5e43a7fea6c4cf9c25ccf269ee113168c19722f87876677c5cb2",
+		},
+
+		{
+			name:               "pinned image labels should not get added on other image",
+			expectedLabel:      map[string]string{labels.ImageLabelKey: labels.ImageLabelValue},
+			configSandboxImage: "k8s.gcr.io/pause:3.9",
+			pullImageName:      "k8s.gcr.io/random:latest",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			criService.config.SandboxImage = tt.configSandboxImage
+			labels := criService.getLabels(context.Background(), tt.pullImageName)
+			assert.Equal(t, tt.expectedLabel, labels)
+
+		})
+	}
+}

pkg/cri/server/image_status.go

@@ -63,12 +63,15 @@ func (c *criService) ImageStatus(ctx context.Context, r *runtime.ImageStatusRequ
 // toCRIImage converts internal image object to CRI runtime.Image.
 func toCRIImage(image imagestore.Image) *runtime.Image {
 	repoTags, repoDigests := parseImageReferences(image.References)
+
 	runtimeImage := &runtime.Image{
 		Id:          image.ID,
 		RepoTags:    repoTags,
 		RepoDigests: repoDigests,
 		Size_:       uint64(image.Size),
+		Pinned:      image.Pinned,
 	}
+
 	uid, username := getUserFromImage(image.ImageSpec.Config.User)
 	if uid != nil {
 		runtimeImage.Uid = &runtime.Int64Value{Value: *uid}