fix: guess correct key with "use" and "alg"

ref: authlib/authlib#771

Author: lepture

src/joserfc/_keys.py

@@ -140,21 +140,32 @@ def as_dict(self, private: bool | None = None, **params: t.Any) -> KeySetSeriali
                 keys.append(key.as_dict(private=private, **params))
         return {"keys": keys}
 
-    def get_by_kid(self, kid: str | None = None) -> Key:
+    def get_by_kid(self, kid: str | None = None, parameters: KeyParameters | None = None) -> Key:
         if kid is None and len(self.keys) == 1:
             return self.keys[0]
 
-        for key in self.keys:
-            if key.kid == kid:
-                return key
-        raise InvalidKeyIdError(f"No key for kid: '{kid}'")
+        keys = [key for key in self.keys if key.kid == kid]
+        if parameters:
+            keys = list(_filter_keys_by_parameters(keys, parameters))
+
+        if len(keys) == 1:
+            return keys[0]
+
+        elif len(keys) == 0:
+            raise InvalidKeyIdError(f"No key for kid: '{kid}'")
+        else:
+            raise InvalidKeyIdError(f"Multiple keys for kid: '{kid}'")
 
-    def pick_random_key(self, algorithm: str) -> t.Optional[Key]:
+    def pick_random_key(self, algorithm: str, parameters: KeyParameters | None = None) -> t.Optional[Key]:
         key_types = self.algorithm_keys.get(algorithm)
         if key_types:
             keys = [k for k in self.keys if k.key_type in key_types]
         else:
             keys = self.keys
+
+        if parameters:
+            keys = list(_filter_keys_by_parameters(keys, parameters))
+
         if keys:
             return random.choice(keys)
         return None
@@ -186,3 +197,19 @@ def generate_key_set(
             keys.append(key)
 
         return cls(keys)
+
+
+def _filter_keys_by_parameters(keys: list[Key], parameters: KeyParameters) -> t.Generator[Key]:
+    _use = parameters.get("use")
+    _alg = parameters.get("alg")
+
+    for key in keys:
+        designed_use = key.get("use")
+        if designed_use and _use and designed_use != _use:
+            continue
+
+        designed_alg = key.get("alg")
+        if designed_alg and _alg and designed_alg != _alg:
+            continue
+
+        yield key

src/joserfc/_rfc7517/models.py

@@ -205,6 +205,15 @@ def as_dict(self, private: t.Optional[bool] = None, **params: t.Any) -> DictKey:
     def check_use(self, use: str) -> None:
         """Check if this key supports the given "use".
 
+        Values defined by this specification are:
+
+        - "sig" (signature)
+        - "enc" (encryption)
+
+        Other values MAY be used.  The "use" value is a case-sensitive
+        string. Use of the "use" member is OPTIONAL, unless the application
+        requires its presence.
+
         :param use: this key is used for, e.g. "sig", "enc"
         :raise: UnsupportedKeyUseError
         """

src/joserfc/jwe.py

@@ -109,7 +109,7 @@ def encrypt_compact(
 
     obj = CompactEncryption(protected, to_bytes(plaintext))
     recipient: Recipient[Key] = Recipient(obj)
-    key = guess_key(public_key, recipient, True)
+    key = guess_key(public_key, recipient, True, use="enc")
     key.check_use("enc")
     recipient.recipient_key = key
     if sender_key:
@@ -155,7 +155,7 @@ def decrypt_compact(
 
     recipient = obj.recipient
     assert recipient is not None
-    key = guess_key(private_key, recipient)
+    key = guess_key(private_key, recipient, use="enc")
     key.check_use("enc")
     recipient.recipient_key = key
     if sender_key:
@@ -228,7 +228,7 @@ def encrypt_json(
             recipient.sender_key = _guess_sender_key(recipient, sender_key, True)
         if not recipient.recipient_key:
             assert public_key is not None
-            key = guess_key(public_key, recipient, True)
+            key = guess_key(public_key, recipient, True, use="enc")
             key.check_use("enc")
             recipient.recipient_key = key
 
@@ -276,7 +276,7 @@ def _attach_recipient_keys(
     recipients: list[Recipient[Key]], private_key: KeyFlexible, sender_key: ECKey | OKPKey | KeySet | None = None
 ) -> None:
     for recipient in recipients:
-        key = guess_key(private_key, recipient)
+        key = guess_key(private_key, recipient, use="enc")
         key.check_use("enc")
         recipient.recipient_key = key
         if sender_key:

src/joserfc/jwk.py

@@ -56,12 +56,18 @@ def set_kid(self, kid: str) -> None: ...
 KeyFlexible = t.Union[KeyBase, KeyCallable]
 
 
-def guess_key(key: KeyFlexible, obj: GuestProtocol, use_random: bool = False) -> Key:
+def guess_key(
+    key: KeyFlexible,
+    obj: GuestProtocol,
+    random: bool = False,
+    use: t.Literal["sig", "enc"] | None = None,
+) -> Key:
     """Guess key from a various sources.
 
     :param key: a very flexible key
     :param obj: a protocol that has ``headers`` and ``set_kid`` methods
-    :param use_random: pick a random key from key set
+    :param random: pick a random key from key set
+    :param use: optional "use" value
     """
     resolved_key: KeyBase
     if callable(key):
@@ -74,15 +80,20 @@ def guess_key(key: KeyFlexible, obj: GuestProtocol, use_random: bool = False) ->
     elif isinstance(resolved_key, KeySet):
         headers = obj.headers()
         kid: str | None = headers.get("kid")
-        if not kid and use_random:
+
+        parameters: KeyParameters = {"alg": headers["alg"]}
+        if use:
+            parameters["use"] = use
+
+        if not kid and random:
             # choose one key by random
-            return_key = resolved_key.pick_random_key(headers["alg"])
+            return_key = resolved_key.pick_random_key(headers["alg"], parameters)
             if return_key is None:
                 raise ValueError("Invalid key")
             return_key.ensure_kid()
             obj.set_kid(t.cast(str, return_key.kid))
         else:
-            return_key = resolved_key.get_by_kid(kid)
+            return_key = resolved_key.get_by_kid(kid, parameters)
         return return_key
     else:
         raise ValueError("Invalid key")

src/joserfc/jws.py

@@ -130,10 +130,10 @@ def serialize_compact(
         if private_key is None:
             raise MissingKeyError()
 
-        key = guess_key(private_key, obj, True)
+        key = guess_key(private_key, obj, True, use="sig")
         key.check_use("sig")
         alg.check_key_type(key)
-        key.check_alg(protected["alg"])
+        key.check_alg(alg.name)
 
     if is_rfc7797:
         out = sign_rfc7515_compact(obj, alg, key)
@@ -170,7 +170,7 @@ def validate_compact(
     if public_key is None:
         raise MissingKeyError()
 
-    key: Key = guess_key(public_key, obj)
+    key: Key = guess_key(public_key, obj, use="sig")
     key.check_use("sig")
     alg.check_key_type(key)
     return verify_compact(obj, alg, key)
@@ -265,8 +265,8 @@ def serialize_json(
     if registry is None:
         registry = construct_registry(algorithms)
 
-    def find_key(obj: Any) -> Key:
-        return guess_key(private_key, obj, True)
+    def find_key(obj: HeaderMember) -> Key:
+        return guess_key(private_key, obj, True, use="sig")
 
     _payload = to_bytes(payload)
     if isinstance(members, list):
@@ -315,8 +315,8 @@ def deserialize_json(
     if registry is None:
         registry = construct_registry(algorithms)
 
-    def find_key(obj: Any) -> Key:
-        return guess_key(public_key, obj)
+    def find_key(obj: HeaderMember) -> Key:
+        return guess_key(public_key, obj, use="sig")
 
     if "signatures" in value:
         general_obj = extract_general_json(value)

tests/jwk/test_key_methods.py

@@ -1,7 +1,6 @@
 from unittest import TestCase
 
-# trigger register_key_set
-import joserfc.jws  # noqa: F401
+from joserfc import jws
 from joserfc.jwk import guess_key, import_key, generate_key, thumbprint_uri
 from joserfc.jwk import KeySet, OctKey, RSAKey, ECKey, OKPKey
 from joserfc.errors import (
@@ -16,7 +15,7 @@
 
 class Guest:
     def __init__(self):
-        self._headers = {}
+        self._headers = {"alg": "HS256"}
 
     def headers(self):
         return self._headers
@@ -138,3 +137,41 @@ def test_thumbprint_uri(self):
         )
         expected = "urn:ietf:params:oauth:jwk-thumbprint:sha-256:w9eYdC6_s_tLQ8lH6PUpc0mddazaqtPgeC2IgWDiqY8"
         self.assertEqual(value, expected)
+
+    def test_find_correct_key_with_use(self):
+        key = OctKey.generate_key()
+        dict_key = key.as_dict()
+
+        key1: OctKey = OctKey.import_key(dict_key, {"use": "enc"})
+        key2: OctKey = OctKey.import_key(dict_key, {"use": "sig"})
+        self.assertEqual(key1.kid, key2.kid)
+
+        key_set = KeySet([key1, key2])
+        # pick randomly
+        jws.serialize_compact({"alg": "HS256"}, "foo", key_set)
+        # get by kid
+        jws.serialize_compact({"alg": "HS256", "kid": key2.kid}, "foo", key_set)
+
+        key_set = KeySet([key1, key2, key2])
+        self.assertRaises(
+            InvalidKeyIdError,
+            jws.serialize_compact,
+            {"alg": "HS256", "kid": key2.kid},
+            "foo",
+            key_set,
+        )
+
+    def test_find_correct_key_with_alg(self):
+        key = OctKey.generate_key()
+        dict_key = key.as_dict()
+
+        key1: OctKey = OctKey.import_key(dict_key, {"alg": "HS256"})
+        key2: OctKey = OctKey.import_key(dict_key, {"alg": "dir"})
+
+        self.assertEqual(key1.kid, key2.kid)
+
+        key_set = KeySet([key1, key2])
+        # pick randomly
+        jws.serialize_compact({"alg": "HS256"}, "foo", key_set)
+        # get by kid
+        jws.serialize_compact({"alg": "HS256", "kid": key2.kid}, "foo", key_set)