annoate pods after mutation

Signed-off-by: Austin Abro <AustinAbro321@gmail.com>
zarf-dev · Jul 24, 2024 · c4c9cc7 · c4c9cc7
1 parent 05ef439
commit c4c9cc7
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/src/internal/agent/hooks/pods.go b/src/internal/agent/hooks/pods.go
@@ -66,6 +66,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 +79,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
 		patches = append(patches, operations.ReplacePatchOperation(path, replacement))
 	}
 
@@ -85,6 +91,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
 		patches = append(patches, operations.ReplacePatchOperation(path, replacement))
 	}
 
@@ -96,11 +103,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[fmt.Sprintf("zarf.dev/original-container-image-%d", idx)] = 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
@@ -85,6 +85,15 @@ func TestPodMutationWebhook(t *testing.T) {
 						"should-be":  "mutated",
 					},
 				),
+				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",
+					},
+				),
 			},
 			code: http.StatusOK,
 		},
@@ -124,6 +133,12 @@ func TestPodMutationWebhook(t *testing.T) {
 					"/metadata/labels",
 					map[string]string{"zarf-agent": "patched"},
 				),
+				operations.ReplacePatchOperation(
+					"/metadata/annotations",
+					map[string]string{
+						"zarf.dev/original-container-image-0": "nginx",
+					},
+				),
 			},
 			code: http.StatusOK,
 		},