test/oidc: extract key type to allow testing different algs

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

test/integration/apiserver/oidc/oidc_test.go

@@ -35,11 +35,13 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"gopkg.in/square/go-jose.v2"
 
 	authenticationv1 "k8s.io/api/authentication/v1"
 	rbacv1 "k8s.io/api/rbac/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	utilrand "k8s.io/apimachinery/pkg/util/rand"
 	"k8s.io/apiserver/pkg/features"
 	utilfeature "k8s.io/apiserver/pkg/util/feature"
 	"k8s.io/client-go/kubernetes"
@@ -116,29 +118,10 @@ func TestStructuredAuthenticationConfig(t *testing.T) {
 }
 
 func runTests(t *testing.T, useAuthenticationConfig bool) {
-	var tests = []struct {
+	var tests = []singleTest[*rsa.PrivateKey, *rsa.PublicKey]{
-		name                    string
-		configureInfrastructure func(t *testing.T, fn authenticationConfigFunc) (
-			oidcServer *utilsoidc.TestServer,
-			apiServer *kubeapiserverapptesting.TestServer,
-			signingPrivateKey *rsa.PrivateKey,
-			caCertContent []byte,
-			caFilePath string,
-		)
-		configureOIDCServerBehaviour func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey)
-		configureClient              func(
-			t *testing.T,
-			restCfg *rest.Config,
-			caCert []byte,
-			certPath,
-			oidcServerURL,
-			oidcServerTokenURL string,
-		) kubernetes.Interface
-		assertErrFn func(t *testing.T, errorToCheck error)
-	}{
 		{
 			name:                    "ID token is ok",
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				idTokenLifetime := time.Second * 1200
 				oidcServer.TokenHandler().EXPECT().Token().Times(1).DoAndReturn(utilsoidc.TokenHandlerBehaviorReturningPredefinedJWT(
@@ -161,7 +144,7 @@ func runTests(t *testing.T, useAuthenticationConfig bool) {
 		},
 		{
 			name:                    "ID token is expired",
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				configureOIDCServerToReturnExpiredIDToken(t, 2, oidcServer, signingPrivateKey)
 			},
@@ -172,7 +155,7 @@ func runTests(t *testing.T, useAuthenticationConfig bool) {
 		},
 		{
 			name:                    "wrong client ID",
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, _ *rsa.PrivateKey) {
 				oidcServer.TokenHandler().EXPECT().Token().Times(2).Return(utilsoidc.Token{}, utilsoidc.ErrBadClientID)
 			},
@@ -189,7 +172,7 @@ func runTests(t *testing.T, useAuthenticationConfig bool) {
 		},
 		{
 			name:                         "client has wrong CA",
-			configureInfrastructure:      configureTestInfrastructure,
+			configureInfrastructure:      configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, _ *rsa.PrivateKey) {},
 			configureClient: func(t *testing.T, restCfg *rest.Config, caCert []byte, _, oidcServerURL, oidcServerTokenURL string) kubernetes.Interface {
 				tempDir := t.TempDir()
@@ -207,7 +190,7 @@ func runTests(t *testing.T, useAuthenticationConfig bool) {
 		},
 		{
 			name:                    "refresh flow does not return ID Token",
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				configureOIDCServerToReturnExpiredIDToken(t, 1, oidcServer, signingPrivateKey)
 				oidcServer.TokenHandler().EXPECT().Token().Times(1).Return(utilsoidc.Token{
@@ -230,7 +213,7 @@ func runTests(t *testing.T, useAuthenticationConfig bool) {
 		},
 		{
 			name: "ID token signature can not be verified due to wrong JWKs",
-			configureInfrastructure: func(t *testing.T, fn authenticationConfigFunc) (
+			configureInfrastructure: func(t *testing.T, fn authenticationConfigFunc, keyFunc func(t *testing.T) (*rsa.PrivateKey, *rsa.PublicKey)) (
 				oidcServer *utilsoidc.TestServer,
 				apiServer *kubeapiserverapptesting.TestServer,
 				signingPrivateKey *rsa.PrivateKey,
@@ -239,8 +222,7 @@ func runTests(t *testing.T, useAuthenticationConfig bool) {
 			) {
 				caCertContent, _, caFilePath, caKeyFilePath := generateCert(t)
 
-				signingPrivateKey, wantErr := rsa.GenerateKey(rand.Reader, rsaKeyBitSize)
+				signingPrivateKey, _ = keyFunc(t)
-				require.NoError(t, wantErr)
 
 				oidcServer = utilsoidc.BuildAndRunTestServer(t, caFilePath, caKeyFilePath)
 
@@ -260,16 +242,16 @@ jwt:
       claim: sub
       prefix: %s
 `, oidcServer.URL(), defaultOIDCClientID, indentCertificateAuthority(string(caCertContent)), defaultOIDCUsernamePrefix)
-					apiServer = startTestAPIServerForOIDC(t, "", "", "", authenticationConfig)
+					apiServer = startTestAPIServerForOIDC(t, "", "", "", authenticationConfig, &signingPrivateKey.PublicKey)
 				} else {
-					apiServer = startTestAPIServerForOIDC(t, oidcServer.URL(), defaultOIDCClientID, caFilePath, "")
+					apiServer = startTestAPIServerForOIDC(t, oidcServer.URL(), defaultOIDCClientID, caFilePath, "", &signingPrivateKey.PublicKey)
 				}
 
 				adminClient := kubernetes.NewForConfigOrDie(apiServer.ClientConfig)
 				configureRBAC(t, adminClient, defaultRole, defaultRoleBinding)
 
-				anotherSigningPrivateKey, wantErr := rsa.GenerateKey(rand.Reader, rsaKeyBitSize)
+				anotherSigningPrivateKey, _ := keyFunc(t)
-				require.NoError(t, wantErr)
+
 				oidcServer.JwksHandler().EXPECT().KeySet().AnyTimes().DoAndReturn(utilsoidc.DefaultJwksHandlerBehavior(t, &anotherSigningPrivateKey.PublicKey))
 
 				return oidcServer, apiServer, signingPrivateKey, caCertContent, caFilePath
@@ -296,7 +278,37 @@ jwt:
 	}
 
 	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
+		t.Run(tt.name, singleTestRunner(useAuthenticationConfig, rsaGenerateKey, tt))
+	}
+}
+
+type singleTest[K utilsoidc.JosePrivateKey, L utilsoidc.JosePublicKey] struct {
+	name                    string
+	configureInfrastructure func(t *testing.T, fn authenticationConfigFunc, keyFunc func(t *testing.T) (K, L)) (
+		oidcServer *utilsoidc.TestServer,
+		apiServer *kubeapiserverapptesting.TestServer,
+		signingPrivateKey K,
+		caCertContent []byte,
+		caFilePath string,
+	)
+	configureOIDCServerBehaviour func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey K)
+	configureClient              func(
+		t *testing.T,
+		restCfg *rest.Config,
+		caCert []byte,
+		certPath,
+		oidcServerURL,
+		oidcServerTokenURL string,
+	) kubernetes.Interface
+	assertErrFn func(t *testing.T, errorToCheck error)
+}
+
+func singleTestRunner[K utilsoidc.JosePrivateKey, L utilsoidc.JosePublicKey](
+	useAuthenticationConfig bool,
+	keyFunc func(t *testing.T) (K, L),
+	tt singleTest[K, L],
+) func(t *testing.T) {
+	return func(t *testing.T) {
 		fn := func(t *testing.T, issuerURL, caCert string) string { return "" }
 		if useAuthenticationConfig {
 			fn = func(t *testing.T, issuerURL, caCert string) string {
@@ -317,7 +329,7 @@ jwt:
 `, issuerURL, defaultOIDCClientID, indentCertificateAuthority(caCert), defaultOIDCUsernamePrefix)
 			}
 		}
-			oidcServer, apiServer, signingPrivateKey, caCert, certPath := tt.configureInfrastructure(t, fn)
+		oidcServer, apiServer, signingPrivateKey, caCert, certPath := tt.configureInfrastructure(t, fn, keyFunc)
 
 		tt.configureOIDCServerBehaviour(t, oidcServer, signingPrivateKey)
 
@@ -330,16 +342,17 @@ jwt:
 		_, err = client.CoreV1().Pods(defaultNamespace).List(ctx, metav1.ListOptions{})
 
 		tt.assertErrFn(t, err)
-		})
 	}
 }
 
 func TestUpdatingRefreshTokenInCaseOfExpiredIDToken(t *testing.T) {
-	var tests = []struct {
+	type testRun[K utilsoidc.JosePrivateKey] struct {
 		name                            string
-		configureUpdatingTokenBehaviour func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey)
+		configureUpdatingTokenBehaviour func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey K)
 		assertErrFn                     func(t *testing.T, errorToCheck error)
-	}{
+	}
+
+	var tests = []testRun[*rsa.PrivateKey]{
 		{
 			name: "cache returns stale client if refresh token is not updated in config",
 			configureUpdatingTokenBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
@@ -369,7 +382,7 @@ func TestUpdatingRefreshTokenInCaseOfExpiredIDToken(t *testing.T) {
 		},
 	}
 
-	oidcServer, apiServer, signingPrivateKey, caCert, certPath := configureTestInfrastructure(t, func(t *testing.T, _, _ string) string { return "" })
+	oidcServer, apiServer, signingPrivateKey, caCert, certPath := configureTestInfrastructure(t, func(t *testing.T, _, _ string) string { return "" }, rsaGenerateKey)
 
 	tokenURL, err := oidcServer.TokenURL()
 	require.NoError(t, err)
@@ -399,17 +412,17 @@ func TestUpdatingRefreshTokenInCaseOfExpiredIDToken(t *testing.T) {
 func TestStructuredAuthenticationConfigCEL(t *testing.T) {
 	defer featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.StructuredAuthenticationConfiguration, true)()
 
-	tests := []struct {
+	type testRun[K utilsoidc.JosePrivateKey, L utilsoidc.JosePublicKey] struct {
 		name                    string
 		authConfigFn            authenticationConfigFunc
-		configureInfrastructure func(t *testing.T, fn authenticationConfigFunc) (
+		configureInfrastructure func(t *testing.T, fn authenticationConfigFunc, keyFunc func(t *testing.T) (K, L)) (
 			oidcServer *utilsoidc.TestServer,
 			apiServer *kubeapiserverapptesting.TestServer,
 			signingPrivateKey *rsa.PrivateKey,
 			caCertContent []byte,
 			caFilePath string,
 		)
-		configureOIDCServerBehaviour func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey)
+		configureOIDCServerBehaviour func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey K)
 		configureClient              func(
 			t *testing.T,
 			restCfg *rest.Config,
@@ -420,7 +433,9 @@ func TestStructuredAuthenticationConfigCEL(t *testing.T) {
 		) kubernetes.Interface
 		assertErrFn func(t *testing.T, errorToCheck error)
 		wantUser    *authenticationv1.UserInfo
-	}{
+	}
+
+	tests := []testRun[*rsa.PrivateKey, *rsa.PublicKey]{
 		{
 			name: "username CEL expression is ok",
 			authConfigFn: func(t *testing.T, issuerURL, caCert string) string {
@@ -439,7 +454,7 @@ jwt:
       expression: "'k8s-' + claims.sub"
 `, issuerURL, defaultOIDCClientID, indentCertificateAuthority(caCert))
 			},
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				idTokenLifetime := time.Second * 1200
 				oidcServer.TokenHandler().EXPECT().Token().Times(1).DoAndReturn(utilsoidc.TokenHandlerBehaviorReturningPredefinedJWT(
@@ -484,7 +499,7 @@ jwt:
       expression: '(claims.roles.split(",") + claims.other_roles.split(",")).map(role, "prefix:" + role)'
 `, issuerURL, defaultOIDCClientID, indentCertificateAuthority(caCert))
 			},
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				idTokenLifetime := time.Second * 1200
 				oidcServer.TokenHandler().EXPECT().Token().Times(1).DoAndReturn(utilsoidc.TokenHandlerBehaviorReturningPredefinedJWT(
@@ -532,7 +547,7 @@ jwt:
     message: "the hd claim must be set to example.com"
 `, issuerURL, defaultOIDCClientID, indentCertificateAuthority(caCert))
 			},
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				idTokenLifetime := time.Second * 1200
 				oidcServer.TokenHandler().EXPECT().Token().Times(1).DoAndReturn(utilsoidc.TokenHandlerBehaviorReturningPredefinedJWT(
@@ -580,7 +595,7 @@ jwt:
     message: "example.org/foo must be bar and example.org/baz must be qux"
 `, issuerURL, defaultOIDCClientID, indentCertificateAuthority(caCert))
 			},
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				idTokenLifetime := time.Second * 1200
 				oidcServer.TokenHandler().EXPECT().Token().Times(1).DoAndReturn(utilsoidc.TokenHandlerBehaviorReturningPredefinedJWT(
@@ -630,7 +645,7 @@ jwt:
       expression: "claims.uid"
 `, issuerURL, defaultOIDCClientID, indentCertificateAuthority(caCert))
 			},
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				idTokenLifetime := time.Second * 1200
 				oidcServer.TokenHandler().EXPECT().Token().Times(1).DoAndReturn(utilsoidc.TokenHandlerBehaviorReturningPredefinedJWT(
@@ -680,7 +695,7 @@ jwt:
     message: "groups cannot used reserved system: prefix"
 `, issuerURL, defaultOIDCClientID, indentCertificateAuthority(caCert))
 			},
-			configureInfrastructure: configureTestInfrastructure,
+			configureInfrastructure: configureTestInfrastructure[*rsa.PrivateKey, *rsa.PublicKey],
 			configureOIDCServerBehaviour: func(t *testing.T, oidcServer *utilsoidc.TestServer, signingPrivateKey *rsa.PrivateKey) {
 				idTokenLifetime := time.Second * 1200
 				oidcServer.TokenHandler().EXPECT().Token().Times(1).DoAndReturn(utilsoidc.TokenHandlerBehaviorReturningPredefinedJWT(
@@ -708,7 +723,7 @@ jwt:
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			oidcServer, apiServer, signingPrivateKey, caCert, certPath := tt.configureInfrastructure(t, tt.authConfigFn)
+			oidcServer, apiServer, signingPrivateKey, caCert, certPath := tt.configureInfrastructure(t, tt.authConfigFn, rsaGenerateKey)
 
 			tt.configureOIDCServerBehaviour(t, oidcServer, signingPrivateKey)
 
@@ -731,10 +746,19 @@ jwt:
 	}
 }
 
-func configureTestInfrastructure(t *testing.T, fn authenticationConfigFunc) (
+func rsaGenerateKey(t *testing.T) (*rsa.PrivateKey, *rsa.PublicKey) {
+	t.Helper()
+
+	privateKey, err := rsa.GenerateKey(rand.Reader, rsaKeyBitSize)
+	require.NoError(t, err)
+
+	return privateKey, &privateKey.PublicKey
+}
+
+func configureTestInfrastructure[K utilsoidc.JosePrivateKey, L utilsoidc.JosePublicKey](t *testing.T, fn authenticationConfigFunc, keyFunc func(t *testing.T) (K, L)) (
 	oidcServer *utilsoidc.TestServer,
 	apiServer *kubeapiserverapptesting.TestServer,
-	signingPrivateKey *rsa.PrivateKey,
+	signingPrivateKey K,
 	caCertContent []byte,
 	caFilePath string,
 ) {
@@ -742,19 +766,18 @@ func configureTestInfrastructure(t *testing.T, fn authenticationConfigFunc) (
 
 	caCertContent, _, caFilePath, caKeyFilePath := generateCert(t)
 
-	signingPrivateKey, err := rsa.GenerateKey(rand.Reader, rsaKeyBitSize)
+	signingPrivateKey, publicKey := keyFunc(t)
-	require.NoError(t, err)
 
 	oidcServer = utilsoidc.BuildAndRunTestServer(t, caFilePath, caKeyFilePath)
 
 	authenticationConfig := fn(t, oidcServer.URL(), string(caCertContent))
 	if len(authenticationConfig) > 0 {
-		apiServer = startTestAPIServerForOIDC(t, "", "", "", authenticationConfig)
+		apiServer = startTestAPIServerForOIDC(t, "", "", "", authenticationConfig, publicKey)
 	} else {
-		apiServer = startTestAPIServerForOIDC(t, oidcServer.URL(), defaultOIDCClientID, caFilePath, "")
+		apiServer = startTestAPIServerForOIDC(t, oidcServer.URL(), defaultOIDCClientID, caFilePath, "", publicKey)
 	}
 
-	oidcServer.JwksHandler().EXPECT().KeySet().AnyTimes().DoAndReturn(utilsoidc.DefaultJwksHandlerBehavior(t, &signingPrivateKey.PublicKey))
+	oidcServer.JwksHandler().EXPECT().KeySet().AnyTimes().DoAndReturn(utilsoidc.DefaultJwksHandlerBehavior(t, publicKey))
 
 	adminClient := kubernetes.NewForConfigOrDie(apiServer.ClientConfig)
 	configureRBAC(t, adminClient, defaultRole, defaultRoleBinding)
@@ -803,7 +826,7 @@ func configureClientConfigForOIDC(t *testing.T, config *rest.Config, clientID, c
 	return cfg
 }
 
-func startTestAPIServerForOIDC(t *testing.T, oidcURL, oidcClientID, oidcCAFilePath, authenticationConfigYAML string) *kubeapiserverapptesting.TestServer {
+func startTestAPIServerForOIDC[L utilsoidc.JosePublicKey](t *testing.T, oidcURL, oidcClientID, oidcCAFilePath, authenticationConfigYAML string, publicKey L) *kubeapiserverapptesting.TestServer {
 	t.Helper()
 
 	var customFlags []string
@@ -816,6 +839,7 @@ func startTestAPIServerForOIDC(t *testing.T, oidcURL, oidcClientID, oidcCAFilePa
 			fmt.Sprintf("--oidc-ca-file=%s", oidcCAFilePath),
 			fmt.Sprintf("--oidc-username-prefix=%s", defaultOIDCUsernamePrefix),
 		}
+		customFlags = append(customFlags, maybeSetSigningAlgs(publicKey)...)
 	}
 	customFlags = append(customFlags, "--authorization-mode=RBAC")
 
@@ -832,6 +856,18 @@ func startTestAPIServerForOIDC(t *testing.T, oidcURL, oidcClientID, oidcCAFilePa
 	return &server
 }
 
+func maybeSetSigningAlgs[K utilsoidc.JoseKey](key K) []string {
+	alg := utilsoidc.GetSignatureAlgorithm(key)
+	if alg == jose.RS256 && randomBool() {
+		return nil // check the default case of RS256 by not always setting the flag
+	}
+	return []string{
+		fmt.Sprintf("--oidc-signing-algs=%s", alg), // all other algs need to be manually set
+	}
+}
+
+func randomBool() bool { return utilrand.Int()%2 == 1 }
+
 func fetchOIDCCredentials(t *testing.T, oidcTokenURL string, caCertContent []byte) (idToken, refreshToken string) {
 	t.Helper()
 

test/utils/oidc/testserver.go

@@ -18,6 +18,7 @@ package oidc
 
 import (
 	"crypto"
+	"crypto/ecdsa"
 	"crypto/rsa"
 	"crypto/tls"
 	"encoding/hex"
@@ -170,17 +171,21 @@ func BuildAndRunTestServer(t *testing.T, caPath, caKeyPath string) *TestServer {
 	return oidcServer
 }
 
+type JosePrivateKey interface {
+	*rsa.PrivateKey | *ecdsa.PrivateKey
+}
+
 // TokenHandlerBehaviorReturningPredefinedJWT describes the scenario when signed JWT token is being created.
 // This behavior should being applied to the MockTokenHandler.
-func TokenHandlerBehaviorReturningPredefinedJWT(
+func TokenHandlerBehaviorReturningPredefinedJWT[K JosePrivateKey](
 	t *testing.T,
-	rsaPrivateKey *rsa.PrivateKey,
+	privateKey K,
 	claims map[string]interface{}, accessToken, refreshToken string,
 ) func() (Token, error) {
 	t.Helper()
 
 	return func() (Token, error) {
-		signer, err := jose.NewSigner(jose.SigningKey{Algorithm: jose.RS256, Key: rsaPrivateKey}, nil)
+		signer, err := jose.NewSigner(jose.SigningKey{Algorithm: GetSignatureAlgorithm(privateKey), Key: privateKey}, nil)
 		require.NoError(t, err)
 
 		payloadJSON, err := json.Marshal(claims)
@@ -199,13 +204,17 @@ func TokenHandlerBehaviorReturningPredefinedJWT(
 	}
 }
 
+type JosePublicKey interface {
+	*rsa.PublicKey | *ecdsa.PublicKey
+}
+
 // DefaultJwksHandlerBehavior describes the scenario when JSON Web Key Set token is being returned.
 // This behavior should being applied to the MockJWKsHandler.
-func DefaultJwksHandlerBehavior(t *testing.T, verificationPublicKey *rsa.PublicKey) func() jose.JSONWebKeySet {
+func DefaultJwksHandlerBehavior[K JosePublicKey](t *testing.T, verificationPublicKey K) func() jose.JSONWebKeySet {
 	t.Helper()
 
 	return func() jose.JSONWebKeySet {
-		key := jose.JSONWebKey{Key: verificationPublicKey, Use: "sig", Algorithm: string(jose.RS256)}
+		key := jose.JSONWebKey{Key: verificationPublicKey, Use: "sig", Algorithm: string(GetSignatureAlgorithm(verificationPublicKey))}
 
 		thumbprint, err := key.Thumbprint(crypto.SHA256)
 		require.NoError(t, err)
@@ -216,3 +225,16 @@ func DefaultJwksHandlerBehavior(t *testing.T, verificationPublicKey *rsa.PublicK
 		}
 	}
 }
+
+type JoseKey interface{ JosePrivateKey | JosePublicKey }
+
+func GetSignatureAlgorithm[K JoseKey](key K) jose.SignatureAlgorithm {
+	switch any(key).(type) {
+	case *rsa.PrivateKey, *rsa.PublicKey:
+		return jose.RS256
+	case *ecdsa.PrivateKey, *ecdsa.PublicKey:
+		return jose.ES256
+	default:
+		panic("unknown key type") // should be impossible
+	}
+}