feat: enable mapping multiple pinned images to a common key

internal/cri/config/config.go

@@ -254,6 +254,24 @@
 	Snapshotter string `toml:"snapshotter" json:"snapshotter"`
 }
 
+// PinnedImageInfo contains the pinned image information for all the requried images
+// of the cluster. This also enables mapping multiple images to a single usage key.
+type PinnedImageInfo struct {
+	// Pinned Image corresponding to the Sandbox Images of the cluster. This is a mandatory
+	// attribute to be set and should be the image that is used for the sandbox containers.
+	// Image name should be full names including domain and tag
+	Sandbox string `toml:"sandbox" json:"sandbox"`
+
+	// AdditionalImages contains the additional images that are required for the cluster.
+	// These are non Sandbox images mapped to the respective usage set as keys
+	// Image names should be full names including domain and tag
+	// Examples:
+	//  base:
+	// 	- "docker.io/library/ubuntu:20.04"
+	// 	- "docker.io/library/ubuntu:22.04"
+	AdditionalImages map[string][]string `toml:"additional_images" json:"additionalImages"`
+}
+
 type ImageConfig struct {
 	// Snapshotter is the snapshotter used by containerd.
 	Snapshotter string `toml:"snapshotter" json:"snapshotter"`
@@ -277,7 +295,7 @@
 	//   "base": "docker.io/library/ubuntu:latest"
 	// Migrated from:
 	// (PluginConfig).SandboxImage string `toml:"sandbox_image" json:"sandboxImage"`
-	PinnedImages map[string]string `toml:"pinned_images" json:"pinned_images"`
+	PinnedImages PinnedImageInfo `toml:"pinned_images" json:"pinned_images"`
 
 	// RuntimePlatforms is map between the runtime and the image platform to
 	// use for that runtime. When resolving an image for a runtime, this

internal/cri/config/config_unix.go

@@ -31,8 +31,8 @@ func DefaultImageConfig() ImageConfig {
 		ImageDecryption: ImageDecryption{
 			KeyModel: KeyModelNode,
 		},
-		PinnedImages: map[string]string{
-			"sandbox": DefaultSandboxImage,
+		PinnedImages: PinnedImageInfo{
+			Sandbox: DefaultSandboxImage,
 		},
 		ImagePullProgressTimeout: defaultImagePullProgressTimeoutDuration.String(),
 		ImagePullWithSyncFs:      false,

internal/cri/config/config_windows.go

@@ -31,8 +31,8 @@ func DefaultImageConfig() ImageConfig {
 		ImageDecryption: ImageDecryption{
 			KeyModel: KeyModelNode,
 		},
-		PinnedImages: map[string]string{
-			"sandbox": DefaultSandboxImage,
+		PinnedImages: PinnedImageInfo{
+			Sandbox: DefaultSandboxImage,
 		},
 		ImagePullProgressTimeout: defaultImagePullProgressTimeoutDuration.String(),
 	}

internal/cri/server/images/image_pull.go

@@ -46,6 +46,7 @@ import (
 	"github.com/containerd/containerd/v2/internal/cri/annotations"
 	criconfig "github.com/containerd/containerd/v2/internal/cri/config"
 	crilabels "github.com/containerd/containerd/v2/internal/cri/labels"
+	"github.com/containerd/containerd/v2/internal/cri/util"
 	snpkg "github.com/containerd/containerd/v2/pkg/snapshotters"
 	"github.com/containerd/containerd/v2/pkg/tracing"
 	"github.com/containerd/errdefs"
@@ -348,8 +349,12 @@ func (c *CRIImageService) createOrUpdateImageReference(ctx context.Context, name
 // getLabels get image labels to be added on CRI image
 func (c *CRIImageService) getLabels(ctx context.Context, name string) map[string]string {
 	labels := map[string]string{crilabels.ImageLabelKey: crilabels.ImageLabelValue}
-	for _, pinned := range c.config.PinnedImages {
-		if pinned == name {
+	if name == c.config.PinnedImages.Sandbox {
+		labels[crilabels.PinnedImageLabelKey] = crilabels.PinnedImageLabelValue
+	}
+
+	for _, pinned := range c.config.PinnedImages.AdditionalImages {
+		if util.ArrayHasElement(pinned, name) {
 			labels[crilabels.PinnedImageLabelKey] = crilabels.PinnedImageLabelValue
 		}
 	}

internal/cri/server/images/image_pull_test.go

@@ -492,42 +492,46 @@ func TestImageGetLabels(t *testing.T) {
 	tests := []struct {
 		name          string
 		expectedLabel map[string]string
-		pinnedImages  map[string]string
+		pinnedImages  criconfig.PinnedImageInfo
 		pullImageName string
 	}{
 		{
 			name:          "pinned image labels should get added on sandbox image",
 			expectedLabel: map[string]string{labels.ImageLabelKey: labels.ImageLabelValue, labels.PinnedImageLabelKey: labels.PinnedImageLabelValue},
-			pinnedImages:  map[string]string{"sandbox": "k8s.gcr.io/pause:3.9"},
+			pinnedImages:  criconfig.PinnedImageInfo{Sandbox: "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},
-			pinnedImages:  map[string]string{"sandboxnotag": "k8s.gcr.io/pause", "sandbox": "k8s.gcr.io/pause:latest"},
+			pinnedImages:  criconfig.PinnedImageInfo{Sandbox: "k8s.gcr.io/pause:latest", AdditionalImages: map[string][]string{"sandboxnotag": {"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},
-			pinnedImages: map[string]string{
-				"sandboxtagdigest": "k8s.gcr.io/pause:3.9@sha256:45b23dee08af5e43a7fea6c4cf9c25ccf269ee113168c19722f87876677c5cb2",
-				"sandbox":          "k8s.gcr.io/pause@sha256:45b23dee08af5e43a7fea6c4cf9c25ccf269ee113168c19722f87876677c5cb2",
+			pinnedImages: criconfig.PinnedImageInfo{
+				Sandbox:          "k8s.gcr.io/pause@sha256:45b23dee08af5e43a7fea6c4cf9c25ccf269ee113168c19722f87876677c5cb2",
+				AdditionalImages: map[string][]string{"sandboxtagdigest": {"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},
-			pinnedImages:  map[string]string{"sandbox": "k8s.gcr.io/pause@sha256:45b23dee08af5e43a7fea6c4cf9c25ccf269ee113168c19722f87876677c5cb2"},
+			pinnedImages: criconfig.PinnedImageInfo{
+				Sandbox: "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},
-			pinnedImages:  map[string]string{"sandbox": "k8s.gcr.io/pause:3.9"},
+			pinnedImages: criconfig.PinnedImageInfo{
+				Sandbox: "k8s.gcr.io/pause:3.9",
+			},
 			pullImageName: "k8s.gcr.io/random:latest",
 		},
 	}

internal/cri/server/images/service.go

@@ -181,10 +181,14 @@
 	return c.imageFSPaths
 }
 
+func (r *CRIImageService) SandboxImage() string {
+	return r.config.PinnedImages.Sandbox
+}
+
 // PinnedImage is used to lookup a pinned image by name.
 // Most often used to get the "sandbox" image.
-func (c *CRIImageService) PinnedImage(name string) string {
-	return c.config.PinnedImages[name]
+func (c *CRIImageService) PinnedImage(name string) []string {
+	return c.config.PinnedImages.AdditionalImages[name]
 }
 
 // GRPCService returns a new CRI Image Service grpc server.

internal/cri/server/images/service_test.go

@@ -50,8 +50,8 @@ func newTestCRIService() (*CRIImageService, *GRPCCRIImageService) {
 }
 
 var testImageConfig = criconfig.ImageConfig{
-	PinnedImages: map[string]string{
-		"sandbox": testSandboxImage,
+	PinnedImages: criconfig.PinnedImageInfo{
+		Sandbox: testSandboxImage,
 	},
 }
 

internal/cri/server/podsandbox/controller.go

@@ -111,7 +111,8 @@ type ImageService interface {
 	GetImage(id string) (imagestore.Image, error)
 	PullImage(ctx context.Context, name string, creds func(string) (string, string, error), sc *runtime.PodSandboxConfig, runtimeHandler string) (string, error)
 	RuntimeSnapshotter(ctx context.Context, ociRuntime criconfig.Runtime) string
-	PinnedImage(string) string
+	SandboxImage() string
+	PinnedImage(string) []string
 }
 
 type Controller struct {

internal/cri/server/podsandbox/sandbox_run.go

@@ -323,7 +323,7 @@ func (c *Controller) getSandboxImageName() string {
 	// returns the name of the sandbox image used to scope pod shared resources used by the pod's containers,
 	// if empty return the default sandbox image.
 	if c.imageService != nil {
-		sandboxImage := c.imageService.PinnedImage("sandbox")
+		sandboxImage := c.imageService.SandboxImage()
 		if sandboxImage != "" {
 			return sandboxImage
 		}

internal/cri/util/util.go

@@ -44,3 +44,12 @@ func NamespacedContext() context.Context {
 func WithNamespace(ctx context.Context) context.Context {
 	return namespaces.WithNamespace(ctx, constants.K8sContainerdNamespace)
 }
+
+func ArrayHasElement(array []string, element string) bool {
+	for _, e := range array {
+		if e == element {
+			return true
+		}
+	}
+	return false
+}