fix(jwa): improve RFC9864 check key logic

Author: lepture

src/joserfc/_rfc9864/jws_eddsa.py

@@ -2,34 +2,30 @@
 from cryptography.exceptions import InvalidSignature
 from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey, Ed25519PrivateKey
 from cryptography.hazmat.primitives.asymmetric.ed448 import Ed448PublicKey, Ed448PrivateKey
-from ..errors import InvalidKeyTypeError
+from ..errors import InvalidKeyCurveError
 from .._rfc7515.model import JWSAlgModel
 from .._rfc8037.okp_key import OKPKey
 
 
-_private_key_mapping = {"Ed25519": Ed25519PrivateKey, "Ed448": Ed448PrivateKey}
-_public_key_mapping = {"Ed25519": Ed25519PublicKey, "Ed448": Ed448PublicKey}
-
-
 class EdDSAAlgorithm(JWSAlgModel):
     key_type = "OKP"
 
     def __init__(self, curve: t.Literal["Ed25519", "Ed448"]):
         self.name = curve
+        self.curve = curve
         self.description = f"EdDSA using the {curve} parameter set"
 
+    def check_key(self, key: OKPKey) -> None:
+        super().check_key(key)
+        if key.curve_name != self.curve:
+            raise InvalidKeyCurveError(f"Key for '{self.name}' not supported, only '{self.curve}' allowed")
+
     def sign(self, msg: bytes, key: OKPKey) -> bytes:
         op_key = t.cast(t.Union[Ed25519PrivateKey, Ed448PrivateKey], key.get_op_key("sign"))
-        private_key_cls = _private_key_mapping[self.name]
-        if not isinstance(op_key, private_key_cls):
-            raise InvalidKeyTypeError(f"Algorithm '{self.name}' requires '{self.name}' OKP key")
         return op_key.sign(msg)
 
     def verify(self, msg: bytes, sig: bytes, key: OKPKey) -> bool:
         op_key = t.cast(t.Union[Ed25519PublicKey, Ed448PublicKey], key.get_op_key("verify"))
-        public_key_cls = _public_key_mapping[self.name]
-        if not isinstance(op_key, public_key_cls):
-            raise InvalidKeyTypeError(f"Algorithm '{self.name}' requires '{self.name}' OKP key")
         try:
             op_key.verify(sig, msg)
             return True

tests/jws/test_eddsa.py

@@ -2,7 +2,7 @@
 from unittest import TestCase
 from joserfc import jwt
 from joserfc.jwk import OKPKey
-from joserfc.errors import InvalidKeyTypeError, BadSignatureError
+from joserfc.errors import InvalidKeyTypeError, InvalidKeyCurveError, BadSignatureError
 from tests.base import load_key
 
 
@@ -23,9 +23,9 @@ def test_Ed25519(self):
         encoded_jwt = jwt.encode({"alg": "Ed25519"}, {}, self.ed25519_key, algorithms=algorithms)
         jwt.decode(encoded_jwt, self.ed25519_key, algorithms=algorithms)
         self.assertRaises(
-            InvalidKeyTypeError, jwt.encode, {"alg": "Ed25519"}, {}, self.ed448_key, algorithms=algorithms
+            InvalidKeyCurveError, jwt.encode, {"alg": "Ed25519"}, {}, self.ed448_key, algorithms=algorithms
         )
-        self.assertRaises(InvalidKeyTypeError, jwt.decode, encoded_jwt, self.ed448_key, algorithms=algorithms)
+        self.assertRaises(InvalidKeyCurveError, jwt.decode, encoded_jwt, self.ed448_key, algorithms=algorithms)
         wrong_key = OKPKey.generate_key("Ed25519", private=False)
         self.assertRaises(BadSignatureError, jwt.decode, encoded_jwt, wrong_key, algorithms=algorithms)
 
@@ -34,8 +34,8 @@ def test_Ed448(self):
         encoded_jwt = jwt.encode({"alg": "Ed448"}, {}, self.ed448_key, algorithms=algorithms)
         jwt.decode(encoded_jwt, self.ed448_key, algorithms=algorithms)
         self.assertRaises(
-            InvalidKeyTypeError, jwt.encode, {"alg": "Ed448"}, {}, self.ed25519_key, algorithms=algorithms
+            InvalidKeyCurveError, jwt.encode, {"alg": "Ed448"}, {}, self.ed25519_key, algorithms=algorithms
         )
-        self.assertRaises(InvalidKeyTypeError, jwt.decode, encoded_jwt, self.ed25519_key, algorithms=algorithms)
+        self.assertRaises(InvalidKeyCurveError, jwt.decode, encoded_jwt, self.ed25519_key, algorithms=algorithms)
         wrong_key = OKPKey.generate_key("Ed448", private=False)
         self.assertRaises(BadSignatureError, jwt.decode, encoded_jwt, wrong_key, algorithms=algorithms)

tests/jws/test_registry.py

@@ -45,3 +45,21 @@ def test_filter_algorithms_default_names(self):
         explicit = JWSRegistry.filter_algorithms(self.rsa_key, all_names)
         default = JWSRegistry.filter_algorithms(self.rsa_key)
         self.assertEqual(explicit, default)
+
+    def test_filter_algorithms_ed25519(self):
+        """Ed25519 keys should only be compatible with EdDSA and Ed25519, not Ed448."""
+        ed25519_key = OKPKey.generate_key("Ed25519")
+        algs = JWSRegistry.filter_algorithms(ed25519_key)
+        names = [alg.name for alg in algs]
+        self.assertIn("EdDSA", names)
+        self.assertIn("Ed25519", names)
+        self.assertNotIn("Ed448", names)
+
+    def test_filter_algorithms_ed448(self):
+        """Ed448 keys should only be compatible with EdDSA and Ed448, not Ed25519."""
+        ed448_key = OKPKey.generate_key("Ed448")
+        algs = JWSRegistry.filter_algorithms(ed448_key)
+        names = [alg.name for alg in algs]
+        self.assertIn("EdDSA", names)
+        self.assertIn("Ed448", names)
+        self.assertNotIn("Ed25519", names)