From c666bd00123b1d636b42914cc46f4c3154a3d89e Mon Sep 17 00:00:00 2001
From: Tim Allclair <tallclair@google.com>
Date: Wed, 24 Apr 2019 15:32:57 -0700
Subject: [PATCH] Drop RuntimeClass from PSP when feature is disabled

---
 pkg/api/podsecuritypolicy/BUILD        |  1 +
 pkg/api/podsecuritypolicy/util.go      |  4 ++
 pkg/api/podsecuritypolicy/util_test.go | 54 ++++++++++++++++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/pkg/api/podsecuritypolicy/BUILD b/pkg/api/podsecuritypolicy/BUILD
index 0219cfc6b9f..06a384e7d3c 100644
--- a/pkg/api/podsecuritypolicy/BUILD
+++ b/pkg/api/podsecuritypolicy/BUILD
@@ -41,5 +41,6 @@ go_test(
         "//staging/src/k8s.io/apimachinery/pkg/util/diff:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/util/feature:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/util/feature/testing:go_default_library",
+        "//vendor/github.com/stretchr/testify/assert:go_default_library",
     ],
 )
diff --git a/pkg/api/podsecuritypolicy/util.go b/pkg/api/podsecuritypolicy/util.go
index 7f28a729c64..f31499a090c 100644
--- a/pkg/api/podsecuritypolicy/util.go
+++ b/pkg/api/podsecuritypolicy/util.go
@@ -38,6 +38,10 @@ func DropDisabledFields(pspSpec, oldPSPSpec *policy.PodSecurityPolicySpec) {
 	if !utilfeature.DefaultFeatureGate.Enabled(features.CSIInlineVolume) {
 		pspSpec.AllowedCSIDrivers = nil
 	}
+	if !utilfeature.DefaultFeatureGate.Enabled(features.RuntimeClass) &&
+		(oldPSPSpec == nil || oldPSPSpec.RuntimeClass == nil) {
+		pspSpec.RuntimeClass = nil
+	}
 }
 
 func allowedProcMountTypesInUse(oldPSPSpec *policy.PodSecurityPolicySpec) bool {
diff --git a/pkg/api/podsecuritypolicy/util_test.go b/pkg/api/podsecuritypolicy/util_test.go
index 1527fdb67a0..8816a2e9d74 100644
--- a/pkg/api/podsecuritypolicy/util_test.go
+++ b/pkg/api/podsecuritypolicy/util_test.go
@@ -21,6 +21,8 @@ import (
 	"reflect"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+
 	"k8s.io/apimachinery/pkg/util/diff"
 	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	utilfeaturetesting "k8s.io/apiserver/pkg/util/feature/testing"
@@ -276,3 +278,55 @@ func TestDropSysctls(t *testing.T) {
 		}
 	}
 }
+
+func TestDropRuntimeClass(t *testing.T) {
+	type testcase struct {
+		name                string
+		featureEnabled      bool
+		pspSpec, oldPSPSpec *policy.PodSecurityPolicySpec
+		expectRuntimeClass  bool
+	}
+	tests := []testcase{}
+	pspGenerator := func(withRuntimeClass bool) *policy.PodSecurityPolicySpec {
+		psp := &policy.PodSecurityPolicySpec{}
+		if withRuntimeClass {
+			psp.RuntimeClass = &policy.RuntimeClassStrategyOptions{
+				AllowedRuntimeClassNames: []string{policy.AllowAllRuntimeClassNames},
+			}
+		}
+		return psp
+	}
+	for _, enabled := range []bool{true, false} {
+		for _, hasRuntimeClass := range []bool{true, false} {
+			tests = append(tests, testcase{
+				name:               fmt.Sprintf("create feature:%t hasRC:%t", enabled, hasRuntimeClass),
+				featureEnabled:     enabled,
+				pspSpec:            pspGenerator(hasRuntimeClass),
+				expectRuntimeClass: enabled && hasRuntimeClass,
+			})
+			for _, hadRuntimeClass := range []bool{true, false} {
+				tests = append(tests, testcase{
+					name:               fmt.Sprintf("update feature:%t hasRC:%t hadRC:%t", enabled, hasRuntimeClass, hadRuntimeClass),
+					featureEnabled:     enabled,
+					pspSpec:            pspGenerator(hasRuntimeClass),
+					oldPSPSpec:         pspGenerator(hadRuntimeClass),
+					expectRuntimeClass: hasRuntimeClass && (enabled || hadRuntimeClass),
+				})
+			}
+		}
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			defer utilfeaturetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.RuntimeClass, test.featureEnabled)()
+
+			DropDisabledFields(test.pspSpec, test.oldPSPSpec)
+
+			if test.expectRuntimeClass {
+				assert.NotNil(t, test.pspSpec.RuntimeClass)
+			} else {
+				assert.Nil(t, test.pspSpec.RuntimeClass)
+			}
+		})
+	}
+}