feat: mutate pods based on pod label (#653)

* feat: mutate pods based on pod label Signed-off-by: Anish Ramasekar <[email protected]> * feat: add warning for pods missing labels Signed-off-by: Anish Ramasekar <[email protected]> * test: add e2e test for validating mutation with labelled pod Signed-off-by: Anish Ramasekar <[email protected]> * add warnings as part of audit annotations Signed-off-by: Anish Ramasekar <[email protected]> * update error message and admission response Signed-off-by: Anish Ramasekar <[email protected]> * update e2e to check for warning Signed-off-by: Anish Ramasekar <[email protected]> Signed-off-by: Anish Ramasekar <[email protected]>
Azure · Dec 7, 2022 · b52c7f9 · b52c7f9
1 parent c4c898f
commit b52c7f9
Show file tree

Hide file tree

Showing 7 changed files with 196 additions and 20 deletions.
diff --git a/pkg/webhook/webhook.go b/pkg/webhook/webhook.go
@@ -30,6 +30,13 @@ var (
 	// ProxyImageVersion is the image version of the proxy init and sidecar.
 	// This is injected via LDFLAGS in the Makefile during the build.
 	ProxyImageVersion string
+
+	PodLabelMissingWarning = fmt.Sprintf(`pod missing label '%s: "true"'. This label will be required in a future release for the webhook to mutate pod. Please add the label to the pod.`, UseWorkloadIdentityLabel)
+)
+
+const (
+	// warningAnnotationKey is the annotation key used to store warning messages in audit annotations
+	warningAnnotationKey = "mutation.azure-workload-identity.io/warning"
 )
 
 // +kubebuilder:webhook:path=/mutate-v1-pod,mutating=true,failurePolicy=ignore,groups="",resources=pods,verbs=create,versions=v1,name=mutation.azure-workload-identity.io,sideEffects=None,admissionReviewVersions=v1;v1beta1,matchPolicy=Equivalent
@@ -81,7 +88,9 @@ func NewPodMutator(client client.Client, reader client.Reader, arcCluster bool,
 }
 
 // PodMutator adds projected service account volume for incoming pods if service account is annotated
-func (m *podMutator) Handle(ctx context.Context, req admission.Request) admission.Response {
+func (m *podMutator) Handle(ctx context.Context, req admission.Request) (response admission.Response) {
+	warnings := []string{}
+	auditAnnotations := make(map[string]string)
 	pod := &corev1.Pod{}
 	timeStart := time.Now()
 
@@ -117,17 +126,32 @@ func (m *podMutator) Handle(ctx context.Context, req admission.Request) admissio
 		}
 	}
 
-	// check if the service account has the annotation
-	if !isServiceAccountAnnotated(serviceAccount) {
-		logger.Info("service account not annotated")
-		return admission.Allowed("service account not annotated")
+	// mutate pod if the service account or pod is annotated
+	// service account annotated with "azure.workload-identity/use" in annotation or label is current behavior
+	// pod labeled with "azure.workload-identity/use" is added for backward compatibility in v0.15.0 release
+	// This check will eventually be removed in a future release (tentatively v1.0.0-alpha) and all pods will be
+	// mutated. (ref: https://github.com/Azure/azure-workload-identity/issues/601)
+	serviceAccountAnnotated := isServiceAccountAnnotated(serviceAccount)
+	podLabeled := isPodLabeled(pod)
+	shouldMutatePod := serviceAccountAnnotated || podLabeled
+	if !shouldMutatePod {
+		logger.Info("pod not labeled or service account not annotated, skipping mutation")
+		return admission.Allowed("pod not labeled or service account not annotated, skipping mutation")
+	}
+	if !podLabeled {
+		warnings = append(warnings, PodLabelMissingWarning)
 	}
 
 	// only log metrics for the pods that are actually mutated
 	defer func() {
 		if m.reporter != nil {
 			m.reporter.ReportRequest(ctx, req.Namespace, time.Since(timeStart))
 		}
+		response.Warnings = warnings
+		if len(warnings) > 0 {
+			auditAnnotations[warningAnnotationKey] = strings.Join(warnings, ";")
+			response.AuditAnnotations = auditAnnotations
+		}
 	}()
 
 	if shouldInjectProxySidecar(pod) {
@@ -279,6 +303,13 @@ func isServiceAccountAnnotated(sa *corev1.ServiceAccount) bool {
 	return isLabeled || isAnnotated
 }
 
+// isPodLabeled checks if the pod has been labelled to use with workload identity
+// "azure.workload.identity/use" needs to be available either in labels
+func isPodLabeled(pod *corev1.Pod) bool {
+	_, isLabeled := pod.Labels[UseWorkloadIdentityLabel]
+	return isLabeled
+}
+
 func shouldInjectProxySidecar(pod *corev1.Pod) bool {
 	if len(pod.Annotations) == 0 {
 		return false

diff --git a/pkg/webhook/webhook_test.go b/pkg/webhook/webhook_test.go
@@ -2,10 +2,12 @@ package webhook
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"path/filepath"
 	"reflect"
 	"strconv"
+	"strings"
 	"testing"
 
 	admissionv1 "k8s.io/api/admission/v1"
@@ -24,6 +26,40 @@ var (
 	serviceAccountTokenExpiry = MinServiceAccountTokenExpiration
 )
 
+func newPod(name, namespace, serviceAccountName string, labels map[string]string) *corev1.Pod {
+	return &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: namespace,
+			Labels:    labels,
+		},
+		Spec: corev1.PodSpec{
+			ServiceAccountName: serviceAccountName,
+			InitContainers: []corev1.Container{
+				{
+					Name:  "init-container",
+					Image: "init-container-image",
+				},
+			},
+			Containers: []corev1.Container{
+				{
+					Name:  "container",
+					Image: "image",
+				},
+			},
+		},
+	}
+}
+
+func newPodRaw(name, namespace, serviceAccountName string, labels map[string]string) []byte {
+	pod := newPod(name, namespace, serviceAccountName, labels)
+	raw, err := json.Marshal(pod)
+	if err != nil {
+		panic(err)
+	}
+	return raw
+}
+
 func TestIsServiceAccountAnnotated(t *testing.T) {
 	tests := []struct {
 		name     string
@@ -744,30 +780,45 @@ func TestHandle(t *testing.T) {
 	tests := []struct {
 		name               string
 		serviceAccountName string
+		podLabels          map[string]string
 		clientObjects      []client.Object
 		readerObjects      []client.Object
+		expectedWarnings   []string
 	}{
 		{
 			name:               "service account in cache",
 			serviceAccountName: "sa",
 			clientObjects:      serviceAccounts,
 			readerObjects:      nil,
+			expectedWarnings:   []string{PodLabelMissingWarning},
 		},
 		{
 			name:               "service account not in cache",
 			serviceAccountName: "sa",
 			clientObjects:      nil,
 			readerObjects:      serviceAccounts,
+			expectedWarnings:   []string{PodLabelMissingWarning},
+		},
+		{
+			name:             "default service account in cache",
+			clientObjects:    serviceAccounts,
+			readerObjects:    nil,
+			expectedWarnings: []string{PodLabelMissingWarning},
 		},
 		{
-			name:          "default service account in cache",
-			clientObjects: serviceAccounts,
-			readerObjects: nil,
+			name:             "default service account not in cache",
+			clientObjects:    nil,
+			readerObjects:    serviceAccounts,
+			expectedWarnings: []string{PodLabelMissingWarning},
 		},
 		{
-			name:          "default service account not in cache",
-			clientObjects: nil,
-			readerObjects: serviceAccounts,
+			name: "pod has the required label, no warnings",
+			podLabels: map[string]string{
+				UseWorkloadIdentityLabel: "true",
+			},
+			clientObjects:    serviceAccounts,
+			readerObjects:    nil,
+			expectedWarnings: nil,
 		},
 	}
 
@@ -780,16 +831,14 @@ func TestHandle(t *testing.T) {
 				decoder: decoder,
 			}
 
-			raw := []byte(fmt.Sprintf(`{"apiVersion":"v1","kind":"Pod","metadata":{"name":"pod","namespace":"ns1"},"spec":{"initContainers":[{"image":"image","name":"cont1"}],"containers":[{"image":"image","name":"cont1"}],"serviceAccountName":"%s"}}`, test.serviceAccountName))
-
 			req := atypes.Request{
 				AdmissionRequest: admissionv1.AdmissionRequest{
 					Kind: metav1.GroupVersionKind{
 						Group:   "",
 						Version: "v1",
 						Kind:    "Pod",
 					},
-					Object:    runtime.RawExtension{Raw: raw},
+					Object:    runtime.RawExtension{Raw: newPodRaw("pod", "ns1", test.serviceAccountName, test.podLabels)},
 					Namespace: "ns1",
 					Operation: admissionv1.Create,
 				},
@@ -799,6 +848,20 @@ func TestHandle(t *testing.T) {
 			if !resp.Allowed {
 				t.Fatalf("expected to be allowed")
 			}
+			if len(resp.Warnings) != len(test.expectedWarnings) {
+				t.Fatalf("expected %d warnings, got %d", len(test.expectedWarnings), len(resp.Warnings))
+			}
+			for i := range resp.Warnings {
+				actual := resp.Warnings[i]
+				actualAuditAnnotations := resp.AuditAnnotations[warningAnnotationKey]
+				expected := test.expectedWarnings[i]
+				if actual != expected {
+					t.Fatalf("expected warning %d to be %s, got %s", i, expected, actual)
+				}
+				if !strings.Contains(actualAuditAnnotations, expected) {
+					t.Fatalf("expected audit annotation to contain %s, got %s", expected, actualAuditAnnotations)
+				}
+			}
 		})
 	}
 }

diff --git a/scripts/ci-e2e.sh b/scripts/ci-e2e.sh
@@ -94,9 +94,8 @@ test_helm_chart() {
     --debug \
     -v=5
   poll_webhook_readiness
-  # TODO(aramase: remove GINKGO_SKIP once helm chart is updated to use v0.12.0 to include the 
-  # new proxy probe feature.
-  GINKGO_SKIP=Proxy make test-e2e-run
+  # TODO(aramase): remove GINKGO_SKIP once helm chart is updated to use v0.15.0
+  GINKGO_SKIP=Proxy\|should.mutate.a.labeled.pod\|should.warn.if.the.pod.is.not.labeled make test-e2e-run
 
   ${HELM} upgrade --install workload-identity-webhook "${REPO_ROOT}/manifest_staging/charts/workload-identity-webhook" \
     --set image.repository="${REGISTRY:-mcr.microsoft.com/oss/azure/workload-identity/webhook}" \

diff --git a/test/e2e/helpers.go b/test/e2e/helpers.go
@@ -55,18 +55,18 @@ func createServiceAccount(c kubernetes.Interface, namespace, name string, labels
 
 // createPodWithServiceAccount creates a pod with two containers, busybox-1 and busybox-2 with customizable
 // namespace, service account, image, command, arguments, environment variables, and annotations.
-func createPodWithServiceAccount(c kubernetes.Interface, namespace, serviceAccount, image string, command, args []string, env []corev1.EnvVar, annotations map[string]string, runAsRoot bool) (*corev1.Pod, error) {
+func createPodWithServiceAccount(c kubernetes.Interface, namespace, serviceAccount, image string, command, args []string, env []corev1.EnvVar, annotations, labels map[string]string, runAsRoot bool) (*corev1.Pod, error) {
 	if arcCluster {
 		createSecretForArcCluster(c, namespace, serviceAccount)
 	}
 
-	pod := generatePodWithServiceAccount(c, namespace, serviceAccount, image, command, args, env, annotations, runAsRoot)
+	pod := generatePodWithServiceAccount(c, namespace, serviceAccount, image, command, args, env, annotations, labels, runAsRoot)
 	return createPod(c, pod)
 }
 
 // generatePodWithServiceAccount generates a pod with two containers, busybox-1 and busybox-2 with customizable
 // namespace, service account, image, command, arguments, environment variables, and annotations.
-func generatePodWithServiceAccount(c kubernetes.Interface, namespace, serviceAccount, image string, command, args []string, env []corev1.EnvVar, annotations map[string]string, runAsRoot bool) *corev1.Pod {
+func generatePodWithServiceAccount(c kubernetes.Interface, namespace, serviceAccount, image string, command, args []string, env []corev1.EnvVar, annotations, labels map[string]string, runAsRoot bool) *corev1.Pod {
 	// this is required for pod to be admitted in kubernetes 1.24+
 	contSecurityContext := &corev1.SecurityContext{
 		AllowPrivilegeEscalation: pointer.BoolPtr(false),
@@ -89,6 +89,7 @@ func generatePodWithServiceAccount(c kubernetes.Interface, namespace, serviceAcc
 			GenerateName: namespace + "-",
 			Namespace:    namespace,
 			Annotations:  annotations,
+			Labels:       labels,
 		},
 		Spec: corev1.PodSpec{
 			TerminationGracePeriodSeconds: pointer.Int64Ptr(0),

diff --git a/test/e2e/proxy_test.go b/test/e2e/proxy_test.go
@@ -44,6 +44,7 @@ var _ = ginkgo.Describe("Proxy [LinuxOnly] [AKSSoakOnly] [Exclude:Arc]", func()
 			[]string{"/bin/sh", "-c", fmt.Sprintf("az login -i -u %s --allow-no-subscriptions --debug; sleep 3600", clientID)},
 			nil,
 			proxyAnnotations,
+			nil,
 			true,
 		)
 
@@ -97,6 +98,7 @@ var _ = ginkgo.Describe("Proxy [LinuxOnly] [AKSSoakOnly] [Exclude:Arc]", func()
 			[]string{"/bin/sh", "-c", "az login -i --allow-no-subscriptions --debug; sleep 3600"},
 			nil,
 			proxyAnnotations,
+			nil,
 			true,
 		)
 

diff --git a/test/e2e/token_exchange.go b/test/e2e/token_exchange.go
@@ -54,6 +54,7 @@ var _ = ginkgo.Describe("TokenExchange [AKSSoakOnly] [Exclude:Arc]", func() {
 				Value: keyvaultSecretName,
 			}},
 			nil,
+			nil,
 			false,
 		)
 		framework.ExpectNoError(err, "failed to create pod %s in %s", pod.Name, namespace)