From 5058a673bef1496195afd9e193a64030df6f28ba Mon Sep 17 00:00:00 2001
From: Austin Abro <37223396+AustinAbro321@users.noreply.github.com>
Date: Thu, 25 Jul 2024 04:44:50 -0400
Subject: [PATCH] feat: annotate image mutation (#2755)

Signed-off-by: Austin Abro <AustinAbro321@gmail.com>
Signed-off-by: Philip Laine <philip.laine@gmail.com>
Co-authored-by: Philip Laine <philip.laine@gmail.com>
---
 src/internal/agent/hooks/pods.go      | 16 ++++++++++++++++
 src/internal/agent/hooks/pods_test.go | 22 +++++++++++++++++++---
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/internal/agent/hooks/pods.go b/src/internal/agent/hooks/pods.go
index 1eaccb5fbb..60362785a5 100644
--- 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 {
@@ -66,6 +72,11 @@ func mutatePod(ctx context.Context, r *v1.AdmissionRequest, cluster *cluster.Clu
 	zarfSecret := []corev1.LocalObjectReference{{Name: config.ZarfImagePullSecretName}}
 	patches = append(patches, operations.ReplacePatchOperation("/spec/imagePullSecrets", zarfSecret))
 
+	updatedAnnotations := pod.Annotations
+	if updatedAnnotations == nil {
+		updatedAnnotations = make(map[string]string)
+	}
+
 	// update the image host for each init container
 	for idx, container := range pod.Spec.InitContainers {
 		path := fmt.Sprintf("/spec/initContainers/%d/image", idx)
@@ -74,6 +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[getImageAnnotationKey(container.Name)] = container.Image
 		patches = append(patches, operations.ReplacePatchOperation(path, replacement))
 	}
 
@@ -85,6 +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[getImageAnnotationKey(container.Name)] = container.Image
 		patches = append(patches, operations.ReplacePatchOperation(path, replacement))
 	}
 
@@ -96,11 +109,14 @@ 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[getImageAnnotationKey(container.Name)] = container.Image
 		patches = append(patches, operations.ReplacePatchOperation(path, replacement))
 	}
 
 	patches = append(patches, getLabelPatch(pod.Labels))
 
+	patches = append(patches, operations.ReplacePatchOperation("/metadata/annotations", updatedAnnotations))
+
 	return &operations.Result{
 		Allowed:  true,
 		PatchOps: patches,
diff --git a/src/internal/agent/hooks/pods_test.go b/src/internal/agent/hooks/pods_test.go
index dafa786f8c..60bd09d0f5 100644
--- 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",
 							},
 						},
@@ -85,6 +86,15 @@ func TestPodMutationWebhook(t *testing.T) {
 						"should-be":  "mutated",
 					},
 				),
+				operations.ReplacePatchOperation(
+					"/metadata/annotations",
+					map[string]string{
+						"zarf.dev/original-image-nginx":     "nginx",
+						"zarf.dev/original-image-alpine":    "alpine",
+						"zarf.dev/original-image-different": "busybox",
+						"should-be":                         "mutated",
+					},
+				),
 			},
 			code: http.StatusOK,
 		},
@@ -108,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{
@@ -124,6 +134,12 @@ func TestPodMutationWebhook(t *testing.T) {
 					"/metadata/labels",
 					map[string]string{"zarf-agent": "patched"},
 				),
+				operations.ReplacePatchOperation(
+					"/metadata/annotations",
+					map[string]string{
+						"zarf.dev/original-image-nginx": "nginx",
+					},
+				),
 			},
 			code: http.StatusOK,
 		},