From e12e0c7d02ff06bee6b5b56da4b604333cef6713 Mon Sep 17 00:00:00 2001 From: kevin9327 Date: Sat, 5 Sep 2026 21:07:13 +0900 Subject: [PATCH] =?UTF-8?q?fix(watchlist):=20=E8=87=AA=E9=80=89=E5=AF=BC?= =?UTF-8?q?=E5=85=A5=E4=B8=8A=E4=BC=A0=E5=88=86=E5=9D=97=E8=AF=BB=E5=8F=96?= =?UTF-8?q?,=20=E8=B6=8A=E9=99=90=E5=8D=B3=E6=8B=92=E7=BB=9D=E8=80=8C?= =?UTF-8?q?=E9=9D=9E=E8=AF=BB=E5=AE=8C=E5=86=8D=E6=8B=92?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit import-csv / import-image 原先 `await file.read()` 之后才比较长度, 上限只在 整个文件进入内存之后生效, 与 issue #204 修复前的扩展数据上传同类。新增 _read_upload_capped 分块读取, 累计超过上限的那一块立即返回 400 并停止读取, 状态码与文案不变; 内存占用不超过上限 + 一块。 测试: 新增 5 例 (含「越限后不再继续读取」的字节计数断言); test_watchlist_csv 的上传 mock 改为像真实 UploadFile 一样读尽返回 b""。 --- backend/app/api/watchlist.py | 30 +++++++-- backend/tests/test_watchlist_csv.py | 4 +- backend/tests/test_watchlist_ocr.py | 4 +- .../tests/test_watchlist_upload_size_limit.py | 67 +++++++++++++++++++ 4 files changed, 97 insertions(+), 8 deletions(-) create mode 100644 backend/tests/test_watchlist_upload_size_limit.py 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""