.

vendor/github.com/MicahParks/keyfunc/v2/multiple.go

@@ -0,0 +1,72 @@
+package keyfunc
+
+import (
+	"errors"
+	"fmt"
+
+	"github.com/golang-jwt/jwt/v5"
+)
+
+// ErrMultipleJWKSSize is returned when the number of JWKS given are not enough to make a MultipleJWKS.
+var ErrMultipleJWKSSize = errors.New("multiple JWKS must have one or more remote JWK Set resources")
+
+// MultipleJWKS manages multiple JWKS and has a field for jwt.Keyfunc.
+type MultipleJWKS struct {
+	keySelector func(multiJWKS *MultipleJWKS, token *jwt.Token) (key interface{}, err error)
+	sets        map[string]*JWKS // No lock is required because this map is read-only after initialization.
+}
+
+// GetMultiple creates a new MultipleJWKS. A map of length one or more JWKS URLs to Options is required.
+//
+// Be careful when choosing Options for each JWKS in the map. If RefreshUnknownKID is set to true for all JWKS in the
+// map then many refresh requests would take place each time a JWT is processed, this should be rate limited by
+// RefreshRateLimit.
+func GetMultiple(multiple map[string]Options, options MultipleOptions) (multiJWKS *MultipleJWKS, err error) {
+	if len(multiple) < 1 {
+		return nil, fmt.Errorf("multiple JWKS must have one or more remote JWK Set resources: %w", ErrMultipleJWKSSize)
+	}
+
+	if options.KeySelector == nil {
+		options.KeySelector = KeySelectorFirst
+	}
+
+	multiJWKS = &MultipleJWKS{
+		sets:        make(map[string]*JWKS, len(multiple)),
+		keySelector: options.KeySelector,
+	}
+
+	for u, opts := range multiple {
+		jwks, err := Get(u, opts)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get JWKS from %q: %w", u, err)
+		}
+		multiJWKS.sets[u] = jwks
+	}
+
+	return multiJWKS, nil
+}
+
+// JWKSets returns a copy of the map of JWK Sets. The map itself is a copy, but the JWKS are not and should be treated
+// as read-only.
+func (m *MultipleJWKS) JWKSets() map[string]*JWKS {
+	sets := make(map[string]*JWKS, len(m.sets))
+	for u, jwks := range m.sets {
+		sets[u] = jwks
+	}
+	return sets
+}
+
+// KeySelectorFirst returns the first key found in the multiple JWK Sets.
+func KeySelectorFirst(multiJWKS *MultipleJWKS, token *jwt.Token) (key interface{}, err error) {
+	kid, alg, err := kidAlg(token)
+	if err != nil {
+		return nil, err
+	}
+	for _, jwks := range multiJWKS.sets {
+		key, err = jwks.getKey(alg, kid)
+		if err == nil {
+			return key, nil
+		}
+	}
+	return nil, fmt.Errorf("failed to find key ID in multiple JWKS: %w", ErrKIDNotFound)
+}