rename image annotation

zarf-dev · Jul 24, 2024 · e7750db · e7750db
1 parent a00f0dd
commit e7750db
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 11 deletions.
diff --git a/src/internal/agent/hooks/pods.go b/src/internal/agent/hooks/pods.go
@@ -20,6 +20,8 @@ import (
 	corev1 "k8s.io/api/core/v1"
 )
 
+const annotationPrefix = "zarf.dev"
+
 // NewPodMutationHook creates a new instance of pods mutation hook.
 func NewPodMutationHook(ctx context.Context, cluster *cluster.Cluster) operations.Hook {
 	return operations.Hook{
@@ -40,6 +42,10 @@ func parsePod(object []byte) (*corev1.Pod, error) {
 	return &pod, nil
 }
 
+func getImageAnnotationKey(containerName string) string {
+	return fmt.Sprintf("%s/original-image-%s", annotationPrefix, containerName)
+}
+
 func mutatePod(ctx context.Context, r *v1.AdmissionRequest, cluster *cluster.Cluster) (*operations.Result, error) {
 	pod, err := parsePod(r.Object.Raw)
 	if err != nil {
@@ -79,7 +85,7 @@ func mutatePod(ctx context.Context, r *v1.AdmissionRequest, cluster *cluster.Clu
 			message.Warnf(lang.AgentErrImageSwap, container.Image)
 			continue // Continue, because we might as well attempt to mutate the other containers for this pod
 		}
-		updatedAnnotations[fmt.Sprintf("zarf.dev/original-init-image-%d", idx)] = container.Image
+		updatedAnnotations[getImageAnnotationKey(container.Name)] = container.Image
 		patches = append(patches, operations.ReplacePatchOperation(path, replacement))
 	}
 
@@ -91,7 +97,7 @@ func mutatePod(ctx context.Context, r *v1.AdmissionRequest, cluster *cluster.Clu
 			message.Warnf(lang.AgentErrImageSwap, container.Image)
 			continue // Continue, because we might as well attempt to mutate the other containers for this pod
 		}
-		updatedAnnotations[fmt.Sprintf("zarf.dev/original-ephemeral-image-%d", idx)] = container.Image
+		updatedAnnotations[getImageAnnotationKey(container.Name)] = container.Image
 		patches = append(patches, operations.ReplacePatchOperation(path, replacement))
 	}
 
@@ -103,7 +109,7 @@ func mutatePod(ctx context.Context, r *v1.AdmissionRequest, cluster *cluster.Clu
 			message.Warnf(lang.AgentErrImageSwap, container.Image)
 			continue // Continue, because we might as well attempt to mutate the other containers for this pod
 		}
-		updatedAnnotations[fmt.Sprintf("zarf.dev/original-container-image-%d", idx)] = container.Image
+		updatedAnnotations[getImageAnnotationKey(container.Name)] = container.Image
 		patches = append(patches, operations.ReplacePatchOperation(path, replacement))
 	}
 

diff --git a/src/internal/agent/hooks/pods_test.go b/src/internal/agent/hooks/pods_test.go
@@ -50,11 +50,12 @@ func TestPodMutationWebhook(t *testing.T) {
 					Annotations: map[string]string{"should-be": "mutated"},
 				},
 				Spec: corev1.PodSpec{
-					Containers:     []corev1.Container{{Image: "nginx"}},
-					InitContainers: []corev1.Container{{Image: "busybox"}},
+					Containers:     []corev1.Container{{Name: "nginx", Image: "nginx"}},
+					InitContainers: []corev1.Container{{Name: "different", Image: "busybox"}},
 					EphemeralContainers: []corev1.EphemeralContainer{
 						{
 							EphemeralContainerCommon: corev1.EphemeralContainerCommon{
+								Name:  "alpine",
 								Image: "alpine",
 							},
 						},
@@ -88,10 +89,10 @@ func TestPodMutationWebhook(t *testing.T) {
 				operations.ReplacePatchOperation(
 					"/metadata/annotations",
 					map[string]string{
-						"zarf.dev/original-container-image-0": "nginx",
-						"zarf.dev/original-ephemeral-image-0": "alpine",
-						"zarf.dev/original-init-image-0":      "busybox",
-						"should-be":                           "mutated",
+						"zarf.dev/original-image-nginx":     "nginx",
+						"zarf.dev/original-image-alpine":    "alpine",
+						"zarf.dev/original-image-different": "busybox",
+						"should-be":                         "mutated",
 					},
 				),
 			},
@@ -117,7 +118,7 @@ func TestPodMutationWebhook(t *testing.T) {
 					Labels: nil,
 				},
 				Spec: corev1.PodSpec{
-					Containers: []corev1.Container{{Image: "nginx"}},
+					Containers: []corev1.Container{{Name: "nginx", Image: "nginx"}},
 				},
 			}),
 			patch: []operations.PatchOperation{
@@ -136,7 +137,7 @@ func TestPodMutationWebhook(t *testing.T) {
 				operations.ReplacePatchOperation(
 					"/metadata/annotations",
 					map[string]string{
-						"zarf.dev/original-container-image-0": "nginx",
+						"zarf.dev/original-image-nginx": "nginx",
 					},
 				),
 			},