kmsv2: run KDF tests in parallel

This change updates the KDF "feature flag" to be per KMS provider
instead of global to the API server.  This allows integration tests
that use distinct provider names to run in parallel.

Locally this change reduced the runtime of
test/integration/controlplane/transformation by 3.5 minutes.

Signed-off-by: Monis Khan <mok@microsoft.com>

Kubernetes-commit: 43740c0def10f22a0ab7f522c1569188913b35a3

pkg/server/options/encryptionconfig/config.go

@@ -107,24 +107,35 @@ const (
 
 var codecs serializer.CodecFactory
 
-// this atomic bool allows us to swap enablement of the KMSv2KDF feature in tests
+// this map allows us to swap enablement of the KMSv2KDF feature in tests
 // as the feature gate is now locked to true starting with v1.29
 // Note: it cannot be set by an end user
-var kdfDisabled atomic.Bool
+// KDF enablement is tracked per KMS provider to allow tests to run in parallel.
+var kdfEnabledPerKMS sync.Map // map[string]bool, KMS name -> KDF enabled
 
 // this function should only be called in tests to swap enablement of the KMSv2KDF feature
-func SetKDFForTests(b bool) func() {
-	kdfDisabled.Store(!b)
-	return func() {
-		kdfDisabled.Store(false)
+// Caller must guarantee that all KMS providers have distinct names across all tests.
+func SetKDFForTests(kmsName string, b bool) func() {
+	if len(kmsName) == 0 { // guarantee that GetKDF("") returns the default value
+		panic("empty KMS name used in test")
 	}
+	if _, loaded := kdfEnabledPerKMS.LoadOrStore(kmsName, b); loaded {
+		panic("duplicate KMS name used in test")
+	}
+	return func() { kdfEnabledPerKMS.Delete(kmsName) }
 }
 
 // this function should be used to determine enablement of the KMSv2KDF feature
 // instead of getting it from DefaultFeatureGate as the feature gate is now locked
 // to true starting with v1.29
-func GetKDF() bool {
-	return !kdfDisabled.Load()
+// to allow integration tests to run in parallel, this "feature flag" can be set
+// per KMS provider as long as all providers use distinct names.
+func GetKDF(kmsName string) bool {
+	kdfEnabled, ok := kdfEnabledPerKMS.Load(kmsName)
+	if !ok {
+		return true // explicit config is missing, but KDF is enabled by default
+	}
+	return kdfEnabled.(bool) // this will panic if a non-bool ever gets stored, which should never happen
 }
 
 func init() {
@@ -390,7 +401,7 @@ func (h *kmsv2PluginProbe) rotateDEKOnKeyIDChange(ctx context.Context, statusKey
 	// this gate can only change during tests, but the check is cheap enough to always make
 	// this allows us to easily exercise both modes without restarting the API server
 	// TODO integration test that this dynamically takes effect
-	useSeed := GetKDF()
+	useSeed := GetKDF(h.name)
 	stateUseSeed := state.EncryptedObject.EncryptedDEKSourceType == kmstypes.EncryptedDEKSourceType_HKDF_SHA256_XNONCE_AES_GCM_SEED
 
 	// state is valid and status keyID is unchanged from when we generated this DEK/seed so there is no need to rotate it

pkg/server/options/encryptionconfig/config_test.go

@@ -1851,7 +1851,7 @@ func TestComputeEncryptionConfigHash(t *testing.T) {
 }
 
 func Test_kmsv2PluginProbe_rotateDEKOnKeyIDChange(t *testing.T) {
-	defaultUseSeed := GetKDF()
+	defaultUseSeed := GetKDF("")
 
 	origNowFunc := envelopekmsv2.NowFunc
 	now := origNowFunc() // freeze time
@@ -2074,9 +2074,10 @@ func Test_kmsv2PluginProbe_rotateDEKOnKeyIDChange(t *testing.T) {
 				`encryptKeyIDHash="", stateKeyIDHash="sha256:d4735e3a265e16eee03f59718b9b5d03019c07d8b6c51f90da3a666eec13ab35", expirationTimestamp=` + now.Format(time.RFC3339),
 		},
 	}
-	for _, tt := range tests {
+	for i, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			defer SetKDFForTests(tt.useSeed)()
+			kmsName := fmt.Sprintf("panda-%d", i)
+			defer SetKDFForTests(kmsName, tt.useSeed)()
 
 			var buf bytes.Buffer
 			klog.SetOutput(&buf)
@@ -2084,7 +2085,7 @@ func Test_kmsv2PluginProbe_rotateDEKOnKeyIDChange(t *testing.T) {
 			ctx := testContext(t)
 
 			h := &kmsv2PluginProbe{
-				name:    "panda",
+				name:    kmsName,
 				service: tt.service,
 			}
 			h.state.Store(&tt.state)
@@ -2133,7 +2134,7 @@ func Test_kmsv2PluginProbe_rotateDEKOnKeyIDChange(t *testing.T) {
 			if _, stateErr := h.getCurrentState(); stateErr == nil || err == nil {
 				transformer := envelopekmsv2.NewEnvelopeTransformer(
 					&testKMSv2EnvelopeService{err: fmt.Errorf("broken")}, // not called
-					"panda",
+					kmsName,
 					h.getCurrentState,
 					"",
 				)