Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
34 commits
Select commit Hold shift + click to select a range
215929d
UN-1924 [FIX] Reject unsupported files by sniffing MIME in API storag…
Deepak-Kesavan Sep 1, 2026
7af4732
UN-1924 [FIX] Terminalise an API execution when every file is rejected
Deepak-Kesavan Sep 1, 2026
63d5237
UN-1924 [FIX] Address review findings on unsupported-file rejection
Deepak-Kesavan Sep 2, 2026
4b436a8
UN-1924 [FIX] Release API rate limit slots by org id, not model instance
Deepak-Kesavan Sep 2, 2026
9929f88
UN-1924 [FIX] Fail one undetectable upload instead of the whole request
Deepak-Kesavan Sep 2, 2026
0efd943
UN-1924 [FIX] Acknowledge results and notify subscribers on an all-re…
Deepak-Kesavan Sep 7, 2026
b27ef4c
UN-1924 [FIX] Keep acknowledgement and notification independent
Deepak-Kesavan Sep 7, 2026
2bf3bb7
UN-1924 [MISC] Drop formatter churn on two untouched assertions
Deepak-Kesavan Sep 7, 2026
cb97bf6
UN-1924 [FIX] Mirror LLMWhisperer's file type gate
Deepak-Kesavan Sep 18, 2026
938a957
Merge remote-tracking branch 'origin/main' into UN-1924-reject-unsupp…
Deepak-Kesavan Sep 18, 2026
6d9f450
UN-1924 [FIX] Recognise a PDF that does not start at offset 0
Deepak-Kesavan Sep 21, 2026
80867ea
Merge remote-tracking branch 'origin/main' into UN-1924-reject-unsupp…
Deepak-Kesavan Sep 21, 2026
5cdce85
UN-1924 [FIX] Harden rejection reporting and the PDF rescue
Deepak-Kesavan Sep 21, 2026
61985aa
UN-1924 [FIX] Require a PDF version header before rescuing a file
Deepak-Kesavan Sep 21, 2026
1b274b5
UN-1924 [FIX] Resolve zip containers, which libmagic will not name fr…
Deepak-Kesavan Sep 21, 2026
048e94f
UN-1924 [REFACTOR] Move the shared MIME gate logic into unstract/core
Deepak-Kesavan Sep 21, 2026
85d3ae5
UN-1924 [FIX] Bound the ODF mimetype read so a zip bomb cannot expand…
Deepak-Kesavan Sep 21, 2026
6c953a9
UN-1924 [FIX] Stop an archive nominating its own file type
Deepak-Kesavan Sep 21, 2026
3b50d5e
UN-1924 [MISC] Leave one throwing call inside the raises block
Deepak-Kesavan Sep 21, 2026
9c8e9ff
UN-1924 [MISC] Pin the zip-carrying-a-PDF-marker case
Deepak-Kesavan Sep 21, 2026
81a9e6a
UN-1924 [FIX] Three small inconsistencies in the new terminal paths
Deepak-Kesavan Sep 21, 2026
f059e50
UN-1924 [FIX] Count and keep files rejected before dispatch
Deepak-Kesavan Sep 21, 2026
2d70dfd
UN-1924 [FIX] Keep two rejections apart when the bytes are identical
Deepak-Kesavan Sep 21, 2026
efcbbf9
UN-1924 [MISC] Drop a line-number cross-reference from a comment
Deepak-Kesavan Oct 12, 2026
188f3e0
UN-1924 [MISC] Shorten the new code comments
Deepak-Kesavan Oct 12, 2026
cbd26fe
UN-1924 [MISC] Clear the two Sonar findings
Deepak-Kesavan Oct 12, 2026
9c619e7
UN-1924 [FIX] Require relationships before naming a zip an Office doc…
Deepak-Kesavan Oct 12, 2026
3afd489
UN-1924 [FIX] Let neither recognised wrapper be renamed by its payload
Deepak-Kesavan Oct 12, 2026
9ce3086
UN-1924 [FIX] Require a declared main part, and one row per rejected …
Deepak-Kesavan Oct 12, 2026
54a5317
UN-1924 [FIX] Correlate the declared main part with what relates it
Deepak-Kesavan Oct 12, 2026
a732034
UN-1924 [FIX] Parse package declarations instead of matching their bytes
Deepak-Kesavan Oct 12, 2026
a5e664f
UN-1924 [MISC] Split the OOXML check to clear the complexity gate
Deepak-Kesavan Oct 12, 2026
6d2b3ab
Merge branch 'main' into UN-1924-reject-unsupported-files
Deepak-Kesavan Oct 12, 2026
f4b0c85
UN-1924 [FIX] Refuse a DTD whatever encoding it is written in
Deepak-Kesavan Oct 12, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions backend/api_v2/deployment_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
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.workflow_v2.dto import ExecutionResponse
from workflow_manager.workflow_v2.enums import ExecutionStatus
Expand Down Expand Up @@ -306,6 +307,26 @@ def execute_workflow(
)
).data

# Staging rejected every file, so there is nothing to dispatch. The worker
# short-circuits an empty file set without writing a status back, which
# would strand this execution in PENDING — terminalise it here instead.
if not hash_values_of_files:
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved
WorkflowExecutionServiceHelper.update_execution_completed(str(execution_id))
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved
APIDeploymentRateLimiter.release_slot(api.organization, str(execution_id))
DestinationConnector.delete_api_storage_dir(
workflow_id=workflow_id, execution_id=execution_id
)
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved
return APIExecutionResponseSerializer(
ExecutionResponse(
workflow_id=workflow_id,
execution_id=execution_id,
execution_status=ExecutionStatus.COMPLETED.value,
result=ResultCacheUtils.get_api_results(
workflow_id=str(workflow_id), execution_id=str(execution_id)
),
)
).data
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved

try:
result = WorkflowHelper.execute_workflow_async(
workflow_id=workflow_id,
Expand Down
67 changes: 62 additions & 5 deletions backend/api_v2/tests/test_deployment_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@ def collaborators():
mocks[
"WorkflowExecutionServiceHelper"
].create_workflow_execution.return_value = execution_row
mocks["SourceConnector"].add_input_file_to_api_storage.side_effect = (
RuntimeError("boom")
mocks["SourceConnector"].add_input_file_to_api_storage.side_effect = RuntimeError(
"boom"
)
yield mocks

Expand Down Expand Up @@ -74,9 +74,9 @@ def test_staging_failure_marks_execution_error(collaborators) -> None:

def test_staging_failure_cleanup_survives_db_marking_error(collaborators) -> None:
"""If marking the row ERROR itself raises, cleanup must still run (not propagate)."""
collaborators["WorkflowExecutionServiceHelper"].update_execution_err.side_effect = (
RuntimeError("db down")
)
collaborators[
"WorkflowExecutionServiceHelper"
].update_execution_err.side_effect = RuntimeError("db down")

# Must NOT raise — a failed error-marking should not break cleanup.
dh.DeploymentHelper.execute_workflow(
Expand All @@ -89,3 +89,60 @@ 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,
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"}
]
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
response = dh.DeploymentHelper.execute_workflow(
organization_name="org",
api=_api(),
file_objs=[],
timeout=-1,
)
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved

# Nothing is dispatched...
mocks["WorkflowHelper"].execute_workflow_async.assert_not_called()
# ...the row is terminalised here instead of being left PENDING...
mocks[
"WorkflowExecutionServiceHelper"
].update_execution_completed.assert_called_once_with("exec-123")
# ...the slot and staging dir are released...
mocks["APIDeploymentRateLimiter"].release_slot.assert_called_once()
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"
61 changes: 39 additions & 22 deletions backend/workflow_manager/endpoint_v2/source.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import logging
import os
import shutil
import uuid
from collections.abc import Collection
from hashlib import sha256
from io import BytesIO
Expand All @@ -25,7 +24,11 @@
SourceConstant,
SourceKey,
)
from workflow_manager.endpoint_v2.dto import FileHash, SourceConfig
from workflow_manager.endpoint_v2.dto import (
FileExecutionResult,
FileHash,
SourceConfig,
)
from workflow_manager.endpoint_v2.enums import AllowedFileTypes
from workflow_manager.endpoint_v2.exceptions import (
InvalidInputDirectory,
Expand All @@ -37,6 +40,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
Expand Down Expand Up @@ -69,6 +73,8 @@ class SourceConnector(BaseConnector):
"""

READ_CHUNK_SIZE = 4194304 # Chunk size for reading files
# libmagic classifies from the leading bytes; reading more only costs memory.
MIME_DETECT_CHUNK_SIZE = 8192
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved

def __init__(
self,
Expand Down Expand Up @@ -1187,6 +1193,22 @@ 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:
"""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.
"""
sample = file.read(cls.MIME_DETECT_CHUNK_SIZE)
file.seek(0)
if not sample:
# libmagic reports "application/x-empty" here, which would reject the
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved
# file as an unsupported type. An empty upload is a distinct failure
# and is reported as such once staging hands off, so let it pass.
return AllowedFileTypes.OCTET_STREAM.value
return magic.from_buffer(sample, mime=True)

@classmethod
def add_input_file_to_api_storage(
cls,
Expand Down Expand Up @@ -1228,30 +1250,25 @@ def add_input_file_to_api_storage(
file_name = file.name
destination_path = os.path.join(api_storage_dir, file_name)

mime_type = file.content_type
mime_type = cls._detect_uploaded_file_mime_type(file)
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved
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}"
)
mime_type = AllowedFileTypes.OCTET_STREAM.value

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,
log_message = (
f"Rejecting file '{file_name}' with unsupported MIME type "
f"'{mime_type}'"
)
workflow_log.log_error(logger=logger, message=log_message)
# Rejected files are never dispatched, so nothing downstream will
# report on them - surface the failure in the API response here.
ResultCacheUtils.update_api_results(
workflow_id=workflow_id,
execution_id=execution_id,
api_result=FileExecutionResult(
file=file_name,
error=log_message,
),
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved
)
file_hashes.update({file_name: file_hash})
continue

file_system = FileSystem(FileStorageType.API_EXECUTION)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
"""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.
"""

from unittest import mock
from unittest.mock import MagicMock

import pytest
from django.core.files.uploadedfile import SimpleUploadedFile

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 PDF header sniffs application/pdf, an HTML document sniffs
# text/html, which is absent from AllowedFileTypes.
PDF_BYTES = b"%PDF-1.7\n1 0 obj\n<< /Type /Catalog >>\nendobj\n"
HTML_BYTES = b"<!DOCTYPE html><html><body>hello</body></html>"


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,
) 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.

An HTML file announced as application/pdf satisfies any header-based check,
so only sniffing the bytes keeps it out.
"""
result = _stage([_upload("evil.pdf", HTML_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", HTML_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 "text/html" in api_result.error
assert api_result.status == ApiDeploymentResultStatus.FAILED


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", HTML_BYTES, "application/pdf"),
]
)

assert set(result) == {"good.pdf"}
assert _staged_names(collaborators["storage"]) == {"good.pdf"}
13 changes: 13 additions & 0 deletions backend/workflow_manager/workflow_v2/execution.py
Original file line number Diff line number Diff line change
Expand Up @@ -387,6 +387,19 @@ 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) -> WorkflowExecution | None:
"""Terminalise an execution that finished without any work to dispatch."""
try:
execution = WorkflowExecution.objects.get(pk=execution_id)
# Same reason as update_execution_err: the model method owns the
# terminal-one-way guard, so this cannot revert an already-final row.
Comment thread
Deepak-Kesavan marked this conversation as resolved.
Resolved
execution.update_execution(status=ExecutionStatus.COMPLETED)
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:
Expand Down
Loading
Loading