Skip to content

Commit

Permalink
feat: parallel generation of JSON web key set (#3561)
Browse files Browse the repository at this point in the history
  • Loading branch information
alnr authored Jul 5, 2023
1 parent 2bdad2c commit 5bd9002
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 19 deletions.
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ require (
golang.org/x/crypto v0.10.0
golang.org/x/exp v0.0.0-20230522175609-2e198f4a06a1
golang.org/x/oauth2 v0.9.0
golang.org/x/sync v0.3.0
golang.org/x/tools v0.10.0
)

Expand Down Expand Up @@ -232,7 +233,6 @@ require (
go.opentelemetry.io/proto/otlp v0.18.0 // indirect
golang.org/x/mod v0.11.0 // indirect
golang.org/x/net v0.11.0 // indirect
golang.org/x/sync v0.3.0 // indirect
golang.org/x/sys v0.9.0 // indirect
golang.org/x/text v0.10.0 // indirect
golang.org/x/xerrors v0.0.0-20220907171357-04be3eba64a2 // indirect
Expand Down
47 changes: 29 additions & 18 deletions jwk/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ import (
"encoding/json"
"net/http"

"golang.org/x/sync/errgroup"

"github.com/ory/x/httprouterx"

"github.com/gofrs/uuid"
Expand Down Expand Up @@ -89,25 +91,34 @@ func (h *Handler) SetRoutes(admin *httprouterx.RouterAdmin, public *httprouterx.
// 200: jsonWebKeySet
// default: errorOAuth2
func (h *Handler) discoverJsonWebKeys(w http.ResponseWriter, r *http.Request) {
var jwks jose.JSONWebKeySet

ctx := r.Context()
for _, set := range stringslice.Unique(h.r.Config().WellKnownKeys(ctx)) {
keys, err := h.r.KeyManager().GetKeySet(ctx, set)
if errors.Is(err, x.ErrNotFound) {
h.r.Logger().Warnf("JSON Web Key Set \"%s\" does not exist yet, generating new key pair...", set)
keys, err = h.r.KeyManager().GenerateAndPersistKeySet(ctx, set, uuid.Must(uuid.NewV4()).String(), string(jose.RS256), "sig")
if err != nil {
h.r.Writer().WriteError(w, r, err)
return
eg, ctx := errgroup.WithContext(r.Context())
wellKnownKeys := stringslice.Unique(h.r.Config().WellKnownKeys(ctx))
keys := make(chan *jose.JSONWebKeySet, len(wellKnownKeys))
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 {
return err
}
} else if err != nil {
h.r.Writer().WriteError(w, r, err)
return
}

keys = ExcludePrivateKeys(keys)
jwks.Keys = append(jwks.Keys, keys.Keys...)
keys <- ExcludePrivateKeys(k)
return nil
})
}
if err := eg.Wait(); err != nil {
h.r.Writer().WriteError(w, r, err)
return
}
close(keys)
var jwks jose.JSONWebKeySet
for k := range keys {
jwks.Keys = append(jwks.Keys, k.Keys...)
}

h.r.Writer().Write(w, r, &jwks)
Expand Down

0 comments on commit 5bd9002

Please sign in to comment.