Repository: jpadilla/pyjwt

Issue:
Error `TypeError: ECPublicKey.verify() takes 3 positional arguments but 4 were given` while using `jwt.decode`
Getting this error `TypeError: ECPublicKey.verify() takes 3 positional arguments but 4 were given` while trying to decode a token with RS256 algorithm.

## Expected Result

Decode some token via jwt.decode

## Actual Result

```
... (omitted) ...
File /opt/venv/lib/python3.12/site-packages/jwt/api_jwt.py:210, in PyJWT.decode(self, jwt, key, algorithms, options, verify, detached_payload, audience, issuer, leeway, **kwargs)
    203 if kwargs:
    204     warnings.warn(
    205         "passing additional kwargs to decode() is deprecated "
    206         "and will be removed in pyjwt version 3. "
    207         f"Unsupported kwargs: {tuple(kwargs.keys())}",
    208         RemovedInPyjwt3Warning,
    209     )
--> 210 decoded = self.decode_complete(
    211     jwt,
    212     key,
    213     algorithms,
    214     options,
    215     verify=verify,
    216     detached_payload=detached_payload,
    217     audience=audience,
    218     issuer=issuer,
    219     leeway=leeway,
    220 )
    221 return decoded["payload"]

File /opt/venv/lib/python3.12/site-packages/jwt/api_jwt.py:151, in PyJWT.decode_complete(self, jwt, key, algorithms, options, verify, detached_payload, audience, issuer, leeway, **kwargs)
    146 if options["verify_signature"] and not algorithms:
    147     raise DecodeError(
    148         'It is required that you pass in a value for the "algorithms" argument when calling decode().'
    149     )
--> 151 decoded = api_jws.decode_complete(
    152     jwt,
    153     key=key,
    154     algorithms=algorithms,
    155     options=options,
    156     detached_payload=detached_payload,
    157 )
    159 payload = self._decode_payload(decoded)
    161 merged_options = {**self.options, **options}

File /opt/venv/lib/python3.12/site-packages/jwt/api_jws.py:209, in PyJWS.decode_complete(self, jwt, key, algorithms, options, detached_payload, **kwargs)
    206     signing_input = b".".join([signing_input.rsplit(b".", 1)[0], payload])
    208 if verify_signature:
--> 209     self._verify_signature(signing_input, header, signature, key, algorithms)
    211 return {
    212     "payload": payload,
    213     "header": header,
    214     "signature": signature,
    215 }

File /opt/venv/lib/python3.12/site-packages/jwt/api_jws.py:309, in PyJWS._verify_signature(self, signing_input, header, signature, key, algorithms)
    306     raise InvalidAlgorithmError("Algorithm not supported") from e
    307 prepared_key = alg_obj.prepare_key(key)
--> 309 if not alg_obj.verify(signing_input, prepared_key, signature):
    310     raise InvalidSignatureError("Signature verification failed")

File /opt/venv/lib/python3.12/site-packages/jwt/algorithms.py:483, in RSAAlgorithm.verify(self, msg, key, sig)
    481 def verify(self, msg: bytes, key: RSAPublicKey, sig: bytes) -> bool:
    482     try:
--> 483         key.verify(sig, msg, padding.PKCS1v15(), self.hash_alg())
    484         return True
    485     except InvalidSignature:

TypeError: ECPublicKey.verify() takes 3 positional arguments but 4 were given
```

## Reproduction Steps

```python
import jwt

token = "..."
key = "..."
jwt.decode(
  jwt=token,
  key=key
  algorithms=["RS256",]
)
```

## System Information

    $ python -m jwt.help

```
{
  "cryptography": {
    "version": "42.0.5"
  },
  "implementation": {
    "name": "CPython",
    "version": "3.12.4"
  },
  "platform": {
    "release": "5.15.0-113-generic",
    "system": "Linux"
  },
  "pyjwt": {
    "version": "2.8.0"
  }
}
```

Discussion:
@RbnRncn: Exact same problem in 2.9.0.

```
{
  "cryptography": {
    "version": "43.0.1"
  },
  "implementation": {
    "name": "CPython",
    "version": "3.12.4"
  },
  "platform": {
    "release": "11",
    "system": "Windows"
  },
  "pyjwt": {
    "version": "2.9.0"
  }
}
```@pachewise: Can you check what  `type(load_pem_public_key(your_key.encode("utf-8"))`is? It seems to be returning an ElliptiveCurvePublicKey. 

```python-console
>>> from cryptography.hazmat.primitives.serialization import load_pem_public_key
>>> your_key = "..."  # unicode string, or you if you already have it as bytes skip the encode() below
>>> type(load_pem_public_key(your_key.encode("utf-8")))
```

(if `load_pem_public_key` throws an error for you, try `load_ssh_public_key` - that is, make sure you are using the loader relevant to how the key is formatted)@RbnRncn: Yes, my problem was caused because I was taking the wrong key from the JWKS endpoint of my openid-configuration. I was taking the first one that it is in fact an elliptic curve key (`{"kty": "EC"}`).

Using a `PyJWTClient` and the method `get_signing_key_from_jwt` as it is showed in the example https://pyjwt.readthedocs.io/en/latest/usage.html#oidc-login-flow it took the correct key and decoded without issues.@pachewise: I do feel like we should catch these issues where the key passed in is from the wrong algo-family (and, ideally, down to ensuring the same algo was used). I can attempt a PR this week if that's the approach we want to take.@RbnRncn: Yes, maybe it should warn you (or throw you) in some way that you are trying to verify with an ECPublicKey when the algorithm you specified is RS256.

But if we are already detecting the key algorithm and we know the algorithm passed by the user, we could catch the error and be more specific.

Only my two cents as an user, I'm by no means an expert in this. I'm glad to help if needed.@pachewise: [Draft PR](https://github.com/jpadilla/pyjwt/pull/985), so we have something to look at while proposing a "fix" for this.@github-actions: This issue is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 7 days

Solution patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index e6ff926..8b19320 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -30,7 +30,7 @@ repos:
         args: [--no-build-isolation]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v1.14.1"
+    rev: "v1.15.0"
     hooks:
       - id: mypy
         additional_dependencies: [cryptography>=3.4.0]
@@ -59,7 +59,7 @@ repos:
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.9.4
+    rev: v0.9.6
     hooks:
       # Run the linter.
       - id: ruff
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 1628e3f..cebc2bb 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,10 @@ This project adheres to `Semantic Versioning <https://semver.org/>`__.
 `Unreleased <https://github.com/jpadilla/pyjwt/compare/2.10.1...HEAD>`__
 ------------------------------------------------------------------------
 
+Fixed
+~~~~~
+- Validate key against allowed types for Algorithm family in `#964 <https://github.com/jpadilla/pyjwt/pull/964>`__
+
 Added
 ~~~~~
 
diff --git a/jwt/algorithms.py b/jwt/algorithms.py
index ccb1500..b5f1615 100644
--- a/jwt/algorithms.py
+++ b/jwt/algorithms.py
@@ -65,23 +65,56 @@ try:
         load_ssh_public_key,
     )
 
+    # pyjwt-964: we use these both for type checking below, as well as for validating the key passed in.
+    # in Py >= 3.10, we can replace this with the Union types below
+    ALLOWED_RSA_KEY_TYPES = (RSAPrivateKey, RSAPublicKey)
+    ALLOWED_EC_KEY_TYPES = (EllipticCurvePrivateKey, EllipticCurvePublicKey)
+    ALLOWED_OKP_KEY_TYPES = (
+        Ed25519PrivateKey,
+        Ed25519PublicKey,
+        Ed448PrivateKey,
+        Ed448PublicKey,
+    )
+    ALLOWED_KEY_TYPES = (
+        ALLOWED_RSA_KEY_TYPES + ALLOWED_EC_KEY_TYPES + ALLOWED_OKP_KEY_TYPES
+    )
+    ALLOWED_PRIVATE_KEY_TYPES = (
+        RSAPrivateKey,
+        EllipticCurvePrivateKey,
+        Ed25519PrivateKey,
+        Ed448PrivateKey,
+    )
+    ALLOWED_PUBLIC_KEY_TYPES = (
+        RSAPublicKey,
+        EllipticCurvePublicKey,
+        Ed25519PublicKey,
+        Ed448PublicKey,
+    )
+
     has_crypto = True
 except ModuleNotFoundError:
     has_crypto = False
 
 
 if TYPE_CHECKING:
+    from typing import TypeAlias
+
+    from cryptography.hazmat.primitives.asymmetric.types import (
+        PrivateKeyTypes,
+        PublicKeyTypes,
+    )
+
     # Type aliases for convenience in algorithms method signatures
-    AllowedRSAKeys = RSAPrivateKey | RSAPublicKey
-    AllowedECKeys = EllipticCurvePrivateKey | EllipticCurvePublicKey
-    AllowedOKPKeys = (
+    AllowedRSAKeys: TypeAlias = RSAPrivateKey | RSAPublicKey
+    AllowedECKeys: TypeAlias = EllipticCurvePrivateKey | EllipticCurvePublicKey
+    AllowedOKPKeys: TypeAlias = (
         Ed25519PrivateKey | Ed25519PublicKey | Ed448PrivateKey | Ed448PublicKey
     )
-    AllowedKeys = AllowedRSAKeys | AllowedECKeys | AllowedOKPKeys
-    AllowedPrivateKeys = (
+    AllowedKeys: TypeAlias = AllowedRSAKeys | AllowedECKeys | AllowedOKPKeys
+    AllowedPrivateKeys: TypeAlias = (
         RSAPrivateKey | EllipticCurvePrivateKey | Ed25519PrivateKey | Ed448PrivateKey
     )
-    AllowedPublicKeys = (
+    AllowedPublicKeys: TypeAlias = (
         RSAPublicKey | EllipticCurvePublicKey | Ed25519PublicKey | Ed448PublicKey
     )
 
@@ -141,6 +174,9 @@ class Algorithm(ABC):
     The interface for an algorithm used to sign and verify tokens.
     """
 
+    # pyjwt-964: Validate to ensure the key passed in was decoded to the correct cryptography key family
+    _crypto_key_types: tuple[type[AllowedKeys], ...] | None = None
+
     def compute_hash_digest(self, bytestr: bytes) -> bytes:
         """
         Compute a hash digest using the specified algorithm's hash algorithm.
@@ -163,6 +199,30 @@ class Algorithm(ABC):
         else:
             return bytes(hash_alg(bytestr).digest())
 
+    def check_crypto_key_type(self, key: PublicKeyTypes | PrivateKeyTypes):
+        """Check that the key belongs to the right cryptographic family.
+
+        Note that this method only works when `cryptography` is installed.
+
+        Args:
+            key (Any): Potentially a cryptography key
+        Raises:
+            ValueError: if `cryptography` is not installed, or this method is called by a non-cryptography algorithm
+            InvalidKeyError: if the key doesn't match the expected key classes
+        """
+        if not has_crypto or self._crypto_key_types is None:
+            raise ValueError(
+                "This method requires the cryptography library, and should only be used by cryptography-based algorithms."
+            )
+
+        if not isinstance(key, self._crypto_key_types):
+            valid_classes = (cls.__name__ for cls in self._crypto_key_types)
+            actual_class = key.__class__.__name__
+            self_class = self.__class__.__name__
+            raise InvalidKeyError(
+                f"Expected one of {valid_classes}, got: {actual_class}. Invalid Key type for {self_class}"
+            )
+
     @abstractmethod
     def prepare_key(self, key: Any) -> Any:
         """
@@ -323,11 +383,13 @@ if has_crypto:
         SHA384: ClassVar[type[hashes.HashAlgorithm]] = hashes.SHA384
         SHA512: ClassVar[type[hashes.HashAlgorithm]] = hashes.SHA512
 
+        _crypto_key_types = ALLOWED_RSA_KEY_TYPES
+
         def __init__(self, hash_alg: type[hashes.HashAlgorithm]) -> None:
             self.hash_alg = hash_alg
 
         def prepare_key(self, key: AllowedRSAKeys | str | bytes) -> AllowedRSAKeys:
-            if isinstance(key, (RSAPrivateKey, RSAPublicKey)):
+            if isinstance(key, self._crypto_key_types):
                 return key
 
             if not isinstance(key, (bytes, str)):
@@ -337,14 +399,20 @@ if has_crypto:
 
             try:
                 if key_bytes.startswith(b"ssh-rsa"):
-                    return cast(RSAPublicKey, load_ssh_public_key(key_bytes))
+                    public_key: PublicKeyTypes = load_ssh_public_key(key_bytes)
+                    self.check_crypto_key_type(public_key)
+                    return cast(RSAPublicKey, public_key)
                 else:
-                    return cast(
-                        RSAPrivateKey, load_pem_private_key(key_bytes, password=None)
+                    private_key: PrivateKeyTypes = load_pem_private_key(
+                        key_bytes, password=None
                     )
+                    self.check_crypto_key_type(private_key)
+                    return cast(RSAPrivateKey, private_key)
             except ValueError:
                 try:
-                    return cast(RSAPublicKey, load_pem_public_key(key_bytes))
+                    public_key = load_pem_public_key(key_bytes)
+                    self.check_crypto_key_type(public_key)
+                    return cast(RSAPublicKey, public_key)
                 except (ValueError, UnsupportedAlgorithm):
                     raise InvalidKeyError(
                         "Could not parse the provided public key."
@@ -493,11 +561,13 @@ if has_crypto:
         SHA384: ClassVar[type[hashes.HashAlgorithm]] = hashes.SHA384
         SHA512: ClassVar[type[hashes.HashAlgorithm]] = hashes.SHA512
 
+        _crypto_key_types = ALLOWED_EC_KEY_TYPES
+
         def __init__(self, hash_alg: type[hashes.HashAlgorithm]) -> None:
             self.hash_alg = hash_alg
 
         def prepare_key(self, key: AllowedECKeys | str | bytes) -> AllowedECKeys:
-            if isinstance(key, (EllipticCurvePrivateKey, EllipticCurvePublicKey)):
+            if isinstance(key, self._crypto_key_types):
                 return key
 
             if not isinstance(key, (bytes, str)):
@@ -510,21 +580,17 @@ if has_crypto:
             # the Verifying Key first.
             try:
                 if key_bytes.startswith(b"ecdsa-sha2-"):
-                    crypto_key = load_ssh_public_key(key_bytes)
+                    public_key: PublicKeyTypes = load_ssh_public_key(key_bytes)
                 else:
-                    crypto_key = load_pem_public_key(key_bytes)  # type: ignore[assignment]
-            except ValueError:
-                crypto_key = load_pem_private_key(key_bytes, password=None)  # type: ignore[assignment]
+                    public_key = load_pem_public_key(key_bytes)
 
-            # Explicit check the key to prevent confusing errors from cryptography
-            if not isinstance(
-                crypto_key, (EllipticCurvePrivateKey, EllipticCurvePublicKey)
-            ):
-                raise InvalidKeyError(
-                    "Expecting a EllipticCurvePrivateKey/EllipticCurvePublicKey. Wrong key provided for ECDSA algorithms"
-                ) from None
-
-            return crypto_key
+                # Explicit check the key to prevent confusing errors from cryptography
+                self.check_crypto_key_type(public_key)
+                return cast(EllipticCurvePublicKey, public_key)
+            except ValueError:
+                private_key = load_pem_private_key(key_bytes, password=None)
+                self.check_crypto_key_type(private_key)
+                return cast(EllipticCurvePrivateKey, private_key)
 
         def sign(self, msg: bytes, key: EllipticCurvePrivateKey) -> bytes:
             der_sig = key.sign(msg, ECDSA(self.hash_alg()))
@@ -715,31 +781,32 @@ if has_crypto:
         This class requires ``cryptography>=2.6`` to be installed.
         """
 
+        _crypto_key_types = ALLOWED_OKP_KEY_TYPES
+
         def __init__(self, **kwargs: Any) -> None:
             pass
 
         def prepare_key(self, key: AllowedOKPKeys | str | bytes) -> AllowedOKPKeys:
-            if isinstance(key, (bytes, str)):
-                key_str = key.decode("utf-8") if isinstance(key, bytes) else key
-                key_bytes = key.encode("utf-8") if isinstance(key, str) else key
-
-                if "-----BEGIN PUBLIC" in key_str:
-                    key = load_pem_public_key(key_bytes)  # type: ignore[assignment]
-                elif "-----BEGIN PRIVATE" in key_str:
-                    key = load_pem_private_key(key_bytes, password=None)  # type: ignore[assignment]
-                elif key_str[0:4] == "ssh-":
-                    key = load_ssh_public_key(key_bytes)  # type: ignore[assignment]
+            if not isinstance(key, (str, bytes)):
+                self.check_crypto_key_type(key)
+                return cast("AllowedOKPKeys", key)
+
+            key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+            key_bytes = key.encode("utf-8") if isinstance(key, str) else key
+
+            loaded_key: PublicKeyTypes | PrivateKeyTypes
+            if "-----BEGIN PUBLIC" in key_str:
+                loaded_key = load_pem_public_key(key_bytes)
+            elif "-----BEGIN PRIVATE" in key_str:
+                loaded_key = load_pem_private_key(key_bytes, password=None)
+            elif key_str[0:4] == "ssh-":
+                loaded_key = load_ssh_public_key(key_bytes)
+            else:
+                raise InvalidKeyError("Not a public or private key")
 
             # Explicit check the key to prevent confusing errors from cryptography
-            if not isinstance(
-                key,
-                (Ed25519PrivateKey, Ed25519PublicKey, Ed448PrivateKey, Ed448PublicKey),
-            ):
-                raise InvalidKeyError(
-                    "Expecting a EllipticCurvePrivateKey/EllipticCurvePublicKey. Wrong key provided for EdDSA algorithms"
-                )
-
-            return key
+            self.check_crypto_key_type(loaded_key)
+            return cast("AllowedOKPKeys", loaded_key)
 
         def sign(
             self, msg: str | bytes, key: Ed25519PrivateKey | Ed448PrivateKey


Test patch:
diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py
index 35f9cc2..0c061d6 100644
--- a/tests/test_algorithms.py
+++ b/tests/test_algorithms.py
@@ -33,6 +33,13 @@ if has_crypto:
 
 
 class TestAlgorithms:
+    def test_check_crypto_key_type_should_fail_when_not_using_crypto(self):
+        """If has_crypto is False, or if _crypto_key_types is None, then this method should throw."""
+
+        algo = NoneAlgorithm()
+        with pytest.raises(ValueError):
+            algo.check_crypto_key_type("key")  # type: ignore[arg-type]
+
     def test_none_algorithm_should_throw_exception_if_key_is_not_none(self):
         algo = NoneAlgorithm()
 
@@ -811,6 +818,7 @@ class TestAlgorithmsRFC7520:
 @crypto_required
 class TestOKPAlgorithms:
     hello_world_sig = b"Qxa47mk/azzUgmY2StAOguAd4P7YBLpyCfU3JdbaiWnXM4o4WibXwmIHvNYgN3frtE2fcyd8OYEaOiD/KiwkCg=="
+    hello_world_sig_pem = b"9ueQE7PT8uudHIQb2zZZ7tB7k1X3jeTnIfOVvGCINZejrqQbru1EXPeuMlGcQEZrGkLVcfMmr99W/+byxfppAg=="
     hello_world = b"Hello World!"
 
     def test_okp_ed25519_should_reject_non_string_key(self):
@@ -825,58 +833,94 @@ class TestOKPAlgorithms:
         with open(key_path("testkey_ed25519.pub")) as keyfile:
             algo.prepare_key(keyfile.read())
 
-    def test_okp_ed25519_sign_should_generate_correct_signature_value(self):
+    @pytest.mark.parametrize(
+        "private_key_file,public_key_file,sig_attr",
+        [
+            ("testkey_ed25519", "testkey_ed25519.pub", "hello_world_sig"),
+            ("testkey_ed25519.pem", "testkey_ed25519.pub.pem", "hello_world_sig_pem"),
+        ],
+    )
+    def test_okp_ed25519_sign_should_generate_correct_signature_value(
+        self, private_key_file, public_key_file, sig_attr
+    ):
         algo = OKPAlgorithm()
 
         jwt_message = self.hello_world
 
-        expected_sig = base64.b64decode(self.hello_world_sig)
+        expected_sig = base64.b64decode(getattr(self, sig_attr))
 
-        with open(key_path("testkey_ed25519")) as keyfile:
+        with open(key_path(private_key_file)) as keyfile:
             jwt_key = cast(Ed25519PrivateKey, algo.prepare_key(keyfile.read()))
 
-        with open(key_path("testkey_ed25519.pub")) as keyfile:
+        with open(key_path(public_key_file)) as keyfile:
             jwt_pub_key = cast(Ed25519PublicKey, algo.prepare_key(keyfile.read()))
 
         algo.sign(jwt_message, jwt_key)
         result = algo.verify(jwt_message, jwt_pub_key, expected_sig)
         assert result
 
-    def test_okp_ed25519_verify_should_return_false_if_signature_invalid(self):
+    @pytest.mark.parametrize(
+        "public_key_file,sig_attr",
+        [
+            ("testkey_ed25519.pub", "hello_world_sig"),
+            ("testkey_ed25519.pub.pem", "hello_world_sig_pem"),
+        ],
+    )
+    def test_okp_ed25519_verify_should_return_false_if_signature_invalid(
+        self, public_key_file, sig_attr
+    ):
         algo = OKPAlgorithm()
 
         jwt_message = self.hello_world
-        jwt_sig = base64.b64decode(self.hello_world_sig)
+        jwt_sig = base64.b64decode(getattr(self, sig_attr))
 
         jwt_sig += b"123"  # Signature is now invalid
 
-        with open(key_path("testkey_ed25519.pub")) as keyfile:
+        with open(key_path(public_key_file)) as keyfile:
             jwt_pub_key = algo.prepare_key(keyfile.read())
 
         result = algo.verify(jwt_message, jwt_pub_key, jwt_sig)
         assert not result
 
-    def test_okp_ed25519_verify_should_return_true_if_signature_valid(self):
+    @pytest.mark.parametrize(
+        "public_key_file,sig_attr",
+        [
+            ("testkey_ed25519.pub", "hello_world_sig"),
+            ("testkey_ed25519.pub.pem", "hello_world_sig_pem"),
+        ],
+    )
+    def test_okp_ed25519_verify_should_return_true_if_signature_valid(
+        self, public_key_file, sig_attr
+    ):
         algo = OKPAlgorithm()
 
         jwt_message = self.hello_world
-        jwt_sig = base64.b64decode(self.hello_world_sig)
+        jwt_sig = base64.b64decode(getattr(self, sig_attr))
 
-        with open(key_path("testkey_ed25519.pub")) as keyfile:
+        with open(key_path(public_key_file)) as keyfile:
             jwt_pub_key = algo.prepare_key(keyfile.read())
 
         result = algo.verify(jwt_message, jwt_pub_key, jwt_sig)
         assert result
 
-    def test_okp_ed25519_prepare_key_should_be_idempotent(self):
+    @pytest.mark.parametrize(
+        "public_key_file", ("testkey_ed25519.pub", "testkey_ed25519.pub.pem")
+    )
+    def test_okp_ed25519_prepare_key_should_be_idempotent(self, public_key_file):
         algo = OKPAlgorithm()
 
-        with open(key_path("testkey_ed25519.pub")) as keyfile:
+        with open(key_path(public_key_file)) as keyfile:
             jwt_pub_key_first = algo.prepare_key(keyfile.read())
             jwt_pub_key_second = algo.prepare_key(jwt_pub_key_first)
 
         assert jwt_pub_key_first == jwt_pub_key_second
 
+    def test_okp_ed25519_prepare_key_should_reject_invalid_key(self):
+        algo = OKPAlgorithm()
+
+        with pytest.raises(InvalidKeyError):
+            algo.prepare_key("not a valid key")
+
     def test_okp_ed25519_jwk_private_key_should_parse_and_verify(self):
         algo = OKPAlgorithm()