fix webhook to deny privileged containers

pkg/scheduler/webhook.go

@@ -19,6 +19,7 @@ package scheduler
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 
 	corev1 "k8s.io/api/core/v1"
@@ -68,15 +69,10 @@ func (h *webhook) Handle(_ context.Context, req admission.Request) admission.Res
 		return admission.Allowed("pod already has different scheduler assigned")
 	}
 	klog.V(5).Infof(template, pod.Namespace, pod.Name, pod.UID)
+	privilegedName, hasPrivileged := privilegedContainerName(pod)
 	hasResource := false
-	for idx, ctr := range pod.Spec.Containers {
+	for idx := range pod.Spec.Containers {
 		c := &pod.Spec.Containers[idx]
-		if ctr.SecurityContext != nil {
-			if ctr.SecurityContext.Privileged != nil && *ctr.SecurityContext.Privileged {
-				klog.Warningf(template+" - Denying admission as container %s is privileged", pod.Namespace, pod.Name, pod.UID, c.Name)
-				continue
-			}
-		}
 		for _, val := range device.GetDevices() {
 			found, err := val.MutateAdmission(c, pod)
 			if err != nil {
@@ -86,6 +82,10 @@ func (h *webhook) Handle(_ context.Context, req admission.Request) admission.Res
 			hasResource = hasResource || found
 		}
 	}
+	if hasPrivileged && hasResource {
+		klog.Warningf(template+" - Denying admission as container %s is privileged", pod.Namespace, pod.Name, pod.UID, privilegedName)
+		return admission.Denied(fmt.Sprintf("container %s is privileged", privilegedName))
+	}
 
 	if !hasResource {
 		klog.V(3).Infof(template+" - Allowing admission: no GPU resource found", pod.Namespace, pod.Name, pod.UID)
@@ -108,6 +108,26 @@ func (h *webhook) Handle(_ context.Context, req admission.Request) admission.Res
 	return admission.PatchResponseFromRaw(req.Object.Raw, marshaledPod)
 }
 
+func privilegedContainerName(pod *corev1.Pod) (string, bool) {
+	for _, ctr := range pod.Spec.InitContainers {
+		if isPrivilegedContainer(&ctr) {
+			return ctr.Name, true
+		}
+	}
+	for _, ctr := range pod.Spec.Containers {
+		if isPrivilegedContainer(&ctr) {
+			return ctr.Name, true
+		}
+	}
+	return "", false
+}
+
+func isPrivilegedContainer(ctr *corev1.Container) bool {
+	return ctr.SecurityContext != nil &&
+		ctr.SecurityContext.Privileged != nil &&
+		*ctr.SecurityContext.Privileged
+}
+
 func fitResourceQuota(pod *corev1.Pod) bool {
 	for deviceName, dev := range device.GetDevices() {
 		// Only supports NVIDIA

pkg/scheduler/webhook_test.go

@@ -18,6 +18,7 @@ package scheduler
 
 import (
 	"context"
+	"strings"
 	"testing"
 
 	admissionv1 "k8s.io/api/admission/v1"
@@ -642,3 +643,171 @@ func TestSchedulerNameEmptyNoOverwrite(t *testing.T) {
 		t.Fatalf("Expected schedulerName patch to %q, got patches: %+v", config.SchedulerName, resp.Patches)
 	}
 }
+
+func TestPrivilegedContainerDenied(t *testing.T) {
+	prevSchedulerName := config.SchedulerName
+	prevDevicesMap := device.DevicesMap
+	prevDevicesToHandle := device.DevicesToHandle
+	t.Cleanup(func() {
+		config.SchedulerName = prevSchedulerName
+		device.DevicesMap = prevDevicesMap
+		device.DevicesToHandle = prevDevicesToHandle
+	})
+
+	config.SchedulerName = "hami-scheduler"
+	sConfig := &config.Config{
+		NvidiaConfig: nvidia.NvidiaConfig{
+			ResourceCountName:            "hami.io/gpu",
+			ResourceMemoryName:           "hami.io/gpumem",
+			ResourceMemoryPercentageName: "hami.io/gpumem-percentage",
+			ResourceCoreName:             "hami.io/gpucores",
+			DefaultMemory:                0,
+			DefaultCores:                 0,
+			DefaultGPUNum:                1,
+		},
+	}
+	if err := config.InitDevicesWithConfig(sConfig); err != nil {
+		t.Fatalf("Failed to initialize devices with config: %v", err)
+	}
+
+	privileged := true
+	testCases := []struct {
+		name    string
+		pod     *corev1.Pod
+		allowed bool
+	}{
+		{
+			name: "privileged container only without gpu",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{Name: "privileged-pod", Namespace: "default"},
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name: "privileged",
+							SecurityContext: &corev1.SecurityContext{
+								Privileged: &privileged,
+							},
+						},
+					},
+				},
+			},
+			allowed: true,
+		},
+		{
+			name: "privileged sidecar with gpu workload",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{Name: "mixed-pod", Namespace: "default"},
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Name: "privileged-sidecar",
+							SecurityContext: &corev1.SecurityContext{
+								Privileged: &privileged,
+							},
+						},
+						{
+							Name: "gpu-workload",
+							Resources: corev1.ResourceRequirements{
+								Limits: corev1.ResourceList{
+									"hami.io/gpu": resource.MustParse("1"),
+								},
+							},
+						},
+					},
+				},
+			},
+			allowed: false,
+		},
+		{
+			name: "privileged init container with gpu workload",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{Name: "init-privileged-pod", Namespace: "default"},
+				Spec: corev1.PodSpec{
+					InitContainers: []corev1.Container{
+						{
+							Name: "privileged-init",
+							SecurityContext: &corev1.SecurityContext{
+								Privileged: &privileged,
+							},
+						},
+					},
+					Containers: []corev1.Container{
+						{
+							Name: "gpu-workload",
+							Resources: corev1.ResourceRequirements{
+								Limits: corev1.ResourceList{
+									"hami.io/gpu": resource.MustParse("1"),
+								},
+							},
+						},
+					},
+				},
+			},
+			allowed: false,
+		},
+		{
+			name: "privileged pod with different scheduler",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{Name: "other-scheduler-pod", Namespace: "default"},
+				Spec: corev1.PodSpec{
+					SchedulerName: "other-scheduler",
+					Containers: []corev1.Container{
+						{
+							Name: "privileged",
+							SecurityContext: &corev1.SecurityContext{
+								Privileged: &privileged,
+							},
+						},
+					},
+				},
+			},
+			allowed: true,
+		},
+	}
+
+	wh, err := NewWebHook()
+	if err != nil {
+		t.Fatalf("Error creating WebHook: %v", err)
+	}
+
+	scheme := runtime.NewScheme()
+	corev1.AddToScheme(scheme)
+	codec := serializer.NewCodecFactory(scheme).LegacyCodec(corev1.SchemeGroupVersion)
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			podBytes, err := runtime.Encode(codec, tc.pod)
+			if err != nil {
+				t.Fatalf("Error encoding pod: %v", err)
+			}
+
+			req := admission.Request{
+				AdmissionRequest: admissionv1.AdmissionRequest{
+					UID:       "test-uid",
+					Namespace: tc.pod.Namespace,
+					Name:      tc.pod.Name,
+					Object: runtime.RawExtension{
+						Raw: podBytes,
+					},
+				},
+			}
+
+			resp := wh.Handle(context.Background(), req)
+			if tc.allowed {
+				if !resp.Allowed {
+					t.Fatalf("Expected allowed response, but got denied: %+v", resp.Result)
+				}
+				return
+			}
+			if resp.Allowed {
+				t.Fatalf("Expected denied response for privileged pod, but got allowed with %d patches", len(resp.Patches))
+			}
+			if len(resp.Patches) != 0 {
+				t.Fatalf("Expected no patches for privileged pod, got %d", len(resp.Patches))
+			}
+			if resp.Result == nil || !strings.Contains(resp.Result.Message, "is privileged") {
+				t.Fatalf("Expected privilege denial message, got: %+v", resp.Result)
+			}
+		})
+	}
+}