From c17a96f02f5513282ffac616f44c34a03a5d97ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?= Date: Mon, 7 Sep 2026 08:59:18 +0000 Subject: [PATCH] refactor(service-api): extract file preview application service (#41732) --- api/.importlinter | 14 + .../service_api/app/file_preview.py | 152 ++----- api/extensions/ext_application_services.py | 7 + .../message_file_preview_repository.py | 66 +++ api/services/message_file_preview_service.py | 73 ++++ .../service_api/app/test_file_preview.py | 384 ++++++------------ .../test_ext_application_services.py | 18 + .../test_message_file_preview_repository.py | 180 ++++++++ .../test_message_file_preview_service.py | 82 ++++ 9 files changed, 609 insertions(+), 367 deletions(-) create mode 100644 api/repositories/message_file_preview_repository.py create mode 100644 api/services/message_file_preview_service.py create mode 100644 api/tests/unit_tests/repositories/test_message_file_preview_repository.py create mode 100644 api/tests/unit_tests/services/test_message_file_preview_service.py diff --git a/api/.importlinter b/api/.importlinter index b29dd9cc4ae..9ec0e118b03 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -279,6 +279,20 @@ forbidden_modules = sqlalchemy werkzeug +[importlinter:contract:message-file-preview-service-boundary] +name = Message file preview application service is framework and persistence neutral +type = forbidden +source_modules = + services.message_file_preview_service +forbidden_modules = + controllers + extensions + flask + models + repositories + sqlalchemy + werkzeug + [importlinter:contract:webapp-access-query-service-boundary] name = Web app access query application service is framework and persistence neutral type = forbidden diff --git a/api/controllers/service_api/app/file_preview.py b/api/controllers/service_api/app/file_preview.py index ad96ac96ce1..561ede3d184 100644 --- a/api/controllers/service_api/app/file_preview.py +++ b/api/controllers/service_api/app/file_preview.py @@ -1,11 +1,9 @@ -import logging from urllib.parse import quote from uuid import UUID from flask import Response from flask_restx import Resource from pydantic import BaseModel, Field -from sqlalchemy import select from controllers.common.fields import BinaryFileResponse from controllers.common.file_response import enforce_download_for_html @@ -18,11 +16,13 @@ from controllers.service_api.app.error import ( ) from controllers.service_api.schema import binary_response from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token -from extensions.ext_database import db -from extensions.ext_storage import storage -from models.model import App, EndUser, Message, MessageFile, UploadFile - -logger = logging.getLogger(__name__) +from extensions.ext_application_services import application_services +from models.model import App, EndUser +from services.message_file_preview_service import ( + MessageFilePreview, + MessageFilePreviewAccessDeniedError, + MessageFilePreviewNotFoundError, +) class FilePreviewQuery(BaseModel): @@ -36,6 +36,20 @@ register_schema_model(service_api_ns, FilePreviewQuery) register_response_schema_model(service_api_ns, BinaryFileResponse) FILE_PREVIEW_RESPONSE_MEDIA_TYPE = "*/*" +_RANGE_MEDIA_TYPES = frozenset( + { + "audio/aac", + "audio/flac", + "audio/mp4", + "audio/mpeg", + "audio/ogg", + "audio/wav", + "audio/x-m4a", + "video/mp4", + "video/quicktime", + "video/webm", + } +) @service_api_ns.route("/files//preview") @@ -88,143 +102,65 @@ class FilePreviewApi(Resource): @service_api_ns.response(200, "File retrieved successfully") @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.QUERY)) @model_validate(FilePreviewQuery) - def get(self, args: FilePreviewQuery, app_model: App, end_user: EndUser, file_id: UUID): + def get(self, args: FilePreviewQuery, app_model: App, end_user: EndUser, file_id: UUID) -> Response: """ Preview/Download a file that was uploaded via Service API. Provides secure file preview/download functionality. Files can only be accessed if they belong to messages within the requesting app's context. """ - file_id_str = str(file_id) - - # Validate file ownership and get file objects - _, upload_file = self._validate_file_ownership(file_id_str, app_model.id) - - # Get file content generator try: - generator = storage.load(upload_file.key, stream=True) - except Exception as e: - raise FileNotFoundError(f"Failed to load file content: {str(e)}") - - # Build response with appropriate headers - response = self._build_file_response(generator, upload_file, args.as_attachment) - - return response - - def _validate_file_ownership(self, file_id: str, app_id: str) -> tuple[MessageFile, UploadFile]: - """ - Validate that the file belongs to a message within the requesting app's context - - Security validations performed: - 1. File exists in MessageFile table (was used in a conversation) - 2. Message belongs to the requesting app - 3. UploadFile record exists and is accessible - 4. File tenant matches app tenant (additional security layer) - - Args: - file_id: UUID of the file to validate - app_id: UUID of the requesting app - - Returns: - Tuple of (MessageFile, UploadFile) if validation passes - - Raises: - FileNotFoundError: File or related records not found - FileAccessDeniedError: File does not belong to the app's context - """ - try: - # Input validation - if not file_id or not app_id: - raise FileAccessDeniedError("Invalid file or app identifier") - - # First, find the MessageFile that references this upload file - message_file = db.session.scalar(select(MessageFile).where(MessageFile.upload_file_id == file_id).limit(1)) - - if not message_file: - raise FileNotFoundError("File not found in message context") - - # Get the message and verify it belongs to the requesting app - message = db.session.scalar( - select(Message).where(Message.id == message_file.message_id, Message.app_id == app_id).limit(1) + preview = application_services().message_file_previews.get_preview( + file_id=str(file_id), + app_id=app_model.id, + tenant_id=app_model.tenant_id, ) + except MessageFilePreviewNotFoundError as error: + raise FileNotFoundError() from error + except MessageFilePreviewAccessDeniedError as error: + raise FileAccessDeniedError() from error - if not message: - raise FileAccessDeniedError("File access denied: not owned by requesting app") + return self._build_file_response(preview=preview, as_attachment=args.as_attachment) - # Get the actual upload file record - upload_file = db.session.get(UploadFile, file_id) - - if not upload_file: - raise FileNotFoundError("Upload file record not found") - - # Additional security: verify tenant isolation - app = db.session.get(App, app_id) - if app and upload_file.tenant_id != app.tenant_id: - raise FileAccessDeniedError("File access denied: tenant mismatch") - - return message_file, upload_file - - except (FileNotFoundError, FileAccessDeniedError): - # Re-raise our custom exceptions - raise - except Exception as e: - # Log unexpected errors for debugging - logger.exception( - "Unexpected error during file ownership validation", - extra={"file_id": file_id, "app_id": app_id, "error": str(e)}, - ) - raise FileAccessDeniedError("File access validation failed") - - def _build_file_response(self, generator, upload_file: UploadFile, as_attachment: bool = False) -> Response: + def _build_file_response(self, *, preview: MessageFilePreview, as_attachment: bool = False) -> Response: """ Build Flask Response object with appropriate headers for file streaming Args: - generator: File content generator from storage - upload_file: UploadFile database record + preview: App-scoped file metadata and content stream as_attachment: Whether to set Content-Disposition as attachment Returns: Flask Response object with streaming file content """ + file = preview.file response = Response( - generator, - mimetype=upload_file.mime_type, + preview.content, + mimetype=file.mime_type, direct_passthrough=True, headers={}, ) # Add Content-Length if known - if upload_file.size and upload_file.size > 0: - response.headers["Content-Length"] = str(upload_file.size) + if file.size > 0: + response.headers["Content-Length"] = str(file.size) # Add Accept-Ranges header for audio/video files to support seeking - if upload_file.mime_type in [ - "audio/mpeg", - "audio/wav", - "audio/mp4", - "audio/ogg", - "audio/flac", - "audio/aac", - "video/mp4", - "video/webm", - "video/quicktime", - "audio/x-m4a", - ]: + if file.mime_type in _RANGE_MEDIA_TYPES: response.headers["Accept-Ranges"] = "bytes" # Set Content-Disposition for downloads - if as_attachment and upload_file.name: - encoded_filename = quote(upload_file.name) + if as_attachment and file.name: + encoded_filename = quote(file.name) response.headers["Content-Disposition"] = f"attachment; filename*=UTF-8''{encoded_filename}" # Override content-type for downloads to force download response.headers["Content-Type"] = "application/octet-stream" enforce_download_for_html( response, - mime_type=upload_file.mime_type, - filename=upload_file.name, - extension=upload_file.extension, + mime_type=file.mime_type, + filename=file.name, + extension=file.extension, ) # Add caching headers for performance diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 9c57bb557f6..6b99f5c7a45 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -49,6 +49,7 @@ from repositories.factory import DifyAPIRepositoryFactory from repositories.file_grant_repository import FileGrantRepository from repositories.human_input_file_upload_repository import SQLAlchemyHumanInputFileUploadRepository from repositories.installation_state_repository import InstallationStateRepository +from repositories.message_file_preview_repository import MessageFilePreviewQueryRepository from repositories.oauth_access_token_repository import SQLAlchemyOAuthAccessTokenRepository from repositories.oauth_server_repository import RedisOAuthServerTokenRepository, SQLAlchemyOAuthServerRepository from repositories.recommended_app_catalog_repository import DatabaseRecommendedAppCatalogRepository @@ -159,6 +160,7 @@ from services.file_service import FileService from services.human_input_file_upload_service import HumanInputFileUploadService from services.init_validation_service import InitValidationService from services.inner_mail_service import InnerMailService +from services.message_file_preview_service import MessageFilePreviewService from services.notification_gateway import BillingNotificationGateway from services.notification_service import NotificationService from services.notion_data_source_gateway import NotionDataSourceGateway @@ -263,6 +265,7 @@ class ApplicationServices: file_grants: FileGrantService files: FileService human_input_file_uploads: HumanInputFileUploadService + message_file_previews: MessageFilePreviewService oauth_server: OAuthServerService init_validation: InitValidationService notifications: NotificationService @@ -662,6 +665,10 @@ def build_application_services( files=file_service, remote_files=remote_file_service, ), + message_file_previews=MessageFilePreviewService( + files=MessageFilePreviewQueryRepository(session_factory=database_client), + storage=storage, + ), oauth_server=_build_oauth_server_service(database_client=database_client, redis=redis), init_validation=InitValidationService( state=installation_state, diff --git a/api/repositories/message_file_preview_repository.py b/api/repositories/message_file_preview_repository.py new file mode 100644 index 00000000000..ad0aa68e219 --- /dev/null +++ b/api/repositories/message_file_preview_repository.py @@ -0,0 +1,66 @@ +"""SQLAlchemy query adapter for app-scoped message file previews.""" + +from typing import override + +from sqlalchemy import case, select +from sqlalchemy.orm import Session, sessionmaker + +from models.model import Message, MessageFile, UploadFile +from services.message_file_preview_service import ( + MessageFilePreviewAccessDeniedError, + MessageFilePreviewNotFoundError, + MessageFilePreviewQuery, + MessageFilePreviewRecord, +) + + +class MessageFilePreviewQueryRepository(MessageFilePreviewQuery): + def __init__(self, *, session_factory: sessionmaker[Session]) -> None: + self._session_factory = session_factory + + @override + def get_for_app( + self, + *, + file_id: str, + app_id: str, + tenant_id: str, + ) -> MessageFilePreviewRecord: + stmt = ( + select( + Message.app_id.label("message_app_id"), + UploadFile.tenant_id.label("file_tenant_id"), + UploadFile.key.label("file_key"), + UploadFile.name.label("file_name"), + UploadFile.size.label("file_size"), + UploadFile.extension.label("file_extension"), + UploadFile.mime_type.label("file_mime_type"), + ) + .select_from(MessageFile) + .outerjoin(Message, Message.id == MessageFile.message_id) + .outerjoin(UploadFile, UploadFile.id == MessageFile.upload_file_id) + .where(MessageFile.upload_file_id == file_id) + # One upload may be linked to multiple messages; prefer a reference owned by this app. + .order_by(case((Message.app_id == app_id, 0), else_=1), MessageFile.id) + .limit(1) + ) + + with self._session_factory() as session: + row = session.execute(stmt).one_or_none() + + if row is None: + raise MessageFilePreviewNotFoundError + if row.message_app_id != app_id: + raise MessageFilePreviewAccessDeniedError + if row.file_key is None: + raise MessageFilePreviewNotFoundError + if row.file_tenant_id != tenant_id: + raise MessageFilePreviewAccessDeniedError + + return MessageFilePreviewRecord( + key=row.file_key, + name=row.file_name, + size=row.file_size, + extension=row.file_extension, + mime_type=row.file_mime_type, + ) diff --git a/api/services/message_file_preview_service.py b/api/services/message_file_preview_service.py new file mode 100644 index 00000000000..ebebc0f312d --- /dev/null +++ b/api/services/message_file_preview_service.py @@ -0,0 +1,73 @@ +"""Application service for previewing files attached to app messages.""" + +from collections.abc import Iterator +from typing import NamedTuple, Protocol + + +class MessageFilePreviewNotFoundError(Exception): + """Raised when a requested file is not attached to a message.""" + + +class MessageFilePreviewAccessDeniedError(Exception): + """Raised when a requested file is outside the authenticated app scope.""" + + +class MessageFilePreviewRecord(NamedTuple): + key: str + name: str + size: int + extension: str + mime_type: str | None + + +class MessageFilePreviewQuery(Protocol): + def get_for_app( + self, + *, + file_id: str, + app_id: str, + tenant_id: str, + ) -> MessageFilePreviewRecord: + """Return preview metadata after enforcing the file ownership chain. + + The upload must be referenced by a MessageFile, its Message must belong + to app_id, and the UploadFile must exist under tenant_id. + """ + ... + + +class FileStreamStorage(Protocol): + def load_stream(self, filename: str) -> Iterator[bytes]: ... + + +class MessageFilePreview(NamedTuple): + content: Iterator[bytes] + file: MessageFilePreviewRecord + + +class MessageFilePreviewService: + def __init__( + self, + *, + files: MessageFilePreviewQuery, + storage: FileStreamStorage, + ) -> None: + self._files = files + self._storage = storage + + def get_preview( + self, + *, + file_id: str, + app_id: str, + tenant_id: str, + ) -> MessageFilePreview: + file = self._files.get_for_app( + file_id=file_id, + app_id=app_id, + tenant_id=tenant_id, + ) + return MessageFilePreview( + content=self._storage.load_stream(file.key), + file=file, + ) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_file_preview.py b/api/tests/unit_tests/controllers/service_api/app/test_file_preview.py index 33b6e83094e..f6e7e9cd1d9 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_file_preview.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_file_preview.py @@ -1,285 +1,151 @@ -"""Unit tests for the Service API file-preview endpoint. +"""Unit tests for the Service API file-preview transport boundary.""" -Ownership checks run against persisted message, file, app, and upload rows so the -tests exercise the same SQLAlchemy statements and tenant boundary as production. -Storage remains mocked because it is the external I/O boundary of the endpoint. -""" - -import logging -from collections.abc import Iterator -from dataclasses import dataclass -from datetime import datetime -from decimal import Decimal -from typing import Protocol, cast +from inspect import unwrap from unittest.mock import Mock, patch -from uuid import uuid4 +from uuid import UUID, uuid4 import pytest -from sqlalchemy import event -from sqlalchemy.engine import Engine -from sqlalchemy.orm import Session +from flask import Response from controllers.service_api.app.error import FileAccessDeniedError, FileNotFoundError -from controllers.service_api.app.file_preview import FilePreviewApi -from extensions.storage.storage_type import StorageType -from graphon.file import FileTransferMethod, FileType -from models.base import TypeBase -from models.enums import ConversationFromSource, CreatorUserRole -from models.model import App, AppMode, Message, MessageFile, UploadFile +from controllers.service_api.app.file_preview import FilePreviewApi, FilePreviewQuery +from models.model import App, EndUser +from services.message_file_preview_service import ( + MessageFilePreview, + MessageFilePreviewAccessDeniedError, + MessageFilePreviewNotFoundError, + MessageFilePreviewRecord, +) -class _FilePreviewLogRecord(Protocol): - file_id: str - app_id: str - error: str - - -@dataclass(frozen=True) -class _Database: - """Expose the real test session through the interface used by the controller.""" - - session: Session - - -@dataclass(frozen=True) -class _PreviewRecords: - app: App - message: Message - message_file: MessageFile - upload_file: UploadFile - - -@pytest.fixture -def database(sqlite_engine: Engine) -> Iterator[_Database]: - """Create only the tables required by file ownership validation.""" - - models = (App, Message, MessageFile, UploadFile) - tables = [TypeBase.metadata.tables[model.__tablename__] for model in models] - TypeBase.metadata.create_all(sqlite_engine, tables=tables) - with Session(sqlite_engine, expire_on_commit=False) as session: - yield _Database(session) - - -@pytest.fixture -def file_preview_api() -> FilePreviewApi: - """Create the resource instance under test.""" - - return FilePreviewApi() - - -def _upload_file(*, tenant_id: str, file_id: str | None = None) -> UploadFile: - upload_file = UploadFile( - tenant_id=tenant_id, - storage_type=StorageType.LOCAL, - key="storage/key/test_file.jpg", - name="test_file.jpg", - size=1024, - extension="jpg", - mime_type="image/jpeg", - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - created_at=datetime(2026, 1, 1), - used=True, - ) - if file_id is not None: - upload_file.id = file_id - return upload_file - - -def _persist_preview_records( - session: Session, +def _preview( *, - app_id: str | None = None, - app_tenant_id: str | None = None, - upload_tenant_id: str | None = None, -) -> _PreviewRecords: - app_id = app_id or str(uuid4()) - app_tenant_id = app_tenant_id or str(uuid4()) - upload_file = _upload_file(tenant_id=upload_tenant_id or app_tenant_id) - app = App( - id=app_id, - tenant_id=app_tenant_id, - name="Preview app", - description="", - mode=AppMode.CHAT, - icon_type=None, - icon="", - icon_background=None, - enable_site=True, - enable_api=True, + mime_type: str | None = "image/jpeg", + name: str = "test_file.jpg", + extension: str = "jpg", + size: int = 1024, +) -> MessageFilePreview: + return MessageFilePreview( + content=iter([b"file content"]), + file=MessageFilePreviewRecord( + key="storage/key/test_file.jpg", + name=name, + size=size, + extension=extension, + mime_type=mime_type, + ), ) - message = Message( - id=str(uuid4()), - app_id=app_id, - conversation_id=str(uuid4()), - _inputs={}, - query="preview", - message={}, - message_unit_price=Decimal(0), - answer="answer", - answer_unit_price=Decimal(0), - currency="USD", - from_source=ConversationFromSource.API, + + +def _app() -> App: + return App(id=str(uuid4()), tenant_id=str(uuid4())) + + +def _get( + *, + api: FilePreviewApi, + args: FilePreviewQuery, + app_model: App, + file_id: UUID, +) -> Response: + return unwrap(api.get)( + api, + args=args, + app_model=app_model, + end_user=Mock(spec=EndUser), + file_id=file_id, ) - message_file = MessageFile( - message_id=message.id, - type=FileType.IMAGE, - transfer_method=FileTransferMethod.LOCAL_FILE, - created_by_role=CreatorUserRole.ACCOUNT, - created_by=str(uuid4()), - upload_file_id=upload_file.id, - ) - session.add_all([app, message, message_file, upload_file]) - session.commit() - return _PreviewRecords(app=app, message=message, message_file=message_file, upload_file=upload_file) class TestFilePreviewApi: - """Exercise ownership validation and response construction.""" + @patch("controllers.service_api.app.file_preview.application_services") + def test_get_uses_authenticated_app_scope(self, mock_application_services: Mock) -> None: + app_model = _app() + file_id = uuid4() + preview = _preview() + service = mock_application_services.return_value.message_file_previews + service.get_preview.return_value = preview - def test_validate_file_ownership_success(self, file_preview_api: FilePreviewApi, database: _Database): - records = _persist_preview_records(database.session) + response = _get( + api=FilePreviewApi(), + args=FilePreviewQuery(), + app_model=app_model, + file_id=file_id, + ) - with patch("controllers.service_api.app.file_preview.db", database): - message_file, upload_file = file_preview_api._validate_file_ownership( - records.upload_file.id, records.app.id - ) - - assert message_file.id == records.message_file.id - assert upload_file.id == records.upload_file.id - assert upload_file.tenant_id == records.app.tenant_id - - def test_validate_file_ownership_file_not_found(self, file_preview_api: FilePreviewApi, database: _Database): - with patch("controllers.service_api.app.file_preview.db", database): - with pytest.raises(FileNotFoundError, match="File not found in message context"): - file_preview_api._validate_file_ownership(str(uuid4()), str(uuid4())) - - def test_validate_file_ownership_access_denied(self, file_preview_api: FilePreviewApi, database: _Database): - records = _persist_preview_records(database.session) - - with patch("controllers.service_api.app.file_preview.db", database): - with pytest.raises(FileAccessDeniedError, match="not owned by requesting app"): - file_preview_api._validate_file_ownership(records.upload_file.id, str(uuid4())) - - def test_validate_file_ownership_upload_file_not_found(self, file_preview_api: FilePreviewApi, database: _Database): - records = _persist_preview_records(database.session) - database.session.delete(records.upload_file) - database.session.commit() - - with patch("controllers.service_api.app.file_preview.db", database): - with pytest.raises(FileNotFoundError, match="Upload file record not found"): - file_preview_api._validate_file_ownership(records.upload_file.id, records.app.id) - - def test_validate_file_ownership_tenant_mismatch(self, file_preview_api: FilePreviewApi, database: _Database): - records = _persist_preview_records(database.session, upload_tenant_id=str(uuid4())) - - with patch("controllers.service_api.app.file_preview.db", database): - with pytest.raises(FileAccessDeniedError, match="tenant mismatch"): - file_preview_api._validate_file_ownership(records.upload_file.id, records.app.id) - - def test_validate_file_ownership_invalid_input(self, file_preview_api: FilePreviewApi): - with pytest.raises(FileAccessDeniedError, match="Invalid file or app identifier"): - file_preview_api._validate_file_ownership("", "app_id") - - with pytest.raises(FileAccessDeniedError, match="Invalid file or app identifier"): - file_preview_api._validate_file_ownership("file_id", "") + service.get_preview.assert_called_once_with( + file_id=str(file_id), + app_id=app_model.id, + tenant_id=app_model.tenant_id, + ) + assert response.response is preview.content + assert response.mimetype == "image/jpeg" @pytest.mark.parametrize( - ("as_attachment", "mime_type", "name", "extension", "size"), + ("service_error", "http_error"), [ - (False, "image/jpeg", "test_file.jpg", "jpg", 1024), - (True, "image/jpeg", "test_file.jpg", "jpg", 1024), - (False, "text/html", "unsafe.html", "html", 1024), - (False, "video/mp4", "test_file.mp4", "mp4", 1024), - (False, "image/jpeg", "test_file.jpg", "jpg", 0), + (MessageFilePreviewNotFoundError(), FileNotFoundError), + (MessageFilePreviewAccessDeniedError(), FileAccessDeniedError), ], ) - def test_build_file_response( + @patch("controllers.service_api.app.file_preview.application_services") + def test_get_maps_expected_errors( self, - file_preview_api: FilePreviewApi, - as_attachment: bool, - mime_type: str, - name: str, - extension: str, - size: int, - ): - upload_file = _upload_file(tenant_id=str(uuid4())) - upload_file.mime_type = mime_type - upload_file.name = name - upload_file.extension = extension - upload_file.size = size + mock_application_services: Mock, + service_error: Exception, + http_error: type[Exception], + ) -> None: + mock_application_services.return_value.message_file_previews.get_preview.side_effect = service_error - response = file_preview_api._build_file_response(Mock(), upload_file, as_attachment) + with pytest.raises(http_error): + _get( + api=FilePreviewApi(), + args=FilePreviewQuery(), + app_model=_app(), + file_id=uuid4(), + ) + + @patch("controllers.service_api.app.file_preview.application_services") + def test_get_does_not_mask_storage_errors(self, mock_application_services: Mock) -> None: + mock_application_services.return_value.message_file_previews.get_preview.side_effect = OSError("storage down") + + with pytest.raises(OSError, match="storage down"): + _get( + api=FilePreviewApi(), + args=FilePreviewQuery(), + app_model=_app(), + file_id=uuid4(), + ) + + @pytest.mark.parametrize( + ("preview", "as_attachment"), + [ + (_preview(), False), + (_preview(name="报告 1.jpg"), True), + (_preview(mime_type="text/html", name="unsafe.html", extension="html"), False), + (_preview(mime_type="video/mp4", name="video.mp4", extension="mp4"), False), + (_preview(size=0), False), + ], + ) + def test_build_file_response(self, preview: MessageFilePreview, as_attachment: bool) -> None: + response = FilePreviewApi()._build_file_response( + preview=preview, + as_attachment=as_attachment, + ) assert response.direct_passthrough is True - assert "Cache-Control" in response.headers - assert ("Content-Length" in response.headers) is bool(size) - if as_attachment or mime_type == "text/html": - assert "attachment" in response.headers["Content-Disposition"] - assert response.headers["Content-Type"] == "application/octet-stream" - else: - assert response.mimetype == mime_type - if mime_type == "text/html": - assert response.headers["X-Content-Type-Options"] == "nosniff" - if mime_type.startswith("video/"): - assert response.headers["Accept-Ranges"] == "bytes" - - @patch("controllers.service_api.app.file_preview.storage") - def test_components_use_validated_file( - self, mock_storage: Mock, file_preview_api: FilePreviewApi, database: _Database - ): - records = _persist_preview_records(database.session) - generator = Mock() - - with patch("controllers.service_api.app.file_preview.db", database): - message_file, upload_file = file_preview_api._validate_file_ownership( - records.upload_file.id, records.app.id + assert response.headers["Cache-Control"] == "public, max-age=3600" + assert ("Content-Length" in response.headers) is (preview.file.size > 0) + if as_attachment: + assert response.headers["Content-Disposition"] == ( + "attachment; filename*=UTF-8''%E6%8A%A5%E5%91%8A%201.jpg" ) - response = file_preview_api._build_file_response(generator, upload_file, False) - - assert message_file.id == records.message_file.id - assert response.mimetype == "image/jpeg" - mock_storage.load.assert_not_called() - - @patch("controllers.service_api.app.file_preview.storage") - def test_storage_error_remains_external( - self, mock_storage: Mock, file_preview_api: FilePreviewApi, database: _Database - ): - records = _persist_preview_records(database.session) - mock_storage.load.side_effect = OSError("Storage error") - - with patch("controllers.service_api.app.file_preview.db", database): - _, upload_file = file_preview_api._validate_file_ownership(records.upload_file.id, records.app.id) - - with pytest.raises(OSError, match="Storage error"): - mock_storage.load(upload_file.key, stream=True) - - def test_validate_file_ownership_unexpected_error_logging( - self, - file_preview_api: FilePreviewApi, - database: _Database, - sqlite_engine: Engine, - caplog: pytest.LogCaptureFixture, - ): - file_id = str(uuid4()) - app_id = str(uuid4()) - - def fail_statement(*_args: object) -> None: - raise RuntimeError("Unexpected database error") - - event.listen(sqlite_engine, "before_cursor_execute", fail_statement) - try: - with patch("controllers.service_api.app.file_preview.db", database): - with caplog.at_level(logging.ERROR, logger="controllers.service_api.app.file_preview"): - with pytest.raises(FileAccessDeniedError, match="File access validation failed"): - file_preview_api._validate_file_ownership(file_id, app_id) - finally: - event.remove(sqlite_engine, "before_cursor_execute", fail_statement) - - assert len(caplog.records) == 1 - log_record = caplog.records[0] - assert log_record.getMessage() == "Unexpected error during file ownership validation" - record = cast(_FilePreviewLogRecord, log_record) - assert record.file_id == file_id - assert record.app_id == app_id - assert record.error == "Unexpected database error" + assert response.headers["Content-Type"] == "application/octet-stream" + elif preview.file.mime_type == "text/html": + assert response.headers["Content-Disposition"] == "attachment; filename*=UTF-8''unsafe.html" + assert response.headers["Content-Type"] == "application/octet-stream" + assert response.headers["X-Content-Type-Options"] == "nosniff" + else: + assert response.mimetype == preview.file.mime_type + if preview.file.mime_type == "video/mp4": + assert response.headers["Accept-Ranges"] == "bytes" diff --git a/api/tests/unit_tests/extensions/test_ext_application_services.py b/api/tests/unit_tests/extensions/test_ext_application_services.py index f0ba570a48c..66bcd463e02 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -32,6 +32,7 @@ from repositories.app_site_command_repository import AppSiteCommandRepository from repositories.app_statistic_query_repository import AppStatisticQueryRepository from repositories.app_tracing_config_repository import SQLAlchemyAppTracingConfigRepository from repositories.human_input_file_upload_repository import SQLAlchemyHumanInputFileUploadRepository +from repositories.message_file_preview_repository import MessageFilePreviewQueryRepository from repositories.sqlalchemy_api_workflow_run_repository import DifyAPISQLAlchemyWorkflowRunRepository from repositories.workflow_app_log_query_repository import WorkflowAppLogQueryRepository from repositories.workflow_run_archive_repository import WorkflowRunArchiveBundleQueryRepository @@ -71,6 +72,7 @@ from services.errors.enterprise import EnterpriseAPIError, EnterpriseAPINotFound from services.file_service import FileService from services.human_input_file_upload_service import HumanInputFileUploadService from services.init_validation_service import InvalidInitializationPasswordError +from services.message_file_preview_service import MessageFilePreviewService from services.partner_tenant_binding_service import PartnerTenantBindingService from services.retention.workflow_run.archive_download_task_cache import WorkflowRunArchiveDownloadTaskCache from services.retention.workflow_run.archive_log_service import WorkflowRunArchiveService @@ -234,6 +236,22 @@ def test_build_application_services_reuses_file_service( assert services.web_app_runtime._file_service is services.files +def test_build_application_services_wires_message_file_previews( + sqlite_session_factory: sessionmaker[Session], +) -> None: + services = ext_application_services.build_application_services( + database_client=sqlite_session_factory, + deployment_edition=DeploymentEdition.COMMUNITY, + initialization_password="", + redis=MagicMock(spec=RedisClientWrapper), + ) + + assert isinstance(services.message_file_previews, MessageFilePreviewService) + assert isinstance(services.message_file_previews._files, MessageFilePreviewQueryRepository) + assert services.message_file_previews._files._session_factory is sqlite_session_factory + assert services.message_file_previews._storage is ext_application_services.storage + + def test_build_application_services_wires_workflow_run_archives( sqlite_session_factory: sessionmaker[Session], ) -> None: diff --git a/api/tests/unit_tests/repositories/test_message_file_preview_repository.py b/api/tests/unit_tests/repositories/test_message_file_preview_repository.py new file mode 100644 index 00000000000..1a098d962ae --- /dev/null +++ b/api/tests/unit_tests/repositories/test_message_file_preview_repository.py @@ -0,0 +1,180 @@ +from datetime import datetime +from decimal import Decimal +from uuid import uuid4 + +import pytest +from sqlalchemy.engine import Engine +from sqlalchemy.orm import Session, sessionmaker + +from extensions.storage.storage_type import StorageType +from graphon.file import FileTransferMethod, FileType +from models.base import TypeBase +from models.enums import ConversationFromSource, CreatorUserRole +from models.model import Message, MessageFile, UploadFile +from repositories.message_file_preview_repository import MessageFilePreviewQueryRepository +from services.message_file_preview_service import ( + MessageFilePreviewAccessDeniedError, + MessageFilePreviewNotFoundError, + MessageFilePreviewRecord, +) + + +@pytest.fixture +def repository( + sqlite_engine: Engine, + sqlite_session_factory: sessionmaker[Session], +) -> MessageFilePreviewQueryRepository: + models = (Message, MessageFile, UploadFile) + tables = [TypeBase.metadata.tables[model.__tablename__] for model in models] + TypeBase.metadata.create_all(sqlite_engine, tables=tables) + return MessageFilePreviewQueryRepository(session_factory=sqlite_session_factory) + + +def _upload_file(*, file_id: str, tenant_id: str) -> UploadFile: + upload_file = UploadFile( + tenant_id=tenant_id, + storage_type=StorageType.LOCAL, + key="upload_files/tenant/file.pdf", + name="file.pdf", + size=42, + extension="pdf", + mime_type="application/pdf", + created_by_role=CreatorUserRole.END_USER, + created_by=str(uuid4()), + created_at=datetime(2026, 1, 1), + used=True, + ) + upload_file.id = file_id + return upload_file + + +def _message(*, app_id: str) -> Message: + return Message( + id=str(uuid4()), + app_id=app_id, + conversation_id=str(uuid4()), + _inputs={}, + query="preview", + message={}, + message_unit_price=Decimal(0), + answer="answer", + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.API, + ) + + +def _message_file(*, message_id: str, file_id: str) -> MessageFile: + return MessageFile( + message_id=message_id, + type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.LOCAL_FILE, + created_by_role=CreatorUserRole.END_USER, + created_by=str(uuid4()), + upload_file_id=file_id, + ) + + +def _persist_message_file( + session: Session, + *, + app_id: str, + file_id: str, +) -> MessageFile: + message = _message(app_id=app_id) + message_file = _message_file(message_id=message.id, file_id=file_id) + session.add_all([message, message_file]) + return message_file + + +def test_get_for_app_returns_detached_file_metadata( + repository: MessageFilePreviewQueryRepository, + sqlite_session: Session, +) -> None: + file_id = str(uuid4()) + app_id = str(uuid4()) + tenant_id = str(uuid4()) + sqlite_session.add(_upload_file(file_id=file_id, tenant_id=tenant_id)) + _persist_message_file(sqlite_session, app_id=app_id, file_id=file_id) + sqlite_session.commit() + + result = repository.get_for_app(file_id=file_id, app_id=app_id, tenant_id=tenant_id) + + assert result == MessageFilePreviewRecord( + key="upload_files/tenant/file.pdf", + name="file.pdf", + size=42, + extension="pdf", + mime_type="application/pdf", + ) + + +def test_get_for_app_rejects_file_without_message_reference( + repository: MessageFilePreviewQueryRepository, + sqlite_session: Session, +) -> None: + file_id = str(uuid4()) + tenant_id = str(uuid4()) + sqlite_session.add(_upload_file(file_id=file_id, tenant_id=tenant_id)) + sqlite_session.commit() + + with pytest.raises(MessageFilePreviewNotFoundError): + repository.get_for_app(file_id=file_id, app_id=str(uuid4()), tenant_id=tenant_id) + + +def test_get_for_app_enforces_app_isolation( + repository: MessageFilePreviewQueryRepository, + sqlite_session: Session, +) -> None: + file_id = str(uuid4()) + tenant_id = str(uuid4()) + sqlite_session.add(_upload_file(file_id=file_id, tenant_id=tenant_id)) + _persist_message_file(sqlite_session, app_id=str(uuid4()), file_id=file_id) + sqlite_session.commit() + + with pytest.raises(MessageFilePreviewAccessDeniedError): + repository.get_for_app(file_id=file_id, app_id=str(uuid4()), tenant_id=tenant_id) + + +def test_get_for_app_rejects_missing_upload_file( + repository: MessageFilePreviewQueryRepository, + sqlite_session: Session, +) -> None: + file_id = str(uuid4()) + app_id = str(uuid4()) + _persist_message_file(sqlite_session, app_id=app_id, file_id=file_id) + sqlite_session.commit() + + with pytest.raises(MessageFilePreviewNotFoundError): + repository.get_for_app(file_id=file_id, app_id=app_id, tenant_id=str(uuid4())) + + +def test_get_for_app_enforces_tenant_isolation( + repository: MessageFilePreviewQueryRepository, + sqlite_session: Session, +) -> None: + file_id = str(uuid4()) + app_id = str(uuid4()) + sqlite_session.add(_upload_file(file_id=file_id, tenant_id=str(uuid4()))) + _persist_message_file(sqlite_session, app_id=app_id, file_id=file_id) + sqlite_session.commit() + + with pytest.raises(MessageFilePreviewAccessDeniedError): + repository.get_for_app(file_id=file_id, app_id=app_id, tenant_id=str(uuid4())) + + +def test_get_for_app_accepts_any_message_reference_owned_by_the_app( + repository: MessageFilePreviewQueryRepository, + sqlite_session: Session, +) -> None: + file_id = str(uuid4()) + app_id = str(uuid4()) + tenant_id = str(uuid4()) + sqlite_session.add(_upload_file(file_id=file_id, tenant_id=tenant_id)) + _persist_message_file(sqlite_session, app_id=str(uuid4()), file_id=file_id) + _persist_message_file(sqlite_session, app_id=app_id, file_id=file_id) + sqlite_session.commit() + + result = repository.get_for_app(file_id=file_id, app_id=app_id, tenant_id=tenant_id) + + assert result.key == "upload_files/tenant/file.pdf" diff --git a/api/tests/unit_tests/services/test_message_file_preview_service.py b/api/tests/unit_tests/services/test_message_file_preview_service.py new file mode 100644 index 00000000000..17599342a05 --- /dev/null +++ b/api/tests/unit_tests/services/test_message_file_preview_service.py @@ -0,0 +1,82 @@ +from unittest.mock import Mock + +import pytest + +from services.message_file_preview_service import ( + FileStreamStorage, + MessageFilePreviewAccessDeniedError, + MessageFilePreviewNotFoundError, + MessageFilePreviewQuery, + MessageFilePreviewRecord, + MessageFilePreviewService, +) + + +def _record() -> MessageFilePreviewRecord: + return MessageFilePreviewRecord( + key="upload_files/tenant/file.pdf", + name="file.pdf", + size=42, + extension="pdf", + mime_type="application/pdf", + ) + + +def test_get_preview_loads_the_authorized_file_stream() -> None: + files = Mock(spec=MessageFilePreviewQuery) + storage = Mock(spec=FileStreamStorage) + file = _record() + content = iter([b"content"]) + files.get_for_app.return_value = file + storage.load_stream.return_value = content + service = MessageFilePreviewService(files=files, storage=storage) + + preview = service.get_preview( + file_id="file-id", + app_id="app-id", + tenant_id="tenant-id", + ) + + files.get_for_app.assert_called_once_with( + file_id="file-id", + app_id="app-id", + tenant_id="tenant-id", + ) + storage.load_stream.assert_called_once_with(file.key) + assert preview.file is file + assert preview.content is content + + +@pytest.mark.parametrize( + "error", + [MessageFilePreviewNotFoundError(), MessageFilePreviewAccessDeniedError()], +) +def test_get_preview_does_not_load_denied_or_missing_files(error: Exception) -> None: + files = Mock(spec=MessageFilePreviewQuery) + storage = Mock(spec=FileStreamStorage) + files.get_for_app.side_effect = error + service = MessageFilePreviewService(files=files, storage=storage) + + with pytest.raises(type(error)): + service.get_preview( + file_id="file-id", + app_id="app-id", + tenant_id="tenant-id", + ) + + storage.load_stream.assert_not_called() + + +def test_get_preview_does_not_mask_storage_errors() -> None: + files = Mock(spec=MessageFilePreviewQuery) + storage = Mock(spec=FileStreamStorage) + files.get_for_app.return_value = _record() + storage.load_stream.side_effect = OSError("storage down") + service = MessageFilePreviewService(files=files, storage=storage) + + with pytest.raises(OSError, match="storage down"): + service.get_preview( + file_id="file-id", + app_id="app-id", + tenant_id="tenant-id", + )