bug fix: make sure cri image is pinned when it is pulled outside cri

henry118 · henry118 · commit 1eaf0c1f04eb · 2024-02-07T16:03:40.000Z
Signed-off-by: Henry Wang &lt;henwang@amazon.com&gt;
diff --git a/integration/containerd_image_test.go b/integration/containerd_image_test.go
@@ -237,6 +237,35 @@ func TestContainerdSandboxImage(t *testing.T) {
 	assert.True(t, pimg.Pinned)
 }
 
+func TestContainerdSandboxImagePulledOutsideCRI(t *testing.T) {
+	var pauseImage = images.Get(images.Pause)
+	ctx := context.Background()
+
+	t.Log("make sure the pause image does not exist")
+	imageService.RemoveImage(&runtime.ImageSpec{Image: pauseImage})
+
+	t.Log("pull pause image")
+	_, err := containerdClient.Pull(ctx, pauseImage)
+	assert.NoError(t, err)
+
+	t.Log("pause image should be seen by cri plugin")
+	var pimg *runtime.Image
+	require.NoError(t, Eventually(func() (bool, error) {
+		pimg, err = imageService.ImageStatus(&runtime.ImageSpec{Image: pauseImage})
+		return pimg != nil, err
+	}, time.Second, 10*time.Second))
+
+	t.Log("verify pinned field is set for pause image")
+	assert.True(t, pimg.Pinned)
+
+	t.Log("make sure the pause image exist")
+	pauseImg, err := containerdClient.GetImage(ctx, pauseImage)
+	require.NoError(t, err)
+
+	t.Log("ensure correct labels are set on pause image")
+	assert.Equal(t, "pinned", pauseImg.Labels()["io.cri-containerd.pinned"])
+}
+
 func TestContainerdImageWithDockerSchema1(t *testing.T) {
 	if goruntime.GOOS == "windows" {
 		t.Skip("Skipped on Windows because the test image is not a multi-platform one.")
diff --git a/internal/cri/server/images/image_pull.go b/internal/cri/server/images/image_pull.go
@@ -319,16 +319,30 @@ func (c *CRIImageService) createImageReference(ctx context.Context, name string,
 	// TODO(random-liu): Figure out which is the more performant sequence create then update or
 	// update then create.
 	// TODO: Call CRIImageService directly
-	oldImg, err := c.images.Create(ctx, img)
+	_, err := c.images.Create(ctx, img)
 	if err == nil {
 		return nil
 	} else if !errdefs.IsAlreadyExists(err) {
 		return err
 	}
-	if oldImg.Target.Digest == img.Target.Digest && oldImg.Labels[crilabels.ImageLabelKey] == labels[crilabels.ImageLabelKey] {
+	// Retrieve oldImg from image store here because Create routine returns an
+	// empty image on ErrAlreadyExists
+	oldImg, err := c.images.Get(ctx, name)
+	if err != nil {
+		return err
+	}
+	fieldpaths := []string{"target"}
+	if oldImg.Labels[crilabels.ImageLabelKey] != labels[crilabels.ImageLabelKey] {
+		fieldpaths = append(fieldpaths, "labels."+crilabels.ImageLabelKey)
+	}
+	if oldImg.Labels[crilabels.PinnedImageLabelKey] != labels[crilabels.PinnedImageLabelKey] &&
+		labels[crilabels.PinnedImageLabelKey] == crilabels.PinnedImageLabelValue {
+		fieldpaths = append(fieldpaths, "labels."+crilabels.PinnedImageLabelKey)
+	}
+	if oldImg.Target.Digest == img.Target.Digest && len(fieldpaths) < 2 {
 		return nil
 	}
-	_, err = c.images.Update(ctx, img, "target", "labels."+crilabels.ImageLabelKey)
+	_, err = c.images.Update(ctx, img, fieldpaths...)
 	return err
 }
 
@@ -360,7 +374,7 @@ func (c *CRIImageService) UpdateImage(ctx context.Context, r string) error {
 			return fmt.Errorf("get image id: %w", err)
 		}
 		id := configDesc.Digest.String()
-		labels := c.getLabels(ctx, id)
+		labels := c.getLabels(ctx, r)
 		if err := c.createImageReference(ctx, id, img.Target(), labels); err != nil {
 			return fmt.Errorf("create image id reference %q: %w", id, err)
 		}
