fix: allow jwt.encode passing encoder_cls and jwt.decode passing decoder_cls

Author: lepture

src/joserfc/jwt.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import json
+from json import JSONEncoder, JSONDecoder
 from .rfc7519.claims import convert_claims
 from .rfc7519.claims import Claims as Claims
 from .rfc7519.claims import check_sensitive_data as check_sensitive_data
@@ -49,18 +50,20 @@ def encode(
         claims: Claims,
         key: KeyFlexible,
         algorithms: list[str] | None = None,
-        registry: JWSRegistry | JWERegistry | None = None) -> str:
+        registry: JWSRegistry | JWERegistry | None = None,
+        encoder_cls: JSONEncoder | None = None) -> str:
     """Encode a JSON Web Token with the given header, and claims.
 
     :param header: A dict of the JWT header
     :param claims: A dict of the JWT claims to be encoded
     :param key: key used to sign the signature
     :param algorithms: a list of allowed algorithms
     :param registry: a ``JWSRegistry`` or ``JWERegistry`` to use
+    :param encoder_cls: A JSONEncoder subclass to use
     """
     # add ``typ`` in header
     _header = {"typ": "JWT", **header}
-    payload = convert_claims(claims)
+    payload = convert_claims(claims, encoder_cls)
     if isinstance(registry, JWERegistry):
         return encrypt_compact(_header, payload, key, algorithms, registry)
     else:
@@ -71,14 +74,16 @@ def decode(
         value: bytes | str,
         key: KeyFlexible,
         algorithms: list[str] | None = None,
-        registry: JWSRegistry | JWERegistry | None = None) -> Token:
+        registry: JWSRegistry | JWERegistry | None = None,
+        decoder_cls: JSONDecoder | None = None) -> Token:
     """Decode the JSON Web Token string with the given key, and validate
     it with the claims requests.
 
     :param value: text of the JWT
     :param key: key used to verify the signature
     :param algorithms: a list of allowed algorithms
     :param registry: a ``JWSRegistry`` or ``JWERegistry`` to use
+    :param decoder_cls: A JSONDecoder subclass to use
     :raise: BadSignatureError
     """
     _value = to_bytes(value)
@@ -90,7 +95,7 @@ def decode(
         header, payload = _decode_jws(_value, key, algorithms, registry)
 
     try:
-        claims: Claims = json.loads(payload)
+        claims: Claims = json.loads(payload, cls=decoder_cls)
     except (TypeError, ValueError):
         raise InvalidPayloadError()
 

src/joserfc/rfc7516/message.py

@@ -56,7 +56,7 @@ def perform_encrypt(obj: EncryptionData, registry: JWERegistry) -> None:
         plaintext = obj.plaintext
 
     # Step 13, Compute the Encoded Protected Header value BASE64URL(UTF8(JWE Protected Header)).
-    aad = json_b64encode(obj.protected, "ascii")
+    aad = json_b64encode(obj.protected)
 
     # Step 14, Let the Additional Authenticated Data encryption parameter be
     # ASCII(Encoded Protected Header).  However, if a JWE AAD value is
@@ -117,7 +117,7 @@ def _perform_decrypt(obj: EncryptionData, registry: JWERegistry) -> None:
     if len(cek) * 8 != enc.cek_size:  # pragma: no cover
         raise InvalidCEKLengthError(f"A key of size {enc.cek_size} bits MUST be used")
 
-    aad = json_b64encode(obj.protected, "ascii")
+    aad = json_b64encode(obj.protected)
     if isinstance(obj, BaseJSONEncryption) and obj.aad:
         aad = aad + b"." + urlsafe_b64encode(obj.aad)
 

src/joserfc/rfc7519/claims.py

@@ -1,8 +1,12 @@
+from __future__ import annotations
+
 import re
+import json
 import datetime
 import calendar
+from json import JSONEncoder
 from typing import Dict, Any
-from ..util import to_bytes, json_dumps
+from ..util import to_bytes
 from ..errors import InsecureClaimError
 
 
@@ -22,13 +26,15 @@
 Claims = Dict[str, Any]
 
 
-def convert_claims(claims: Claims) -> bytes:
+def convert_claims(claims: Claims, encoder_cls: JSONEncoder | None = None) -> bytes:
     """Turn claims into bytes payload."""
     for k in ["exp", "iat", "nbf"]:
         claim = claims.get(k)
         if isinstance(claim, datetime.datetime):
             claims[k] = calendar.timegm(claim.utctimetuple())
-    return to_bytes(json_dumps(claims))
+
+    content = json.dumps(claims, ensure_ascii=False, separators=(",", ":"), cls=encoder_cls)
+    return to_bytes(content)
 
 
 def check_sensitive_data(claims: Claims) -> None:

src/joserfc/rfc7638/__init__.py

@@ -1,7 +1,8 @@
 import typing as t
+import json
 import hashlib
 from collections import OrderedDict
-from ..util import to_bytes, json_dumps, urlsafe_b64encode
+from ..util import to_bytes, urlsafe_b64encode
 
 
 def thumbprint(
@@ -15,7 +16,7 @@ def thumbprint(
     for k in sorted_fields:
         data[k] = dict_value[k]
 
-    json_data = json_dumps(data)
+    json_data = json.dumps(data, ensure_ascii=True, separators=(",", ":"))
     hash_value = hashlib.new(digest_method, to_bytes(json_data))
     digest_data = hash_value.digest()
     return urlsafe_b64encode(digest_data).decode("utf-8")

src/joserfc/util.py

@@ -22,10 +22,6 @@ def to_str(x: bytes | str, charset: str = "utf-8") -> str:
     return x
 
 
-def json_dumps(data: Any, ensure_ascii: bool = False) -> str:
-    return json.dumps(data, ensure_ascii=ensure_ascii, separators=(",", ":"))
-
-
 def urlsafe_b64decode(s: bytes) -> bytes:
     if b"+" in s or b"/" in s:
         raise binascii.Error
@@ -51,11 +47,11 @@ def int_to_base64(num: int) -> str:
     return urlsafe_b64encode(s).decode("utf-8", "strict")
 
 
-def json_b64encode(text: Any, charset: str = "utf-8") -> bytes:
+def json_b64encode(text: Any) -> bytes:
     if isinstance(text, dict):
-        text = json_dumps(text)
-    return urlsafe_b64encode(to_bytes(text, charset))
+        text = json.dumps(text, ensure_ascii=True, separators=(",", ":"))
+    return urlsafe_b64encode(to_bytes(text, "ascii"))
 
 
-def json_b64decode(text: Any, charset: str = "utf-8") -> Any:
-    return json.loads(urlsafe_b64decode(to_bytes(text, charset)))
+def json_b64decode(text: Any) -> Any:
+    return json.loads(urlsafe_b64decode(to_bytes(text, "ascii")))

tests/jwt/test_claims.py

@@ -1,5 +1,7 @@
 import time
 import datetime
+import json
+import uuid
 from unittest import TestCase
 from joserfc import jwt
 from joserfc.jwk import OctKey
@@ -12,6 +14,13 @@
 )
 
 
+class UUIDEncoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, uuid.UUID):
+            return str(o)
+        return super().default(o)
+
+
 class TestJWTClaims(TestCase):
     def test_check_sensitive_data(self):
         jwt.check_sensitive_data({})
@@ -135,3 +144,11 @@ def test_validate_nbf(self):
         claims_requests = jwt.JWTClaimsRegistry(now=now, leeway=500)
         claims_requests.validate({"nbf": now})
         self.assertRaises(InvalidTokenError, claims_requests.validate, {"nbf": now + 1000})
+
+    def test_claims_with_uuid_field(self):
+        value = uuid.uuid4()
+        claims = {"uuid": value}
+        key = OctKey.import_key("secret")
+        encoded_text = jwt.encode({"alg": "HS256"}, claims, key)
+        decoded_data = jwt.decode(encoded_text, key)
+        self.assertEqual(decoded_data, {"uuid": str(value)})