diff --git a/api/.importlinter b/api/.importlinter index 5f0c8ea0915..0c91b3c82f8 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -515,6 +515,20 @@ forbidden_modules = sqlalchemy werkzeug +[importlinter:contract:tool-file-download-service-boundary] +name = Tool file download application service is framework and persistence neutral +type = forbidden +source_modules = + services.tool_file_download_service +forbidden_modules = + controllers + extensions + flask + models + repositories + sqlalchemy + werkzeug + [importlinter:contract:recommended-app-query-service-boundary] name = Recommended app query application service is framework and persistence neutral type = forbidden diff --git a/api/controllers/files/tool_files.py b/api/controllers/files/tool_files.py index 03f635627de..ce20480bbd6 100644 --- a/api/controllers/files/tool_files.py +++ b/api/controllers/files/tool_files.py @@ -6,12 +6,14 @@ from flask_restx import Resource from pydantic import BaseModel, Field from werkzeug.exceptions import Forbidden, NotFound -from controllers.common.errors import UnsupportedFileTypeError from controllers.common.file_response import enforce_download_for_html -from controllers.common.schema import register_schema_models +from controllers.common.schema import query_params_from_model, register_schema_models from controllers.files import files_ns -from core.tools.signature import verify_tool_file_signature -from core.tools.tool_file_manager import ToolFileManager +from extensions.ext_application_services import application_services +from services.tool_file_download_service import ( + ToolFileDownloadAccessDeniedError, + ToolFileDownloadNotFoundError, +) class ToolFileQuery(BaseModel): @@ -32,10 +34,7 @@ class ToolFileApi(Resource): params={ "file_id": "Tool file identifier", "extension": "Expected file extension", - "timestamp": "Unix timestamp used in the signature", - "nonce": "Random string used in the signature", - "sign": "HMAC signature verifying the request", - "as_attachment": "Whether to download the file as an attachment", + **query_params_from_model(ToolFileQuery), } ) @files_ns.doc( @@ -43,52 +42,38 @@ class ToolFileApi(Resource): 200: "Tool file stream returned successfully", 403: "Forbidden - invalid signature", 404: "File not found", - 415: "Unsupported file type", } ) - def get(self, file_id: UUID, extension: str): - file_id_str = str(file_id) - - args = ToolFileQuery.model_validate(request.args.to_dict()) - if not verify_tool_file_signature( - file_id=file_id_str, timestamp=args.timestamp, nonce=args.nonce, sign=args.sign - ): - raise Forbidden("Invalid request.") - + def get(self, file_id: UUID, extension: str) -> Response: + args = ToolFileQuery.model_validate(request.args.to_dict(flat=True)) try: - tool_file_manager = ToolFileManager() - stream, tool_file = tool_file_manager.get_file_generator_by_tool_file_id( - file_id_str, + download = application_services().tool_file_downloads.get_signed_file( + file_id=str(file_id), + timestamp=args.timestamp, + nonce=args.nonce, + sign=args.sign, ) - - if not stream or not tool_file: - raise NotFound("file is not found") - - except NotFound: - raise - - except Exception as e: - raise UnsupportedFileTypeError() from e - - mime_type = tool_file.mime_type - filename = tool_file.filename + except ToolFileDownloadAccessDeniedError as error: + raise Forbidden("Invalid request.") from error + except ToolFileDownloadNotFoundError as error: + raise NotFound("file is not found") from error response = Response( - stream, - mimetype=mime_type, + download.content, + mimetype=download.mime_type, direct_passthrough=True, headers={}, ) - if tool_file.size > 0: - response.headers["Content-Length"] = str(tool_file.size) - if args.as_attachment and filename: - encoded_filename = quote(filename) + if download.size > 0: + response.headers["Content-Length"] = str(download.size) + if args.as_attachment and download.filename: + encoded_filename = quote(download.filename) response.headers["Content-Disposition"] = f"attachment; filename*=UTF-8''{encoded_filename}" enforce_download_for_html( response, - mime_type=mime_type, - filename=filename, + mime_type=download.mime_type, + filename=download.filename, extension=extension, ) diff --git a/api/extensions/ext_application_services.py b/api/extensions/ext_application_services.py index 338ed8b0a48..666313a2a30 100644 --- a/api/extensions/ext_application_services.py +++ b/api/extensions/ext_application_services.py @@ -186,6 +186,7 @@ from services.setup_service import SetupService from services.step_by_step_tour_service import StepByStepTourService from services.system_feature_service import SystemFeatureService from services.tag_application_service import TagApplicationService +from services.tool_file_download_service import ToolFileDownloadService from services.trial_app_usage import TrialAppUsageRecorder from services.upload_file_delivery_service import UploadFileDeliveryService from services.web_app_runtime_query_service import WebAppRuntimeQueryService @@ -268,6 +269,7 @@ class ApplicationServices: files: FileService human_input_file_uploads: HumanInputFileUploadService message_file_previews: MessageFilePreviewService + tool_file_downloads: ToolFileDownloadService upload_file_delivery: UploadFileDeliveryService oauth_server: OAuthServerService init_validation: InitValidationService @@ -672,6 +674,7 @@ def build_application_services( files=MessageFilePreviewQueryRepository(session_factory=database_client), storage=storage, ), + tool_file_downloads=ToolFileDownloadService(tool_files=ToolFileManager()), upload_file_delivery=UploadFileDeliveryService( files=UploadFileDeliveryQueryRepository(session_factory=database_client), storage=storage, diff --git a/api/services/tool_file_download_service.py b/api/services/tool_file_download_service.py new file mode 100644 index 00000000000..bfee566abde --- /dev/null +++ b/api/services/tool_file_download_service.py @@ -0,0 +1,61 @@ +"""Application service for signed ToolFile downloads.""" + +from collections.abc import Iterator +from typing import NamedTuple, Protocol + +from core.tools.signature import verify_tool_file_signature +from graphon.file import File + + +class ToolFileDownloadAccessDeniedError(PermissionError): + pass + + +class ToolFileDownloadNotFoundError(LookupError): + pass + + +class ToolFileDownloadSource(Protocol): + def get_file_generator_by_tool_file_id( + self, + tool_file_id: str, + ) -> tuple[Iterator[bytes] | None, File | None]: ... + + +class ToolFileDownload(NamedTuple): + content: Iterator[bytes] + mime_type: str | None + filename: str | None + size: int + + +class ToolFileDownloadService: + def __init__(self, *, tool_files: ToolFileDownloadSource) -> None: + self._tool_files = tool_files + + def get_signed_file( + self, + *, + file_id: str, + timestamp: str, + nonce: str, + sign: str, + ) -> ToolFileDownload: + if not verify_tool_file_signature( + file_id=file_id, + timestamp=timestamp, + nonce=nonce, + sign=sign, + ): + raise ToolFileDownloadAccessDeniedError + + content, file = self._tool_files.get_file_generator_by_tool_file_id(tool_file_id=file_id) + if content is None or file is None: + raise ToolFileDownloadNotFoundError + + return ToolFileDownload( + content=content, + mime_type=file.mime_type, + filename=file.filename, + size=file.size, + ) diff --git a/api/tests/unit_tests/controllers/files/test_tool_files.py b/api/tests/unit_tests/controllers/files/test_tool_files.py index a41e2a056b1..3a7544dcbcf 100644 --- a/api/tests/unit_tests/controllers/files/test_tool_files.py +++ b/api/tests/unit_tests/controllers/files/test_tool_files.py @@ -1,168 +1,222 @@ import types +from collections.abc import Iterator from inspect import unwrap -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from werkzeug.exceptions import Forbidden, NotFound import controllers.files.tool_files as module +from services.tool_file_download_service import ( + ToolFileDownload, + ToolFileDownloadAccessDeniedError, + ToolFileDownloadNotFoundError, +) -def fake_request(args: dict): - return types.SimpleNamespace(args=types.SimpleNamespace(to_dict=lambda flat=True: args)) +def _fake_request(args: dict[str, object]) -> types.SimpleNamespace: + return types.SimpleNamespace(args=types.SimpleNamespace(to_dict=lambda **_kwargs: args)) -class DummyToolFile: - def __init__(self, mime_type="text/plain", size=10, filename="tool.txt"): - self.mime_type = mime_type - self.size = size - self.filename = filename +def _set_request(monkeypatch: pytest.MonkeyPatch, args: dict[str, object]) -> None: + monkeypatch.setattr(module, "request", _fake_request(args)) -@pytest.fixture(autouse=True) -def mock_global_db(): - fake_db = types.SimpleNamespace(engine=object()) - module.global_db = fake_db +def _download( + *, + content: Iterator[bytes] | None = None, + mime_type: str | None = "text/plain", + filename: str | None = "tool.txt", + size: int = 10, +) -> ToolFileDownload: + return ToolFileDownload( + content=content if content is not None else iter([b"data"]), + mime_type=mime_type, + filename=filename, + size=size, + ) class TestToolFileApi: - @patch.object(module, "verify_tool_file_signature", return_value=True) - @patch.object(module, "ToolFileManager") + @patch.object(module, "application_services") def test_success_stream( self, - mock_tool_file_manager, - mock_verify, - ): - module.request = fake_request( + mock_application_services: MagicMock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _set_request( + monkeypatch, { "timestamp": "123", "nonce": "abc", "sign": "sig", "as_attachment": False, - } + }, ) - stream = iter([b"data"]) - tool_file = DummyToolFile(size=100) + service = mock_application_services.return_value.tool_file_downloads + service.get_signed_file.return_value = _download(content=stream, size=100) - mock_tool_file_manager.return_value.get_file_generator_by_tool_file_id.return_value = ( - stream, - tool_file, - ) - - api = module.ToolFileApi() - get_fn = unwrap(api.get) - - response = get_fn("file-id", "txt") + response = unwrap(module.ToolFileApi().get)("file-id", "txt") + assert response.response is stream assert response.mimetype == "text/plain" assert response.headers["Content-Length"] == "100" - mock_verify.assert_called_once_with( + assert response.direct_passthrough is True + service.get_signed_file.assert_called_once_with( file_id="file-id", timestamp="123", nonce="abc", sign="sig", ) - @patch.object(module, "verify_tool_file_signature", return_value=True) - @patch.object(module, "ToolFileManager") - def test_as_attachment( + @patch.object(module, "application_services") + def test_zero_size_omits_content_length( self, - mock_tool_file_manager, - mock_verify, - ): - module.request = fake_request( + mock_application_services: MagicMock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _set_request( + monkeypatch, + { + "timestamp": "123", + "nonce": "abc", + "sign": "sig", + "as_attachment": False, + }, + ) + mock_application_services.return_value.tool_file_downloads.get_signed_file.return_value = _download(size=0) + + response = unwrap(module.ToolFileApi().get)("file-id", "txt") + + assert "Content-Length" not in response.headers + + @patch.object(module, "application_services") + def test_as_attachment_preserves_mime_type( + self, + mock_application_services: MagicMock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _set_request( + monkeypatch, { "timestamp": "123", "nonce": "abc", "sign": "sig", "as_attachment": True, - } + }, ) - - stream = iter([b"data"]) - tool_file = DummyToolFile( + mock_application_services.return_value.tool_file_downloads.get_signed_file.return_value = _download( mime_type="application/pdf", - filename="doc.pdf", + filename="报告.pdf", ) - mock_tool_file_manager.return_value.get_file_generator_by_tool_file_id.return_value = ( - stream, - tool_file, + response = unwrap(module.ToolFileApi().get)("file-id", "pdf") + + assert response.headers["Content-Disposition"] == "attachment; filename*=UTF-8''%E6%8A%A5%E5%91%8A.pdf" + assert response.headers["Content-Type"] == "application/pdf" + + @pytest.mark.parametrize( + ("mime_type", "filename", "route_extension"), + [ + pytest.param("text/html", "file.txt", "txt", id="mime-type"), + pytest.param("text/plain", "file.HTML", "txt", id="filename"), + pytest.param("text/plain", "file.txt", "html", id="route-extension"), + ], + ) + @patch.object(module, "application_services") + def test_html_forces_download( + self, + mock_application_services: MagicMock, + monkeypatch: pytest.MonkeyPatch, + mime_type: str, + filename: str, + route_extension: str, + ) -> None: + _set_request( + monkeypatch, + { + "timestamp": "123", + "nonce": "abc", + "sign": "sig", + "as_attachment": False, + }, + ) + mock_application_services.return_value.tool_file_downloads.get_signed_file.return_value = _download( + mime_type=mime_type, + filename=filename, ) - api = module.ToolFileApi() - get_fn = unwrap(api.get) - - response = get_fn("file-id", "pdf") + response = unwrap(module.ToolFileApi().get)("file-id", route_extension) assert response.headers["Content-Disposition"].startswith("attachment") - mock_verify.assert_called_once() + assert response.headers["Content-Type"] == "application/octet-stream" + assert response.headers["X-Content-Type-Options"] == "nosniff" - @patch.object(module, "verify_tool_file_signature", return_value=False) - def test_invalid_signature(self, mock_verify): - module.request = fake_request( + @patch.object(module, "application_services") + def test_invalid_signature( + self, + mock_application_services: MagicMock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _set_request( + monkeypatch, { "timestamp": "123", "nonce": "abc", "sign": "bad-sig", "as_attachment": False, - } + }, + ) + mock_application_services.return_value.tool_file_downloads.get_signed_file.side_effect = ( + ToolFileDownloadAccessDeniedError() ) - api = module.ToolFileApi() - get_fn = unwrap(api.get) + with pytest.raises(Forbidden, match=r"Invalid request\."): + unwrap(module.ToolFileApi().get)("file-id", "txt") - with pytest.raises(Forbidden): - get_fn("file-id", "txt") - - @patch.object(module, "verify_tool_file_signature", return_value=True) - @patch.object(module, "ToolFileManager") + @patch.object(module, "application_services") def test_file_not_found( self, - mock_tool_file_manager, - mock_verify, - ): - module.request = fake_request( + mock_application_services: MagicMock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + _set_request( + monkeypatch, { "timestamp": "123", "nonce": "abc", "sign": "sig", "as_attachment": False, - } + }, + ) + mock_application_services.return_value.tool_file_downloads.get_signed_file.side_effect = ( + ToolFileDownloadNotFoundError() ) - mock_tool_file_manager.return_value.get_file_generator_by_tool_file_id.return_value = ( - None, - None, - ) + with pytest.raises(NotFound, match="file is not found"): + unwrap(module.ToolFileApi().get)("file-id", "txt") - api = module.ToolFileApi() - get_fn = unwrap(api.get) - - with pytest.raises(NotFound): - get_fn("file-id", "txt") - - @patch.object(module, "verify_tool_file_signature", return_value=True) - @patch.object(module, "ToolFileManager") - def test_unsupported_file_type( + @pytest.mark.parametrize("service_error", [RuntimeError("database unavailable"), OSError("storage unavailable")]) + @patch.object(module, "application_services") + def test_unexpected_error_is_not_converted( self, - mock_tool_file_manager, - mock_verify, - ): - module.request = fake_request( + mock_application_services: MagicMock, + monkeypatch: pytest.MonkeyPatch, + service_error: Exception, + ) -> None: + _set_request( + monkeypatch, { "timestamp": "123", "nonce": "abc", "sign": "sig", "as_attachment": False, - } + }, ) + mock_application_services.return_value.tool_file_downloads.get_signed_file.side_effect = service_error - mock_tool_file_manager.return_value.get_file_generator_by_tool_file_id.side_effect = Exception("boom") + with pytest.raises(type(service_error)) as error_info: + unwrap(module.ToolFileApi().get)("file-id", "txt") - api = module.ToolFileApi() - get_fn = unwrap(api.get) - - with pytest.raises(module.UnsupportedFileTypeError): - get_fn("file-id", "txt") + assert error_info.value is service_error 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 10fb414d88e..c54de8aaec4 100644 --- a/api/tests/unit_tests/extensions/test_ext_application_services.py +++ b/api/tests/unit_tests/extensions/test_ext_application_services.py @@ -13,6 +13,7 @@ from pydantic import ValidationError from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker +from core.tools.tool_file_manager import ToolFileManager from enums import DeploymentEdition, WebAppAccessMode from extensions import ext_application_services from extensions.ext_redis import RedisClientWrapper @@ -78,6 +79,7 @@ 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 from services.tag_application_service import TagApplicationService +from services.tool_file_download_service import ToolFileDownloadService from services.upload_file_delivery_service import UploadFileDeliveryService from services.webapp_access_query_service import WebAppAccessUnavailableError from services.workflow_app_log_query_service import WorkflowAppLogQueryService @@ -254,6 +256,20 @@ def test_build_application_services_wires_message_file_previews( assert services.message_file_previews._storage is ext_application_services.storage +def test_build_application_services_wires_tool_file_downloads( + 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.tool_file_downloads, ToolFileDownloadService) + assert isinstance(services.tool_file_downloads._tool_files, ToolFileManager) + + def test_build_application_services_wires_upload_file_delivery( sqlite_session_factory: sessionmaker[Session], ) -> None: diff --git a/api/tests/unit_tests/services/test_tool_file_download_service.py b/api/tests/unit_tests/services/test_tool_file_download_service.py new file mode 100644 index 00000000000..43f6110e769 --- /dev/null +++ b/api/tests/unit_tests/services/test_tool_file_download_service.py @@ -0,0 +1,94 @@ +from collections.abc import Iterator +from unittest.mock import Mock, patch + +import pytest + +from graphon.file import File, FileTransferMethod, FileType +from services.tool_file_download_service import ( + ToolFileDownload, + ToolFileDownloadAccessDeniedError, + ToolFileDownloadNotFoundError, + ToolFileDownloadService, + ToolFileDownloadSource, +) + + +def _file() -> File: + return File( + file_id="file-id", + tenant_id="tenant-id", + file_type=FileType.DOCUMENT, + transfer_method=FileTransferMethod.TOOL_FILE, + related_id="file-id", + filename="tool.txt", + extension=".txt", + mime_type="text/plain", + size=12, + storage_key="tools/tenant-id/file.txt", + ) + + +@pytest.fixture +def tool_files() -> Mock: + return Mock(spec=ToolFileDownloadSource) + + +@pytest.fixture +def service(tool_files: Mock) -> ToolFileDownloadService: + return ToolFileDownloadService(tool_files=tool_files) + + +def test_invalid_signature_does_not_load_file(service: ToolFileDownloadService, tool_files: Mock) -> None: + with patch("services.tool_file_download_service.verify_tool_file_signature", return_value=False) as verify: + with pytest.raises(ToolFileDownloadAccessDeniedError): + service.get_signed_file(file_id="file-id", timestamp="1", nonce="nonce", sign="invalid") + + verify.assert_called_once_with(file_id="file-id", timestamp="1", nonce="nonce", sign="invalid") + tool_files.get_file_generator_by_tool_file_id.assert_not_called() + + +def test_missing_file_is_reported_after_signature_validation( + service: ToolFileDownloadService, + tool_files: Mock, +) -> None: + tool_files.get_file_generator_by_tool_file_id.return_value = (None, None) + + with patch("services.tool_file_download_service.verify_tool_file_signature", return_value=True): + with pytest.raises(ToolFileDownloadNotFoundError): + service.get_signed_file(file_id="missing", timestamp="1", nonce="nonce", sign="valid") + + tool_files.get_file_generator_by_tool_file_id.assert_called_once_with(tool_file_id="missing") + + +def test_signed_file_returns_stream_and_metadata( + service: ToolFileDownloadService, + tool_files: Mock, +) -> None: + stream: Iterator[bytes] = iter([b"a", b"b"]) + tool_files.get_file_generator_by_tool_file_id.return_value = (stream, _file()) + + with patch("services.tool_file_download_service.verify_tool_file_signature", return_value=True): + result = service.get_signed_file(file_id="file-id", timestamp="1", nonce="nonce", sign="valid") + + assert result == ToolFileDownload( + content=stream, + mime_type="text/plain", + filename="tool.txt", + size=12, + ) + tool_files.get_file_generator_by_tool_file_id.assert_called_once_with(tool_file_id="file-id") + + +@pytest.mark.parametrize("source_error", [RuntimeError("database unavailable"), OSError("storage unavailable")]) +def test_source_error_is_not_converted( + service: ToolFileDownloadService, + tool_files: Mock, + source_error: Exception, +) -> None: + tool_files.get_file_generator_by_tool_file_id.side_effect = source_error + + with patch("services.tool_file_download_service.verify_tool_file_signature", return_value=True): + with pytest.raises(type(source_error)) as error_info: + service.get_signed_file(file_id="file-id", timestamp="1", nonce="nonce", sign="valid") + + assert error_info.value is source_error