Merge pull request #128136 from enj/enj/t/non_global_kms_kdf_via_name

kmsv2: run KDF tests in parallel

Kubernetes-commit: a8fd407d2f982267d847eeedc2d8097beedccc9e

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,
 					"",
 				)