host-ctr: add label flag to host-ctr pull-image.

Sandbox container image on k8s 1.29 isn't pinned and causea the Sandbox container image would be GC. we add label flag to host-ctr pull-image and pass "io.cri-containerd.pinned=pinned" label to all 1.29 varaints. Then the Sandbox contianer image will be pinned and avoid GC issue.
bottlerocket-os · Feb 6, 2024 · c73afaa · c73afaa
1 parent 6057f8e
commit c73afaa
Show file tree

Hide file tree

Showing 9 changed files with 153 additions and 20 deletions.
diff --git a/packages/kubernetes-1.23/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.23/prestart-pull-pause-ctr.conf
@@ -6,4 +6,5 @@ ExecStartPre=/usr/bin/host-ctr \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
     --registry-config=/etc/host-containers/host-ctr.toml \
-    --skip-if-image-exists=true
+    --skip-if-image-exists=true \
+    --label="io.cri-containerd.pinned=pinned"
diff --git a/packages/kubernetes-1.24/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.24/prestart-pull-pause-ctr.conf
@@ -6,4 +6,5 @@ ExecStartPre=/usr/bin/host-ctr \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
     --registry-config=/etc/host-containers/host-ctr.toml \
-    --skip-if-image-exists=true
+    --skip-if-image-exists=true \
+    --label="io.cri-containerd.pinned=pinned"
diff --git a/packages/kubernetes-1.25/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.25/prestart-pull-pause-ctr.conf
@@ -6,4 +6,5 @@ ExecStartPre=/usr/bin/host-ctr \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
     --registry-config=/etc/host-containers/host-ctr.toml \
-    --skip-if-image-exists=true
+    --skip-if-image-exists=true \
+    --label="io.cri-containerd.pinned=pinned"
diff --git a/packages/kubernetes-1.26/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.26/prestart-pull-pause-ctr.conf
@@ -6,4 +6,5 @@ ExecStartPre=/usr/bin/host-ctr \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
     --registry-config=/etc/host-containers/host-ctr.toml \
-    --skip-if-image-exists=true
+    --skip-if-image-exists=true \
+    --label="io.cri-containerd.pinned=pinned"
diff --git a/packages/kubernetes-1.27/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.27/prestart-pull-pause-ctr.conf
@@ -6,4 +6,5 @@ ExecStartPre=/usr/bin/host-ctr \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
     --registry-config=/etc/host-containers/host-ctr.toml \
-    --skip-if-image-exists=true
+    --skip-if-image-exists=true \
+    --label="io.cri-containerd.pinned=pinned"
diff --git a/packages/kubernetes-1.28/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.28/prestart-pull-pause-ctr.conf
@@ -6,4 +6,5 @@ ExecStartPre=/usr/bin/host-ctr \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
     --registry-config=/etc/host-containers/host-ctr.toml \
-    --skip-if-image-exists=true
+    --skip-if-image-exists=true \
+    --label="io.cri-containerd.pinned=pinned"
diff --git a/packages/kubernetes-1.29/prestart-pull-pause-ctr.conf b/packages/kubernetes-1.29/prestart-pull-pause-ctr.conf
@@ -6,4 +6,5 @@ ExecStartPre=/usr/bin/host-ctr \
     pull-image \
     --source=${POD_INFRA_CONTAINER_IMAGE} \
     --registry-config=/etc/host-containers/host-ctr.toml \
-    --skip-if-image-exists=true
+    --skip-if-image-exists=true \
+    --label="io.cri-containerd.pinned=pinned"
diff --git a/sources/host-ctr/cmd/host-ctr/main.go b/sources/host-ctr/cmd/host-ctr/main.go
@@ -39,6 +39,13 @@ import (
 // Example 2: 777777777777.dkr.ecr.cn-north-1.amazonaws.com.cn/my_image:latest
 var ecrRegex = regexp.MustCompile(`(^[a-zA-Z0-9][a-zA-Z0-9-_]*)\.dkr\.ecr\.([a-zA-Z0-9][a-zA-Z0-9-_]*)\.amazonaws\.com(\.cn)?.*`)
 
+// derived from: /~https://github.com/containerd/containerd/blob/f5e7fe0cb6286333a597a407029fe7c760626669/pkg/labels/validate.go#L25-L28
+const (
+	maxSize = 4096
+	// maximum length of key portion of error message if len of key + len of value > maxSize
+	keyMaxLen = 64
+)
+
 func init() {
 	rand.New(rand.NewSource(time.Now().UnixNano()))
 	// Dispatch logging output instead of writing all levels' messages to
@@ -161,9 +168,14 @@ func App() *cli.App {
 					Destination: &useCachedImage,
 					Value:       false,
 				},
+				&cli.StringSliceFlag{
+					Name:  "label",
+					Usage: "add multiple labels to the pulled image. The label should be in the format of `key=value`",
+				},
 			},
 			Action: func(c *cli.Context) error {
-				return pullImageOnly(containerdSocket, namespace, source, registryConfig, useCachedImage)
+				labels := strings.Join(c.StringSlice("label"), `, `)
+				return pullImageOnly(containerdSocket, namespace, source, registryConfig, useCachedImage, labels)
 			},
 		},
 		{
@@ -274,13 +286,14 @@ func runCtr(containerdSocket string, namespace string, containerID string, sourc
 	// Check if the image source is an ECR image. If it is, then we need to handle it with the ECR resolver.
 	isECRImage := ecrRegex.MatchString(source)
 	var img containerd.Image
+
 	if isECRImage {
-		img, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage)
+		img, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage, "")
 		if err != nil {
 			return err
 		}
 	} else {
-		img, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage)
+		img, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage, "")
 		if err != nil {
 			log.G(ctx).WithField("ref", source).Error(err)
 			return err
@@ -494,7 +507,7 @@ func runCtr(containerdSocket string, namespace string, containerID string, sourc
 }
 
 // pullImageOnly pulls the specified container image
-func pullImageOnly(containerdSocket string, namespace string, source string, registryConfigPath string, useCachedImage bool) error {
+func pullImageOnly(containerdSocket string, namespace string, source string, registryConfigPath string, useCachedImage bool, labels string) error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	ctx = namespaces.WithNamespace(ctx, namespace)
@@ -508,12 +521,12 @@ func pullImageOnly(containerdSocket string, namespace string, source string, reg
 	// Check if the image source is an ECR image. If it is, then we need to handle it with the ECR resolver.
 	isECRImage := ecrRegex.MatchString(source)
 	if isECRImage {
-		_, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage)
+		_, err = fetchECRImage(ctx, source, client, registryConfigPath, useCachedImage, labels)
 		if err != nil {
 			return err
 		}
 	} else {
-		_, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage)
+		_, err = fetchImage(ctx, source, client, registryConfigPath, useCachedImage, labels)
 		if err != nil {
 			log.G(ctx).WithField("ref", source).Error(err)
 			return err
@@ -646,7 +659,7 @@ func fetchECRRef(ctx context.Context, input string) (ecr.ECRSpec, error) {
 }
 
 // fetchECRImage does some additional conversions before resolving the image reference and fetches the image.
-func fetchECRImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, fetchCachedImageIfExist bool) (containerd.Image, error) {
+func fetchECRImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, fetchCachedImageIfExist bool, labels string) (containerd.Image, error) {
 	ecrRef, err := fetchECRRef(ctx, source)
 	if err != nil {
 		return nil, err
@@ -658,7 +671,7 @@ func fetchECRImage(ctx context.Context, source string, client *containerd.Client
 		WithField("source", source).
 		Debug("parsed ECR reference from URI")
 
-	img, err := fetchImage(ctx, ref, client, registryConfigPath, fetchCachedImageIfExist)
+	img, err := fetchImage(ctx, ref, client, registryConfigPath, fetchCachedImageIfExist, labels)
 	if err != nil {
 		log.G(ctx).WithField("ref", ref).Error(err)
 		return nil, err
@@ -971,7 +984,7 @@ func withProxyEnv() oci.SpecOpts {
 }
 
 // fetchImage returns a `containerd.Image` given an image source.
-func fetchImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, useCachedImage bool) (containerd.Image, error) {
+func fetchImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, useCachedImage bool, labels string) (containerd.Image, error) {
 	// Check the containerd image store to see if image exists
 	img, err := client.GetImage(ctx, source)
 	if err != nil {
@@ -986,11 +999,11 @@ func fetchImage(ctx context.Context, source string, client *containerd.Client, r
 		log.G(ctx).WithField("ref", source).Info("Image exists, fetching cached image from image store")
 		return img, err
 	}
-	return pullImage(ctx, source, client, registryConfigPath)
+	return pullImage(ctx, source, client, registryConfigPath, labels)
 }
 
 // pullImage pulls an image from the specified source.
-func pullImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string) (containerd.Image, error) {
+func pullImage(ctx context.Context, source string, client *containerd.Client, registryConfigPath string, labels string) (containerd.Image, error) {
 	// Handle registry config
 	var registryConfig *RegistryConfig
 	if registryConfigPath != "" {
@@ -1017,9 +1030,22 @@ func pullImage(ctx context.Context, source string, client *containerd.Client, re
 	var img containerd.Image
 	for {
 		var err error
-		img, err = client.Pull(ctx, source,
+
+		pullOpts := []containerd.RemoteOpt{
 			withDynamicResolver(ctx, source, registryConfig),
-			containerd.WithSchema1Conversion)
+			containerd.WithSchema1Conversion,
+		}
+
+		if len(labels) != 0 {
+			labelsMap, err := convertLabels(labels)
+			if err != nil {
+				return nil, err
+			}
+			pullOpts = append(pullOpts, containerd.WithPullLabels(labelsMap))
+		}
+
+		img, err = client.Pull(ctx, source, pullOpts...)
+
 		if err == nil {
 			log.G(ctx).WithField("img", img.Name()).Info("pulled image successfully")
 			break
@@ -1184,3 +1210,38 @@ func withUnmaskedPaths(unmaskPaths []string) oci.SpecOpts {
 		return nil
 	}
 }
+
+// Convert label to map[string]string for containerd.WithPullLabels.
+// Label are in the format of "key=value".
+func convertLabels(labels string) (map[string]string, error) {
+	labelsMap := make(map[string]string)
+	labelsList := strings.Split(labels, ",")
+	for _, label := range labelsList {
+		labelKeyValue := strings.Split(label, "=")
+
+		if len(labelKeyValue) != 2 {
+			return labelsMap, fmt.Errorf("invalid label format: %s. It should be key=value", label)
+		}
+		if err := Validate(labelKeyValue[0], labelKeyValue[1]); err != nil {
+			return labelsMap, err
+		}
+
+		labelsMap[labelKeyValue[0]] = labelKeyValue[1]
+	}
+	return labelsMap, nil
+}
+
+// Validate a label's key and value are under 4096 bytes
+//
+// validation logic derived from:
+// /~https://github.com/containerd/containerd/blob/f5e7fe0cb6286333a597a407029fe7c760626669/pkg/labels/validate.go#L32-L41
+func Validate(k, v string) error {
+	total := len(k) + len(v)
+	if total > maxSize {
+		if len(k) > keyMaxLen {
+			k = k[:keyMaxLen]
+		}
+		return fmt.Errorf("label key and value length (%d bytes) greater than maximum size (%d bytes), key: %s: %w", total, maxSize, k, errdefs.ErrInvalidArgument)
+	}
+	return nil
+}
diff --git a/sources/host-ctr/cmd/host-ctr/main_test.go b/sources/host-ctr/cmd/host-ctr/main_test.go
@@ -222,3 +222,68 @@ func TestFetchECRRef(t *testing.T) {
 		})
 	}
 }
+
+func TestConvertLabel(t *testing.T) {
+	tests := []struct {
+		name             string
+		labels           string
+		expectedErr      bool
+		expectedLabelMap map[string]string
+	}{
+		{
+			"valid single label",
+			"io.cri-containerd.pinned=pinned",
+			false,
+			map[string]string{
+				"io.cri-containerd.pinned": "pinned",
+			},
+		},
+		{
+			"Invalid single label",
+			"io.cri-containerd.pinned,pinned",
+			true,
+			map[string]string{
+				"io.cri-containerd.pinned": "pinned",
+			},
+		},
+		{
+			"empty label",
+			" ",
+			true,
+			map[string]string{},
+		},
+		{
+			"valid multiple labels",
+			"io.cri-containerd.pinned=pinned,io.cri-containerd.test=test",
+			false,
+			map[string]string{
+				"io.cri-containerd.pinned": "pinned",
+				"io.cri-containerd.test":   "test",
+			},
+		},
+		{
+			"invalid multiple labels",
+			"io.cri-containerd.pinned=pinned,io.cri-containerd.test,test",
+			true,
+			map[string]string{
+				"io.cri-containerd.pinned": "pinned",
+				"io.cri-containerd.test":   "test",
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.labels, func(t *testing.T) {
+			result, err := convertLabels(tc.labels)
+			if tc.expectedErr {
+				// handle error cases
+				if err == nil {
+					t.Fail()
+				}
+			} else {
+				// handle happy paths
+				assert.Equal(t, tc.expectedLabelMap, result)
+			}
+		})
+	}
+}