fix(jws): validating content size to avoid DoS

Author: lepture

.pre-commit-config.yaml

@@ -7,11 +7,6 @@ repos:
         args: [--fix, --exit-non-zero-on-fix]
       - id: ruff-format
 
-  - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: 'v1.18.1'
-    hooks:
-      - id: mypy
-
   - repo: https://github.com/codespell-project/codespell
     rev: 'v2.4.1'
     hooks:

src/joserfc/_rfc7515/json.py

@@ -14,6 +14,7 @@
 from .registry import JWSRegistry
 from ..registry import reject_unprotected_crit_header
 from ..util import (
+    to_bytes,
     json_b64encode,
     json_b64decode,
     urlsafe_b64encode,
@@ -86,25 +87,27 @@ def sign_json_member(
     return rv
 
 
-def extract_general_json(value: GeneralJSONSerialization) -> GeneralJSONSignature:
+def extract_general_json(value: GeneralJSONSerialization, registry: JWSRegistry) -> GeneralJSONSignature:
     payload_segment: bytes = value["payload"].encode("utf-8")
+    registry.validate_payload_size(payload_segment)
     try:
         payload = urlsafe_b64decode(payload_segment)
     except (TypeError, ValueError):
         raise DecodeError("Invalid payload")
 
     signatures: list[JSONSignatureDict] = value["signatures"]
-    members = [__signature_to_member(sig) for sig in signatures]
+    members = [__signature_to_member(sig, registry) for sig in signatures]
     obj = GeneralJSONSignature(members, payload)
     obj.signatures = signatures
     obj.segments = {"payload": payload_segment}
     return obj
 
 
-def __signature_to_member(sig: JSONSignatureDict) -> HeaderMember:
+def __signature_to_member(sig: JSONSignatureDict, registry: JWSRegistry) -> HeaderMember:
     member = HeaderMember()
     if "protected" in sig:
-        protected_segment = sig["protected"]
+        protected_segment = to_bytes(sig["protected"])
+        registry.validate_header_size(protected_segment)
         member.protected = json_b64decode(protected_segment)
     if "header" in sig:
         member.header = sig["header"]
@@ -139,11 +142,16 @@ def verify_signature(
     alg = registry.get_alg(headers["alg"])
     key = find_key(member)
     alg.check_key(key)
+
     if "protected" in signature:
-        protected_segment = signature["protected"].encode("utf-8")
+        protected_segment = to_bytes(signature["protected"])
     else:
         protected_segment = b""
-    sig = urlsafe_b64decode(signature["signature"].encode("utf-8"))
+
+    signature_segment = to_bytes(signature["signature"])
+    registry.validate_signature_size(signature_segment)
+
+    sig = urlsafe_b64decode(signature_segment)
     signing_input = b".".join([protected_segment, payload_segment])
     return alg.verify(signing_input, sig, key)
 

src/joserfc/_rfc7515/registry.py

@@ -3,7 +3,12 @@
 from typing import Any
 from enum import Enum
 from .model import JWSAlgModel
-from ..errors import UnsupportedAlgorithmError, SecurityWarning, JoseError
+from ..errors import (
+    UnsupportedAlgorithmError,
+    SecurityWarning,
+    JoseError,
+    ExceededSizeError,
+)
 from ..registry import (
     JWS_HEADER_REGISTRY,
     Header,
@@ -40,6 +45,13 @@ class Strategy(Enum):
     algorithms: dict[str, JWSAlgModel] = {}
     recommended: list[str] = []
 
+    #: max header content's size in bytes
+    max_header_length: int = 512
+    #: max payload content's size in bytes
+    max_payload_length: int = 8000
+    #: max signature's size in bytes
+    max_signature_length: int = 1024
+
     def __init__(
         self,
         header_registry: HeaderRegistryDict | None = None,
@@ -87,6 +99,18 @@ def check_header(self, header: Header) -> None:
         if self.strict_check_header:
             check_supported_header(self.header_registry, header)
 
+    def validate_header_size(self, header: bytes) -> None:
+        if header and len(header) > self.max_header_length:
+            raise ExceededSizeError(f"Header size of '{header!r}' exceeds {self.max_header_length} bytes.")
+
+    def validate_payload_size(self, payload: bytes) -> None:
+        if payload and len(payload) > self.max_payload_length:
+            raise ExceededSizeError(f"Payload size of '{payload!r}' exceeds {self.max_payload_length} bytes.")
+
+    def validate_signature_size(self, signature: bytes) -> None:
+        if len(signature) > self.max_signature_length:
+            raise ExceededSizeError(f"Signature of '{signature!r}' exceeds {self.max_signature_length} bytes.")
+
     @classmethod
     def guess_alg(cls, key: Any, strategy: Strategy) -> str | None:
         """Guess the JWS algorithm for a given key.

src/joserfc/_rfc7797/compact.py

@@ -10,6 +10,7 @@
 )
 from .._rfc7515.model import JWSAlgModel, CompactSignature
 from .._rfc7515.compact import decode_header
+from .._rfc7515.registry import JWSRegistry, default_registry
 from ..errors import DecodeError
 from .util import is_rfc7797_enabled
 
@@ -27,18 +28,28 @@ def sign_rfc7515_compact(obj: CompactSignature, alg: JWSAlgModel, key: Any) -> b
     return out
 
 
-def extract_rfc7515_compact(value: bytes, payload: bytes | str | None = None) -> CompactSignature:
+def extract_rfc7515_compact(
+    value: bytes, payload: bytes | str | None = None, registry: JWSRegistry | None = None
+) -> CompactSignature:
     """Extract the JWS Compact Serialization from bytes to object.
 
     :param value: JWS in bytes
     :param payload: optional payload, required with detached content
+    :param registry: optional JWSRegistry instance
     :raise: DecodeError
     """
     parts = value.split(b".")
     if len(parts) != 3:
         raise DecodeError("Invalid JSON Web Signature")
 
+    if registry is None:
+        registry = default_registry
+
     header_segment, payload_segment, signature_segment = parts
+
+    registry.validate_header_size(header_segment)
+    registry.validate_signature_size(signature_segment)
+
     protected = decode_header(header_segment)
 
     if is_rfc7797_enabled(protected):
@@ -50,6 +61,7 @@ def extract_rfc7515_compact(value: bytes, payload: bytes | str | None = None) ->
             payload = to_bytes(payload)
             payload_segment = urlsafe_b64encode(payload)
         else:
+            registry.validate_payload_size(payload_segment)
             try:
                 payload = urlsafe_b64decode(payload_segment)
             except (TypeError, ValueError):

src/joserfc/_rfc7797/json.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from .._rfc7515.types import FlattenedJSONSerialization, JSONSignatureDict
 from .._rfc7515.model import HeaderMember, FlattenedJSONSignature
 from .._rfc7515.registry import JWSRegistry
@@ -18,9 +19,10 @@ def sign_rfc7797_json(
     return data
 
 
-def extract_rfc7797_json(value: FlattenedJSONSerialization) -> FlattenedJSONSignature:
+def extract_rfc7797_json(value: FlattenedJSONSerialization, registry: JWSRegistry) -> FlattenedJSONSignature:
     if "protected" in value:
         protected_segment = to_bytes(value["protected"])
+        registry.validate_header_size(protected_segment)
         protected = json_b64decode(protected_segment)
     else:
         protected = None
@@ -32,6 +34,7 @@ def extract_rfc7797_json(value: FlattenedJSONSerialization) -> FlattenedJSONSign
     if is_rfc7797_enabled(member.headers()):
         payload = payload_segment
     else:
+        registry.validate_payload_size(payload_segment)
         try:
             payload = urlsafe_b64decode(payload_segment)
         except (TypeError, ValueError):

src/joserfc/errors.py

@@ -135,8 +135,8 @@ class BadSignatureError(JoseError):
 
 
 class ExceededSizeError(JoseError):
-    """This error is designed for DEF zip algorithm. It raised when the
-    compressed data exceeds the maximum allowed length."""
+    """This error is designed for validating the token's content size.
+    It raised when the data exceeds the maximum allowed length."""
 
     error = "exceeded_size"
 

src/joserfc/jws.py

@@ -180,7 +180,7 @@ def deserialize_compact(
     :param payload: optional payload, required with detached content
     :return: object of the ``CompactSignature``
     """
-    obj = extract_compact(to_bytes(value), payload)
+    obj = extract_compact(to_bytes(value), payload, registry)
     if not validate_compact(obj, public_key, algorithms, registry):
         raise BadSignatureError()
     return obj
@@ -296,12 +296,12 @@ def find_key(obj: HeaderMember) -> Key:
         return guess_key(public_key, obj, use="sig")
 
     if "signatures" in value:
-        general_obj = extract_general_json(value)
+        general_obj = extract_general_json(value, registry)
         if not verify_general_json(general_obj, registry, find_key):
             raise BadSignatureError()
         return general_obj
     else:
-        flattened_obj = extract_flattened_json(value)
+        flattened_obj = extract_flattened_json(value, registry)
         if not verify_flattened_json(flattened_obj, registry, find_key):
             raise BadSignatureError()
         return flattened_obj

tests/jws/test_compact.py

@@ -12,17 +12,20 @@
     MissingAlgorithmError,
     UnsupportedAlgorithmError,
     UnsupportedHeaderError,
+    ExceededSizeError,
 )
+from joserfc.util import urlsafe_b64encode, json_b64encode
 
 
 class TestCompact(TestCase):
+    key = OctKey.import_key("secret")
+
     def test_registry_is_none(self):
-        key = OctKey.import_key("secret")
-        value = serialize_compact({"alg": "HS256"}, b"foo", key)
+        value = serialize_compact({"alg": "HS256"}, b"foo", self.key)
         expected = "eyJhbGciOiJIUzI1NiJ9.Zm9v.0pehoi-RMZM1jl-4TP_C4Y6BJ-bcmsuzfDyQpkpJkh0"
         self.assertEqual(value, expected)
 
-        obj = deserialize_compact(value, key)
+        obj = deserialize_compact(value, self.key)
         self.assertEqual(obj.payload, b"foo")
 
     def test_bad_signature_error(self):
@@ -31,28 +34,42 @@ def test_bad_signature_error(self):
         self.assertRaises(BadSignatureError, deserialize_compact, value, key)
 
     def test_raise_unsupported_algorithm_error(self):
-        key = OctKey.import_key("secret")
-        self.assertRaises(UnsupportedAlgorithmError, serialize_compact, {"alg": "HS512"}, b"foo", key)
-        self.assertRaises(UnsupportedAlgorithmError, serialize_compact, {"alg": "NOT"}, b"foo", key)
+        self.assertRaises(UnsupportedAlgorithmError, serialize_compact, {"alg": "HS512"}, b"foo", self.key)
+        self.assertRaises(UnsupportedAlgorithmError, serialize_compact, {"alg": "NOT"}, b"foo", self.key)
 
     def test_invalid_length(self):
-        key = OctKey.import_key("secret")
-        self.assertRaises(DecodeError, deserialize_compact, b"a.b.c.d", key)
+        self.assertRaises(DecodeError, deserialize_compact, b"a.b.c.d", self.key)
 
     def test_no_invalid_header(self):
         # invalid base64
         value = b"abc.Zm9v.0pehoi"
-        key = OctKey.import_key("secret")
-        self.assertRaises(DecodeError, deserialize_compact, value, key)
+        self.assertRaises(DecodeError, deserialize_compact, value, self.key)
 
         # no alg value
         value = b"eyJhIjoiYiJ9.Zm9v.0pehoi"
-        self.assertRaises(MissingAlgorithmError, deserialize_compact, value, key)
+        self.assertRaises(MissingAlgorithmError, deserialize_compact, value, self.key)
 
     def test_invalid_payload(self):
         value = b"eyJhbGciOiJIUzI1NiJ9.a$b.0pehoi"
-        key = OctKey.import_key("secret")
-        self.assertRaises(DecodeError, deserialize_compact, value, key)
+        self.assertRaises(DecodeError, deserialize_compact, value, self.key)
+
+    def test_header_exceeded_size_error(self):
+        exceeded_header = json_b64encode({f"a{i}": f"a{i}" for i in range(1000)})
+        other = urlsafe_b64encode(b"o")
+        fake_jws = exceeded_header + b"." + other + b"." + other
+        self.assertRaises(ExceededSizeError, deserialize_compact, fake_jws, self.key)
+
+    def test_payload_exceeded_size_error(self):
+        header = json_b64encode({"alg": "HS256"})
+        exceeded_payload = urlsafe_b64encode(("o" * 10000).encode("utf8"))
+        fake_jws = header + b"." + exceeded_payload + b"." + urlsafe_b64encode(b"o")
+        self.assertRaises(ExceededSizeError, deserialize_compact, fake_jws, self.key)
+
+    def test_signature_exceeded_size_error(self):
+        header = json_b64encode({"alg": "HS256"})
+        exceeded_signature = urlsafe_b64encode(("o" * 1000).encode("utf8"))
+        fake_jws = header + b"." + urlsafe_b64encode(b"o") + b"." + exceeded_signature
+        self.assertRaises(ExceededSizeError, deserialize_compact, fake_jws, self.key)
 
     def test_with_key_set(self):
         keys = KeySet(
@@ -73,20 +90,18 @@ def test_with_key_set(self):
 
     def test_strict_check_header(self):
         header = {"alg": "HS256", "custom": "hi"}
-        key = OctKey.import_key("secret")
-        self.assertRaises(UnsupportedHeaderError, serialize_compact, header, b"hi", key)
+        self.assertRaises(UnsupportedHeaderError, serialize_compact, header, b"hi", self.key)
 
         registry = JWSRegistry(strict_check_header=False)
-        serialize_compact(header, b"hi", key, registry=registry)
+        serialize_compact(header, b"hi", self.key, registry=registry)
 
     def test_non_canonical_signature_encoding(self):
         text = "eyJhbGciOiJIUzI1NiJ9.eyJ1c2VyIjoiYWRtaW4ifQ.VI29GgHzuh2xfF0bkRYvZIsSuQnbTXSIvuRyt7RDrwo"[:-1] + "p"
         self.assertRaises(BadSignatureError, deserialize_compact, text, OctKey.import_key("secret"))
 
     def test_detached_content(self):
-        key = OctKey.import_key("secret")
-        value = detach_content(serialize_compact({"alg": "HS256"}, b"foo", key))
+        value = detach_content(serialize_compact({"alg": "HS256"}, b"foo", self.key))
         expected = "eyJhbGciOiJIUzI1NiJ9..0pehoi-RMZM1jl-4TP_C4Y6BJ-bcmsuzfDyQpkpJkh0"
         self.assertEqual(value, expected)
-        obj = deserialize_compact(value, key, payload=b"foo")
+        obj = deserialize_compact(value, self.key, payload=b"foo")
         self.assertEqual(obj.payload, b"foo")