diff --git a/backend/app/api/watchlist.py b/backend/app/api/watchlist.py index 771c617..0164f8a 100644 --- a/backend/app/api/watchlist.py +++ b/backend/app/api/watchlist.py @@ -40,6 +40,28 @@ _IMPORT_CSV_TYPES = { "text/plain", "application/csv", } +# 上传分块读取粒度 (与 ext_data 上传一致) +_UPLOAD_CHUNK_BYTES = 1024 * 1024 + + +async def _read_upload_capped(file: UploadFile, max_bytes: int, too_large: str) -> bytes: + """分块读取上传内容, 累计超过 max_bytes 立即拒绝(400), 返回完整字节。 + + 与 ext_data._write_upload_capped 同类保护: 一次性 `await file.read()` 会先把整个 + 文件读入内存再比较长度, 上限在那之后才生效, 一个远超上限的上传照样把进程内存 + 顶满; 分块读取在越过上限的那一块就停止, 内存占用不超过上限 + 一块。 + """ + chunks: list[bytes] = [] + total = 0 + while True: + chunk = await file.read(_UPLOAD_CHUNK_BYTES) + if not chunk: + break + total += len(chunk) + if total > max_bytes: + raise HTTPException(400, too_large) + chunks.append(chunk) + return b"".join(chunks) class AddRequest(BaseModel): @@ -186,11 +208,9 @@ async def import_from_image(request: Request, file: UploadFile = File(...)): if not ok_type and not ok_ext: raise HTTPException(400, "仅支持 JPG / PNG / WebP / BMP / GIF 图片") - data = await file.read() + data = await _read_upload_capped(file, _MAX_IMPORT_IMAGE_BYTES, "图片过大(上限 12MB)") if not data: raise HTTPException(400, "空文件") - if len(data) > _MAX_IMPORT_IMAGE_BYTES: - raise HTTPException(400, "图片过大(上限 12MB)") existing = {r["symbol"] for r in watchlist.list_symbols()} data_dir = request.app.state.repo.store.data_dir @@ -242,11 +262,9 @@ async def import_from_csv(request: Request, file: UploadFile = File(...)): if not ok_type and not ok_ext: raise HTTPException(400, "仅支持 CSV / TXT 文件") - data = await file.read() + data = await _read_upload_capped(file, _MAX_IMPORT_CSV_BYTES, "文件过大(上限 5MB)") if not data: raise HTTPException(400, "空文件") - if len(data) > _MAX_IMPORT_CSV_BYTES: - raise HTTPException(400, "文件过大(上限 5MB)") data_dir = request.app.state.repo.store.data_dir # 解码与自选/instruments parquet 读取为同步 CPU/IO,挪线程池避免卡事件循环 diff --git a/backend/tests/test_watchlist_csv.py b/backend/tests/test_watchlist_csv.py index 1477638..c3a51f1 100644 --- a/backend/tests/test_watchlist_csv.py +++ b/backend/tests/test_watchlist_csv.py @@ -40,7 +40,9 @@ def _mock_upload(*, content: bytes, content_type: str = "text/csv", filename: st file = MagicMock() file.content_type = content_type file.filename = filename - file.read = AsyncMock(return_value=content) + # 像真实 UploadFile 一样: 第一次 read 返回全部内容, 之后返回 b"" 表示读尽 + # (端点已改为分块读取, 一直返回同一段内容的 mock 会被当成无限大的文件) + file.read = AsyncMock(side_effect=[content, b""]) return file diff --git a/backend/tests/test_watchlist_ocr.py b/backend/tests/test_watchlist_ocr.py index 843043b..0e4e6df 100644 --- a/backend/tests/test_watchlist_ocr.py +++ b/backend/tests/test_watchlist_ocr.py @@ -54,7 +54,9 @@ def _mock_upload( file = MagicMock() file.content_type = content_type file.filename = filename - file.read = AsyncMock(return_value=content) + # 像真实 UploadFile 一样: 第一次 read 返回全部内容, 之后返回 b"" 表示读尽 + # (端点已改为分块读取, 一直返回同一段内容的 mock 会被当成无限大的文件) + file.read = AsyncMock(side_effect=[content, b""]) return file diff --git a/backend/tests/test_watchlist_upload_size_limit.py b/backend/tests/test_watchlist_upload_size_limit.py new file mode 100644 index 0000000..ca0d5e8 --- /dev/null +++ b/backend/tests/test_watchlist_upload_size_limit.py @@ -0,0 +1,67 @@ +"""自选导入上传体积上限回归测试 (import-csv / import-image)。 + +两个端点原先 `data = await file.read()` 之后才比较长度: 上限只在整个文件已经进入 +内存之后生效, 与 ext_data 上传修复 (issue #204) 前的问题同类。_read_upload_capped +分块读取, 越过上限即拒绝(400), 内存占用不超过上限 + 一块。纯逻辑, 不需真实数据源。 +""" +from __future__ import annotations + +import io + +import pytest +from fastapi import HTTPException +from starlette.datastructures import UploadFile + +from app.api import watchlist as watchlist_api +from app.api.watchlist import _read_upload_capped + + +class _CountingStream(io.BytesIO): + """记录被读取的字节数, 用来证明越限后不再继续读。""" + + def __init__(self, data: bytes) -> None: + super().__init__(data) + self.consumed = 0 + + def read(self, size: int = -1) -> bytes: # type: ignore[override] + chunk = super().read(size) + self.consumed += len(chunk) + return chunk + + +def _upload(stream: io.BytesIO) -> UploadFile: + return UploadFile(stream, filename="x.csv") + + +async def test_under_limit_returns_full_content(): + data = "code,name\n600519.SH,贵州茅台\n".encode() + got = await _read_upload_capped(_upload(io.BytesIO(data)), 1024, "too large") + assert got == data + + +async def test_at_limit_is_allowed(): + data = b"x" * 64 + got = await _read_upload_capped(_upload(io.BytesIO(data)), 64, "too large") + assert got == data + + +async def test_over_limit_raises_400_with_the_message_the_endpoint_passes(): + data = b"x" * 200 + with pytest.raises(HTTPException) as exc: + await _read_upload_capped(_upload(io.BytesIO(data)), 64, "文件过大") + assert exc.value.status_code == 400 + assert exc.value.detail == "文件过大" + + +async def test_over_limit_stops_reading_at_the_first_chunk_past_the_cap(monkeypatch): + # 4 块的文件, 上限 1.5 块: 第 2 块越限即拒绝, 第 3/4 块不应被读取。 + monkeypatch.setattr(watchlist_api, "_UPLOAD_CHUNK_BYTES", 16) + stream = _CountingStream(b"x" * 64) + with pytest.raises(HTTPException): + await _read_upload_capped(_upload(stream), 24, "too large") + assert stream.consumed == 32 + + +async def test_empty_upload_returns_empty_bytes(): + got = await _read_upload_capped(_upload(io.BytesIO(b"")), 64, "too large") + assert got == b""