test: pod agent unit tests

packages/zarf-agent/manifests/deployment.yaml

@@ -20,6 +20,7 @@ spec:
       imagePullSecrets:
         - name: private-registry
       priorityClassName: system-node-critical
+      serviceAccountName: zarf
       containers:
         - name: server
           image: "###ZARF_REGISTRY###/###ZARF_CONST_AGENT_IMAGE###:###ZARF_CONST_AGENT_IMAGE_TAG###"

packages/zarf-agent/manifests/role.yaml

@@ -0,0 +1,11 @@
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+  name: zarf-agent
+rules:
+- apiGroups:
+  - ""
+  resources:
+  - secrets
+  verbs:
+  - get

packages/zarf-agent/manifests/rolebinding.yaml

@@ -0,0 +1,12 @@
+apiVersion: rbac.authorization.k8s.io/v1
+kind: RoleBinding
+metadata:
+  name: zarf-agent-binding
+roleRef:
+  apiGroup: rbac.authorization.k8s.io
+  kind: Role
+  name: zarf-agent
+subjects:
+- kind: ServiceAccount
+  name: zarf
+  namespace: zarf

packages/zarf-agent/manifests/serviceaccount.yaml

@@ -0,0 +1,5 @@
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: zarf
+  namespace: zarf

packages/zarf-agent/zarf.yaml

@@ -27,6 +27,9 @@ components:
           - manifests/secret.yaml
           - manifests/deployment.yaml
           - manifests/webhook.yaml
+          - manifests/role.yaml
+          - manifests/rolebinding.yaml
+          - manifests/serviceaccount.yaml
     actions:
       onCreate:
         before:

src/config/lang/english.go

@@ -664,7 +664,7 @@ const (
 	AgentErrBadRequest             = "could not read request body: %s"
 	AgentErrBindHandler            = "Unable to bind the webhook handler"
 	AgentErrCouldNotDeserializeReq = "could not deserialize request: %s"
-	AgentErrGetState               = "failed to load zarf state from file: %w"
+	AgentErrGetState               = "failed to load zarf state: %w"
 	AgentErrHostnameMatch          = "failed to complete hostname matching: %w"
 	AgentErrImageSwap              = "Unable to swap the host for (%s)"
 	AgentErrInvalidMethod          = "invalid method only POST requests are allowed"

src/internal/agent/hooks/pods.go

@@ -5,13 +5,14 @@
 package hooks
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 
 	"github.com/defenseunicorns/zarf/src/config"
 	"github.com/defenseunicorns/zarf/src/config/lang"
 	"github.com/defenseunicorns/zarf/src/internal/agent/operations"
-	"github.com/defenseunicorns/zarf/src/internal/agent/state"
+	"github.com/defenseunicorns/zarf/src/pkg/cluster"
 	"github.com/defenseunicorns/zarf/src/pkg/message"
 	"github.com/defenseunicorns/zarf/src/pkg/transform"
 	v1 "k8s.io/api/admission/v1"
@@ -20,11 +21,15 @@ import (
 )
 
 // NewPodMutationHook creates a new instance of pods mutation hook.
-func NewPodMutationHook() operations.Hook {
+func NewPodMutationHook(ctx context.Context, cluster *cluster.Cluster) operations.Hook {
 	message.Debug("hooks.NewMutationHook()")
 	return operations.Hook{
-		Create: mutatePod,
-		Update: mutatePod,
+		Create: func(r *v1.AdmissionRequest) (*operations.Result, error) {
+			return mutatePod(ctx, r, cluster)
+		},
+		Update: func(r *v1.AdmissionRequest) (*operations.Result, error) {
+			return mutatePod(ctx, r, cluster)
+		},
 	}
 }
 
@@ -38,10 +43,9 @@ func parsePod(object []byte) (*corev1.Pod, error) {
 	return &pod, nil
 }
 
-func mutatePod(r *v1.AdmissionRequest) (*operations.Result, error) {
+func mutatePod(ctx context.Context, r *v1.AdmissionRequest, cluster *cluster.Cluster) (*operations.Result, error) {
 	message.Debugf("hooks.mutatePod()(*v1.AdmissionRequest) - %#v , %s/%s: %#v", r.Kind, r.Namespace, r.Name, r.Operation)
 
-	var patchOperations []operations.PatchOperation
 	pod, err := parsePod(r.Object.Raw)
 	if err != nil {
 		return &operations.Result{Msg: err.Error()}, nil
@@ -51,24 +55,26 @@ func mutatePod(r *v1.AdmissionRequest) (*operations.Result, error) {
 		// We've already played with this pod, just keep swimming 🐟
 		return &operations.Result{
 			Allowed:  true,
-			PatchOps: patchOperations,
+			PatchOps: []operations.PatchOperation{},
 		}, nil
 	}
 
-	// Add the zarf secret to the podspec
-	zarfSecret := []corev1.LocalObjectReference{{Name: config.ZarfImagePullSecretName}}
-	patchOperations = append(patchOperations, operations.ReplacePatchOperation("/spec/imagePullSecrets", zarfSecret))
-
-	zarfState, err := state.GetZarfStateFromAgentPod()
+	state, err := cluster.LoadZarfState(ctx)
 	if err != nil {
 		return nil, fmt.Errorf(lang.AgentErrGetState, err)
 	}
-	containerRegistryURL := zarfState.RegistryInfo.Address
+	registryURL := state.RegistryInfo.Address
+
+	var patchOperations []operations.PatchOperation
+
+	// Add the zarf secret to the podspec
+	zarfSecret := []corev1.LocalObjectReference{{Name: config.ZarfImagePullSecretName}}
+	patchOperations = append(patchOperations, operations.ReplacePatchOperation("/spec/imagePullSecrets", zarfSecret))
 
 	// update the image host for each init container
 	for idx, container := range pod.Spec.InitContainers {
 		path := fmt.Sprintf("/spec/initContainers/%d/image", idx)
-		replacement, err := transform.ImageTransformHost(containerRegistryURL, container.Image)
+		replacement, err := transform.ImageTransformHost(registryURL, container.Image)
 		if err != nil {
 			message.Warnf(lang.AgentErrImageSwap, container.Image)
 			continue // Continue, because we might as well attempt to mutate the other containers for this pod
@@ -79,7 +85,7 @@ func mutatePod(r *v1.AdmissionRequest) (*operations.Result, error) {
 	// update the image host for each ephemeral container
 	for idx, container := range pod.Spec.EphemeralContainers {
 		path := fmt.Sprintf("/spec/ephemeralContainers/%d/image", idx)
-		replacement, err := transform.ImageTransformHost(containerRegistryURL, container.Image)
+		replacement, err := transform.ImageTransformHost(registryURL, container.Image)
 		if err != nil {
 			message.Warnf(lang.AgentErrImageSwap, container.Image)
 			continue // Continue, because we might as well attempt to mutate the other containers for this pod
@@ -90,7 +96,7 @@ func mutatePod(r *v1.AdmissionRequest) (*operations.Result, error) {
 	// update the image host for each normal container
 	for idx, container := range pod.Spec.Containers {
 		path := fmt.Sprintf("/spec/containers/%d/image", idx)
-		replacement, err := transform.ImageTransformHost(containerRegistryURL, container.Image)
+		replacement, err := transform.ImageTransformHost(registryURL, container.Image)
 		if err != nil {
 			message.Warnf(lang.AgentErrImageSwap, container.Image)
 			continue // Continue, because we might as well attempt to mutate the other containers for this pod

src/internal/agent/hooks/pods_test.go

@@ -0,0 +1,149 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2021-Present The Zarf Authors
+
+package hooks
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"testing"
+
+	"github.com/defenseunicorns/zarf/src/config"
+	"github.com/defenseunicorns/zarf/src/internal/agent/http/admission"
+	"github.com/defenseunicorns/zarf/src/internal/agent/operations"
+	"github.com/defenseunicorns/zarf/src/types"
+	"github.com/stretchr/testify/require"
+	v1 "k8s.io/api/admission/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+)
+
+func createPodAdmissionRequest(t *testing.T, op v1.Operation, pod *corev1.Pod) *v1.AdmissionRequest {
+	t.Helper()
+	raw, err := json.Marshal(pod)
+	require.NoError(t, err)
+	return &v1.AdmissionRequest{
+		Operation: op,
+		Object: runtime.RawExtension{
+			Raw: raw,
+		},
+	}
+}
+
+func TestPodMutationWebhook(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	state := &types.ZarfState{RegistryInfo: types.RegistryInfo{Address: "127.0.0.1:31999"}}
+	c := createTestClientWithZarfState(ctx, t, state)
+	handler := admission.NewHandler().Serve(NewPodMutationHook(ctx, c))
+
+	tests := []struct {
+		name          string
+		admissionReq  *v1.AdmissionRequest
+		expectedPatch []operations.PatchOperation
+		code          int
+	}{
+		{
+			name: "pod with label should be mutated",
+			admissionReq: createPodAdmissionRequest(t, v1.Create, &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{"should-be": "mutated"},
+				},
+				Spec: corev1.PodSpec{
+					Containers:     []corev1.Container{{Image: "nginx"}},
+					InitContainers: []corev1.Container{{Image: "busybox"}},
+					EphemeralContainers: []corev1.EphemeralContainer{
+						{
+							EphemeralContainerCommon: corev1.EphemeralContainerCommon{
+								Image: "alpine",
+							},
+						},
+					},
+				},
+			}),
+			expectedPatch: []operations.PatchOperation{
+				operations.ReplacePatchOperation(
+					"/spec/imagePullSecrets",
+					[]corev1.LocalObjectReference{{Name: config.ZarfImagePullSecretName}},
+				),
+				operations.ReplacePatchOperation(
+					"/spec/initContainers/0/image",
+					"127.0.0.1:31999/library/busybox:latest-zarf-2140033595",
+				),
+				operations.ReplacePatchOperation(
+					"/spec/ephemeralContainers/0/image",
+					"127.0.0.1:31999/library/alpine:latest-zarf-1117969859",
+				),
+				operations.ReplacePatchOperation(
+					"/spec/containers/0/image",
+					"127.0.0.1:31999/library/nginx:latest-zarf-3793515731",
+				),
+				operations.ReplacePatchOperation(
+					"/metadata/labels/zarf-agent",
+					"patched",
+				),
+			},
+			code: http.StatusOK,
+		},
+		{
+			name: "pod with zarf-agent patched label should not be mutated",
+			admissionReq: createPodAdmissionRequest(t, v1.Create, &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{"zarf-agent": "patched"},
+				},
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{{Image: "nginx"}},
+				},
+			}),
+			expectedPatch: nil,
+			code:          http.StatusOK,
+		},
+		{
+			name: "pod with no labels should not error",
+			admissionReq: createPodAdmissionRequest(t, v1.Create, &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: nil,
+				},
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{{Image: "nginx"}},
+				},
+			}),
+			expectedPatch: []operations.PatchOperation{
+				operations.ReplacePatchOperation(
+					"/spec/imagePullSecrets",
+					[]corev1.LocalObjectReference{{Name: config.ZarfImagePullSecretName}},
+				),
+				operations.ReplacePatchOperation(
+					"/spec/containers/0/image",
+					"127.0.0.1:31999/library/nginx:latest-zarf-3793515731",
+				),
+				operations.AddPatchOperation(
+					"/metadata/labels",
+					map[string]string{"zarf-agent": "patched"},
+				),
+			},
+			code: http.StatusOK,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			resp := sendAdmissionRequest(t, tt.admissionReq, handler, tt.code)
+			if tt.expectedPatch != nil {
+				expectedPatchJSON, err := json.Marshal(tt.expectedPatch)
+				require.NoError(t, err)
+				require.NotNil(t, resp)
+				require.True(t, resp.Allowed)
+				require.JSONEq(t, string(expectedPatchJSON), string(resp.Patch))
+			} else {
+				require.Empty(t, string(resp.Patch))
+			}
+		})
+	}
+}

src/internal/agent/hooks/utils_test.go

@@ -0,0 +1,70 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: 2021-Present The Zarf Authors
+
+package hooks
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/defenseunicorns/zarf/src/pkg/cluster"
+	"github.com/defenseunicorns/zarf/src/pkg/k8s"
+	"github.com/defenseunicorns/zarf/src/types"
+	"github.com/stretchr/testify/require"
+	v1 "k8s.io/api/admission/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/kubernetes/fake"
+)
+
+func createTestClientWithZarfState(ctx context.Context, t *testing.T, state *types.ZarfState) *cluster.Cluster {
+	t.Helper()
+	c := &cluster.Cluster{K8s: &k8s.K8s{Clientset: fake.NewSimpleClientset()}}
+	stateData, err := json.Marshal(state)
+	require.NoError(t, err)
+
+	secret := &corev1.Secret{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      cluster.ZarfStateSecretName,
+			Namespace: cluster.ZarfNamespaceName,
+		},
+		Data: map[string][]byte{
+			cluster.ZarfStateDataKey: stateData,
+		},
+	}
+	_, err = c.Clientset.CoreV1().Secrets(cluster.ZarfNamespaceName).Create(ctx, secret, metav1.CreateOptions{})
+	require.NoError(t, err)
+	return c
+}
+
+// sendAdmissionRequest sends an admission request to the handler and returns the response.
+func sendAdmissionRequest(t *testing.T, admissionReq *v1.AdmissionRequest, handler http.HandlerFunc, code int) *v1.AdmissionResponse {
+	t.Helper()
+
+	b, err := json.Marshal(&v1.AdmissionReview{
+		Request: admissionReq,
+	})
+	require.NoError(t, err)
+
+	// Note: The URL ("/test") doesn't matter here because we are directly invoking the handler.
+	// The handler processes the request based on the HTTP method and body content, not the URL path.
+	req := httptest.NewRequest(http.MethodPost, "/test", bytes.NewReader(b))
+	req.Header.Set("Content-Type", "application/json")
+
+	rr := httptest.NewRecorder()
+	handler.ServeHTTP(rr, req)
+
+	require.Equal(t, code, rr.Code)
+
+	var admissionReview v1.AdmissionReview
+	if rr.Code == http.StatusOK {
+		err = json.NewDecoder(rr.Body).Decode(&admissionReview)
+		require.NoError(t, err)
+	}
+
+	return admissionReview.Response
+}