fix(jwe): limit DEF decompress size to 250k bytes

Author: lepture

src/joserfc/errors.py

@@ -72,6 +72,12 @@ class BadSignatureError(JoseError):
     error = "bad_signature"
 
 
+class ExceededSizeError(JoseError):
+    """This error is designed for DEF zip algorithm. It raised when the
+    compressed data exceeds the maximum allowed length."""
+    error = "exceeded_size"
+
+
 class InvalidEncryptionAlgorithmError(JoseError):
     """This error is designed for JWE. It is raised when "enc" value
     does not work together with "alg" value.

src/joserfc/rfc7518/jwe_zips.py

@@ -1,6 +1,10 @@
 import zlib
 from typing import List
 from ..rfc7516.models import JWEZipModel
+from ..errors import ExceededSizeError
+
+GZIP_HEAD = bytes([120, 156])
+MAX_SIZE = 250 * 1024
 
 
 class DeflateZipModel(JWEZipModel):
@@ -10,12 +14,20 @@ class DeflateZipModel(JWEZipModel):
     def compress(self, s: bytes) -> bytes:
         """Compress bytes data with DEFLATE algorithm."""
         data = zlib.compress(s)
-        # drop gzip headers and tail
+        # https://datatracker.ietf.org/doc/html/rfc1951
+        # since DEF is always gzip, we can drop gzip headers and tail
         return data[2:-4]
 
     def decompress(self, s: bytes) -> bytes:
         """Decompress DEFLATE bytes data."""
-        return zlib.decompress(s, -zlib.MAX_WBITS)
+        if s.startswith(GZIP_HEAD):
+            decompressor = zlib.decompressobj()
+        else:
+            decompressor = zlib.decompressobj(-zlib.MAX_WBITS)
+        value = decompressor.decompress(s, MAX_SIZE)
+        if decompressor.unconsumed_tail:
+            raise ExceededSizeError(f"Decompressed string exceeds {MAX_SIZE} bytes")
+        return value
 
 
 JWE_ZIP_MODELS: List[JWEZipModel] = [DeflateZipModel()]

tests/jwe/test_compact.py

@@ -7,6 +7,7 @@
     MissingAlgorithmError,
     MissingEncryptionError,
     DecodeError,
+    ExceededSizeError,
 )
 from joserfc.util import json_b64encode
 from tests.base import load_key
@@ -176,6 +177,25 @@ def test_with_zip_header(self):
         obj = decrypt_compact(result, private_key)
         self.assertEqual(obj.plaintext, plaintext)
 
+    def test_decompress_zip_with_gzip_head(self):
+        key = OctKey.import_key({'k': 'pyL42ncDFSYnenl-GiZjRw', 'kty': 'oct'})
+        s = (
+            'eyJhbGciOiJkaXIiLCJlbmMiOiJBMTI4R0NNIiwiemlwIjoiREVGIn0..'
+            'YbDfdYa6p-wAEFul.YK7j0MsH-Dko6ifsEg.wES6-QAOEbErZqXiS0JHRw'
+        )
+        result = decrypt_compact(s, key)
+        self.assertEqual(result.plaintext, b'hello')
+        self.assertEqual(result.headers().get('zip'), 'DEF')
+
+    def test_decompress_zip_exceeds_size(self):
+        key = OctKey.import_key({'k': 'pyL42ncDFSYnenl-GiZjRw', 'kty': 'oct'})
+        result = encrypt_compact(
+            {"alg": "dir", "enc": "A128GCM", "zip": "DEF"},
+            b'h' * 300000,
+            key
+        )
+        self.assertRaises(ExceededSizeError, decrypt_compact, result, key)
+
     def test_invalid_compact_data(self):
         private_key: RSAKey = load_key('rsa-openssl-private.pem')
         value = b"a.b.c.d.e.f.g"