diff --git a/backend/package/yuxi/services/conversation_service.py b/backend/package/yuxi/services/conversation_service.py index e71dfa0f..b24c16b9 100644 --- a/backend/package/yuxi/services/conversation_service.py +++ b/backend/package/yuxi/services/conversation_service.py @@ -2,7 +2,6 @@ import uuid from dataclasses import dataclass from pathlib import Path -import aiofiles from fastapi import HTTPException, UploadFile from sqlalchemy.ext.asyncio import AsyncSession from yuxi.agents.backends.sandbox import ( @@ -13,6 +12,7 @@ from yuxi.agents.buildin import agent_manager from yuxi.config import config as app_config from yuxi.plugins.parser import Parser from yuxi.repositories.conversation_repository import ConversationRepository +from yuxi.services.upload_utils import write_upload_to_path from yuxi.utils.datetime_utils import utc_isoformat from yuxi.utils.logging_config import logger from yuxi.utils.paths import VIRTUAL_PATH_UPLOADS @@ -41,21 +41,12 @@ def _ensure_workdir() -> Path: async def _write_upload_to_disk(upload: UploadFile, dest: Path) -> int: - await upload.seek(0) - written = 0 - chunk_size = 1024 * 1024 - - async with aiofiles.open(dest, "wb") as buffer: - while True: - chunk = await upload.read(chunk_size) - if not chunk: - break - written += len(chunk) - if written > MAX_ATTACHMENT_SIZE_BYTES: - raise ValueError("附件过大，当前仅支持 5 MB 以内的文件") - await buffer.write(chunk) - - return written + return await write_upload_to_path( + upload, + dest, + max_size_bytes=MAX_ATTACHMENT_SIZE_BYTES, + too_large_message="附件过大，当前仅支持 5 MB 以内的文件", + ) def _truncate_markdown(markdown: str) -> tuple[str, bool]: diff --git a/backend/package/yuxi/services/upload_utils.py b/backend/package/yuxi/services/upload_utils.py new file mode 100644 index 00000000..6a1b7902 --- /dev/null +++ b/backend/package/yuxi/services/upload_utils.py @@ -0,0 +1,43 @@ +from pathlib import Path + +import aiofiles +from fastapi import UploadFile + + +async def write_upload_to_buffer( + upload: UploadFile, + buffer, + *, + max_size_bytes: int, + too_large_message: str, + chunk_size: int = 1024 * 1024, +) -> int: + await upload.seek(0) + written = 0 + + while chunk := await upload.read(chunk_size): + written += len(chunk) + if written > max_size_bytes: + raise ValueError(too_large_message) + await buffer.write(chunk) + + return written + + +async def write_upload_to_path( + upload: UploadFile, + dest: Path, + *, + max_size_bytes: int, + too_large_message: str, + mode: str = "wb", + chunk_size: int = 1024 * 1024, +) -> int: + async with aiofiles.open(dest, mode) as buffer: + return await write_upload_to_buffer( + upload, + buffer, + max_size_bytes=max_size_bytes, + too_large_message=too_large_message, + chunk_size=chunk_size, + ) diff --git a/backend/package/yuxi/services/workspace_service.py b/backend/package/yuxi/services/workspace_service.py index c34aa3a0..509a0d3d 100644 --- a/backend/package/yuxi/services/workspace_service.py +++ b/backend/package/yuxi/services/workspace_service.py @@ -12,12 +12,14 @@ import aiofiles from fastapi import HTTPException, UploadFile from fastapi.responses import FileResponse, StreamingResponse from yuxi.agents.backends.sandbox.paths import _global_user_data_dir, ensure_workspace_default_files +from yuxi.services.upload_utils import write_upload_to_buffer from yuxi.services.viewer_filesystem_service import _detect_preview_type from yuxi.storage.postgres.models_business import User from yuxi.utils.datetime_utils import utc_isoformat_from_timestamp from yuxi.utils.paths import WORKSPACE_DIR_NAME EDITABLE_WORKSPACE_SUFFIXES = {".md", ".markdown", ".mdx", ".txt"} +MAX_WORKSPACE_UPLOAD_SIZE_BYTES = 100 * 1024 * 1024 def _workspace_root(user: User) -> Path: @@ -237,11 +239,17 @@ async def upload_workspace_file(*, parent_path: str, file: UploadFile, current_u try: async with aiofiles.open(target, "xb") as buffer: created_file = True - while chunk := await file.read(1024 * 1024): - await buffer.write(chunk) + await write_upload_to_buffer( + file, + buffer, + max_size_bytes=MAX_WORKSPACE_UPLOAD_SIZE_BYTES, + too_large_message="文件过大，当前仅支持 100 MB 以内的文件", + ) upload_completed = True except FileExistsError as exc: raise HTTPException(status_code=400, detail="同名文件或文件夹已存在") from exc + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc except PermissionError as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc finally: diff --git a/backend/test/unit/services/test_workspace_service.py b/backend/test/unit/services/test_workspace_service.py index 24233ce8..55de5bfd 100644 --- a/backend/test/unit/services/test_workspace_service.py +++ b/backend/test/unit/services/test_workspace_service.py @@ -1,10 +1,11 @@ from __future__ import annotations +from io import BytesIO from pathlib import Path from types import SimpleNamespace import pytest -from fastapi import HTTPException +from fastapi import HTTPException, UploadFile from yuxi.agents.backends.sandbox import paths as workspace_paths from yuxi.services import workspace_service as svc @@ -146,3 +147,37 @@ async def test_write_workspace_file_content_blocks_path_traversal(tmp_path: Path ) assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_upload_workspace_file_writes_file(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(workspace_paths.conf, "save_dir", str(tmp_path)) + user = SimpleNamespace(id="user-1") + root = svc._workspace_root(user) + upload = UploadFile(filename="demo.txt", file=BytesIO(b"hello")) + + result = await svc.upload_workspace_file(parent_path="/", file=upload, current_user=user) + + assert result["success"] is True + assert result["entry"]["path"] == "/demo.txt" + assert result["entry"]["size"] == 5 + assert (root / "demo.txt").read_bytes() == b"hello" + + +@pytest.mark.asyncio +async def test_upload_workspace_file_rejects_oversized_file_and_cleans_partial_file( + tmp_path: Path, + monkeypatch, +) -> None: + monkeypatch.setattr(workspace_paths.conf, "save_dir", str(tmp_path)) + monkeypatch.setattr(svc, "MAX_WORKSPACE_UPLOAD_SIZE_BYTES", 5) + user = SimpleNamespace(id="user-1") + root = svc._workspace_root(user) + upload = UploadFile(filename="large.txt", file=BytesIO(b"123456")) + + with pytest.raises(HTTPException) as exc_info: + await svc.upload_workspace_file(parent_path="/", file=upload, current_user=user) + + assert exc_info.value.status_code == 400 + assert "100 MB" in exc_info.value.detail + assert not (root / "large.txt").exists()