Repository: jpadilla/pyjwt

Issue:
Support exposure of hash algorithm digest to handle OIDC at_hash, potentially other spec extensions
#296 was closed with the suggestion that rather than special-case support for OIDC, PyJWT could support pluggable verifiers which handle additional, specialized claim validation behavior.

If there's a general market for verifiers -- or if they're going to be prepackaged and provided in this library as optional extensions -- that's a good avenue, but at present #295 is the _only_ such issue.
Additionally, #258 runs directly counter to any extensible family of verifiers being hooked into PyJWT.

"Pluggable verifiers" seem therefore, at first blush, to be overwrought as a solution. I want PyJWT to support external verification of a claim based on the following information:
- content of an arbitrary claim in the payload (trivially easy)
- a digest from the hash algorithm specified in the JWT header (right now, hard, should be easy)

Client code should be able to extract this info without having to worry about how PyJWT behaves internally when `cryptography` is present.

The heart of the #296 implementation is really quite short:
```python
        alg_obj = self.get_algo_by_name(algorithm)  # self is a PyJWT object
        hash_alg = alg_obj.hash_alg

        def get_digest(bytestr):
            if has_crypto and (
                    isinstance(hash_alg, type) and
                    issubclass(hash_alg, hashes.HashAlgorithm)):
                digest = hashes.Hash(hash_alg(), backend=default_backend())
                digest.update(bytestr)
                return digest.finalize()
            else:
                return hash_alg(bytestr).digest()
```
where `get_algo_by_name` is [this segment](https://github.com/jpadilla/pyjwt/blob/7f7d524b78ab98e42e9064f4f2bab5b6e126f766/jwt/api_jws.py#L99-L111) abstracted into a function.


Having this hooked into the `PyJWT.decode()` call is not in any way necessary -- an external method like `PyJWT.compute_hash_digest(<string>)` would satisfy this just fine.

OIDC Client code would then look like this:
```python
def verify_at_hash(decoded_id_token, access_token):
    digest = decoded_id_token.compute_hash_digest(access_token.encode('utf-8'))
    computed_at_hash = base64.urlsafe_b64encode(digest[:(len(digest) // 2)]).rstrip('=').decode('utf-8')
    assert decoded_id_token['at_hash'] == computed_at_hash

decoded = jwt.decode(...)
verify_at_hash(decoded, access_token)
```

However, there's a significant, nasty wrinkle: `PyJWT.decode()` only returns the payload and discards the header.

One of my initial inclinations is to make the payload a subclass of `dict` with additional attributes for `signing_input, header, signature` -- the other things returned from `PyJWT._load()` -- and any useful helper methods (in this case, `compute_hash_digest`).
This solution fixes the decision by `PyJWT.decode()` to discard potentially useful information. In fact, this is the critical information which makes it presently impossible to verify the OIDC `at_hash` with PyJWT.

If that's unattractive, the big alternative is to add `jwt.compute_hash_digest(header, <input_string>)`, and add an argument to `PyJWT.decode` like `verification_callbacks=...` which is an iterable of callables which consume `payload, header, signing_input, signature` as inputs.


@mark-adams, which of these directions for PyJWT is more appealing to you? Do you see issues with exposing `compute_hash_digest` in one form or another?

Subclassing `dict` may sound messier at first, but it results in less code for PyJWT to maintain and doesn't expand the public API surface significantly vs `verification_callbacks`.
It also has a nice benefit over `verification_callbacks` in that adding attributes to an object if necessary in the future won't break most sane usage, but if people write their `verification_callbacks` without accepting `**kwargs` (which would be documented as incorrect, but people would still do it) then adding more information into those callbacks would break those clients' usage.

Discussion:

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 8ec7ede..4d56207 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -18,6 +18,10 @@ Fixed
 Added
 ~~~~~
 
+- Add ``compute_hash_digest`` as a method of ``Algorithm`` objects, which uses
+  the underlying hash algorithm to compute a digest. If there is no appropriate
+  hash algorithm, a ``NotImplementedError`` will be raised
+
 `v2.6.0 <https://github.com/jpadilla/pyjwt/compare/2.5.0...2.6.0>`__
 -----------------------------------------------------------------------
 
diff --git a/docs/usage.rst b/docs/usage.rst
index 91d9679..a85fa18 100644
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -297,3 +297,74 @@ Retrieve RSA signing keys from a JWKS endpoint
     ... )
     >>> print(data)
     {'iss': 'https://dev-87evx9ru.auth0.com/', 'sub': 'aW4Cca79xReLWUz0aE2H6kD0O3cXBVtC@clients', 'aud': 'https://expenses-api', 'iat': 1572006954, 'exp': 1572006964, 'azp': 'aW4Cca79xReLWUz0aE2H6kD0O3cXBVtC', 'gty': 'client-credentials'}
+
+OIDC Login Flow
+---------------
+
+The following usage demonstrates an OIDC login flow using pyjwt. Further
+reading about the OIDC spec is recommended for implementers.
+
+In particular, this demonstrates validation of the ``at_hash`` claim.
+This claim relies on data from outside of the the JWT for validation. Methods
+are provided which support computation and validation of this claim, but it
+is not built into pyjwt.
+
+.. code-block:: python
+
+    import base64
+    import jwt
+    import requests
+
+
+    # Part 1: setup
+    # get the OIDC config and JWKs to use
+
+    # in OIDC, you must know your client_id (this is the OAuth 2.0 client_id)
+    client_id = ...
+
+    # example of fetching data from your OIDC server
+    # see: https://openid.net/specs/openid-connect-discovery-1_0.html#ProviderConfig
+    oidc_server = ...
+    oidc_config = requests.get(
+        f"https://{oidc_server}/.well-known/openid-configuration"
+    ).json()
+    signing_algos = oidc_config["id_token_signing_alg_values_supported"]
+
+    # setup a PyJWKClient to get the appropriate signing key
+    jwks_client = jwt.PyJWKClient(oidc_config["jwks_uri"])
+
+
+    # Part 2: login / authorization
+    # when a user completes an OIDC login flow, there will be a well-formed
+    # response object to parse/handle
+
+    # data from the login flow
+    # see: https://openid.net/specs/openid-connect-core-1_0.html#TokenResponse
+    token_response = ...
+    id_token = token_response["id_token"]
+    access_token = token_response["access_token"]
+
+
+    # Part 3: decode and validate at_hash
+    # after the login is complete, the id_token needs to be decoded
+    # this is the stage at which an OIDC client must verify the at_hash
+
+    # get signing_key from id_token
+    signing_key = jwks_client.get_signing_key_from_jwt(id_token)
+
+    # now, decode_complete to get payload + header
+    data = jwt.decode_complete(
+        id_token,
+        key=signing_key.key,
+        algorithms=signing_algos,
+        audience=client_id,
+    )
+    payload, header = data["payload"], data["header"]
+
+    # get the pyjwt algorithm object
+    alg_obj = jwt.get_algorithm_by_name(header["alg"])
+
+    # compute at_hash, then validate / assert
+    digest = alg_obj.compute_hash_digest(access_token)
+    at_hash = base64.urlsafe_b64encode(digest[: (len(digest) // 2)]).rstrip("=")
+    assert at_hash == payload["at_hash"]
diff --git a/jwt/algorithms.py b/jwt/algorithms.py
index 93fadf4..4fae441 100644
--- a/jwt/algorithms.py
+++ b/jwt/algorithms.py
@@ -18,6 +18,7 @@ from .utils import (
 try:
     import cryptography.exceptions
     from cryptography.exceptions import InvalidSignature
+    from cryptography.hazmat.backends import default_backend
     from cryptography.hazmat.primitives import hashes
     from cryptography.hazmat.primitives.asymmetric import ec, padding
     from cryptography.hazmat.primitives.asymmetric.ec import (
@@ -111,6 +112,28 @@ class Algorithm:
     The interface for an algorithm used to sign and verify tokens.
     """
 
+    def compute_hash_digest(self, bytestr: bytes) -> bytes:
+        """
+        Compute a hash digest using the specified algorithm's hash algorithm.
+
+        If there is no hash algorithm, raises a NotImplementedError.
+        """
+        # lookup self.hash_alg if defined in a way that mypy can understand
+        hash_alg = getattr(self, "hash_alg", None)
+        if hash_alg is None:
+            raise NotImplementedError
+
+        if (
+            has_crypto
+            and isinstance(hash_alg, type)
+            and issubclass(hash_alg, hashes.HashAlgorithm)
+        ):
+            digest = hashes.Hash(hash_alg(), backend=default_backend())
+            digest.update(bytestr)
+            return digest.finalize()
+        else:
+            return hash_alg(bytestr).digest()
+
     def prepare_key(self, key):
         """
         Performs necessary validation and conversions on the key and returns


Test patch:
diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py
index 538078a..894ce28 100644
--- a/tests/test_algorithms.py
+++ b/tests/test_algorithms.py
@@ -45,6 +45,12 @@ class TestAlgorithms:
         with pytest.raises(NotImplementedError):
             algo.to_jwk("value")
 
+    def test_algorithm_should_throw_exception_if_compute_hash_digest_not_impl(self):
+        algo = Algorithm()
+
+        with pytest.raises(NotImplementedError):
+            algo.compute_hash_digest(b"value")
+
     def test_none_algorithm_should_throw_exception_if_key_is_not_none(self):
         algo = NoneAlgorithm()
 
@@ -1054,3 +1060,20 @@ class TestOKPAlgorithms:
         signature_2 = algo.sign(b"Hello World!", priv_key_2)
         assert algo.verify(b"Hello World!", pub_key_2, signature_1)
         assert algo.verify(b"Hello World!", pub_key_2, signature_2)
+
+    @crypto_required
+    def test_rsa_can_compute_digest(self):
+        # this is the well-known sha256 hash of "foo"
+        foo_hash = base64.b64decode(b"LCa0a2j/xo/5m0U8HTBBNBNCLXBkg7+g+YpeiGJm564=")
+
+        algo = RSAAlgorithm(RSAAlgorithm.SHA256)
+        computed_hash = algo.compute_hash_digest(b"foo")
+        assert computed_hash == foo_hash
+
+    def test_hmac_can_compute_digest(self):
+        # this is the well-known sha256 hash of "foo"
+        foo_hash = base64.b64decode(b"LCa0a2j/xo/5m0U8HTBBNBNCLXBkg7+g+YpeiGJm564=")
+
+        algo = HMACAlgorithm(HMACAlgorithm.SHA256)
+        computed_hash = algo.compute_hash_digest(b"foo")
+        assert computed_hash == foo_hash