Merge pull request #106097 from liggitt/feature/pod-security/unique-controller-pods-validation

Update pods validation based on uniqueness of controller

staging/src/k8s.io/pod-security-admission/admission/admission.go

@@ -34,6 +34,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 	admissionapi "k8s.io/pod-security-admission/admission/api"
 	"k8s.io/pod-security-admission/admission/api/validation"
@@ -534,18 +535,16 @@ func (a *Admission) EvaluatePodsInNamespace(ctx context.Context, namespace strin
 
 		podWarnings        []string
 		podWarningsToCount = make(map[string]podCount)
+		prioritizedPods    = a.prioritizePods(pods)
 	)
-	totalPods := len(pods)
+
-	if len(pods) > a.namespaceMaxPodsToCheck {
+	totalPods := len(prioritizedPods)
-		pods = pods[0:a.namespaceMaxPodsToCheck]
+	if len(prioritizedPods) > a.namespaceMaxPodsToCheck {
+		prioritizedPods = prioritizedPods[0:a.namespaceMaxPodsToCheck]
 	}
 
-	checkedPods := len(pods)
+	checkedPods := len(prioritizedPods)
-	for i, pod := range pods {
+	for i, pod := range prioritizedPods {
-		// short-circuit on exempt runtimeclass
-		if a.exemptRuntimeClass(pod.Spec.RuntimeClassName) {
-			continue
-		}
 		r := policy.AggregateCheckResults(a.Evaluator.EvaluatePod(enforce, &pod.ObjectMeta, &pod.Spec))
 		if !r.Allowed {
 			warning := r.ForbiddenReason()
@@ -743,6 +742,36 @@ func (a *Admission) exemptRuntimeClass(runtimeClass *string) bool {
 	// TODO: consider optimizing to O(1) lookup
 	return containsString(*runtimeClass, a.Configuration.Exemptions.RuntimeClasses)
 }
+
+// Filter and prioritize pods based on runtimeclass and uniqueness of the controller respectively for evaluation.
+// The input slice is modified in place and should not be reused.
+func (a *Admission) prioritizePods(pods []*corev1.Pod) []*corev1.Pod {
+	// accumulate the list of prioritized pods in-place to avoid double-allocating
+	prioritizedPods := pods[:0]
+	// accumulate any additional replicated pods after the first one encountered for a given controller uid
+	var duplicateReplicatedPods []*corev1.Pod
+	evaluatedControllers := make(map[types.UID]bool)
+	for _, pod := range pods {
+		// short-circuit on exempt runtimeclass
+		if a.exemptRuntimeClass(pod.Spec.RuntimeClassName) {
+			continue
+		}
+		// short-circuit if pod from the same controller is evaluated
+		podOwnerControllerRef := metav1.GetControllerOfNoCopy(pod)
+		if podOwnerControllerRef == nil {
+			prioritizedPods = append(prioritizedPods, pod)
+			continue
+		}
+		if evaluatedControllers[podOwnerControllerRef.UID] {
+			duplicateReplicatedPods = append(duplicateReplicatedPods, pod)
+			continue
+		}
+		prioritizedPods = append(prioritizedPods, pod)
+		evaluatedControllers[podOwnerControllerRef.UID] = true
+	}
+	return append(prioritizedPods, duplicateReplicatedPods...)
+}
+
 func containsString(needle string, haystack []string) bool {
 	for _, s := range haystack {
 		if s == needle {

staging/src/k8s.io/pod-security-admission/admission/admission_test.go

@@ -18,6 +18,7 @@ package admission
 
 import (
 	"context"
+	"math/rand"
 	"reflect"
 	"strings"
 	"testing"
@@ -32,6 +33,8 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/runtime/schema"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/uuid"
 	admissionapi "k8s.io/pod-security-admission/admission/api"
 	"k8s.io/pod-security-admission/api"
 	"k8s.io/pod-security-admission/metrics"
@@ -434,7 +437,44 @@ func TestValidateNamespace(t *testing.T) {
 				`pod1 (and 3 other pods): message`,
 			},
 		},
-		// TODO: test for prioritizing evaluating pods from unique controllers
+		{
+			name:                 "prioritized pods",
+			exemptRuntimeClasses: []string{"runtimeclass1"},
+			newLabels:            map[string]string{api.EnforceLevelLabel: string(api.LevelRestricted)},
+			oldLabels:            map[string]string{api.EnforceLevelLabel: string(api.LevelBaseline)},
+			pods: []*corev1.Pod{
+				// ensure exempt pods don't use up the limit of evaluated pods
+				{ObjectMeta: metav1.ObjectMeta{Name: "exemptpod1", Annotations: map[string]string{"error": "message1"}}, Spec: corev1.PodSpec{RuntimeClassName: pointer.String("runtimeclass1")}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "exemptpod2", Annotations: map[string]string{"error": "message1"}}, Spec: corev1.PodSpec{RuntimeClassName: pointer.String("runtimeclass1")}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "exemptpod3", Annotations: map[string]string{"error": "message1"}}, Spec: corev1.PodSpec{RuntimeClassName: pointer.String("runtimeclass1")}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "exemptpod4", Annotations: map[string]string{"error": "message1"}}, Spec: corev1.PodSpec{RuntimeClassName: pointer.String("runtimeclass1")}},
+				// ensure replicas from the same controller don't use up limit of evaluated pods
+				{ObjectMeta: metav1.ObjectMeta{Name: "replicaset1pod1", Annotations: map[string]string{"error": "replicaset1error"}, OwnerReferences: []metav1.OwnerReference{{UID: types.UID("1"), Controller: pointer.Bool(true)}}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "replicaset1pod2", Annotations: map[string]string{"error": "replicaset1error"}, OwnerReferences: []metav1.OwnerReference{{UID: types.UID("1"), Controller: pointer.Bool(true)}}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "replicaset1pod3", Annotations: map[string]string{"error": "replicaset1error"}, OwnerReferences: []metav1.OwnerReference{{UID: types.UID("1"), Controller: pointer.Bool(true)}}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "replicaset1pod4", Annotations: map[string]string{"error": "replicaset1error"}, OwnerReferences: []metav1.OwnerReference{{UID: types.UID("1"), Controller: pointer.Bool(true)}}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "replicaset2pod1", Annotations: map[string]string{"error": "replicaset2error"}, OwnerReferences: []metav1.OwnerReference{{UID: types.UID("2"), Controller: pointer.Bool(true)}}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "replicaset2pod2", Annotations: map[string]string{"error": "replicaset2error"}, OwnerReferences: []metav1.OwnerReference{{UID: types.UID("2"), Controller: pointer.Bool(true)}}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "replicaset2pod3", Annotations: map[string]string{"error": "replicaset2error"}, OwnerReferences: []metav1.OwnerReference{{UID: types.UID("2"), Controller: pointer.Bool(true)}}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "replicaset2pod4", Annotations: map[string]string{"error": "replicaset2error"}, OwnerReferences: []metav1.OwnerReference{{UID: types.UID("2"), Controller: pointer.Bool(true)}}}},
+				// ensure unique pods are prioritized before additional replicas
+				{ObjectMeta: metav1.ObjectMeta{Name: "uniquepod1", Annotations: map[string]string{"error": "uniquemessage1"}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "uniquepod2", Annotations: map[string]string{"error": "uniquemessage2"}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "uniquepod3", Annotations: map[string]string{"error": "uniquemessage3"}}},
+				{ObjectMeta: metav1.ObjectMeta{Name: "uniquepod4", Annotations: map[string]string{"error": "uniquemessage4"}}},
+			},
+			expectAllowed:  true,
+			expectListPods: true,
+			expectEvaluate: api.LevelVersion{Level: api.LevelRestricted, Version: api.LatestVersion()},
+			expectWarnings: []string{
+				`new PodSecurity enforce level only checked against the first 4 of 12 existing pods`,
+				`existing pods in namespace "test" violate the new PodSecurity enforce level "restricted:latest"`,
+				`replicaset1pod1: replicaset1error`,
+				`replicaset2pod1: replicaset2error`,
+				`uniquepod1: uniquemessage1`,
+				`uniquepod2: uniquemessage2`,
+			},
+		},
 	}
 
 	for _, tc := range testcases {
@@ -536,7 +576,7 @@ func TestValidateNamespace(t *testing.T) {
 				t.Errorf("expected to evaluate %v, got %v", tc.expectEvaluate, evaluator.lv)
 			}
 			if !reflect.DeepEqual(result.Warnings, tc.expectWarnings) {
-				t.Errorf("expected warnings:\n%v\ngot\n%v", tc.expectWarnings, result.Warnings)
+				t.Errorf("expected warnings:\n%v\ngot\n%v", strings.Join(tc.expectWarnings, "\n"), strings.Join(result.Warnings, "\n"))
 			}
 		})
 	}
@@ -765,3 +805,60 @@ func (r *FakeRecorder) ExpectEvaluations(t *testing.T, expected []EvaluationReco
 	t.Helper()
 	assert.ElementsMatch(t, expected, r.evaluations)
 }
+
+func TestPrioritizePods(t *testing.T) {
+	isController := true
+	sampleOwnerReferences := []struct {
+		ownerRefs []metav1.OwnerReference
+	}{
+		{
+			ownerRefs: []metav1.OwnerReference{
+				{
+					UID:        uuid.NewUUID(),
+					Controller: &isController,
+				},
+			},
+		}, {
+			ownerRefs: []metav1.OwnerReference{
+				{
+					UID:        uuid.NewUUID(),
+					Controller: &isController,
+				},
+			},
+		}, {
+			ownerRefs: []metav1.OwnerReference{
+				{
+					UID:        uuid.NewUUID(),
+					Controller: &isController,
+				},
+			},
+		},
+	}
+
+	var pods []*corev1.Pod
+	randomSource := rand.NewSource(time.Now().Unix())
+	for _, sampleOwnerRef := range sampleOwnerReferences {
+		// Generate multiple pods for a controller
+		for i := 0; i < rand.New(randomSource).Intn(5)+len(sampleOwnerReferences); i++ {
+			pods = append(pods, &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					OwnerReferences: sampleOwnerRef.ownerRefs,
+				},
+				Spec: corev1.PodSpec{},
+			})
+		}
+	}
+	a := &Admission{}
+	prioritizedPods := a.prioritizePods(pods)
+	controllerRef := make(map[types.UID]bool)
+
+	for i := 0; i < len(sampleOwnerReferences); i++ {
+		if controllerRef[metav1.GetControllerOfNoCopy(prioritizedPods[i]).UID] {
+			assert.Fail(t, "Pods are not prioritized based on uniqueness of the controller")
+		}
+		controllerRef[metav1.GetControllerOfNoCopy(prioritizedPods[i]).UID] = true
+	}
+	if len(prioritizedPods) != len(pods) {
+		assert.Fail(t, "Pod count is not the same after prioritization")
+	}
+}