Deflake supported key ack

pkg/serviceaccount/externaljwt/plugin/testing/v1alpha1/externalsigner_mock.go

@@ -28,7 +28,6 @@ import (
 	"net"
 	"os"
 	"sync"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -48,10 +47,12 @@ type MockSigner struct {
 	SigningKeyID              string
 	SigningAlg                string
 	TokenType                 string
-	SupportedKeys             atomic.Pointer[map[string]KeyT]
-	AckKeyFetch               chan bool
 	MaxTokenExpirationSeconds int64
 
+	supportedKeys        map[string]KeyT
+	supportedKeysLock    sync.RWMutex
+	supportedKeysFetched *sync.Cond
+
 	FetchError    error
 	MetadataError error
 	errorLock     sync.RWMutex
@@ -70,9 +71,9 @@ func NewMockSigner(t *testing.T, socketPath string) *MockSigner {
 	m := &MockSigner{
 		socketPath:                socketPath,
 		server:                    server,
-		AckKeyFetch:               make(chan bool),
 		MaxTokenExpirationSeconds: 10 * 60, // 10m
 	}
+	m.supportedKeysFetched = sync.NewCond(&m.supportedKeysLock)
 
 	if err := m.Reset(); err != nil {
 		t.Fatalf("failed to load keys for mock signer: %v", err)
@@ -91,6 +92,22 @@ func NewMockSigner(t *testing.T, socketPath string) *MockSigner {
 	return m
 }
 
+func (m *MockSigner) GetSupportedKeys() map[string]KeyT {
+	m.supportedKeysLock.RLock()
+	defer m.supportedKeysLock.RUnlock()
+	return m.supportedKeys
+}
+func (m *MockSigner) SetSupportedKeys(keys map[string]KeyT) {
+	m.supportedKeysLock.Lock()
+	defer m.supportedKeysLock.Unlock()
+	m.supportedKeys = keys
+}
+func (m *MockSigner) WaitForSupportedKeysFetch() {
+	m.supportedKeysLock.Lock()
+	defer m.supportedKeysLock.Unlock()
+	m.supportedKeysFetched.Wait()
+}
+
 func (m *MockSigner) Sign(ctx context.Context, req *v1alpha1.SignJWTRequest) (*v1alpha1.SignJWTResponse, error) {
 
 	header := &struct {
@@ -132,18 +149,16 @@ func (m *MockSigner) FetchKeys(ctx context.Context, req *v1alpha1.FetchKeysReque
 
 	keys := []*v1alpha1.Key{}
 
-	for id, k := range *m.SupportedKeys.Load() {
+	m.supportedKeysLock.RLock()
+	for id, k := range m.supportedKeys {
 		keys = append(keys, &v1alpha1.Key{
 			KeyId:                    id,
 			Key:                      k.Key,
 			ExcludeFromOidcDiscovery: k.ExcludeFromOidcDiscovery,
 		})
 	}
-
+	m.supportedKeysFetched.Broadcast()
-	select {
+	m.supportedKeysLock.RUnlock()
-	case <-m.AckKeyFetch:
-	default:
-	}
 
 	return &v1alpha1.FetchKeysResponse{
 		RefreshHintSeconds: 5,
@@ -185,7 +200,7 @@ func (m *MockSigner) Reset() error {
 	m.SigningKeyID = "kid-1"
 	m.SigningAlg = "RS256"
 	m.TokenType = "JWT"
-	m.SupportedKeys.Store(&map[string]KeyT{
+	m.SetSupportedKeys(map[string]KeyT{
 		"kid-1": {Key: pub1},
 		"kid-2": {Key: pub2},
 		"kid-3": {Key: pub3},

test/integration/serviceaccount/external_jwt_signer_test.go

@@ -121,14 +121,14 @@ func TestExternalJWTSigningAndAuth(t *testing.T) {
 				mockSigner.SigningKeyID = "updated-kid-1"
 
 				cpy := make(map[string]v1alpha1testing.KeyT)
-				for key, value := range *mockSigner.SupportedKeys.Load() {
+				for key, value := range mockSigner.GetSupportedKeys() {
 					cpy[key] = value
 				}
 				cpy["updated-kid-1"] = v1alpha1testing.KeyT{
 					Key:                      pubKey1Bytes,
 					ExcludeFromOidcDiscovery: true,
 				}
-				mockSigner.SupportedKeys.Store(&cpy)
+				mockSigner.SetSupportedKeys(cpy)
 			},
 			preValidationSignerUpdate: func() { /*no-op*/ },
 			wantTokenReqErr:           fmt.Errorf("failed to generate token: while validating header: key used for signing JWT (kid: updated-kid-1) is excluded from OIDC discovery docs"),
@@ -163,7 +163,7 @@ func TestExternalJWTSigningAndAuth(t *testing.T) {
 				mockSigner.SigningKey = key1
 			},
 			preValidationSignerUpdate: func() {
-				mockSigner.SupportedKeys.Store(&map[string]v1alpha1testing.KeyT{})
+				mockSigner.SetSupportedKeys(map[string]v1alpha1testing.KeyT{})
 			},
 			shouldPassAuth: false,
 		},
@@ -174,12 +174,12 @@ func TestExternalJWTSigningAndAuth(t *testing.T) {
 			},
 			preValidationSignerUpdate: func() {
 				cpy := make(map[string]v1alpha1testing.KeyT)
-				for key, value := range *mockSigner.SupportedKeys.Load() {
+				for key, value := range mockSigner.GetSupportedKeys() {
 					cpy[key] = value
 				}
 				cpy["kid-1"] = v1alpha1testing.KeyT{Key: pubKey1Bytes}
-				mockSigner.SupportedKeys.Store(&cpy)
+				mockSigner.SetSupportedKeys(cpy)
-				mockSigner.AckKeyFetch <- true
+				mockSigner.WaitForSupportedKeysFetch()
 			},
 			shouldPassAuth: true,
 		},
@@ -192,7 +192,7 @@ func TestExternalJWTSigningAndAuth(t *testing.T) {
 			if err != nil {
 				t.Fatalf("failed to reset signer for the test %q: %v", tc.desc, err)
 			}
-			mockSigner.AckKeyFetch <- true
+			mockSigner.WaitForSupportedKeysFetch()
 
 			// Adjust parameters on mock signer for the test.
 			tc.preTestSignerUpdate()