Repository: jpadilla/pyjwt
Pull request #832 (2022-11-28 14:37:57)

Add `sort_headers` parameter to `api_jwt.encode`
This allows you to not sort headers, which prevents a breaking change between v2.4.0 and v2.5.0

fixes https://github.com/jpadilla/pyjwt/issues/811

Should I create a unittest for this?

evroon: @auvipy I have added a unittest, but it only makes sure that encoding and decoding is the same when using `sort_headers=True` and `sort_headers=False`, as well as testing that the encoded payload is not the same between those two options. But I don't see a way to check that the headers are properly sorted before the payload is encoded.

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_jws.py b/jwt/api_jws.py
index ab8490f..c914db1 100644
--- a/jwt/api_jws.py
+++ b/jwt/api_jws.py
@@ -101,6 +101,7 @@ class PyJWS:
         headers: dict[str, Any] | None = None,
         json_encoder: Type[json.JSONEncoder] | None = None,
         is_payload_detached: bool = False,
+        sort_headers: bool = True,
     ) -> str:
         segments = []
 
@@ -133,9 +134,8 @@ class PyJWS:
             # True is the standard value for b64, so no need for it
             del header["b64"]
 
-        # Fix for headers misorder - issue #715
         json_header = json.dumps(
-            header, separators=(",", ":"), cls=json_encoder, sort_keys=True
+            header, separators=(",", ":"), cls=json_encoder, sort_keys=sort_headers
         ).encode()
 
         segments.append(base64url_encode(json_header))
diff --git a/jwt/api_jwt.py b/jwt/api_jwt.py
index 77857d2..5d21afc 100644
--- a/jwt/api_jwt.py
+++ b/jwt/api_jwt.py
@@ -45,6 +45,7 @@ class PyJWT:
         algorithm: Optional[str] = "HS256",
         headers: Optional[Dict[str, Any]] = None,
         json_encoder: Optional[Type[json.JSONEncoder]] = None,
+        sort_headers: bool = True,
     ) -> str:
         # Check that we get a mapping
         if not isinstance(payload, Mapping):
@@ -60,11 +61,38 @@ 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,
+            sort_headers=sort_headers,
+        )
 
-        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,
@@ -125,12 +153,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 +163,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_api_jws.py b/tests/test_api_jws.py
index cfbbe21..fe6c2d4 100644
--- a/tests/test_api_jws.py
+++ b/tests/test_api_jws.py
@@ -414,6 +414,17 @@ class TestJWS:
 
         assert decoded_payload == payload
 
+    @pytest.mark.parametrize("sort_headers", (False, True))
+    def test_sorting_of_headers(self, jws, payload, sort_headers):
+        jws_message = jws.encode(
+            payload,
+            key="\xc2",
+            headers={"b": "1", "a": "2"},
+            sort_headers=sort_headers,
+        )
+        header_json = base64url_decode(jws_message.split(".")[0])
+        assert sort_headers == (header_json.index(b'"a"') < header_json.index(b'"b"'))
+
     def test_decode_invalid_header_padding(self, jws):
         example_jws = (
             "aeyJhbGciOiAiSFMyNTYiLCAidHlwIjogIkpXVCJ9"
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,"
+        ),
+    }