Skip to content

Commit 9f23fee

Browse files
terevterev
committed
fix: cpu contention when reading JWKs and suppress generating duplicate JWKs
Previously each concurrent caller would need to lock a shared mutex when reading or writing a given JWK set. The read path now doesn't require locking a mutex at all and instead returns valid query results directly. The write path is now protected by a concurrency control mechanism (using x/sync/singleflight) to ensure only one JWK set is generated and persisted. Note: Duplicate JWK sets may still be improperly generated if running more than one Hydra instance in a high traffic environment.
1 parent 825c24d commit 9f23fee

17 files changed

Lines changed: 108 additions & 116 deletions

‎Makefile‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ format: .bin/goimports .bin/ory node_modules
109109
mocks: .bin/mockgen
110110
mockgen -package oauth2_test -destination oauth2/oauth2_provider_mock_test.go github.com/ory/fosite OAuth2Provider
111111
mockgen -package jwk_test -destination jwk/registry_mock_test.go -source=jwk/registry.go
112+
mockgen -package jwk_test -destination jwk/manager_mock_test.go -source=jwk/manager.go
112113
go generate ./...
113114

114115
# Generates the SDKs

‎cmd/server/helper_cert.go‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,6 @@ import (
1212
"encoding/pem"
1313
"sync"
1414

15-
"github.com/gofrs/uuid"
16-
1715
"github.com/go-jose/go-jose/v3"
1816

1917
"github.com/ory/hydra/v2/driver"
@@ -58,7 +56,7 @@ func GetOrCreateTLSCertificate(ctx context.Context, d driver.Registry, iface con
5856
}
5957

6058
// no certificates configured: self-sign a new cert
61-
priv, err := jwk.GetOrGenerateKeys(ctx, d, d.SoftwareKeyManager(), TlsKeyName, uuid.Must(uuid.NewV4()).String(), "RS256")
59+
priv, err := jwk.GetOrGenerateKeySetPrivateKey(ctx, d.SoftwareKeyManager(), TlsKeyName, "", "RS256")
6260
if err != nil {
6361
d.Logger().WithError(err).Fatal("Unable to fetch or generate HTTPS TLS key pair")
6462
return nil // in case Fatal is hooked

‎hsm/manager_hsm.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,16 @@ import (
1616
"net/http"
1717
"sync"
1818

19-
"github.com/ory/hydra/v2/driver/config"
2019
"github.com/ory/x/otelx"
2120

21+
"github.com/ory/hydra/v2/driver/config"
22+
2223
"github.com/pkg/errors"
2324

2425
"github.com/pborman/uuid"
2526

2627
"github.com/ory/fosite"
28+
2729
"github.com/ory/hydra/v2/jwk"
2830

2931
"github.com/miekg/pkcs11"

‎hsm/manager_nohsm.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,10 @@ import (
1010
"context"
1111
"sync"
1212

13-
"github.com/ory/hydra/v2/driver/config"
1413
"github.com/ory/x/logrusx"
1514

15+
"github.com/ory/hydra/v2/driver/config"
16+
1617
"github.com/pkg/errors"
1718

1819
"github.com/ory/hydra/v2/jwk"

‎internal/mock/config_cookie.go‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎jwk/handler.go‎

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,6 @@ import (
1212
"github.com/ory/herodot"
1313
"github.com/ory/x/httprouterx"
1414

15-
"github.com/gofrs/uuid"
16-
"github.com/pkg/errors"
17-
1815
"github.com/ory/x/urlx"
1916

2017
"github.com/ory/x/errorsx"
@@ -101,17 +98,11 @@ func (h *Handler) discoverJsonWebKeys(w http.ResponseWriter, r *http.Request) {
10198
for _, set := range wellKnownKeys {
10299
set := set
103100
eg.Go(func() error {
104-
k, err := h.r.KeyManager().GetKeySet(ctx, set)
105-
if errors.Is(err, x.ErrNotFound) {
106-
h.r.Logger().Warnf("JSON Web Key Set %q does not exist yet, generating new key pair...", set)
107-
k, err = h.r.KeyManager().GenerateAndPersistKeySet(ctx, set, uuid.Must(uuid.NewV4()).String(), string(jose.RS256), "sig")
108-
if err != nil {
109-
return err
110-
}
111-
} else if err != nil {
101+
keySet, err := GetOrGenerateKeySet(ctx, h.r.KeyManager(), set, "", string(jose.RS256))
102+
if err != nil {
112103
return err
113104
}
114-
keys <- ExcludePrivateKeys(k)
105+
keys <- ExcludePrivateKeys(keySet)
115106
return nil
116107
})
117108
}

‎jwk/helper.go‎

Lines changed: 25 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -12,69 +12,51 @@ import (
1212
"crypto/x509"
1313
"encoding/json"
1414
"encoding/pem"
15-
"sync"
16-
17-
hydra "github.com/ory/hydra-client-go/v2"
1815

1916
"github.com/ory/x/josex"
2017

21-
"github.com/ory/x/errorsx"
22-
18+
hydra "github.com/ory/hydra-client-go/v2"
2319
"github.com/ory/hydra/v2/x"
2420

21+
"github.com/ory/x/errorsx"
22+
2523
jose "github.com/go-jose/go-jose/v3"
2624
"github.com/pkg/errors"
2725
)
2826

29-
var mapLock sync.RWMutex
30-
var locks = map[string]*sync.RWMutex{}
31-
32-
func getLock(set string) *sync.RWMutex {
33-
mapLock.Lock()
34-
defer mapLock.Unlock()
35-
if _, ok := locks[set]; !ok {
36-
locks[set] = new(sync.RWMutex)
37-
}
38-
return locks[set]
39-
}
40-
4127
func EnsureAsymmetricKeypairExists(ctx context.Context, r InternalRegistry, alg, set string) error {
42-
_, err := GetOrGenerateKeys(ctx, r, r.KeyManager(), set, set, alg)
28+
_, err := GetOrGenerateKeySetPrivateKey(ctx, r.KeyManager(), set, set, alg)
4329
return err
4430
}
4531

46-
func GetOrGenerateKeys(ctx context.Context, r InternalRegistry, m Manager, set, kid, alg string) (private *jose.JSONWebKey, err error) {
47-
getLock(set).Lock()
48-
defer getLock(set).Unlock()
49-
50-
keys, err := m.GetKeySet(ctx, set)
51-
if errors.Is(err, x.ErrNotFound) || keys != nil && len(keys.Keys) == 0 {
52-
r.Logger().Warnf("JSON Web Key Set \"%s\" does not exist yet, generating new key pair...", set)
53-
keys, err = m.GenerateAndPersistKeySet(ctx, set, kid, alg, "sig")
54-
if err != nil {
55-
return nil, err
56-
}
57-
} else if err != nil {
32+
func GetOrGenerateKeySetPrivateKey(ctx context.Context, m Manager, set, kid, alg string) (*jose.JSONWebKey, error) {
33+
keySet, err := GetOrGenerateKeySet(ctx, m, set, kid, alg)
34+
if err != nil {
5835
return nil, err
5936
}
6037

61-
privKey, privKeyErr := FindPrivateKey(keys)
62-
if privKeyErr == nil {
38+
privKey, err := FindPrivateKey(keySet)
39+
if err == nil {
6340
return privKey, nil
64-
} else {
65-
r.Logger().WithField("jwks", set).Warnf("JSON Web Key not found in JSON Web Key Set %s, generating new key pair...", set)
41+
}
6642

67-
keys, err = m.GenerateAndPersistKeySet(ctx, set, kid, alg, "sig")
68-
if err != nil {
69-
return nil, err
70-
}
43+
keySet, err = m.GenerateAndPersistKeySet(ctx, set, kid, alg, "sig")
44+
if err != nil {
45+
return nil, err
46+
}
7147

72-
privKey, err := FindPrivateKey(keys)
73-
if err != nil {
74-
return nil, err
75-
}
76-
return privKey, nil
48+
return FindPrivateKey(keySet)
49+
}
50+
51+
func GetOrGenerateKeySet(ctx context.Context, m Manager, set, kid, alg string) (*jose.JSONWebKeySet, error) {
52+
keys, err := m.GetKeySet(ctx, set)
53+
if err != nil && !errors.Is(err, x.ErrNotFound) {
54+
return nil, err
55+
} else if keys != nil && len(keys.Keys) > 0 {
56+
return keys, nil
7757
}
58+
59+
return m.GenerateAndPersistKeySet(ctx, set, kid, alg, "sig")
7860
}
7961

8062
func First(keys []jose.JSONWebKey) *jose.JSONWebKey {

‎jwk/helper_test.go‎

Lines changed: 14 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,20 +17,19 @@ import (
1717
"strings"
1818
"testing"
1919

20+
gomock "github.com/golang/mock/gomock"
21+
"github.com/pborman/uuid"
22+
"github.com/pkg/errors"
23+
2024
hydra "github.com/ory/hydra-client-go/v2"
2125

2226
"github.com/go-jose/go-jose/v3"
2327
"github.com/go-jose/go-jose/v3/cryptosigner"
24-
"github.com/golang/mock/gomock"
25-
"github.com/pborman/uuid"
26-
"github.com/pkg/errors"
2728
"github.com/stretchr/testify/assert"
2829
"github.com/stretchr/testify/require"
2930

30-
"github.com/ory/hydra/v2/internal"
3131
"github.com/ory/hydra/v2/jwk"
3232
"github.com/ory/hydra/v2/x"
33-
"github.com/ory/x/contextx"
3433
)
3534

3635
type fakeSigner struct {
@@ -210,7 +209,6 @@ func TestExcludeOpaquePrivateKeys(t *testing.T) {
210209

211210
func TestGetOrGenerateKeys(t *testing.T) {
212211
t.Parallel()
213-
reg := internal.NewMockedRegistry(t, &contextx.Default{})
214212

215213
setId := uuid.NewUUID().String()
216214
keyId := uuid.NewUUID().String()
@@ -226,46 +224,46 @@ func TestGetOrGenerateKeys(t *testing.T) {
226224
return NewMockManager(ctrl)
227225
}
228226

229-
t.Run("Test_Helper/Run_GetOrGenerateKeys_With_GetKeySetError", func(t *testing.T) {
227+
t.Run("Test_Helper/Run_GetOrGenerateKeySetPrivateKey_With_GetKeySetError", func(t *testing.T) {
230228
keyManager := km(t)
231229
keyManager.EXPECT().GetKeySet(gomock.Any(), gomock.Eq(setId)).Return(nil, errors.New("GetKeySetError"))
232-
privKey, err := jwk.GetOrGenerateKeys(context.TODO(), reg, keyManager, setId, keyId, "RS256")
230+
privKey, err := jwk.GetOrGenerateKeySetPrivateKey(context.TODO(), keyManager, setId, keyId, "RS256")
233231
assert.Nil(t, privKey)
234232
assert.EqualError(t, err, "GetKeySetError")
235233
})
236234

237-
t.Run("Test_Helper/Run_GetOrGenerateKeys_With_GenerateAndPersistKeySetError", func(t *testing.T) {
235+
t.Run("Test_Helper/Run_GetOrGenerateKeySetPrivateKey_With_GenerateAndPersistKeySetError", func(t *testing.T) {
238236
keyManager := km(t)
239237
keyManager.EXPECT().GetKeySet(gomock.Any(), gomock.Eq(setId)).Return(nil, errors.Wrap(x.ErrNotFound, ""))
240238
keyManager.EXPECT().GenerateAndPersistKeySet(gomock.Any(), gomock.Eq(setId), gomock.Eq(keyId), gomock.Eq("RS256"), gomock.Eq("sig")).Return(nil, errors.New("GetKeySetError"))
241-
privKey, err := jwk.GetOrGenerateKeys(context.TODO(), reg, keyManager, setId, keyId, "RS256")
239+
privKey, err := jwk.GetOrGenerateKeySetPrivateKey(context.TODO(), keyManager, setId, keyId, "RS256")
242240
assert.Nil(t, privKey)
243241
assert.EqualError(t, err, "GetKeySetError")
244242
})
245243

246-
t.Run("Test_Helper/Run_GetOrGenerateKeys_With_GenerateAndPersistKeySetError", func(t *testing.T) {
244+
t.Run("Test_Helper/Run_GetOrGenerateKeySetPrivateKey_With_GenerateAndPersistKeySetError", func(t *testing.T) {
247245
keyManager := km(t)
248246
keyManager.EXPECT().GetKeySet(gomock.Any(), gomock.Eq(setId)).Return(keySetWithoutPrivateKey, nil)
249247
keyManager.EXPECT().GenerateAndPersistKeySet(gomock.Any(), gomock.Eq(setId), gomock.Eq(keyId), gomock.Eq("RS256"), gomock.Eq("sig")).Return(nil, errors.New("GetKeySetError"))
250-
privKey, err := jwk.GetOrGenerateKeys(context.TODO(), reg, keyManager, setId, keyId, "RS256")
248+
privKey, err := jwk.GetOrGenerateKeySetPrivateKey(context.TODO(), keyManager, setId, keyId, "RS256")
251249
assert.Nil(t, privKey)
252250
assert.EqualError(t, err, "GetKeySetError")
253251
})
254252

255-
t.Run("Test_Helper/Run_GetOrGenerateKeys_With_GetKeySet_ContainsMissingPrivateKey", func(t *testing.T) {
253+
t.Run("Test_Helper/Run_GetOrGenerateKeySetPrivateKey_With_GetKeySet_ContainsMissingPrivateKey", func(t *testing.T) {
256254
keyManager := km(t)
257255
keyManager.EXPECT().GetKeySet(gomock.Any(), gomock.Eq(setId)).Return(keySetWithoutPrivateKey, nil)
258256
keyManager.EXPECT().GenerateAndPersistKeySet(gomock.Any(), gomock.Eq(setId), gomock.Eq(keyId), gomock.Eq("RS256"), gomock.Eq("sig")).Return(keySet, nil)
259-
privKey, err := jwk.GetOrGenerateKeys(context.TODO(), reg, keyManager, setId, keyId, "RS256")
257+
privKey, err := jwk.GetOrGenerateKeySetPrivateKey(context.TODO(), keyManager, setId, keyId, "RS256")
260258
assert.NoError(t, err)
261259
assert.Equal(t, privKey, &keySet.Keys[0])
262260
})
263261

264-
t.Run("Test_Helper/Run_GetOrGenerateKeys_With_GenerateAndPersistKeySet_ContainsMissingPrivateKey", func(t *testing.T) {
262+
t.Run("Test_Helper/Run_GetOrGenerateKeySetPrivateKey_With_GenerateAndPersistKeySet_ContainsMissingPrivateKey", func(t *testing.T) {
265263
keyManager := km(t)
266264
keyManager.EXPECT().GetKeySet(gomock.Any(), gomock.Eq(setId)).Return(keySetWithoutPrivateKey, nil)
267265
keyManager.EXPECT().GenerateAndPersistKeySet(gomock.Any(), gomock.Eq(setId), gomock.Eq(keyId), gomock.Eq("RS256"), gomock.Eq("sig")).Return(keySetWithoutPrivateKey, nil).Times(1)
268-
privKey, err := jwk.GetOrGenerateKeys(context.TODO(), reg, keyManager, setId, keyId, "RS256")
266+
privKey, err := jwk.GetOrGenerateKeySetPrivateKey(context.TODO(), keyManager, setId, keyId, "RS256")
269267
assert.Nil(t, privKey)
270268
assert.EqualError(t, err, "key not found")
271269
})

‎jwk/jwt_strategy.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@ import (
1010
"github.com/ory/x/josex"
1111

1212
"github.com/go-jose/go-jose/v3"
13-
"github.com/gofrs/uuid"
1413

1514
"github.com/ory/fosite"
15+
1616
"github.com/ory/hydra/v2/driver/config"
1717

1818
"github.com/pkg/errors"
@@ -40,7 +40,7 @@ func NewDefaultJWTSigner(c *config.DefaultProvider, r InternalRegistry, setID st
4040
}
4141

4242
func (j *DefaultJWTSigner) getKeys(ctx context.Context) (private *jose.JSONWebKey, err error) {
43-
private, err = GetOrGenerateKeys(ctx, j.r, j.r.KeyManager(), j.setID, uuid.Must(uuid.NewV4()).String(), string(jose.RS256))
43+
private, err = GetOrGenerateKeySetPrivateKey(ctx, j.r.KeyManager(), j.setID, "", string(jose.RS256))
4444
if err == nil {
4545
return private, nil
4646
}

‎jwk/manager.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,10 @@ import (
1111

1212
"github.com/pkg/errors"
1313

14+
"github.com/ory/x/errorsx"
15+
1416
"github.com/ory/hydra/v2/aead"
1517
"github.com/ory/hydra/v2/x"
16-
"github.com/ory/x/errorsx"
1718

1819
jose "github.com/go-jose/go-jose/v3"
1920
"github.com/gofrs/uuid"

0 commit comments

Comments
 (0)