Repository: jpadilla/pyjwt
Pull request #829 (2022-11-22 11:54:35)

Add PyJWT._{de,en}code_payload hooks
This enables custom subclasses of `PyJWT` to customize decoding of the `payload` data of the JWS.

There's an example on how to subclass the function in the test case.

Refs #666 (reimplementation of it)
Refs #753 (reimplementation of it, borrows test case)


russoale: Alternatively, this could be implemented as well by passing a callable 
```python 
    def decode(
        self,
        jwt: str,
        ...
        # hook for custom payload decompression 
        decompress_payload: Optional[Callable[Dict]] = None,
        ...
        **kwargs,
    ) -> Dict[str, Any]:

   # passing to 

    def decode_complete(
        self,
        jwt: str,
                ...
        # hook for custom payload decompression 
        decompress_payload: Optional[Callable[Dict]] = None,
        ...
        # kwargs
        **kwargs,
    ) -> Dict[str, Any]:
```

Then checking if the function has been passed and use it or proceed with default
akx: @russoale Possibly, but those are pretty arg-heavy functions as-is. Also, "decompress" is not a good name – the function can do something else than just decompress.
simon-sk: @akx Would it be possible to add a similar hook for the `encode` function as well?
akx: @simon-sk Good idea. Added.

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 25d8d51..16edcb3 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -12,7 +12,7 @@ repos:
         args: ["--target-version=py37"]
 
   - repo: https://github.com/PyCQA/flake8
-    rev: 5.0.4
+    rev: 6.0.0
     hooks:
       - id: flake8
         language_version: python3.8
@@ -23,14 +23,14 @@ repos:
       - id: isort
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v4.4.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
       - id: debug-statements
 
   - repo: https://github.com/mgedmin/check-manifest
-    rev: "0.48"
+    rev: "0.49"
     hooks:
       - id: check-manifest
         args: [--no-build-isolation]
diff --git a/jwt/api_jwt.py b/jwt/api_jwt.py
index 77857d2..e1df3c7 100644
--- a/jwt/api_jwt.py
+++ b/jwt/api_jwt.py
@@ -60,12 +60,32 @@ class PyJWT:
             if isinstance(payload.get(time_claim), datetime):
                 payload[time_claim] = timegm(payload[time_claim].utctimetuple())
 
-        json_payload = json.dumps(
-            payload, separators=(",", ":"), cls=json_encoder
-        ).encode("utf-8")
+        json_payload = self._encode_payload(
+            payload,
+            headers=headers,
+            json_encoder=json_encoder,
+        )
 
         return api_jws.encode(json_payload, key, algorithm, headers, json_encoder)
 
+    def _encode_payload(
+        self,
+        payload: Dict[str, Any],
+        headers: Optional[Dict[str, Any]] = None,
+        json_encoder: Optional[Type[json.JSONEncoder]] = None,
+    ) -> bytes:
+        """
+        Encode a given payload to the bytes to be signed.
+
+        This method is intended to be overridden by subclasses that need to
+        encode the payload in a different way, e.g. compress the payload.
+        """
+        return json.dumps(
+            payload,
+            separators=(",", ":"),
+            cls=json_encoder,
+        ).encode("utf-8")
+
     def decode_complete(
         self,
         jwt: str,
@@ -125,12 +145,7 @@ class PyJWT:
             detached_payload=detached_payload,
         )
 
-        try:
-            payload = json.loads(decoded["payload"])
-        except ValueError as e:
-            raise DecodeError(f"Invalid payload string: {e}")
-        if not isinstance(payload, dict):
-            raise DecodeError("Invalid payload string: must be a json object")
+        payload = self._decode_payload(decoded)
 
         merged_options = {**self.options, **options}
         self._validate_claims(
@@ -140,6 +155,22 @@ class PyJWT:
         decoded["payload"] = payload
         return decoded
 
+    def _decode_payload(self, decoded: Dict[str, Any]) -> Any:
+        """
+        Decode the payload from a JWS dictionary (payload, signature, header).
+
+        This method is intended to be overridden by subclasses that need to
+        decode the payload in a different way, e.g. decompress compressed
+        payloads.
+        """
+        try:
+            payload = json.loads(decoded["payload"])
+        except ValueError as e:
+            raise DecodeError(f"Invalid payload string: {e}")
+        if not isinstance(payload, dict):
+            raise DecodeError("Invalid payload string: must be a json object")
+        return payload
+
     def decode(
         self,
         jwt: str,


Test patch:
diff --git a/tests/test_compressed_jwt.py b/tests/test_compressed_jwt.py
new file mode 100644
index 0000000..21fac3f
--- /dev/null
+++ b/tests/test_compressed_jwt.py
@@ -0,0 +1,37 @@
+import json
+import zlib
+
+from jwt import PyJWT
+
+
+class CompressedPyJWT(PyJWT):
+    def _decode_payload(self, decoded):
+        return json.loads(
+            # wbits=-15 has zlib not worry about headers of crc's
+            zlib.decompress(decoded["payload"], wbits=-15).decode("utf-8")
+        )
+
+
+def test_decodes_complete_valid_jwt_with_compressed_payload():
+    # Test case from https://github.com/jpadilla/pyjwt/pull/753/files
+    example_payload = {"hello": "world"}
+    example_secret = "secret"
+    # payload made with the pako (https://nodeca.github.io/pako/) library in Javascript:
+    # Buffer.from(pako.deflateRaw('{"hello": "world"}')).toString('base64')
+    example_jwt = (
+        b"eyJhbGciOiAiSFMyNTYiLCAidHlwIjogIkpXVCJ9"
+        b".q1bKSM3JyVeyUlAqzy/KSVGqBQA="
+        b".08wHYeuh1rJXmcBcMrz6NxmbxAnCQp2rGTKfRNIkxiw="
+    )
+    decoded = CompressedPyJWT().decode_complete(
+        example_jwt, example_secret, algorithms=["HS256"]
+    )
+
+    assert decoded == {
+        "header": {"alg": "HS256", "typ": "JWT"},
+        "payload": example_payload,
+        "signature": (
+            b"\xd3\xcc\x07a\xeb\xa1\xd6\xb2W\x99\xc0\\2\xbc\xfa7"
+            b"\x19\x9b\xc4\t\xc2B\x9d\xab\x192\x9fD\xd2$\xc6,"
+        ),
+    }