Merge branch 'main' into feat/language-detection-custom-fallback

Author: PastelStorm

CHANGELOG.md

@@ -1,5 +1,11 @@
+## 0.20.7
+
+### Fixes
+- **Cap size when decompressing elements JSON file**: Prevents situations where decompression can consume an arbitrarily large portion in memory and on the filesystem.
+
 ## 0.20.6
 
+### Fixes
 - fix: remap parent id after hashing to preserve right reference
 
 ## 0.20.5

test_unstructured/staging/test_base.py

@@ -4,6 +4,8 @@
 import pathlib
 import platform
 import tempfile
+import zlib
+from unittest.mock import patch
 
 import pandas as pd
 import pytest
@@ -28,6 +30,7 @@
     Text,
     Title,
 )
+from unstructured.errors import DecompressedSizeExceededError
 from unstructured.partition.email import partition_email
 from unstructured.partition.text import partition_text
 from unstructured.staging import base
@@ -45,6 +48,31 @@ def test_base64_gzipped_json_to_elements_can_deserialize_compressed_elements_fro
     assert elements == [Title("Lorem"), Text("Lorem Ipsum")]
 
 
+def test_elements_from_base64_gzipped_json_raises_error_if_decompression_is_incomplete():
+    base64_elements_str = (
+        "eJyFzcsKwjAQheFXKVm7yDS3xjcQXNaViKTJjBR6o46glr67zVI3Lmf4Dv95EdhhjwNf2yT2hYDGUaWtJVm5WDoq"
+        "NUL0UoJrqtLHJHaF6JFDChw2v6zbzfjkvD2OM/YZ8GvC/Khb7lBs5LcilUwRyCsblQYTiBQpZRxYZcCA/1spDtP9"
+        "8dU6DTEw3sa5fWOqs10vH0cL="
+    )
+
+    with pytest.raises(zlib.error):
+        base.elements_from_base64_gzipped_json(base64_elements_str)
+
+
+def test_elements_from_base64_gzipped_json_raises_error_if_decompression_exceeds_max_size():
+    base64_elements_str = (
+        "eJyFzcsKwjAQheFXKVm7yDS3xjcQXNaViKTJjBR6o46glr67zVI3Lmf4Dv95EdhhjwNf2yT2hYDGUaWtJVm5WDoq"
+        "NUL0UoJrqtLHJHaF6JFDChw2v6zbzfjkvD2OM/YZ8GvC/Khb7lBs5LcilUwRyCsblQYTiBQpZRxYZcCA/1spDtP9"
+        "8dU6DTEw3sa5fWOqs10vH0cLQn0="
+    )
+
+    with (
+        patch("unstructured.staging.base.MAX_DECOMPRESSED_SIZE", 32),
+        pytest.raises(DecompressedSizeExceededError),
+    ):
+        base.elements_from_base64_gzipped_json(base64_elements_str)
+
+
 def test_elements_to_base64_gzipped_json_can_serialize_elements_to_a_base64_str():
     elements = assign_hash_ids([Title("Lorem"), Text("Lorem Ipsum")])
 

unstructured/__version__.py

@@ -1 +1 @@
-__version__ = "0.20.6"  # pragma: no cover
+__version__ = "0.20.7"  # pragma: no cover

unstructured/errors.py

@@ -13,3 +13,15 @@ def __init__(self, document_pages: int, pdf_hi_res_max_pages: int):
 
 class UnprocessableEntityError(Exception):
     """Error raised when a file is not valid."""
+
+
+class DecompressedSizeExceededError(ValueError):
+    """Error raised when decompressed data exceeds the maximum size limit."""
+
+    def __init__(self, max_size: int):
+        self.max_size = max_size
+        self.message = (
+            f"Decompressed data exceeds maximum allowed size of {max_size} bytes "
+            f"({max_size / (1024 * 1024):.1f} MB)."
+        )
+        super().__init__(self.message)

unstructured/staging/base.py

@@ -19,6 +19,7 @@
     Table,
     Title,
 )
+from unstructured.errors import DecompressedSizeExceededError
 from unstructured.file_utils.ndjson import dumps as ndjson_dumps
 from unstructured.partition.common.common import exactly_one
 from unstructured.utils import Point, dependency_exists, requires_dependencies
@@ -35,6 +36,8 @@
 
 # == DESERIALIZERS ===============================
 
+MAX_DECOMPRESSED_SIZE = 200 * 1024 * 1024  # 200MB
+
 
 def elements_from_base64_gzipped_json(b64_encoded_elements: str) -> list[Element]:
     """Restore Base64-encoded gzipped JSON elements to element objects.
@@ -45,7 +48,17 @@ def elements_from_base64_gzipped_json(b64_encoded_elements: str) -> list[Element
     # -- Base64 str -> gzip-encoded (JSON) bytes --
     decoded_b64_bytes = base64.b64decode(b64_encoded_elements)
     # -- undo gzip compression --
-    elements_json_bytes = zlib.decompress(decoded_b64_bytes)
+    dobj = zlib.decompressobj()
+    elements_json_bytes = dobj.decompress(decoded_b64_bytes, max_length=MAX_DECOMPRESSED_SIZE)
+    # -- Check if decompression completed successfully --
+    if not dobj.eof:
+        # Check if we hit the size limit or if data is actually incomplete
+        if len(elements_json_bytes) >= MAX_DECOMPRESSED_SIZE:
+            raise DecompressedSizeExceededError(
+                max_size=MAX_DECOMPRESSED_SIZE,
+            )
+        else:
+            raise zlib.error("Incomplete or corrupted compressed data")
     # -- JSON (bytes) to JSON (str) --
     elements_json_str = elements_json_bytes.decode("utf-8")
     # -- JSON (str) -> dicts --