diff --git a/backend/api_v2/deployment_helper.py b/backend/api_v2/deployment_helper.py
index 5fbec7999c..38b1a8795f 100644
--- a/backend/api_v2/deployment_helper.py
+++ b/backend/api_v2/deployment_helper.py
@@ -19,7 +19,9 @@
from utils.constants import Account, CeleryQueue
from utils.local_context import StateStore
from workflow_manager.endpoint_v2.destination import DestinationConnector
+from workflow_manager.endpoint_v2.result_cache_utils import ResultCacheUtils
from workflow_manager.endpoint_v2.source import SourceConnector
+from workflow_manager.utils.pipeline_utils import PipelineUtils
from workflow_manager.workflow_v2.dto import ExecutionResponse
from workflow_manager.workflow_v2.enums import ExecutionStatus
from workflow_manager.workflow_v2.execution import WorkflowExecutionServiceHelper
@@ -293,7 +295,12 @@ def execute_workflow(
logger.exception(f"Failed to mark execution {execution_id} as ERROR")
# Async job never started — release the rate limit slot and clean up.
- APIDeploymentRateLimiter.release_slot(api.organization, str(execution_id))
+ # The id string, never the model instance: release_slot builds its
+ # Redis key from what it is passed, so an instance silently removes
+ # nothing and leaves the slot held for its full TTL.
+ APIDeploymentRateLimiter.release_slot(
+ str(api.organization.organization_id), str(execution_id)
+ )
DestinationConnector.delete_api_storage_dir(
workflow_id=workflow_id, execution_id=execution_id
)
@@ -306,6 +313,69 @@ def execute_workflow(
)
).data
+ # Nothing was staged, so nothing will be dispatched to write a status
+ # back. Terminalise here or the execution is stranded in PENDING.
+ if not hash_values_of_files:
+ # Isolated so the slot and staging dir are released even if it raises.
+ execution = None
+ terminalise_error: str | None = None
+ try:
+ execution = WorkflowExecutionServiceHelper.update_execution_completed(
+ str(execution_id),
+ total_files=len(file_objs),
+ failed_files=len(file_objs),
+ )
+ except Exception as error:
+ logger.exception(f"Failed to mark execution {execution_id} as COMPLETED")
+ # Reported, so the caller is not given a bare failure with no reason.
+ terminalise_error = str(error)
+
+ APIDeploymentRateLimiter.release_slot(
+ str(api.organization.organization_id), str(execution_id)
+ )
+ DestinationConnector.delete_api_storage_dir(
+ workflow_id=workflow_id, execution_id=execution_id
+ )
+ api_results = ResultCacheUtils.get_api_results(
+ workflow_id=str(workflow_id), execution_id=str(execution_id)
+ )
+ if execution is not None:
+ # Separate blocks: independent obligations, so a failure in one
+ # must not silence the other.
+ try:
+ # This response carries the results, so mark them consumed or a
+ # later GET /status serves them again.
+ WorkflowHelper.set_result_acknowledge(execution)
+ except Exception:
+ logger.exception(
+ f"Failed to acknowledge results for execution {execution_id}"
+ )
+
+ try:
+ # Terminalising here bypasses the usual notifier, so subscribers
+ # would otherwise hear nothing.
+ PipelineUtils.update_pipeline_status(
+ pipeline_id=pipeline_id, workflow_execution=execution
+ )
+ except Exception:
+ logger.exception(
+ f"Failed to notify subscribers for execution {execution_id}"
+ )
+
+ # The stored status, not an assumed COMPLETED: the row may be missing
+ # or the guard may have refused the change.
+ return APIExecutionResponseSerializer(
+ ExecutionResponse(
+ workflow_id=workflow_id,
+ execution_id=execution_id,
+ execution_status=(
+ execution.status if execution else ExecutionStatus.ERROR.value
+ ),
+ error=terminalise_error,
+ result=api_results,
+ )
+ ).data
+
try:
result = WorkflowHelper.execute_workflow_async(
workflow_id=workflow_id,
@@ -352,7 +422,9 @@ def execute_workflow(
# Dispatch failures are marked ERROR internally by execute_workflow_async;
# post-dispatch failures (enrichment/config) must not overwrite a running
# execution's status, so only release the slot and clean up storage here.
- APIDeploymentRateLimiter.release_slot(api.organization, str(execution_id))
+ APIDeploymentRateLimiter.release_slot(
+ str(api.organization.organization_id), str(execution_id)
+ )
# Clean up storage
DestinationConnector.delete_api_storage_dir(
diff --git a/backend/api_v2/tests/test_deployment_helper.py b/backend/api_v2/tests/test_deployment_helper.py
index 39b23e5b16..2e8c3ec7d3 100644
--- a/backend/api_v2/tests/test_deployment_helper.py
+++ b/backend/api_v2/tests/test_deployment_helper.py
@@ -48,6 +48,7 @@ def _api() -> MagicMock:
api = MagicMock()
api.workflow.id = "wf-1"
api.id = "pipe-1"
+ api.organization.organization_id = "org-uuid-1"
return api
@@ -89,3 +90,195 @@ def test_staging_failure_cleanup_survives_db_marking_error(collaborators) -> Non
# Cleanup still runs even though error-marking raised.
collaborators["APIDeploymentRateLimiter"].release_slot.assert_called_once()
collaborators["DestinationConnector"].delete_api_storage_dir.assert_called_once()
+
+
+@pytest.fixture
+def staging_rejects_everything():
+ """Patch execute_workflow's collaborators; staging returns no dispatchable files."""
+ with mock.patch.multiple(
+ dh,
+ WorkflowExecutionServiceHelper=mock.DEFAULT,
+ SourceConnector=mock.DEFAULT,
+ DestinationConnector=mock.DEFAULT,
+ APIDeploymentRateLimiter=mock.DEFAULT,
+ WorkflowHelper=mock.DEFAULT,
+ ResultCacheUtils=mock.DEFAULT,
+ PipelineUtils=mock.DEFAULT,
+ Tag=mock.DEFAULT,
+ logger=mock.DEFAULT,
+ ) as mocks:
+ execution_row = MagicMock()
+ execution_row.id = "exec-123"
+ mocks[
+ "WorkflowExecutionServiceHelper"
+ ].create_workflow_execution.return_value = execution_row
+ mocks["SourceConnector"].add_input_file_to_api_storage.return_value = {}
+ mocks["ResultCacheUtils"].get_api_results.return_value = [
+ {"file": "evil.pdf", "status": "Failed", "error": "unsupported MIME type"}
+ ]
+ completed_row = MagicMock()
+ completed_row.status = "COMPLETED"
+ mocks[
+ "WorkflowExecutionServiceHelper"
+ ].update_execution_completed.return_value = completed_row
+ yield mocks
+
+
+def test_all_files_rejected_completes_without_dispatch(
+ staging_rejects_everything,
+) -> None:
+ """A request whose every file is rejected must reach a terminal status.
+
+ The worker short-circuits an empty file set without writing a status back, so
+ dispatching one strands the execution in PENDING and the caller polls forever.
+ """
+ mocks = staging_rejects_everything
+ # A non-empty upload whose staging result is empty. Passing [] instead would
+ # leave the branch satisfied by `not file_objs` too, and the original bug -
+ # dispatching a request whose files were all rejected - would pass this test.
+ response = dh.DeploymentHelper.execute_workflow(
+ organization_name="org",
+ api=_api(),
+ file_objs=[MagicMock()],
+ timeout=-1,
+ )
+
+ # Nothing is dispatched...
+ mocks["WorkflowHelper"].execute_workflow_async.assert_not_called()
+ # ...the row is terminalised here instead of being left PENDING, and the
+ # counters are written so the run does not read back as a clean success...
+ mocks[
+ "WorkflowExecutionServiceHelper"
+ ].update_execution_completed.assert_called_once_with(
+ "exec-123", total_files=1, failed_files=1
+ )
+ # ...the slot and staging dir are released. The slot must be released by org
+ # id string: release_slot formats its argument into the Redis key, so passing
+ # the model instance removes a non-member and silently holds the slot.
+ mocks["APIDeploymentRateLimiter"].release_slot.assert_called_once_with(
+ "org-uuid-1", "exec-123"
+ )
+ mocks["DestinationConnector"].delete_api_storage_dir.assert_called_once()
+ # ...and the caller still sees why each file failed.
+ assert response["execution_status"] == "COMPLETED"
+ assert response["result"][0]["file"] == "evil.pdf"
+ assert response["result"][0]["status"] == "Failed"
+
+
+def test_all_files_rejected_acknowledges_and_notifies(
+ staging_rejects_everything,
+) -> None:
+ """The early return owes the caller what the dispatch path would have done.
+
+ It hands back the results in its own response and reaches a terminal status
+ without going through WorkflowHelper, so both the acknowledgement and the
+ subscriber notification have to happen here or they happen nowhere.
+ """
+ mocks = staging_rejects_everything
+ completed_row = mocks[
+ "WorkflowExecutionServiceHelper"
+ ].update_execution_completed.return_value
+
+ dh.DeploymentHelper.execute_workflow(
+ organization_name="org",
+ api=_api(),
+ file_objs=[MagicMock()],
+ timeout=-1,
+ )
+
+ # Results were served in this response, so a later GET /status must 406
+ # rather than serve them again.
+ mocks["WorkflowHelper"].set_result_acknowledge.assert_called_once_with(completed_row)
+ # PipelineUtils is the only dispatcher of API deployment notifications.
+ mocks["PipelineUtils"].update_pipeline_status.assert_called_once_with(
+ pipeline_id="pipe-1", workflow_execution=completed_row
+ )
+
+
+def test_notification_survives_a_failed_acknowledgement(
+ staging_rejects_everything,
+) -> None:
+ """Acknowledgement and notification are independent obligations.
+
+ Sharing one try block would let a failed acknowledgement silence the
+ notification that subscribers depend on.
+ """
+ mocks = staging_rejects_everything
+ mocks["WorkflowHelper"].set_result_acknowledge.side_effect = Exception("db is down")
+
+ response = dh.DeploymentHelper.execute_workflow(
+ organization_name="org",
+ api=_api(),
+ file_objs=[MagicMock()],
+ timeout=-1,
+ )
+
+ mocks["PipelineUtils"].update_pipeline_status.assert_called_once()
+ assert response["execution_status"] == "COMPLETED"
+
+
+def test_all_files_rejected_still_responds_if_notification_fails(
+ staging_rejects_everything,
+) -> None:
+ """A failing webhook must not turn a handled rejection into a 500."""
+ mocks = staging_rejects_everything
+ mocks["PipelineUtils"].update_pipeline_status.side_effect = Exception("webhook down")
+
+ response = dh.DeploymentHelper.execute_workflow(
+ organization_name="org",
+ api=_api(),
+ file_objs=[MagicMock()],
+ timeout=-1,
+ )
+
+ assert response["execution_status"] == "COMPLETED"
+ assert response["result"][0]["file"] == "evil.pdf"
+
+
+def test_files_staged_successfully_are_dispatched(staging_rejects_everything) -> None:
+ """The short-circuit must not fire when staging did return files.
+
+ Sibling to the test above: together they pin the branch to the staging result
+ rather than to the upload list.
+ """
+ mocks = staging_rejects_everything
+ mocks["SourceConnector"].add_input_file_to_api_storage.return_value = {
+ "good.pdf": MagicMock()
+ }
+
+ dh.DeploymentHelper.execute_workflow(
+ organization_name="org",
+ api=_api(),
+ file_objs=[MagicMock()],
+ timeout=-1,
+ )
+
+ mocks["WorkflowHelper"].execute_workflow_async.assert_called_once()
+ mocks["WorkflowExecutionServiceHelper"].update_execution_completed.assert_not_called()
+
+
+def test_all_files_rejected_cleanup_survives_db_marking_error(
+ staging_rejects_everything,
+) -> None:
+ """A failing status write must not strand the slot or the staging dir.
+
+ update_execution_completed only catches DoesNotExist, so a lock timeout or a
+ dropped connection propagates; without isolation the org's rate limit slot
+ stays held for its full TTL and throttles every other call for that org.
+ """
+ mocks = staging_rejects_everything
+ mocks[
+ "WorkflowExecutionServiceHelper"
+ ].update_execution_completed.side_effect = Exception("db is down")
+
+ response = dh.DeploymentHelper.execute_workflow(
+ organization_name="org",
+ api=_api(),
+ file_objs=[MagicMock()],
+ timeout=-1,
+ )
+
+ mocks["APIDeploymentRateLimiter"].release_slot.assert_called_once()
+ mocks["DestinationConnector"].delete_api_storage_dir.assert_called_once()
+ # The row never reached COMPLETED, so the response must not claim it did.
+ assert response["execution_status"] == "ERROR"
diff --git a/backend/workflow_manager/endpoint_v2/enums.py b/backend/workflow_manager/endpoint_v2/enums.py
index 79725ea005..7fa01dd1b3 100644
--- a/backend/workflow_manager/endpoint_v2/enums.py
+++ b/backend/workflow_manager/endpoint_v2/enums.py
@@ -1,5 +1,10 @@
from enum import Enum
+from unstract.core.mime_gate import ( # noqa: F401 re-exported
+ INCONCLUSIVE_MIME_TYPES,
+ resolve_inconclusive_mime_type,
+)
+
class FileProcessingStatus(Enum):
SUCCESS = "SUCCESS"
@@ -7,6 +12,12 @@ class FileProcessingStatus(Enum):
class AllowedFileTypes(Enum):
+ """MIME types accepted into a workflow.
+
+ Keep in sync with LLMWhisperer's supported types, and with the copy in the
+ workers package.
+ """
+
PLAIN_TEXT = "text/plain"
PDF = "application/pdf"
JPEG = "image/jpeg"
@@ -19,7 +30,6 @@ class AllowedFileTypes(Enum):
DOC = "application/msword"
XLSX = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
XLS = "application/vnd.ms-excel"
- XLSM = "application/vnd.ms-excel.sheet.macroenabled.12"
PPTX = "application/vnd.openxmlformats-officedocument.presentationml.presentation"
PPT = "application/vnd.ms-powerpoint"
ODT = "application/vnd.oasis.opendocument.text"
@@ -28,8 +38,14 @@ class AllowedFileTypes(Enum):
CDFV2 = "application/CDFV2"
JSON = "application/json"
CSV = "text/csv"
- OCTET_STREAM = "application/octet-stream"
@classmethod
def is_allowed(cls, mime_type: str) -> bool:
+ """Whether LLMWhisperer would accept this type.
+
+ Any `text/*` passes; html, xml, tsv and the like reach the extractor as
+ text rather than as listed formats.
+ """
+ if mime_type.startswith("text/"):
+ return True
return mime_type in cls._value2member_map_
diff --git a/backend/workflow_manager/endpoint_v2/source.py b/backend/workflow_manager/endpoint_v2/source.py
index 23b0d04f2c..4de2edd045 100644
--- a/backend/workflow_manager/endpoint_v2/source.py
+++ b/backend/workflow_manager/endpoint_v2/source.py
@@ -2,7 +2,7 @@
import logging
import os
import shutil
-import uuid
+import tempfile
from collections.abc import Collection
from hashlib import sha256
from io import BytesIO
@@ -25,8 +25,16 @@
SourceConstant,
SourceKey,
)
-from workflow_manager.endpoint_v2.dto import FileHash, SourceConfig
-from workflow_manager.endpoint_v2.enums import AllowedFileTypes
+from workflow_manager.endpoint_v2.dto import (
+ FileExecutionResult,
+ FileHash,
+ SourceConfig,
+)
+from workflow_manager.endpoint_v2.enums import (
+ INCONCLUSIVE_MIME_TYPES,
+ AllowedFileTypes,
+ resolve_inconclusive_mime_type,
+)
from workflow_manager.endpoint_v2.exceptions import (
InvalidInputDirectory,
InvalidSourceConnectionType,
@@ -37,6 +45,7 @@
UnsupportedMimeTypeError,
)
from workflow_manager.endpoint_v2.models import WorkflowEndpoint
+from workflow_manager.endpoint_v2.result_cache_utils import ResultCacheUtils
from workflow_manager.file_execution.models import WorkflowFileExecution
from workflow_manager.utils.workflow_log import WorkflowLog
from workflow_manager.workflow_v2.enums import ExecutionStatus
@@ -46,6 +55,7 @@
from workflow_manager.workflow_v2.models.workflow import Workflow
from unstract.connectors.filesystems.unstract_file_system import UnstractFileSystem
+from unstract.core.mime_gate import identify_zip_container
from unstract.filesystem import FileStorageType, FileSystem
from unstract.sdk1.file_storage import FileStorage
from unstract.workflow_execution.enums import LogLevel, LogStage, LogState
@@ -69,6 +79,10 @@ class SourceConnector(BaseConnector):
"""
READ_CHUNK_SIZE = 4194304 # Chunk size for reading files
+ # Most formats are identifiable from their leading bytes.
+ MIME_DETECT_CHUNK_SIZE = 8192
+ # Shared, so the API and connector paths agree on what counts as undecided.
+ CONTAINER_MIME_TYPES = INCONCLUSIVE_MIME_TYPES
def __init__(
self,
@@ -1000,7 +1014,9 @@ def add_input_from_connector_to_volume(self, file_hash: FileHash) -> str:
file_content_hash.update(chunk)
if first_iteration:
# Detect MIME type using first chunk
- mime_type = magic.from_buffer(chunk, mime=True)
+ mime_type = resolve_inconclusive_mime_type(
+ magic.from_buffer(chunk, mime=True), chunk
+ )
logger.info(
f"Detected MIME type: {mime_type} for file {input_file_path}"
)
@@ -1187,6 +1203,143 @@ def load_file(self, input_file_path: str) -> tuple[str, BytesIO]:
return os.path.basename(input_file_path), file_stream
+ @classmethod
+ def _detect_uploaded_file_mime_type(cls, file: UploadedFile) -> str | None:
+ """Detect an uploaded file's MIME type from its own bytes.
+
+ The multipart Content-Type is supplied by the caller and never verified,
+ so it cannot be used to decide what is allowed into API storage.
+
+ Returns None for an empty upload, which has no type to judge.
+ """
+ sample = file.read(cls.MIME_DETECT_CHUNK_SIZE)
+ file.seek(0)
+ if not sample:
+ # An empty upload is its own failure downstream, not a type error.
+ return None
+
+ mime_type = magic.from_buffer(sample, mime=True)
+ if mime_type in cls.CONTAINER_MIME_TYPES:
+ mime_type = cls._detect_container_mime_type(file, fallback=mime_type)
+ return resolve_inconclusive_mime_type(mime_type, sample)
+
+ @classmethod
+ def _detect_container_mime_type(cls, file: UploadedFile, fallback: str) -> str:
+ """Resolve a wrapper by classifying the file from a path.
+
+ A path rather than a buffer, because libmagic will not name a zip
+ container from one; an in-memory upload is spilled to a temp file.
+ """
+ temporary_file_path = getattr(file, "temporary_file_path", None)
+ if temporary_file_path is not None:
+ return cls._classify_file_path(temporary_file_path(), fallback)
+
+ with tempfile.NamedTemporaryFile(suffix=".upload") as spill:
+ for chunk in file.chunks(chunk_size=cls.READ_CHUNK_SIZE):
+ spill.write(chunk)
+ spill.flush()
+ file.seek(0)
+ return cls._classify_file_path(spill.name, fallback)
+
+ @classmethod
+ def _classify_file_path(cls, path: str, fallback: str) -> str:
+ """Classify a staged path, looking inside a zip when that is all we get."""
+ mime_type = magic.from_file(path, mime=True)
+ if mime_type == "application/zip":
+ inside = identify_zip_container(path)
+ return inside or mime_type
+ return mime_type or fallback
+
+ @classmethod
+ def _report_rejected_file(
+ cls,
+ workflow_log: WorkflowLog,
+ workflow_id: str,
+ execution_id: str,
+ file: UploadedFile,
+ file_name: str,
+ mime_type: str | None,
+ log_message: str,
+ position: int,
+ ) -> None:
+ """Record a rejection for the caller without risking the whole request.
+
+ Runs per file inside the staging loop, so nothing here may escape and
+ cost the caller the files already staged.
+ """
+ workflow_log.log_error(logger=logger, message=log_message)
+ try:
+ ResultCacheUtils.update_api_results(
+ workflow_id=workflow_id,
+ execution_id=execution_id,
+ api_result=FileExecutionResult(file=file_name, error=log_message),
+ )
+ except Exception:
+ logger.exception(
+ f"Could not record the rejection of '{file_name}' for the caller"
+ )
+
+ try:
+ cls._persist_rejected_file_execution(
+ execution_id=execution_id,
+ file=file,
+ file_name=file_name,
+ mime_type=mime_type,
+ log_message=log_message,
+ position=position,
+ )
+ except Exception:
+ logger.exception(
+ f"Could not persist a file execution row for rejected '{file_name}'"
+ )
+
+ @classmethod
+ def _persist_rejected_file_execution(
+ cls,
+ execution_id: str,
+ file: UploadedFile,
+ file_name: str,
+ mime_type: str | None,
+ log_message: str,
+ position: int,
+ ) -> None:
+ """Record a rejected file as a terminal ERROR file execution.
+
+ The cached result is short-lived, so the row is what keeps the rejection
+ both auditable and countable.
+ """
+ workflow_execution = WorkflowExecution.objects.get(pk=execution_id)
+
+ # The row's identity rests on this alone, since the API path has no
+ # file_path to tell two rejections apart. Content, name and the upload's
+ # position in the request, so every rejected part gets its own record
+ # even when two are byte-identical and share a name. Not a plain content
+ # hash on purpose: this file was never processed and must not look to
+ # file history as though it were.
+ digest = sha256()
+ for chunk in file.chunks(chunk_size=cls.READ_CHUNK_SIZE):
+ digest.update(chunk)
+ file.seek(0)
+ row_identity = sha256(
+ digest.digest() + file_name.encode("utf-8") + str(position).encode("ascii")
+ ).hexdigest()
+
+ file_execution = WorkflowFileExecution.objects.get_or_create_file_execution(
+ workflow_execution=workflow_execution,
+ file_hash=FileHash(
+ file_path=None,
+ source_connection_type=WorkflowEndpoint.ConnectionType.API,
+ file_name=file_name,
+ file_hash=row_identity,
+ file_size=file.size,
+ mime_type=mime_type,
+ ),
+ is_api=True,
+ )
+ file_execution.update_status(
+ status=ExecutionStatus.ERROR, execution_error=log_message
+ )
+
@classmethod
def add_input_file_to_api_storage(
cls,
@@ -1224,34 +1377,53 @@ def add_input_file_to_api_storage(
file_hashes: dict[str, FileHash] = {}
unique_file_hashes: set[str] = set()
connection_type = WorkflowEndpoint.ConnectionType.API
- for file in file_objs:
+ for position, file in enumerate(file_objs):
file_name = file.name
destination_path = os.path.join(api_storage_dir, file_name)
- mime_type = file.content_type
- logger.info(f"Detected MIME type: {mime_type} for file {file_name}")
- if not mime_type:
- logger.info(
- f"MIME type not found for file {file_name}, using default MIME type: {AllowedFileTypes.OCTET_STREAM.value}"
+ try:
+ mime_type = cls._detect_uploaded_file_mime_type(file)
+ except (OSError, ValueError, magic.MagicException):
+ # Narrow on purpose: these are faults in THIS upload's bytes, so
+ # failing the one file is right. A broken libmagic database or an
+ # unreadable temp dir hits every file in every request, and
+ # swallowing that would tell each caller their files are invalid
+ # while the platform is down - it must propagate and fail loudly.
+ log_message = (
+ f"Rejecting file '{file_name}': could not determine its type"
+ )
+ logger.exception(log_message)
+ cls._report_rejected_file(
+ workflow_log,
+ workflow_id,
+ execution_id,
+ file,
+ file_name,
+ None,
+ log_message,
+ position,
)
- mime_type = AllowedFileTypes.OCTET_STREAM.value
+ continue
- if not AllowedFileTypes.is_allowed(mime_type):
- log_message = f"Skipping file '{file_name}' to stage due to unsupported MIME type '{mime_type}'"
- workflow_log.log_info(logger=logger, message=log_message)
- # Generate a clearly marked temporary hash to avoid reading the file content
- # Helps to prevent duplicate entries in file executions
- fake_hash = f"temp-hash-{uuid.uuid4().hex}"
- file_hash = FileHash(
- file_path=destination_path,
- source_connection_type=connection_type,
- file_name=file_name,
- file_hash=fake_hash,
- is_executed=True,
- file_size=file.size,
- mime_type=mime_type,
+ logger.info(f"Detected MIME type: {mime_type} for file {file_name}")
+
+ if mime_type is not None and not AllowedFileTypes.is_allowed(mime_type):
+ log_message = (
+ f"Rejecting file '{file_name}' with unsupported MIME type "
+ f"'{mime_type}'"
+ )
+ # Rejected files are never dispatched, so nothing downstream will
+ # report on them - surface the failure in the API response here.
+ cls._report_rejected_file(
+ workflow_log,
+ workflow_id,
+ execution_id,
+ file,
+ file_name,
+ mime_type,
+ log_message,
+ position,
)
- file_hashes.update({file_name: file_hash})
continue
file_system = FileSystem(FileStorageType.API_EXECUTION)
diff --git a/backend/workflow_manager/endpoint_v2/tests/test_api_storage_mime_validation.py b/backend/workflow_manager/endpoint_v2/tests/test_api_storage_mime_validation.py
new file mode 100644
index 0000000000..7d466ff867
--- /dev/null
+++ b/backend/workflow_manager/endpoint_v2/tests/test_api_storage_mime_validation.py
@@ -0,0 +1,970 @@
+"""MIME validation for files staged into API storage.
+
+``SourceConnector.add_input_file_to_api_storage`` is the single funnel through
+which API-deployment uploads reach the API storage bucket, so an unsupported
+file has to be rejected here or it reaches the extraction step and fails there
+with an error that does not name the real cause.
+
+Unit tests: the real classmethod runs with its DB/storage-touching
+collaborators patched on the imported module, so no database is needed. MIME
+detection itself is deliberately *not* patched — sniffing the bytes with
+libmagic is the behaviour under test.
+"""
+
+import io
+import zipfile
+from unittest import mock
+from unittest.mock import MagicMock
+
+import pytest
+from django.core.files.uploadedfile import SimpleUploadedFile, TemporaryUploadedFile
+
+import workflow_manager.endpoint_v2.source as src_mod
+from workflow_manager.endpoint_v2.constants import ApiDeploymentResultStatus
+from workflow_manager.endpoint_v2.source import SourceConnector
+
+# Bytes chosen from what libmagic actually reports (verified against the pinned
+# python-magic). A WAV header is the unsupported case: audio/x-wav is neither
+# text/* nor a listed type, and RIFF classifies stably on every libmagic build.
+# HTML is the opposite case - it looks like something to reject, but it reaches
+# the extractor as text/*, so the gate has to let it through.
+PDF_BYTES = b"%PDF-1.7\n1 0 obj\n<< /Type /Catalog >>\nendobj\n"
+WAV_BYTES = b"RIFF\x24\x08\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x02\x00"
+HTML_BYTES = b"hello"
+
+
+# A package declares its main part and relates it from the root; an archive
+# carrying only the member names is what the extractor cannot open.
+_OOXML_MAIN = (
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+ ".main+xml"
+)
+_OOXML_CONTENT_TYPES = (
+ ''
+ f''
+ ""
+)
+_OOXML_RELS = (
+ ''
+ ''
+ ""
+)
+
+API_STORAGE_DIR = "/api-storage/exec-1"
+
+
+@pytest.fixture
+def collaborators():
+ """Patch everything the staging loop touches except MIME detection."""
+ with (
+ mock.patch.multiple(
+ src_mod,
+ UserContext=mock.DEFAULT,
+ WorkflowLog=mock.DEFAULT,
+ Workflow=mock.DEFAULT,
+ FileSystem=mock.DEFAULT,
+ FileHistoryHelper=mock.DEFAULT,
+ ResultCacheUtils=mock.DEFAULT,
+ WorkflowExecution=mock.DEFAULT,
+ WorkflowFileExecution=mock.DEFAULT,
+ ) as mocks,
+ mock.patch.object(
+ SourceConnector,
+ "get_api_storage_dir_path",
+ return_value=API_STORAGE_DIR,
+ ),
+ ):
+ storage = MagicMock()
+ mocks["FileSystem"].return_value.get_file_storage.return_value = storage
+ mocks["storage"] = storage
+ yield mocks
+
+
+def _upload(name: str, content: bytes, declared: str) -> SimpleUploadedFile:
+ """An uploaded file whose declared Content-Type may not match its bytes."""
+ return SimpleUploadedFile(name, content, content_type=declared)
+
+
+def _stage(files):
+ return SourceConnector.add_input_file_to_api_storage(
+ pipeline_id="pipe-1",
+ workflow_id="wf-1",
+ execution_id="exec-1",
+ file_objs=files,
+ )
+
+
+def _staged_names(storage: MagicMock) -> set[str]:
+ """File names that actually had bytes written to API storage."""
+ return {
+ call.kwargs["path"].rsplit("/", 1)[-1] for call in storage.write.call_args_list
+ }
+
+
+def test_supported_file_is_staged(collaborators) -> None:
+ """A real PDF is staged and returned for dispatch."""
+ result = _stage([_upload("doc.pdf", PDF_BYTES, "application/pdf")])
+
+ assert set(result) == {"doc.pdf"}
+ assert result["doc.pdf"].mime_type == "application/pdf"
+ assert _staged_names(collaborators["storage"]) == {"doc.pdf"}
+
+
+def test_unsupported_bytes_rejected_despite_supported_declared_type(
+ collaborators,
+) -> None:
+ """The declared Content-Type must not decide what reaches the bucket.
+
+ A WAV file announced as application/pdf satisfies any header-based check,
+ so only sniffing the bytes keeps it out.
+ """
+ result = _stage([_upload("evil.pdf", WAV_BYTES, "application/pdf")])
+
+ # Never dispatched...
+ assert result == {}
+ # ...and never written to the bucket.
+ collaborators["storage"].write.assert_not_called()
+
+
+def test_rejection_is_reported_to_the_caller(collaborators) -> None:
+ """A rejected file gets its own failed entry in the API response."""
+ _stage([_upload("evil.pdf", WAV_BYTES, "application/pdf")])
+
+ collaborators["ResultCacheUtils"].update_api_results.assert_called_once()
+ api_result = collaborators["ResultCacheUtils"].update_api_results.call_args.kwargs[
+ "api_result"
+ ]
+ assert api_result.file == "evil.pdf"
+ # The message has to name the offending type, not a downstream symptom.
+ assert "audio/x-wav" in api_result.error
+ assert api_result.status == ApiDeploymentResultStatus.FAILED
+
+
+def test_rejection_is_persisted_as_a_failed_file_execution(collaborators) -> None:
+ """A rejection has to outlive the result cache.
+
+ The cache entry is deleted on the first status poll and expires with the
+ result TTL, so without a row nothing can answer "why was my file not
+ processed?" afterwards. The row is also what makes the rejection countable:
+ the execution serializer derives successful/failed by counting these, so a
+ partially-rejected run stops reading as a clean success.
+ """
+ _stage(
+ [
+ _upload("good.pdf", PDF_BYTES, "application/pdf"),
+ _upload("evil.pdf", WAV_BYTES, "application/pdf"),
+ ]
+ )
+
+ manager = collaborators["WorkflowFileExecution"].objects
+ manager.get_or_create_file_execution.assert_called_once()
+ kwargs = manager.get_or_create_file_execution.call_args.kwargs
+ assert kwargs["is_api"] is True
+ recorded = kwargs["file_hash"]
+ assert recorded.file_name == "evil.pdf"
+ # The real type is what a support question needs, not the declared one.
+ assert recorded.mime_type == "audio/x-wav"
+ # An opaque identity, not the raw name: the API path has no file_path, so
+ # this value is the only thing keeping two rejections apart.
+ assert recorded.file_hash
+ assert recorded.file_hash != "evil.pdf"
+
+ created = manager.get_or_create_file_execution.return_value
+ created.update_status.assert_called_once()
+ status_kwargs = created.update_status.call_args.kwargs
+ assert status_kwargs["status"] is src_mod.ExecutionStatus.ERROR
+ assert "audio/x-wav" in status_kwargs["execution_error"]
+
+
+def test_two_rejected_files_get_two_rows(collaborators) -> None:
+ """Two rejections must never fold into one row.
+
+ Identical bytes under different names is the case that breaks a
+ content-only identity: the second upload would reuse the first row, whose
+ name was set at creation, so one rejection would be recorded instead of two
+ and the counts would be short by one.
+ """
+ _stage(
+ [
+ _upload("a.pdf", WAV_BYTES, "application/pdf"),
+ _upload("b.pdf", WAV_BYTES, "application/pdf"),
+ ]
+ )
+
+ manager = collaborators["WorkflowFileExecution"].objects
+ hashes = {
+ call.kwargs["file_hash"].file_hash
+ for call in manager.get_or_create_file_execution.call_args_list
+ }
+ assert len(hashes) == 2
+
+
+def test_two_identical_same_named_rejections_get_two_rows(collaborators) -> None:
+ """Every rejected part is its own record.
+
+ Same bytes under the same name is the case content and name alone cannot
+ separate, so the counts would be short by one for each repeat.
+ """
+ _stage(
+ [
+ _upload("same.pdf", WAV_BYTES, "application/pdf"),
+ _upload("same.pdf", WAV_BYTES, "application/pdf"),
+ ]
+ )
+
+ manager = collaborators["WorkflowFileExecution"].objects
+ identities = {
+ call.kwargs["file_hash"].file_hash
+ for call in manager.get_or_create_file_execution.call_args_list
+ }
+ assert len(identities) == 2
+
+
+def test_a_failed_row_write_does_not_fail_the_request(collaborators) -> None:
+ """Bookkeeping must not cost the caller their good files."""
+ collaborators[
+ "WorkflowFileExecution"
+ ].objects.get_or_create_file_execution.side_effect = Exception("db is down")
+
+ result = _stage(
+ [
+ _upload("good.pdf", PDF_BYTES, "application/pdf"),
+ _upload("evil.pdf", WAV_BYTES, "application/pdf"),
+ ]
+ )
+
+ assert set(result) == {"good.pdf"}
+
+
+def test_missing_declared_type_falls_back_to_sniffed_type(collaborators) -> None:
+ """A supported file with no declared Content-Type is still staged.
+
+ The recorded type comes from the bytes, so an absent header neither blocks
+ the file nor degrades it to application/octet-stream.
+ """
+ result = _stage([_upload("doc.pdf", PDF_BYTES, "")])
+
+ assert result["doc.pdf"].mime_type == "application/pdf"
+
+
+def test_supported_files_survive_a_rejected_sibling(collaborators) -> None:
+ """One bad file does not fail the whole request."""
+ result = _stage(
+ [
+ _upload("good.pdf", PDF_BYTES, "application/pdf"),
+ _upload("evil.pdf", WAV_BYTES, "application/pdf"),
+ ]
+ )
+
+ assert set(result) == {"good.pdf"}
+ assert _staged_names(collaborators["storage"]) == {"good.pdf"}
+
+
+def _ole2_like(total_size: int) -> bytes:
+ """An OLE2 compound file whose format markers sit past the sample window.
+
+ libmagic resolves .doc/.xls/.ppt through the OLE2 directory sector, which
+ lives at the end of the file. Only the container signature is visible in the
+ leading bytes, which is exactly the shape that made these files unstageable.
+ """
+ header = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1" + b"\x00" * 504
+ return header + b"\x00" * (total_size - len(header))
+
+
+def _docx_bytes(pad_first_member: bool) -> bytes:
+ """A minimal but genuine OOXML package.
+
+ With `pad_first_member`, `[Content_Types].xml` is no longer the first entry
+ — what re-zipping or a streaming writer produces. libmagic then declines to
+ name it beyond `application/zip`, which is the shape that has to be resolved
+ by looking inside rather than by asking libmagic again.
+ """
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
+ if pad_first_member:
+ archive.writestr("junk.bin", bytes(range(256)) * 4)
+ archive.writestr("[Content_Types].xml", _OOXML_CONTENT_TYPES)
+ archive.writestr("_rels/.rels", _OOXML_RELS)
+ archive.writestr("word/document.xml", '')
+ return buf.getvalue()
+
+
+@pytest.mark.parametrize("repackaged", [False, True], ids=["normal", "repackaged"])
+def test_real_ooxml_bytes_are_accepted(collaborators, repackaged: bool) -> None:
+ """Real container bytes, with libmagic unmocked.
+
+ Every other container test scripts libmagic's answers, so none of them would
+ notice the detector asking a question libmagic cannot answer. A repackaged
+ OOXML package reads as a bare zip however many bytes it is given, so it is
+ the case that proves the resolution works rather than the mocks agreeing
+ with each other.
+ """
+ data = _docx_bytes(pad_first_member=repackaged)
+
+ result = _stage([_upload("report.docx", data, "application/pdf")])
+
+ assert set(result) == {"report.docx"}
+ assert result["report.docx"].mime_type == (
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+ )
+ # Detection copies the upload to classify it; staging must still see it all.
+ written = b"".join(
+ call.kwargs["data"] for call in collaborators["storage"].write.call_args_list
+ )
+ assert written == data
+
+
+def test_an_oversized_mimetype_member_is_not_decompressed(collaborators) -> None:
+ """A hostile archive must not turn the type check into a decompression bomb.
+
+ `mimetype` holds one media type string, so a member declaring megabytes is
+ not the thing being looked for. Reading it unbounded would materialise
+ whatever the archive declares, during synchronous staging, failing the whole
+ request before anything is dispatched.
+ """
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
+ # Highly compressible, so the archive stays small while the member does not.
+ archive.writestr("mimetype", "a" * (8 * 1024 * 1024))
+ payload = buf.getvalue()
+ assert len(payload) < 100_000 # small on the wire, large on the way out
+
+ result = _stage([_upload("bomb.pdf", payload, "application/pdf")])
+
+ # Not recognised, and never expanded to find that out.
+ assert result == {}
+
+
+def _odf_bytes(declared: str, *, stored: bool = True, first: bool = True) -> bytes:
+ """An ODF-shaped archive declaring `declared` in its mimetype member."""
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
+ if not first:
+ archive.writestr("junk.bin", "x" * 64)
+ archive.writestr(
+ zipfile.ZipInfo("mimetype"),
+ declared,
+ compress_type=zipfile.ZIP_STORED if stored else zipfile.ZIP_DEFLATED,
+ )
+ archive.writestr("content.xml", "")
+ return buf.getvalue()
+
+
+def test_a_genuine_odf_document_is_recognised(collaborators) -> None:
+ """The mimetype member identifies which ODF document this is."""
+ data = _odf_bytes("application/vnd.oasis.opendocument.text")
+
+ result = _stage([_upload("notes.pdf", data, "application/pdf")])
+
+ assert set(result) == {"notes.pdf"}
+ assert result["notes.pdf"].mime_type == "application/vnd.oasis.opendocument.text"
+
+
+@pytest.mark.parametrize(
+ "declared, stored, first, why",
+ [
+ ("application/pdf", True, True, "names an unrelated format"),
+ ("text/plain", True, True, "names a format it cannot vouch for"),
+ ("application/vnd.oasis.opendocument.text", False, True, "is compressed"),
+ ("application/vnd.oasis.opendocument.text", True, False, "is not first"),
+ ],
+)
+def test_a_mimetype_member_cannot_nominate_a_format(
+ collaborators, declared: str, stored: bool, first: bool, why: str
+) -> None:
+ """The member is archive-controlled, so it may identify but never choose.
+
+ Without this an arbitrary zip declaring an allow-listed type would be waved
+ straight through — recreating the deferred extraction failure the whole gate
+ exists to prevent. The spec's own conditions are what make it evidence.
+ """
+ data = _odf_bytes(declared, stored=stored, first=first)
+
+ result = _stage([_upload("spoof.pdf", data, "application/pdf")])
+
+ assert result == {}, f"accepted an archive whose mimetype member {why}"
+
+
+@pytest.mark.parametrize(
+ "stored", [True, False], ids=["stored-entry", "deflated-entry"]
+)
+def test_a_zip_carrying_a_pdf_marker_is_still_a_zip(collaborators, stored: bool) -> None:
+ """A zip whose contents mention a PDF header must not be staged as a PDF.
+
+ An entry stored uncompressed puts `%PDF-1.7` in the archive's literal leading
+ bytes, so a scan of the head sees it exactly as it would in a real document.
+ Resolving the wrapper before the marker is consulted is what keeps that from
+ becoming a way past the gate.
+ """
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w") as archive:
+ archive.writestr(
+ zipfile.ZipInfo("a.txt"),
+ "%PDF-1.7 pretending to be a pdf",
+ compress_type=zipfile.ZIP_STORED if stored else zipfile.ZIP_DEFLATED,
+ )
+ archive.writestr("b.bin", bytes(range(256)) * 4)
+ data = buf.getvalue()
+ if stored:
+ assert b"%PDF-1.7" in data[:1024], "fixture must put the marker in the head"
+
+ result = _stage([_upload("archive.pdf", data, "application/pdf")])
+
+ assert result == {}
+ collaborators["storage"].write.assert_not_called()
+
+
+def test_an_incomplete_ooxml_package_is_rejected(collaborators) -> None:
+ """Marker directories alone are not a document.
+
+ Scoped to the repackaged shape, which is the one resolved by reading the
+ archive: with `[Content_Types].xml` first, libmagic names the package itself
+ and this code never runs. An archive with the markers but no relationships
+ is not loadable, so naming it a document would only move the failure to
+ extraction — verified live, where one of exactly this shape was accepted
+ here and then rejected by the extractor.
+ """
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
+ # Ahead of the markers, so libmagic will only call this a zip.
+ archive.writestr("junk.bin", bytes(range(256)) * 4)
+ # Every member name a package has, but nothing declaring a main part.
+ archive.writestr("[Content_Types].xml", '')
+ archive.writestr("_rels/.rels", '')
+ archive.writestr("word/document.xml", '')
+ data = buf.getvalue()
+ # Which wrapper libmagic names here differs between builds - octet-stream on
+ # some, application/zip on others - so assert only the property the test
+ # depends on: it is not named a document, so the archive gets inspected.
+ assert src_mod.magic.from_buffer(data, mime=True) in src_mod.INCONCLUSIVE_MIME_TYPES
+
+ result = _stage([_upload("partial.pdf", data, "application/pdf")])
+
+ assert result == {}
+
+
+@pytest.mark.parametrize(
+ "wrapper", ["application/zip", "application/x-ole-storage"]
+)
+def test_a_recognised_wrapper_is_not_promoted_by_its_contents(
+ collaborators, wrapper: str
+) -> None:
+ """A wrapper libmagic named keeps its own answer.
+
+ Both wrappers, because excluding only one leaves the other able to be
+ reclassified as a PDF by a marker sitting in its payload.
+ """
+ payload = b"\x00" * 64 + b"%PDF-1.7 not a pdf" + b"\x00" * 512
+
+ # First call classifies the sample; a second would only happen if the rescue
+ # ran, and it answers application/pdf so a promotion would be visible as an
+ # accepted file rather than hidden behind the same wrapper verdict.
+ with mock.patch.object(
+ src_mod.magic, "from_buffer", side_effect=[wrapper, "application/pdf"]
+ ), mock.patch.object(src_mod.magic, "from_file", return_value=wrapper):
+ result = _stage([_upload("wrapped.pdf", payload, "application/pdf")])
+
+ assert result == {}
+
+
+def _ooxml_archive(content_types: str, rels: str, part: str) -> bytes:
+ """A repackaged archive with the three members set as given."""
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
+ # Ahead of the markers, so libmagic will only call this a zip and the
+ # archive has to be read to resolve it.
+ archive.writestr("junk.bin", bytes(range(256)) * 4)
+ archive.writestr("[Content_Types].xml", content_types)
+ archive.writestr("_rels/.rels", rels)
+ archive.writestr(part, "")
+ return buf.getvalue()
+
+
+@pytest.mark.parametrize(
+ "content_types, rels, part, why",
+ [
+ (
+ '',
+ '',
+ "word/document.xml",
+ "carries the marker as stray text and relates nothing",
+ ),
+ (
+ _OOXML_CONTENT_TYPES,
+ '',
+ "word/document.xml",
+ "declares a main part that nothing relates",
+ ),
+ (
+ _OOXML_CONTENT_TYPES,
+ _OOXML_RELS.replace("word/document.xml", "word/other.xml"),
+ "word/document.xml",
+ "relates a different part than it declares",
+ ),
+ (
+ _OOXML_CONTENT_TYPES,
+ _OOXML_RELS,
+ "word/styles.xml",
+ "declares and relates a part the archive does not contain",
+ ),
+ ],
+ ids=["stray-marker", "unrelated", "disagrees", "part-absent"],
+)
+def test_an_archive_wearing_the_names_is_not_a_document(
+ collaborators, content_types: str, rels: str, part: str, why: str
+) -> None:
+ """Names and markers are not evidence; the declaration must hold together.
+
+ The package has to declare its main part, relate that same part from the
+ root, and actually contain it. Each case here satisfies some of that and is
+ still not openable, so accepting it would only move the failure to
+ extraction — which the live run of one of these shapes confirmed.
+ """
+ data = _ooxml_archive(content_types, rels, part)
+
+ result = _stage([_upload("spoof.pdf", data, "application/pdf")])
+
+ assert result == {}, f"accepted an archive that {why}"
+
+
+def test_attribute_order_and_quoting_do_not_matter(collaborators) -> None:
+ """A real package is still recognised however its XML is written."""
+ main = (
+ "application/vnd.openxmlformats-officedocument.wordprocessingml"
+ ".document.main+xml"
+ )
+ rel = (
+ "http://schemas.openxmlformats.org/officeDocument/2006/"
+ "relationships/officeDocument"
+ )
+ data = _ooxml_archive(
+ f"",
+ f"",
+ "word/document.xml",
+ )
+
+ result = _stage([_upload("report.docx", data, "application/pdf")])
+
+ assert result["report.docx"].mime_type == (
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+ )
+
+
+_ODT = "application/vnd.oasis.opendocument.text"
+
+
+def _odf_archive(declared: str, manifest: str | None) -> bytes:
+ """An ODF-shaped archive: mimetype first and stored, as the spec requires."""
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
+ archive.writestr(
+ zipfile.ZipInfo("mimetype"), declared, compress_type=zipfile.ZIP_STORED
+ )
+ if manifest is not None:
+ archive.writestr("META-INF/manifest.xml", manifest)
+ archive.writestr("content.xml", "")
+ return buf.getvalue()
+
+
+def _odf_manifest(media_type: str) -> str:
+ return (
+ ''
+ f''
+ )
+
+
+@pytest.mark.parametrize(
+ "manifest, expected, why",
+ [
+ (_odf_manifest(_ODT), _ODT, "manifest agrees with the mimetype member"),
+ (None, None, "there is no manifest"),
+ (
+ "",
+ None,
+ "the manifest declares nothing",
+ ),
+ (_odf_manifest("application/zip"), None, "the manifest disagrees"),
+ ],
+ ids=["agrees", "no-manifest", "empty-manifest", "disagrees"],
+)
+def test_odf_is_recognised_only_when_its_manifest_agrees(
+ tmp_path, manifest: str | None, expected: str | None, why: str
+) -> None:
+ """The mimetype member is archive-controlled, so it needs corroborating.
+
+ Exercised against the helper directly: libmagic names this archive shape
+ itself, so staging never escalates it and the branch would otherwise go
+ untested — and libmagic's container verdicts differ between builds, so the
+ branch is reachable even though it is not on this one.
+ """
+ archive = tmp_path / "candidate.zip"
+ archive.write_bytes(_odf_archive(_ODT, manifest))
+
+ assert src_mod.identify_zip_container(str(archive)) == expected, why
+
+
+@pytest.mark.parametrize(
+ "content_types, why",
+ [
+ (
+ '",
+ "declares its main part only inside a comment",
+ ),
+ (
+ ']>&a;',
+ "carries a DTD, which is never parsed",
+ ),
+ ],
+ ids=["commented-out", "dtd"],
+)
+def test_a_declaration_must_survive_parsing(
+ collaborators, content_types: str, why: str
+) -> None:
+ """Matching raw bytes accepted things parsing does not.
+
+ A comment satisfied a pattern looking for the Override tag, and a DTD is an
+ expansion risk rather than a declaration, so both are read by parsing the
+ member instead.
+ """
+ data = _ooxml_archive(content_types, _OOXML_RELS, "word/document.xml")
+
+ result = _stage([_upload("spoof.pdf", data, "application/pdf")])
+
+ assert result == {}, f"accepted an archive that {why}"
+
+
+def test_a_namespaced_declaration_is_still_recognised(collaborators) -> None:
+ """Prefixes are a presentation detail, not a different document."""
+ main = (
+ "application/vnd.openxmlformats-officedocument.wordprocessingml"
+ ".document.main+xml"
+ )
+ data = _ooxml_archive(
+ '',
+ _OOXML_RELS,
+ "word/document.xml",
+ )
+
+ result = _stage([_upload("report.docx", data, "application/pdf")])
+
+ assert result["report.docx"].mime_type == (
+ "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+ )
+
+
+@pytest.mark.parametrize(
+ "encoding", ["utf-8", "utf-16", "utf-32"], ids=["utf-8", "utf-16", "utf-32"]
+)
+def test_a_dtd_is_refused_whatever_it_is_encoded_in(
+ collaborators, encoding: str
+) -> None:
+ """The parser detects the encoding, so the guard has to as well.
+
+ A DTD written in UTF-16 is invisible to a search for ASCII markers but is
+ still expanded once parsing starts, which is a compact upload turning into
+ real work during staging.
+ """
+ main = (
+ "application/vnd.openxmlformats-officedocument.wordprocessingml"
+ ".document.main+xml"
+ )
+ declaration = (
+ f''
+ ''
+ ']>'
+ f''
+ "&b;"
+ )
+
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
+ archive.writestr("junk.bin", bytes(range(256)) * 4)
+ archive.writestr("[Content_Types].xml", declaration.encode(encoding))
+ archive.writestr("_rels/.rels", _OOXML_RELS)
+ archive.writestr("word/document.xml", "")
+
+ result = _stage([_upload("bomb.pdf", buf.getvalue(), "application/pdf")])
+
+ assert result == {}
+
+
+def test_a_plain_zip_stays_rejected(collaborators) -> None:
+ """Looking inside a zip widens recognition, not the allow-list."""
+ buf = io.BytesIO()
+ with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
+ archive.writestr("notes.txt", "just a zip of files")
+ archive.writestr("data.bin", bytes(range(256)) * 8)
+
+ result = _stage([_upload("archive.pdf", buf.getvalue(), "application/pdf")])
+
+ assert result == {}
+ collaborators["storage"].write.assert_not_called()
+
+
+def test_container_prefix_triggers_a_full_file_sniff(collaborators) -> None:
+ """A container type seen in the sample must not decide the verdict alone.
+
+ Pins the regression directly: an OLE2 upload sniffs application/x-ole-storage
+ from its first bytes, which is absent from AllowedFileTypes, so resolving from
+ the sample alone rejects every legacy Office file bigger than the window.
+
+ The sniff results are stubbed because libmagic's container reporting differs
+ between builds; what must hold everywhere is that an inconclusive sample is
+ escalated to the full file instead of being treated as a verdict.
+ """
+ ole_bytes = _ole2_like(SourceConnector.MIME_DETECT_CHUNK_SIZE * 4)
+
+ with mock.patch.object(
+ src_mod.magic, "from_buffer", return_value="application/x-ole-storage"
+ ) as sniff, mock.patch.object(
+ src_mod.magic, "from_file", return_value="application/msword"
+ ) as sniff_path:
+ result = _stage([_upload("legacy.doc", ole_bytes, "application/msword")])
+
+ # Only the cheap sample came from a buffer...
+ assert sniff.call_count == 1
+ assert len(sniff.call_args_list[0].args[0]) == SourceConnector.MIME_DETECT_CHUNK_SIZE
+ # ...and the verdict came from a path, which is the only way libmagic will
+ # name a container.
+ assert sniff_path.call_count == 1
+ # ...and the answer from the full file is what decides.
+ assert set(result) == {"legacy.doc"}
+ assert result["legacy.doc"].mime_type == "application/msword"
+
+
+def test_container_still_rejected_when_the_full_file_is_unsupported(
+ collaborators,
+) -> None:
+ """The full-file re-sniff widens the evidence, not the allow-list."""
+ ole_bytes = _ole2_like(SourceConnector.MIME_DETECT_CHUNK_SIZE * 4)
+ with mock.patch.object(
+ src_mod.magic, "from_buffer", return_value="application/x-ole-storage"
+ ), mock.patch.object(
+ src_mod.magic, "from_file", return_value="application/x-dosexec"
+ ):
+ result = _stage([_upload("legacy.doc", ole_bytes, "application/msword")])
+
+ assert result == {}
+ collaborators["storage"].write.assert_not_called()
+
+
+def test_container_upload_is_not_consumed_by_detection(collaborators) -> None:
+ """Reading the whole file to classify it must still leave it stageable."""
+ ole_bytes = _ole2_like(SourceConnector.MIME_DETECT_CHUNK_SIZE * 4)
+ with mock.patch.object(
+ src_mod.magic, "from_buffer", return_value="application/x-ole-storage"
+ ), mock.patch.object(
+ src_mod.magic, "from_file", return_value="application/msword"
+ ):
+ _stage([_upload("legacy.doc", ole_bytes, "application/msword")])
+
+ written = b"".join(
+ call.kwargs["data"] for call in collaborators["storage"].write.call_args_list
+ )
+ assert written == ole_bytes
+
+
+def test_undetectable_file_fails_alone(collaborators) -> None:
+ """A stream that cannot be read fails its own file, not the whole request."""
+ with mock.patch.object(
+ SourceConnector,
+ "_detect_uploaded_file_mime_type",
+ side_effect=[OSError("stream is gone"), "application/pdf"],
+ ):
+ result = _stage(
+ [
+ _upload("broken.pdf", PDF_BYTES, "application/pdf"),
+ _upload("good.pdf", PDF_BYTES, "application/pdf"),
+ ]
+ )
+
+ assert set(result) == {"good.pdf"}
+ api_result = collaborators["ResultCacheUtils"].update_api_results.call_args.kwargs[
+ "api_result"
+ ]
+ assert api_result.file == "broken.pdf"
+ # An I/O fault and an unsupported format need different follow-ups, so the
+ # message must not blame the file's type.
+ assert "could not determine its type" in api_result.error
+
+
+def test_disk_backed_upload_is_classified_from_its_temp_file(collaborators) -> None:
+ """The branch real legacy-Office uploads actually take must be exercised.
+
+ Django spills anything over FILE_UPLOAD_MAX_MEMORY_SIZE to a
+ TemporaryUploadedFile, and .doc/.xls/.ppt are usually over it. Every other
+ test here builds a SimpleUploadedFile, which has no temporary_file_path, so
+ the in-memory branch is the only one they reach.
+ """
+ ole_bytes = _ole2_like(SourceConnector.MIME_DETECT_CHUNK_SIZE * 4)
+ upload = TemporaryUploadedFile(
+ "legacy.doc", "application/msword", len(ole_bytes), None
+ )
+ upload.write(ole_bytes)
+ upload.seek(0)
+ assert upload.temporary_file_path() # the branch under test
+
+ with mock.patch.object(
+ src_mod.magic, "from_buffer", return_value="application/x-ole-storage"
+ ), mock.patch.object(
+ src_mod.magic, "from_file", return_value="application/msword"
+ ) as from_file:
+ result = _stage([upload])
+
+ # Classified from the path, so a large upload is never buffered whole...
+ from_file.assert_called_once_with(upload.temporary_file_path(), mime=True)
+ # ...and it is staged with the type the full file resolved to.
+ assert set(result) == {"legacy.doc"}
+ assert result["legacy.doc"].mime_type == "application/msword"
+
+
+def test_one_bad_file_survives_a_failing_result_cache(collaborators) -> None:
+ """A cache fault while reporting a rejection must not fail the request.
+
+ update_api_results runs inside the staging loop and its pipeline execute is
+ unguarded, so letting it escape would mark the whole execution ERROR and
+ discard the good files already written.
+ """
+ collaborators["ResultCacheUtils"].update_api_results.side_effect = Exception(
+ "redis is down"
+ )
+
+ result = _stage(
+ [
+ _upload("good.pdf", PDF_BYTES, "application/pdf"),
+ _upload("evil.pdf", WAV_BYTES, "application/pdf"),
+ ]
+ )
+
+ assert set(result) == {"good.pdf"}
+ assert _staged_names(collaborators["storage"]) == {"good.pdf"}
+
+
+def test_systemic_detection_failure_is_not_reported_as_bad_files(
+ collaborators,
+) -> None:
+ """A fault that is not about this file's bytes must fail the request loudly.
+
+ A broken libmagic database or an unreadable temp dir hits every file in
+ every request. Swallowing it per-file would answer 200 COMPLETED with every
+ file marked invalid, hiding a platform outage behind a clean success.
+ """
+ # Built outside the raises block so only _stage can satisfy it: otherwise
+ # the test would still pass if the upload itself were what raised.
+ uploads = [_upload("doc.pdf", PDF_BYTES, "application/pdf")]
+
+ with mock.patch.object(
+ src_mod.magic, "from_buffer", side_effect=RuntimeError("magic db is broken")
+ ):
+ with pytest.raises(RuntimeError, match="magic db is broken"):
+ _stage(uploads)
+
+
+def test_empty_upload_is_staged_rather_than_called_unsupported(collaborators) -> None:
+ """An empty file must reach the downstream empty-file error, not a type error.
+
+ libmagic calls zero bytes application/x-empty, which is absent from
+ AllowedFileTypes; without the short-circuit an empty upload would be reported
+ as an unsupported type, which names the wrong cause.
+ """
+ result = _stage([_upload("empty.pdf", b"", "application/pdf")])
+
+ assert set(result) == {"empty.pdf"}
+ # No bytes means no type to judge, so the gate is skipped rather than passed.
+ assert result["empty.pdf"].mime_type is None
+ collaborators["ResultCacheUtils"].update_api_results.assert_not_called()
+
+
+def test_types_llmwhisperer_supports_are_accepted(collaborators) -> None:
+ """The allow-list mirrors what LLMWhisperer can extract.
+
+ HTML and XML look like things to reject and were rejected until the two sets
+ were reconciled, but LLMWhisperer extracts both - so rejecting them loses a
+ file that would have worked. Accepting what it cannot read and rejecting what
+ it can are the same bug in opposite directions.
+ """
+ result = _stage(
+ [
+ _upload("page.html", HTML_BYTES, "text/html"),
+ _upload("data.xml", b'1', "text/xml"),
+ ]
+ )
+
+ assert set(result) == {"page.html", "data.xml"}
+ assert result["page.html"].mime_type == "text/html"
+ collaborators["ResultCacheUtils"].update_api_results.assert_not_called()
+
+
+def test_pdf_behind_leading_bytes_is_accepted(collaborators) -> None:
+ """A PDF that does not start at offset 0 must still be recognised.
+
+ libmagic's PDF rule only matches at offset 0, so eight bytes of padding are
+ enough to make a real PDF sniff as application/octet-stream. Extractors
+ tolerate the prefix, and LLMWhisperer scans the same window for exactly this
+ reason, so rejecting it here would fail a file that works everywhere else -
+ including in ETL runs, which share this allow-list.
+ """
+ wrapped = b"\x00" * 8 + PDF_BYTES
+
+ result = _stage([_upload("wrapped.pdf", wrapped, "application/pdf")])
+
+ assert set(result) == {"wrapped.pdf"}
+ assert result["wrapped.pdf"].mime_type == "application/pdf"
+ collaborators["ResultCacheUtils"].update_api_results.assert_not_called()
+
+
+def test_pdf_marker_inside_other_content_does_not_smuggle_a_file_through(
+ collaborators,
+) -> None:
+ """The `%PDF-` rescue must not become a way past the gate.
+
+ Promoting on a bare substring would let any unidentifiable blob carrying
+ that text near its start be staged as a PDF, which is the opposite of what
+ this check exists for. Re-classifying from the marker's own offset is what
+ keeps the rescue narrow.
+ """
+ # These bytes must actually sniff as application/octet-stream, or the test
+ # would pass without ever reaching the rescue it is meant to constrain.
+ blob = b"\x00\x01\x02\x03" * 30 + b"%PDF- not really a pdf" + b"\x00\x01\x02\x03" * 100
+ assert src_mod.magic.from_buffer(blob, mime=True) == "application/octet-stream"
+
+ result = _stage([_upload("smuggled.pdf", blob, "application/pdf")])
+
+ assert result == {}
+ collaborators["storage"].write.assert_not_called()
+
+
+def test_a_recognised_zip_is_never_promoted_to_pdf(collaborators) -> None:
+ """A zip is a zip, whatever its entries happen to contain."""
+ with mock.patch.object(
+ src_mod.magic, "from_buffer", return_value="application/zip"
+ ), mock.patch.object(
+ src_mod.magic, "from_file", return_value="application/zip"
+ ):
+ result = _stage(
+ [_upload("archive.pdf", b"PK\x03\x04" + b"%PDF-1.7" * 8, "application/pdf")]
+ )
+
+ assert result == {}
+
+
+def test_unidentifiable_binary_is_rejected(collaborators) -> None:
+ """octet-stream is no longer a free pass.
+
+ Anything libmagic cannot name is something LLMWhisperer cannot extract; it
+ used to be allow-listed, which is how a zip renamed .pdf reached the
+ extractor and came back as a 415.
+ """
+ with mock.patch.object(
+ src_mod.magic, "from_buffer", return_value="application/octet-stream"
+ ):
+ result = _stage([_upload("mystery.pdf", b"\x00\x01\x02\x03" * 64, "application/pdf")])
+
+ assert result == {}
+ collaborators["storage"].write.assert_not_called()
diff --git a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py
index d9cc5ba4da..25770e8c6e 100644
--- a/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py
+++ b/backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py
@@ -633,7 +633,7 @@ def test_result_acknowledge_does_not_touch_status_or_counters(self):
stale = self._stale(
ex, ExecutionStatus.EXECUTING, successful_files=None, failed_files=None
)
- WorkflowHelper._set_result_acknowledge(stale)
+ WorkflowHelper.set_result_acknowledge(stale)
ex.refresh_from_db()
assert ex.status == ExecutionStatus.COMPLETED.value
assert ex.successful_files == 1
diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py
index 66b22d3216..6d6c969022 100644
--- a/backend/workflow_manager/internal_views.py
+++ b/backend/workflow_manager/internal_views.py
@@ -934,14 +934,53 @@ def _restamp_execution_time_from_files(execution) -> None:
@staticmethod
def _update_file_aggregates(execution, validated_data) -> None:
- """Persist total / successful / failed file counts when present."""
+ """Persist total / failed / successful file counts when present.
+
+ Floored at what the file rows say, because the worker's aggregate only
+ counts files it was dispatched and misses any rejected before that.
+ """
update_fields: list[str] = []
for field in ("total_files", "successful_files", "failed_files"):
if validated_data.get(field) is not None:
setattr(execution, field, validated_data[field])
update_fields.append(field)
- if update_fields:
- execution.save(update_fields=update_fields)
+ if not update_fields:
+ return
+
+ WorkflowExecutionInternalViewSet._floor_counters_at_file_rows(
+ execution, update_fields
+ )
+ execution.save(update_fields=update_fields)
+
+ @staticmethod
+ def _floor_counters_at_file_rows(execution, update_fields: "list[str]") -> None:
+ """Raise the counters being written to match the file rows on disk.
+
+ Only ever corrects upward, and only to a count already true, so a run
+ with nothing rejected is left alone.
+ """
+ wants_failed = "failed_files" in update_fields
+ wants_total = "total_files" in update_fields
+ if not (wants_failed or wants_total):
+ return
+
+ from workflow_manager.workflow_v2.enums import ExecutionStatus
+
+ try:
+ rows = execution.file_executions
+ error_rows = rows.filter(status=ExecutionStatus.ERROR.value).count()
+ row_total = rows.count()
+ except Exception:
+ logger.exception(
+ "Could not read file rows to floor the counters for execution %s",
+ execution.id,
+ )
+ return
+
+ if wants_failed:
+ execution.failed_files = max(execution.failed_files or 0, error_rows)
+ if wants_total:
+ execution.total_files = max(execution.total_files or 0, row_total)
class FileBatchCreateAPIView(APIView):
diff --git a/backend/workflow_manager/workflow_v2/execution.py b/backend/workflow_manager/workflow_v2/execution.py
index 213c81302b..050639fa06 100644
--- a/backend/workflow_manager/workflow_v2/execution.py
+++ b/backend/workflow_manager/workflow_v2/execution.py
@@ -387,6 +387,46 @@ def update_execution_err(execution_id: str, err_msg: str = "") -> WorkflowExecut
except WorkflowExecution.DoesNotExist:
logger.error(f"execution doesn't exist {execution_id}")
+ @staticmethod
+ def update_execution_completed(
+ execution_id: str, total_files: int = 0, failed_files: int = 0
+ ) -> WorkflowExecution | None:
+ """Terminalise an execution that finished without any work to dispatch.
+
+ The counters must be written alongside the status: a terminal row whose
+ failed_files is NULL reads as a clean success to is_failure_run() and to
+ run history, which would hide a run whose files were all rejected.
+
+ Returns the row as persisted, so callers can see whether the status
+ actually changed rather than assuming it did.
+ """
+ try:
+ execution = WorkflowExecution.objects.get(pk=execution_id)
+ # Routed through the model method so the PG terminal-one-way guard
+ # applies; the legacy transport has no such guard.
+ execution.update_execution(status=ExecutionStatus.COMPLETED)
+ execution.total_files = total_files
+ execution.successful_files = 0
+ execution.failed_files = failed_files
+ # Field-scoped for the columns, but save() still republishes this
+ # object's status to the execution cache. Safe only because the row is
+ # undispatched here, so no worker can have moved it on.
+ execution.save(
+ update_fields=[
+ "total_files",
+ "successful_files",
+ "failed_files",
+ "modified_at",
+ ]
+ )
+ # The guard may have refused the status change without raising; re-read
+ # so the returned row reflects what is actually stored.
+ execution.refresh_from_db()
+ return execution
+ except WorkflowExecution.DoesNotExist:
+ logger.error(f"execution doesn't exist {execution_id}")
+ return None
+
@staticmethod
def update_execution_task(execution_id: str, task_id: str) -> None:
try:
@@ -459,7 +499,7 @@ def update_execution_queue_message_id(
# update_fields) re-runs _handle_execution_cache(), which would republish
# this method's stale in-memory status to the Redis execution cache and can
# clobber a status the worker has since advanced — the same reason
- # _set_result_acknowledge uses a queryset .update(). This marker write must
+ # set_result_acknowledge uses a queryset .update(). This marker write must
# touch ONLY the handle column and never the status/counters or the cache.
updated = WorkflowExecution.objects.filter(pk=execution_id).update(
queue_message_id=queue_message_id
diff --git a/backend/workflow_manager/workflow_v2/tests/test_record_dispatch_handle.py b/backend/workflow_manager/workflow_v2/tests/test_record_dispatch_handle.py
index dcba6ac655..7611837f5b 100644
--- a/backend/workflow_manager/workflow_v2/tests/test_record_dispatch_handle.py
+++ b/backend/workflow_manager/workflow_v2/tests/test_record_dispatch_handle.py
@@ -84,7 +84,7 @@ class TestUpdateQueueMessageIdWriteShape:
has advanced status, it reverts the cache to the stale value with no later
corrector and the API-deployment sync-poll blocks to its full timeout. The
codebase already uses the queryset-``.update()`` pattern for the same reason
- in ``_set_result_acknowledge``; this pins it for the marker write too.
+ in ``set_result_acknowledge``; this pins it for the marker write too.
"""
def test_write_uses_queryset_update_not_save(self):
diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py
index 2c926d0dcd..f50ba77713 100644
--- a/backend/workflow_manager/workflow_v2/workflow_helper.py
+++ b/backend/workflow_manager/workflow_v2/workflow_helper.py
@@ -420,7 +420,7 @@ def get_status_of_async_task(
task_result = ResultCacheUtils.get_api_results(
workflow_id=str(execution.workflow.id), execution_id=execution_id
)
- cls._set_result_acknowledge(execution)
+ cls.set_result_acknowledge(execution)
result_response = ExecutionResponse(
workflow_id=str(execution.workflow.id),
@@ -432,7 +432,7 @@ def get_status_of_async_task(
return result_response
@staticmethod
- def _set_result_acknowledge(execution: WorkflowExecution) -> None:
+ def set_result_acknowledge(execution: WorkflowExecution) -> None:
"""Mark the result as acknowledged and update the database.
This method is called once the task has completed and its result is forgotten.
@@ -701,7 +701,7 @@ def execute_workflow_async(
task_result = ResultCacheUtils.get_api_results(
workflow_id=workflow_id, execution_id=execution_id
)
- cls._set_result_acknowledge(workflow_execution)
+ cls.set_result_acknowledge(workflow_execution)
else:
task_result = None
return ExecutionResponse(
diff --git a/unstract/core/src/unstract/core/mime_gate.py b/unstract/core/src/unstract/core/mime_gate.py
new file mode 100644
index 0000000000..d40b0b13e5
--- /dev/null
+++ b/unstract/core/src/unstract/core/mime_gate.py
@@ -0,0 +1,233 @@
+"""Shared rules for deciding whether a file may enter a workflow.
+
+Lives here because the backend and the workers must answer the same question
+identically. Keep in sync with LLMWhisperer's supported types.
+"""
+
+import re
+import xml.etree.ElementTree as ElementTree
+import zipfile
+
+import magic
+
+# libmagic reports these when it recognises a wrapper but not its contents, so
+# they are never a verdict on their own.
+RECOGNISED_WRAPPERS = frozenset({"application/zip", "application/x-ole-storage"})
+INCONCLUSIVE_MIME_TYPES = RECOGNISED_WRAPPERS | {"application/octet-stream"}
+
+# libmagic matches the PDF header only at offset 0, while extractors tolerate
+# bytes before it, so a wrapped PDF needs the header looked for.
+PDF_HEADER_SCAN_BYTES = 1024
+
+# `%PDF-` alone appears in plenty of content that is not a PDF; a real header
+# carries a version, so require one before promoting anything on its strength.
+PDF_HEADER_PATTERN = re.compile(rb"%PDF-\d\.\d")
+
+# A package's declaration members are small; anything larger is not one.
+PACKAGE_MEMBER_LIMIT = 65536
+
+# Those members never carry a DTD, so refusing one keeps entity expansion away
+# from the parser rather than trusting the parser to be safe about it.
+DTD_MARKERS = (" str:
+ """A tag or attribute name without its namespace."""
+ return tag.rsplit("}", 1)[-1]
+
+
+def _attribute(element, name: str) -> str | None:
+ """An attribute by local name, whatever prefix it carries."""
+ for key, value in element.attrib.items():
+ if _local_name(key) == name:
+ return value.strip()
+ return None
+
+
+def _read_member(archive: zipfile.ZipFile, name: str) -> bytes | None:
+ """A declaration member's bytes, bounded, or None if unusable."""
+ try:
+ info = archive.getinfo(name)
+ if info.file_size > PACKAGE_MEMBER_LIMIT:
+ return None
+ with archive.open(info) as member:
+ return member.read(PACKAGE_MEMBER_LIMIT)
+ except (KeyError, OSError, zipfile.BadZipFile, ValueError):
+ return None
+
+
+def _declares_dtd(raw: bytes) -> bool:
+ """Whether these bytes carry a DTD under any encoding the parser accepts."""
+ for encoding in DTD_SEARCH_ENCODINGS:
+ text = raw.decode(encoding, "ignore").upper()
+ if any(marker in text for marker in DTD_MARKERS):
+ return True
+ return False
+
+
+def _parse_member(archive: zipfile.ZipFile, name: str):
+ """Parse a declaration member, or None if it is unusable.
+
+ Parsed rather than pattern-matched because the archive is untrusted: a
+ comment, a namespace prefix or an escaped attribute each defeat matching on
+ raw bytes, and each was a way past this check.
+ """
+ raw = _read_member(archive, name)
+ if raw is None or _declares_dtd(raw):
+ return None
+ try:
+ return ElementTree.fromstring(raw)
+ except ElementTree.ParseError:
+ return None
+
+
+def _normalise_part(part: str | None) -> str | None:
+ return part.lstrip("/") or None if part else None
+
+
+def _declared_main_part(content_types) -> "tuple[str, str] | None":
+ """The part a package declares as its main one, and the type that implies."""
+ for element in content_types.iter():
+ if _local_name(element.tag) != "Override":
+ continue
+ mime_type = OOXML_MAIN_CONTENT_TYPES.get(_attribute(element, "ContentType") or "")
+ if not mime_type:
+ continue
+ part = _normalise_part(_attribute(element, "PartName"))
+ return (part, mime_type) if part else None
+ return None
+
+
+def _relates_main_part(relationships, part: str) -> bool:
+ """Whether the package root relates `part` as its office document."""
+ for element in relationships.iter():
+ if _local_name(element.tag) != "Relationship":
+ continue
+ relation = _attribute(element, "Type") or ""
+ if not relation.rstrip("/").endswith(OFFICE_DOCUMENT_RELATION):
+ continue
+ if _normalise_part(_attribute(element, "Target")) == part:
+ return True
+ return False
+
+
+def _identify_ooxml(archive: zipfile.ZipFile, names: "list[str]") -> str | None:
+ """Name an OOXML package, or None if the archive only looks like one.
+
+ The package must declare its main part, relate that same part from the root,
+ and contain it. Any of those alone is satisfied by an archive the extractor
+ cannot open.
+ """
+ content_types = _parse_member(archive, "[Content_Types].xml")
+ relationships = _parse_member(archive, "_rels/.rels")
+ if content_types is None or relationships is None:
+ return None
+
+ declared = _declared_main_part(content_types)
+ if declared is None:
+ return None
+ part, mime_type = declared
+ if part not in names:
+ return None
+ return mime_type if _relates_main_part(relationships, part) else None
+
+
+def _identify_odf(archive: zipfile.ZipFile, declared: str) -> str | None:
+ """Confirm the manifest agrees with what the mimetype member claims."""
+ manifest = _parse_member(archive, "META-INF/manifest.xml")
+ if manifest is None:
+ return None
+ for element in manifest.iter():
+ if _local_name(element.tag) != "file-entry":
+ continue
+ if (
+ _attribute(element, "full-path") == "/"
+ and _attribute(element, "media-type") == declared
+ ):
+ return declared
+ return None
+
+
+def identify_zip_container(path: str) -> str | None:
+ """Name the document inside a zip, or None if it is just a zip.
+
+ Only ever returns a constant from this module, never a string read out of
+ the archive, so it widens what is recognised and not what is permitted.
+ """
+ try:
+ with zipfile.ZipFile(path) as archive:
+ names = archive.namelist()
+ # ODF requires this member to be first and stored uncompressed, and
+ # its manifest to declare the same type for the package root.
+ info = archive.infolist()[0] if names else None
+ if (
+ info is not None
+ and info.filename == "mimetype"
+ and info.compress_type == zipfile.ZIP_STORED
+ and info.file_size <= PACKAGE_MEMBER_LIMIT
+ ):
+ raw = _read_member(archive, "mimetype") or b""
+ declared = raw.decode("ascii", "ignore").strip()
+ if declared in ODF_MIME_TYPES:
+ return _identify_odf(archive, declared)
+ if {"[Content_Types].xml", "_rels/.rels"} <= set(names):
+ return _identify_ooxml(archive, names)
+ except (zipfile.BadZipFile, OSError, KeyError, ValueError):
+ return None
+ return None
+
+
+def resolve_inconclusive_mime_type(mime_type: str, head: bytes) -> str:
+ """Give a wrapper-only classification a second chance from the leading bytes.
+
+ This promotes a file into the allow-list, so it stays narrow: a recognised
+ wrapper keeps its own answer, and the marker must look like a real header.
+ """
+ if mime_type not in INCONCLUSIVE_MIME_TYPES:
+ return mime_type
+ # A wrapper libmagic positively identified keeps its own answer; only bytes
+ # it could not name at all are open to being read this way.
+ if mime_type in RECOGNISED_WRAPPERS:
+ return mime_type
+ match = PDF_HEADER_PATTERN.search(head[:PDF_HEADER_SCAN_BYTES])
+ if match is None:
+ return mime_type
+ # Classified from the header's own offset rather than trusted as a substring.
+ return magic.from_buffer(head[match.start() :], mime=True)
diff --git a/workers/api-deployment/tasks.py b/workers/api-deployment/tasks.py
index 4d4b56b3b5..fa7a218d1a 100644
--- a/workers/api-deployment/tasks.py
+++ b/workers/api-deployment/tasks.py
@@ -218,7 +218,52 @@ def _unified_api_execution(
converted_files = FileProcessingUtils.convert_file_hash_data(hash_values_of_files)
if not converted_files:
- logger.warning("No valid files to process after conversion")
+ # An empty result means either nothing was dispatched or every file
+ # failed to convert; only the first is a success.
+ if hash_values_of_files:
+ error_message = (
+ f"None of the {len(hash_values_of_files)} dispatched files could "
+ "be converted for processing"
+ )
+ logger.error(error_message)
+ api_client.update_workflow_execution_status(
+ execution_id=execution_id,
+ status=ExecutionStatus.ERROR.value,
+ error_message=error_message,
+ total_files=len(hash_values_of_files),
+ successful_files=0,
+ failed_files=len(hash_values_of_files),
+ )
+ if pipeline_id:
+ api_client.update_pipeline_status(
+ pipeline_id=pipeline_id,
+ status=PipelineStatus.FAILURE.value,
+ )
+ return {
+ "execution_id": execution_id,
+ "status": "ERROR",
+ # "error", matching this function's other ERROR returns.
+ "error": error_message,
+ "files_processed": 0,
+ }
+
+ logger.warning("No files dispatched for this execution")
+ # Without this write the row keeps its dispatch status and the caller
+ # polls forever.
+ api_client.update_workflow_execution_status(
+ execution_id=execution_id,
+ status=ExecutionStatus.COMPLETED.value,
+ total_files=0,
+ # Explicit, because NULL counters on a terminal row read as a
+ # clean success.
+ successful_files=0,
+ failed_files=0,
+ )
+ if pipeline_id:
+ api_client.update_pipeline_status(
+ pipeline_id=pipeline_id,
+ status=PipelineStatus.SUCCESS.value,
+ )
return {
"execution_id": execution_id,
"status": "COMPLETED",
diff --git a/workers/shared/enums/file_types.py b/workers/shared/enums/file_types.py
index f49b7880e3..b29ee8e57d 100644
--- a/workers/shared/enums/file_types.py
+++ b/workers/shared/enums/file_types.py
@@ -7,6 +7,11 @@
from enum import Enum
from typing import Any
+from unstract.core.mime_gate import ( # noqa: F401 re-exported
+ INCONCLUSIVE_MIME_TYPES,
+ resolve_inconclusive_mime_type,
+)
+
class AllowedFileTypes(Enum):
"""Allowed MIME types for file processing.
@@ -47,18 +52,22 @@ class AllowedFileTypes(Enum):
# Other formats
CDFV2 = "application/CDFV2"
- OCTET_STREAM = "application/octet-stream"
@classmethod
def is_allowed(cls, mime_type: str) -> bool:
"""Check if a MIME type is allowed for processing.
+ Any `text/*` passes, which is how html, xml, tsv, rtf and markdown reach
+ the extractor as text rather than as a listed format.
+
Args:
mime_type: The MIME type string to validate
Returns:
bool: True if the MIME type is allowed, False otherwise
"""
+ if mime_type.startswith("text/"):
+ return True
return mime_type in cls._value2member_map_
diff --git a/workers/shared/workflow/execution/service.py b/workers/shared/workflow/execution/service.py
index a62253d51c..ebf8bce17c 100644
--- a/workers/shared/workflow/execution/service.py
+++ b/workers/shared/workflow/execution/service.py
@@ -10,7 +10,7 @@
from typing import Any
import magic
-from shared.enums.file_types import AllowedFileTypes
+from shared.enums.file_types import AllowedFileTypes, resolve_inconclusive_mime_type
from shared.exceptions.execution_exceptions import (
NotFoundDestinationConfiguration,
NotFoundSourceConfiguration,
@@ -1219,7 +1219,9 @@ def _copy_filesystem_file(
while chunk := source_file.read(self.READ_CHUNK_SIZE):
# MIME type detection and validation on first chunk
if first_chunk:
- mime_type = magic.from_buffer(chunk, mime=True)
+ mime_type = resolve_inconclusive_mime_type(
+ magic.from_buffer(chunk, mime=True), chunk
+ )
logger.info(f"Detected MIME type: {mime_type} for file {file_path}")
if not AllowedFileTypes.is_allowed(mime_type):