feat(jwe): add content size validation to avoid DoS

Author: lepture

src/joserfc/_rfc7515/registry.py

@@ -4,9 +4,9 @@
 from enum import Enum
 from .model import JWSAlgModel
 from ..errors import (
+    JoseError,
     UnsupportedAlgorithmError,
     SecurityWarning,
-    JoseError,
     ExceededSizeError,
 )
 from ..registry import (

src/joserfc/_rfc7516/compact.py

@@ -1,4 +1,5 @@
 from .models import CompactEncryption, Recipient
+from .registry import JWERegistry
 from .._keys import Key
 from ..errors import (
     MissingAlgorithmError,
@@ -32,12 +33,17 @@ def represent_compact(obj: CompactEncryption) -> bytes:
     )
 
 
-def extract_compact(value: bytes) -> CompactEncryption:
+def extract_compact(value: bytes, registry: JWERegistry) -> CompactEncryption:
     parts = value.split(b".")
     if len(parts) != 5:
         raise ValueError("Invalid JSON Web Encryption")
 
     header_segment, ek_segment, iv_segment, ciphertext_segment, tag_segment = parts
+    registry.validate_protected_header_size(header_segment)
+    registry.validate_encrypted_key_size(ek_segment)
+    registry.validate_initialization_vector_size(iv_segment)
+    registry.validate_ciphertext_size(ciphertext_segment)
+    registry.validate_auth_tag_size(tag_segment)
     try:
         protected = json_b64decode(header_segment)
         if "alg" not in protected:

src/joserfc/_rfc7516/json.py

@@ -1,10 +1,12 @@
+from __future__ import annotations
 import typing as t
 from .models import (
     BaseJSONEncryption,
     GeneralJSONEncryption,
     FlattenedJSONEncryption,
     Recipient,
 )
+from .registry import JWERegistry
 from .types import (
     JSONRecipientDict,
     GeneralJSONSerialization,
@@ -70,44 +72,49 @@ def __represent_json_serialization(obj: BaseJSONEncryption):  # type: ignore[no-
     return data
 
 
-def extract_general_json(data: GeneralJSONSerialization) -> GeneralJSONEncryption:
-    protected = json_b64decode(data["protected"])
+def extract_general_json(data: GeneralJSONSerialization, registry: JWERegistry) -> GeneralJSONEncryption:
+    protected_segment = to_bytes(data["protected"])
+    registry.validate_protected_header_size(protected_segment)
+    protected = json_b64decode(protected_segment)
+
     unprotected = data.get("unprotected")
-    base64_segments, bytes_segments, aad = __extract_segments(data)
+    base64_segments, bytes_segments, aad = __extract_segments(data, registry)
+
     obj = GeneralJSONEncryption(protected, None, unprotected, aad)
     obj.base64_segments = base64_segments
     obj.bytes_segments = bytes_segments
     for item in data["recipients"]:
-        recipient: Recipient[Key] = Recipient(obj, item.get("header"))
-        if "encrypted_key" in item:
-            recipient.encrypted_key = urlsafe_b64decode(to_bytes(item["encrypted_key"]))
+        recipient = __extract_recipient(obj, item, registry)
         obj.recipients.append(recipient)
     return obj
 
 
-def extract_flattened_json(data: FlattenedJSONSerialization) -> FlattenedJSONEncryption:
-    protected = json_b64decode(data["protected"])
+def extract_flattened_json(data: FlattenedJSONSerialization, registry: JWERegistry) -> FlattenedJSONEncryption:
+    protected_segment = to_bytes(data["protected"])
+    registry.validate_protected_header_size(protected_segment)
+    protected = json_b64decode(protected_segment)
     unprotected = data.get("unprotected")
-    base64_segments, bytes_segments, aad = __extract_segments(data)
+    base64_segments, bytes_segments, aad = __extract_segments(data, registry)
     obj = FlattenedJSONEncryption(protected, None, unprotected, aad)
     obj.base64_segments = base64_segments
     obj.bytes_segments = bytes_segments
-
-    recipient: Recipient[Key] = Recipient(obj, data.get("header"))
-    if "encrypted_key" in data:
-        recipient.encrypted_key = urlsafe_b64decode(to_bytes(data["encrypted_key"]))
+    recipient = __extract_recipient(obj, data, registry)
     obj.recipients.append(recipient)
     return obj
 
 
 def __extract_segments(
     data: t.Union[GeneralJSONSerialization, FlattenedJSONSerialization],
+    registry: JWERegistry,
 ) -> tuple[dict[str, bytes], dict[str, bytes], t.Optional[bytes]]:
     base64_segments: dict[str, bytes] = {
         "iv": to_bytes(data["iv"]),
         "ciphertext": to_bytes(data["ciphertext"]),
         "tag": to_bytes(data["tag"]),
     }
+    registry.validate_initialization_vector_size(base64_segments["iv"])
+    registry.validate_ciphertext_size(base64_segments["ciphertext"])
+    registry.validate_auth_tag_size(base64_segments["tag"])
     bytes_segments: dict[str, bytes] = {
         "iv": urlsafe_b64decode(base64_segments["iv"]),
         "ciphertext": urlsafe_b64decode(base64_segments["ciphertext"]),
@@ -118,3 +125,16 @@ def __extract_segments(
     else:
         aad = None
     return base64_segments, bytes_segments, aad
+
+
+def __extract_recipient(
+    obj: FlattenedJSONEncryption | GeneralJSONEncryption,
+    data: FlattenedJSONSerialization | JSONRecipientDict,
+    registry: JWERegistry,
+) -> Recipient[Key]:
+    recipient: Recipient[Key] = Recipient(obj, data.get("header"))
+    if "encrypted_key" in data:
+        ek_segment = to_bytes(data["encrypted_key"])
+        registry.validate_encrypted_key_size(ek_segment)
+        recipient.encrypted_key = urlsafe_b64decode(ek_segment)
+    return recipient

src/joserfc/_rfc7516/registry.py

@@ -2,7 +2,11 @@
 import warnings
 import typing as t
 from .models import JWEAlgModel, JWEEncModel, JWEZipModel
-from ..errors import UnsupportedAlgorithmError, SecurityWarning
+from ..errors import (
+    UnsupportedAlgorithmError,
+    SecurityWarning,
+    ExceededSizeError,
+)
 from ..registry import (
     Header,
     HeaderRegistryDict,
@@ -41,13 +45,24 @@ class JWERegistry:
     :param strict_check_header: only allow header key in the registry to be used
     """
 
-    algorithms: t.ClassVar[AlgorithmsDict] = {
+    algorithms: AlgorithmsDict = {
         "alg": {},
         "enc": {},
         "zip": {},
     }
     recommended: t.ClassVar[list[str]] = []
 
+    #: max protected header content's size in bytes
+    max_protected_header_length: int = 1024
+    #: max encrypted key's size in bytes
+    max_encrypted_key_length: int = 1024
+    #: max initialization vector's size in bytes
+    max_initialization_vector_length: int = 64
+    #: max ciphertext's size in bytes
+    max_ciphertext_length: int = 65536  # 64KB
+    #: max auth tag's size in bytes
+    max_auth_tag_length: int = 64
+
     def __init__(
         self,
         header_registry: t.Optional[HeaderRegistryDict] = None,
@@ -85,6 +100,28 @@ def check_header(self, header: Header, check_more: bool = False) -> None:
         elif self.strict_check_header:
             check_supported_header(self.header_registry, header)
 
+    def validate_protected_header_size(self, header: bytes) -> None:
+        if header and len(header) > self.max_protected_header_length:
+            raise ExceededSizeError(f"Header size of '{header!r}' exceeds {self.max_protected_header_length} bytes.")
+
+    def validate_encrypted_key_size(self, ek: bytes) -> None:
+        if ek and len(ek) > self.max_encrypted_key_length:
+            raise ExceededSizeError(f"Encrypted key size of '{ek!r}' exceeds {self.max_encrypted_key_length} bytes.")
+
+    def validate_initialization_vector_size(self, iv: bytes) -> None:
+        if iv and len(iv) > self.max_initialization_vector_length:
+            raise ExceededSizeError(
+                f"Initialization vector size of '{iv!r}' exceeds {self.max_initialization_vector_length} bytes."
+            )
+
+    def validate_ciphertext_size(self, ciphertext: bytes) -> None:
+        if ciphertext and len(ciphertext) > self.max_ciphertext_length:
+            raise ExceededSizeError(f"Ciphertext size of '{ciphertext!r}' exceeds {self.max_ciphertext_length} bytes.")
+
+    def validate_auth_tag_size(self, tag: bytes) -> None:
+        if tag and len(tag) > self.max_auth_tag_length:
+            raise ExceededSizeError(f"Auth tag size of '{tag!r}' exceeds {self.max_auth_tag_length} bytes.")
+
     def get_alg(self, name: str) -> JWEAlgModel:
         """Get the allowed ("alg") algorithm instance of the given name.
 

src/joserfc/jwe.py

@@ -119,12 +119,12 @@ def decrypt_compact(
     :param sender_key: only required when using ECDH-1PU
     :return: object of the ``CompactEncryption``
     """
-    obj = extract_compact(to_bytes(value))
     if algorithms:
         registry = JWERegistry(algorithms=algorithms)
     elif registry is None:
         registry = default_registry
 
+    obj = extract_compact(to_bytes(value), registry)
     recipient = obj.recipient
     assert recipient is not None
     key = guess_key(private_key, recipient, use="enc")
@@ -235,12 +235,12 @@ def decrypt_json(
 
     reject_unprotected_crit_header(data.get("unprotected"))
     if "recipients" in data:
-        general_obj = extract_general_json(data)  # type: ignore[arg-type]
+        general_obj = extract_general_json(data, registry)  # type: ignore[arg-type]
         _attach_recipient_keys(general_obj.recipients, private_key, sender_key)
         perform_decrypt(general_obj, registry)
         return general_obj
     else:
-        flattened_obj = extract_flattened_json(data)  # type: ignore[arg-type]
+        flattened_obj = extract_flattened_json(data, registry)  # type: ignore[arg-type]
         _attach_recipient_keys(flattened_obj.recipients, private_key, sender_key)
         perform_decrypt(flattened_obj, registry)
         return flattened_obj

tests/jwe/test_compact.py

@@ -15,7 +15,7 @@
     ExceededSizeError,
     InvalidHeaderValueError,
 )
-from joserfc.util import json_b64encode
+from joserfc.util import json_b64encode, urlsafe_b64encode
 from tests.base import load_key
 
 
@@ -204,6 +204,37 @@ def test_decompress_zip_exceeds_size(self):
         result = encrypt_compact({"alg": "dir", "enc": "A128GCM", "zip": "DEF"}, b"h" * 300000, key)
         self.assertRaises(ExceededSizeError, decrypt_compact, result, key)
 
+    def test_header_exceeds_size(self):
+        header = json_b64encode({f"a{i}": "a" * i for i in range(1000)}).decode("utf-8")
+        s = header + "..YbDfdYa6p-wAEFul.YK7j0MsH-Dko6ifsEg.wES6-QAOEbErZqXiS0JHRw"
+        self.assertRaises(ExceededSizeError, decrypt_compact, s, OctKey.import_key("secret"))
+
+    def test_encrypted_key_exceeds_size(self):
+        header = json_b64encode({"alg": "dir", "enc": "A128GCM"}).decode("utf-8")
+        ek = urlsafe_b64encode(("a" * 1000).encode("utf-8")).decode("utf-8")
+        s = header + "." + ek + ".YbDfdYa6p-wAEFul.YK7j0MsH-Dko6ifsEg.wES6-QAOEbErZqXiS0JHRw"
+        key = OctKey.import_key({"k": "pyL42ncDFSYnenl-GiZjRw", "kty": "oct"})
+        self.assertRaises(ExceededSizeError, decrypt_compact, s, key)
+
+    def test_initialization_vector_size(self):
+        header = json_b64encode({"alg": "dir", "enc": "A128GCM"}).decode("utf-8")
+        iv = urlsafe_b64encode(("a" * 1000).encode("utf-8")).decode("utf-8")
+        s = header + ".." + iv + ".YK7j0MsH-Dko6ifsEg.wES6-QAOEbErZqXiS0JHRw"
+        key = OctKey.import_key({"k": "pyL42ncDFSYnenl-GiZjRw", "kty": "oct"})
+        self.assertRaises(ExceededSizeError, decrypt_compact, s, key)
+
+    def test_ciphertext_exceeds_size(self):
+        header = json_b64encode({"alg": "dir", "enc": "A128GCM"}).decode("utf-8")
+        ciphertext = urlsafe_b64encode(("a" * 70000).encode("utf-8")).decode("utf-8")
+        s = header + "..YbDfdYa6p-wAEFul." + ciphertext + ".wES6-QAOEbErZqXiS0JHRw"
+        self.assertRaises(ExceededSizeError, decrypt_compact, s, OctKey.import_key("secret"))
+
+    def test_auth_tag_exceeds_size(self):
+        header = json_b64encode({"alg": "dir", "enc": "A128GCM"}).decode("utf-8")
+        tag = urlsafe_b64encode(("a" * 80).encode("utf-8")).decode("utf-8")
+        s = header + "..YbDfdYa6p-wAEFul.YK7j0MsH-Dko6ifsEg." + tag
+        self.assertRaises(ExceededSizeError, decrypt_compact, s, OctKey.import_key("secret"))
+
     def test_invalid_compact_data(self):
         private_key: RSAKey = load_key("rsa-openssl-private.pem")
         value = b"a.b.c.d.e.f.g"