Accept multiple keys given a key ID (lestrrat-go#758)

* Allow matching of JWKS keys with not unique kid for signature verification

This changes the matching of JWKS keys with not unique kid for signature verification.
In case a JWKS store contains keys with not unique kid, all matching keys are selected. The selected keys are then tried to see if one produces a correct signature.
Not unique kid are allowed by the standard
https://datatracker.ietf.org/doc/html/rfc7517#section-4.5

* Refactor tests, add a couple of extra cases

* refactor code

* Make the new feature an option

* Update Changes

Co-authored-by: Michael Mueller <[email protected]>

Changes

@@ -4,6 +4,13 @@ Changes
 v2 has many incompatibilities with v1. To see the full list of differences between
 v1 and v2, please read the Changes-v2.md file (https://github.com/lestrrat-go/jwx/blob/develop/v2/Changes-v2.md)
 
+v2.0.4 - UNRELEASED
+[New Features]
+  * [jws] Add `jws.WithMultipleKeysPerKeyID()` sub-option to allow non-unique
+    key IDs in a given JWK set. By default we assume that a key ID is unique
+    within a key set, but enabling this option allows you to handle JWK sets
+    that contain multiple keys that contain the same key ID.
+
 v2.0.3 - 13 Jun 2022
 [Bug Fixes]
   * [jwk] Update dependency on github.com/lestrrat-go/httprc to v1.0.2 to

jws/jws_test.go

@@ -30,6 +30,7 @@ import (
 	"github.com/lestrrat-go/jwx/v2/jws"
 	"github.com/lestrrat-go/jwx/v2/x25519"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 const examplePayload = `{"iss":"joe",` + "\r\n" + ` "exp":1300819380,` + "\r\n" + ` "http://example.com/is_root":true}`
@@ -1065,6 +1066,92 @@ func TestReadFile(t *testing.T) {
 	}
 }
 
+func TestVerifyNonUniqueKid(t *testing.T) {
+	t.Parallel()
+	const payload = "Lorem ipsum"
+	const kid = "notUniqueKid"
+	privateKey, err := jwxtest.GenerateRsaJwk()
+	if !assert.NoError(t, err, "jwxtest.GenerateJwk should succeed") {
+		return
+	}
+	_ = privateKey.Set(jwk.KeyIDKey, kid)
+	signed, err := jws.Sign([]byte(payload), jws.WithKey(jwa.RS256, privateKey))
+	if !assert.NoError(t, err, `jws.Sign should succeed`) {
+		return
+	}
+	correctKey, _ := jwk.PublicKeyOf(privateKey)
+	_ = correctKey.Set(jwk.AlgorithmKey, jwa.RS256)
+
+	makeSet := func(keys ...jwk.Key) jwk.Set {
+		set := jwk.NewSet()
+		for _, key := range keys {
+			_ = set.AddKey(key)
+		}
+		return set
+	}
+
+	testcases := []struct {
+		Name string
+		Key  func() jwk.Key // Generates the "wrong" key
+	}{
+		{
+			Name: `match 2 keys via same "kid"`,
+			Key: func() jwk.Key {
+				privateKey, _ := jwxtest.GenerateRsaJwk()
+				wrongKey, _ := jwk.PublicKeyOf(privateKey)
+				_ = wrongKey.Set(jwk.KeyIDKey, kid)
+				_ = wrongKey.Set(jwk.AlgorithmKey, jwa.RS256)
+				return wrongKey
+			},
+		},
+		{
+			Name: `match 2 keys via same "kid", same key value but different alg`,
+			Key: func() jwk.Key {
+				wrongKey, _ := correctKey.Clone()
+				_ = wrongKey.Set(jwk.KeyIDKey, kid)
+				_ = wrongKey.Set(jwk.AlgorithmKey, jwa.RS512)
+				return wrongKey
+			},
+		},
+		{
+			Name: `match 2 keys via same "kid", same key type but different alg`,
+			Key: func() jwk.Key {
+				privateKey, _ := jwxtest.GenerateRsaJwk()
+				wrongKey, _ := jwk.PublicKeyOf(privateKey)
+				_ = wrongKey.Set(jwk.KeyIDKey, kid)
+				_ = wrongKey.Set(jwk.AlgorithmKey, jwa.RS512)
+				return wrongKey
+			},
+		},
+		{
+			Name: `match 2 keys via same "kid" and different key type / alg`,
+			Key: func() jwk.Key {
+				privateKey, _ := jwxtest.GenerateEcdsaKey(jwa.P256)
+				wrongKey, _ := jwk.PublicKeyOf(privateKey)
+				_ = wrongKey.Set(jwk.KeyIDKey, kid)
+				_ = wrongKey.Set(jwk.AlgorithmKey, jwa.ES256K)
+				return wrongKey
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.Name, func(t *testing.T) {
+			t.Parallel()
+			wrongKey := tc.Key()
+			// Try matching in different orders
+			for _, set := range []jwk.Set{makeSet(wrongKey, correctKey), makeSet(correctKey, wrongKey)} {
+				var usedKey jwk.Key
+				_, err = jws.Verify(signed, jws.WithKeySet(set, jws.WithMultipleKeysPerKeyID(true)), jws.WithKeyUsed(&usedKey))
+				if !assert.NoError(t, err, `jws.Verify should succeed`) {
+					return
+				}
+				require.Equal(t, usedKey, correctKey)
+			}
+		})
+	}
+}
+
 func TestVerifySet(t *testing.T) {
 	t.Parallel()
 	const payload = "Lorem ipsum"

jws/key_provider.go

@@ -102,10 +102,11 @@ func (kp *staticKeyProvider) FetchKeys(_ context.Context, sink KeySink, _ *Signa
 }
 
 type keySetProvider struct {
-	set            jwk.Set
-	requireKid     bool // true if `kid` must be specified
-	useDefault     bool // true if the first key should be used iff there's exactly one key in set
-	inferAlgorithm bool // true if the algorithm should be inferred from key type
+	set                  jwk.Set
+	requireKid           bool // true if `kid` must be specified
+	useDefault           bool // true if the first key should be used iff there's exactly one key in set
+	inferAlgorithm       bool // true if the algorithm should be inferred from key type
+	multipleKeysPerKeyID bool // true if we should attempt to match multiple keys per key ID. if false we assume that only one key exists for a given key ID
 }
 
 func (kp *keySetProvider) selectKey(sink KeySink, key jwk.Key, sig *Signature, _ *Message) error {
@@ -151,8 +152,6 @@ func (kp *keySetProvider) selectKey(sink KeySink, key jwk.Key, sig *Signature, _
 
 func (kp *keySetProvider) FetchKeys(_ context.Context, sink KeySink, sig *Signature, msg *Message) error {
 	if kp.requireKid {
-		var key jwk.Key
-
 		wantedKid := sig.ProtectedHeaders().KeyID()
 		if wantedKid == "" {
 			// If the kid is NOT specified... kp.useDefault needs to be true, and the
@@ -165,19 +164,43 @@ func (kp *keySetProvider) FetchKeys(_ context.Context, sink KeySink, sig *Signat
 
 			// if we got here, then useDefault == true AND there is exactly
 			// one key in the set.
-			key, _ = kp.set.Key(0)
-		} else {
-			// Otherwise we better be able to look up the key, baby.
-			v, ok := kp.set.LookupKeyID(wantedKid)
+			key, _ := kp.set.Key(0)
+			return kp.selectKey(sink, key, sig, msg)
+		}
+
+		// Otherwise we better be able to look up the key.
+		// <= v2.0.3 backwards compatible case: only match a single key
+		// whose key ID matches `wantedKid`
+		if !kp.multipleKeysPerKeyID {
+			key, ok := kp.set.LookupKeyID(wantedKid)
 			if !ok {
 				return fmt.Errorf(`failed to find key with key ID %q in key set`, wantedKid)
 			}
-			key = v
+			return kp.selectKey(sink, key, sig, msg)
 		}
 
-		return kp.selectKey(sink, key, sig, msg)
+		// if multipleKeysPerKeyID is true, we attempt all keys whose key ID matches
+		// the wantedKey
+		var ok bool
+		for i := 0; i < kp.set.Len(); i++ {
+			key, _ := kp.set.Key(i)
+			if key.KeyID() != wantedKid {
+				continue
+			}
+
+			if err := kp.selectKey(sink, key, sig, msg); err != nil {
+				continue
+			}
+			ok = true
+			// continue processing so that we try all keys with the same key ID
+		}
+		if !ok {
+			return fmt.Errorf(`failed to find key with key ID %q in key set`, wantedKid)
+		}
+		return nil
 	}
 
+	// Otherwise just try all keys
 	for i := 0; i < kp.set.Len(); i++ {
 		key, _ := kp.set.Key(i)
 		if err := kp.selectKey(sink, key, sig, msg); err != nil {

jws/options.go

@@ -119,24 +119,27 @@ func WithKey(alg jwa.KeyAlgorithm, key interface{}, options ...WithKeySuboption)
 // suboption types.
 func WithKeySet(set jwk.Set, options ...WithKeySetSuboption) VerifyOption {
 	requireKid := true
-	var useDefault, inferAlgorithm bool
+	var useDefault, inferAlgorithm, multipleKeysPerKeyID bool
 	for _, option := range options {
 		//nolint:forcetypeassert
 		switch option.Ident() {
 		case identRequireKid{}:
 			requireKid = option.Value().(bool)
 		case identUseDefault{}:
 			useDefault = option.Value().(bool)
+		case identMultipleKeysPerKeyID{}:
+			multipleKeysPerKeyID = option.Value().(bool)
 		case identInferAlgorithmFromKey{}:
 			inferAlgorithm = option.Value().(bool)
 		}
 	}
 
 	return WithKeyProvider(&keySetProvider{
-		set:            set,
-		requireKid:     requireKid,
-		useDefault:     useDefault,
-		inferAlgorithm: inferAlgorithm,
+		set:                  set,
+		requireKid:           requireKid,
+		useDefault:           useDefault,
+		multipleKeysPerKeyID: multipleKeysPerKeyID,
+		inferAlgorithm:       inferAlgorithm,
 	})
 }
 

jws/options.yaml

@@ -117,9 +117,18 @@ options:
     interface: WithKeySetSuboption
     argument_type: bool
     comment: |
-      WithrequiredKid specifies whether the keys in the jwk.Set should
+      WithRequiredKid specifies whether the keys in the jwk.Set should
       only be matched if the target JWS message's Key ID and the Key ID
       in the given key matches.
+  - ident: MultipleKeysPerKeyID
+    interface: WithKeySetSuboption
+    argument_type: bool
+    comment: |
+      WithMultipleKeysPerKeyID specifies if we should expect multiple keys
+      to match against a key ID. By default it is assumed that key IDs are
+      unique, i.e. for a given key ID, the key set only contains a single
+      key that has the matching ID. When this option is set to true,
+      multiple keys that match the same key ID in the set can be tried.
   - ident: Pretty
     interface: WithJSONSuboption
     argument_type: bool

jws/options_gen.go

@@ -132,6 +132,7 @@ type identKey struct{}
 type identKeyProvider struct{}
 type identKeyUsed struct{}
 type identMessage struct{}
+type identMultipleKeysPerKeyID struct{}
 type identPretty struct{}
 type identProtectedHeaders struct{}
 type identPublicHeaders struct{}
@@ -175,6 +176,10 @@ func (identMessage) String() string {
 	return "WithMessage"
 }
 
+func (identMultipleKeysPerKeyID) String() string {
+	return "WithMultipleKeysPerKeyID"
+}
+
 func (identPretty) String() string {
 	return "WithPretty"
 }
@@ -268,6 +273,15 @@ func WithMessage(v *Message) VerifyOption {
 	return &verifyOption{option.New(identMessage{}, v)}
 }
 
+// WithMultipleKeysPerKeyID specifies if we should expect multiple keys
+// to match against a key ID. By default it is assumed that key IDs are
+// unique, i.e. for a given key ID, the key set only contains a single
+// key that has the matching ID. When this option is set to true,
+// multiple keys that match the same key ID in the set can be tried.
+func WithMultipleKeysPerKeyID(v bool) WithKeySetSuboption {
+	return &withKeySetSuboption{option.New(identMultipleKeysPerKeyID{}, v)}
+}
+
 // WithPretty specifies whether the JSON output should be formatted and
 // indented
 func WithPretty(v bool) WithJSONSuboption {
@@ -293,7 +307,7 @@ func WithPublicHeaders(v Headers) WithKeySuboption {
 	return &withKeySuboption{option.New(identPublicHeaders{}, v)}
 }
 
-// WithrequiredKid specifies whether the keys in the jwk.Set should
+// WithRequiredKid specifies whether the keys in the jwk.Set should
 // only be matched if the target JWS message's Key ID and the Key ID
 // in the given key matches.
 func WithRequireKid(v bool) WithKeySetSuboption {

jws/options_gen_test.go

@@ -18,6 +18,7 @@ func TestOptionIdent(t *testing.T) {
 	require.Equal(t, "WithKeyProvider", identKeyProvider{}.String())
 	require.Equal(t, "WithKeyUsed", identKeyUsed{}.String())
 	require.Equal(t, "WithMessage", identMessage{}.String())
+	require.Equal(t, "WithMultipleKeysPerKeyID", identMultipleKeysPerKeyID{}.String())
 	require.Equal(t, "WithPretty", identPretty{}.String())
 	require.Equal(t, "WithProtectedHeaders", identProtectedHeaders{}.String())
 	require.Equal(t, "WithPublicHeaders", identPublicHeaders{}.String())

tools/cmd/genoptions/main.go

@@ -121,6 +121,7 @@ func genOptions(objects *Objects) error {
 	o.LL(`package %s`, objects.PackageName)
 
 	imports := append(objects.Imports, []string{
+		`io/fs`, // for some reason without this the goimports in my environment tries to import a differnet package
 		`github.com/lestrrat-go/jwx/v2/jwa`,
 		`github.com/lestrrat-go/jwx/v2/jwe`,
 		`github.com/lestrrat-go/jwx/v2/jwk`,