From 169b731b4e4da047037e9d65f1db4bc216a793e3 Mon Sep 17 00:00:00 2001
From: Trevor Foster <trevorfoster19@gmail.com>
Date: Fri, 1 Nov 2024 03:21:24 -0400
Subject: [PATCH] Suppress duplicate jwk generation for the well-known jwks
 route too

---
 jwk/handler.go | 14 +++-----------
 jwk/helper.go  | 28 +++++++++++++---------------
 2 files changed, 16 insertions(+), 26 deletions(-)

diff --git a/jwk/handler.go b/jwk/handler.go
index 7d48445321e..c7f24f6589d 100644
--- a/jwk/handler.go
+++ b/jwk/handler.go
@@ -13,8 +13,6 @@ import (
 	"github.com/ory/x/httprouterx"
 
 	"github.com/gofrs/uuid"
-	"github.com/pkg/errors"
-
 	"github.com/ory/x/urlx"
 
 	"github.com/ory/x/errorsx"
@@ -101,17 +99,11 @@ func (h *Handler) discoverJsonWebKeys(w http.ResponseWriter, r *http.Request) {
 	for _, set := range wellKnownKeys {
 		set := set
 		eg.Go(func() error {
-			k, err := h.r.KeyManager().GetKeySet(ctx, set)
-			if errors.Is(err, x.ErrNotFound) {
-				h.r.Logger().Warnf("JSON Web Key Set %q does not exist yet, generating new key pair...", set)
-				k, err = h.r.KeyManager().GenerateAndPersistKeySet(ctx, set, uuid.Must(uuid.NewV4()).String(), string(jose.RS256), "sig")
-				if err != nil {
-					return err
-				}
-			} else if err != nil {
+			keySet, err := GetOrGenerateKeySet(ctx, h.r, h.r.KeyManager(), set, uuid.Must(uuid.NewV4()).String(), string(jose.RS256))
+			if err != nil {
 				return err
 			}
-			keys <- ExcludePrivateKeys(k)
+			keys <- ExcludePrivateKeys(keySet)
 			return nil
 		})
 	}
diff --git a/jwk/helper.go b/jwk/helper.go
index 574db10557c..e555738d44d 100644
--- a/jwk/helper.go
+++ b/jwk/helper.go
@@ -21,8 +21,6 @@ import (
 
 	"github.com/ory/x/errorsx"
 
-	"github.com/ory/hydra/v2/x"
-
 	jose "github.com/go-jose/go-jose/v3"
 	"github.com/pkg/errors"
 )
@@ -35,18 +33,23 @@ func EnsureAsymmetricKeypairExists(ctx context.Context, r InternalRegistry, alg,
 }
 
 func GetOrGenerateKeys(ctx context.Context, r InternalRegistry, m Manager, set, kid, alg string) (private *jose.JSONWebKey, err error) {
-	keys, err := m.GetKeySet(ctx, set)
-	if err != nil && !errors.Is(err, x.ErrNotFound) {
+	keySet, err := GetOrGenerateKeySet(ctx, r, m, set, kid, alg)
+	if err != nil {
 		return nil, err
 	}
 
-	if keys != nil && len(keys.Keys) > 0 {
-		privKey, privKeyErr := FindPrivateKey(keys)
-		if privKeyErr == nil {
-			return privKey, nil
-		}
+	privKey, err := FindPrivateKey(keySet)
+	if err != nil {
+		return nil, err
 	}
+	return privKey, nil
+}
 
+func GetOrGenerateKeySet(ctx context.Context, r InternalRegistry, m Manager, set, kid, alg string) (*jose.JSONWebKeySet, error) {
+	keys, err := m.GetKeySet(ctx, set)
+	if err == nil && (keys != nil && len(keys.Keys) > 0) {
+		return keys, nil
+	}
 	// Suppress duplicate key set generation jobs where the set+alg match.
 	keysResult, err, _ := jwkGenFlightGroup.Do(set+alg, func() (any, error) {
 		r.Logger().WithField("jwks", set).Warnf("JSON Web Key not found in JSON Web Key Set %s, generating new key pair...", set)
@@ -55,12 +58,7 @@ func GetOrGenerateKeys(ctx context.Context, r InternalRegistry, m Manager, set,
 	if err != nil {
 		return nil, err
 	}
-
-	privKey, err := FindPrivateKey(keysResult.(*jose.JSONWebKeySet))
-	if err != nil {
-		return nil, err
-	}
-	return privKey, nil
+	return keysResult.(*jose.JSONWebKeySet), nil
 }
 
 func First(keys []jose.JSONWebKey) *jose.JSONWebKey {