opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-09 05:41:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(api): extract tool file download service (#41790)

Browse Source
This commit is contained in:
非法操作 2026-09-08 02:51:11 +00:00 committed by GitHub
parent 22189f69f6
commit d43a23d8f5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 359 additions and 132 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
api/.importlinter
View File

@ -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

67
api/controllers/files/tool_files.py
View File

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

3
api/extensions/ext_application_services.py
View File

@ -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,

61
api/services/tool_file_download_service.py Normal file
View File

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

236
api/tests/unit_tests/controllers/files/test_tool_files.py
View File

@ -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

16
api/tests/unit_tests/extensions/test_ext_application_services.py
View File

@ -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:

94
api/tests/unit_tests/services/test_tool_file_download_service.py Normal file
View File

@ -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