Просмотр исходного кода

feat(backend): validate upload magic bytes and image-only helper

Co-authored-by: Cursor <cursoragent@cursor.com>
wuzilong 2 месяцев назад
Родитель
Сommit
90860d4e7a
2 измененных файлов с 125 добавлено и 0 удалено
  1. 86 0
      backend/app/core/upload_validate.py
  2. 39 0
      backend/tests/test_upload_validate.py

+ 86 - 0
backend/app/core/upload_validate.py

@@ -0,0 +1,86 @@
+"""上传文件校验:扩展名白名单与大小限制。"""
+from __future__ import annotations
+
+import os
+from typing import Callable, Optional
+
+from app.core.response import AppHTTPError
+
+# 对齐原型:图片、视频、常见办公文档与压缩包
+ALLOWED_EXTENSIONS = {
+    ".jpg",
+    ".jpeg",
+    ".png",
+    ".gif",
+    ".webp",
+    ".bmp",
+    ".mp4",
+    ".mov",
+    ".avi",
+    ".webm",
+    ".pdf",
+    ".doc",
+    ".docx",
+    ".xls",
+    ".xlsx",
+    ".zip",
+}
+
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
+
+MAX_UPLOAD_BYTES = 20 * 1024 * 1024  # 20MB
+
+
+def _starts_with(prefix: bytes) -> Callable[[bytes], bool]:
+    def check(content: bytes) -> bool:
+        return content.startswith(prefix)
+
+    return check
+
+
+def _starts_with_any(*prefixes: bytes) -> Callable[[bytes], bool]:
+    def check(content: bytes) -> bool:
+        return any(content.startswith(p) for p in prefixes)
+
+    return check
+
+
+def _webp(content: bytes) -> bool:
+    return len(content) >= 12 and content[:4] == b"RIFF" and content[8:12] == b"WEBP"
+
+
+_MAGIC_CHECKS: dict[str, Callable[[bytes], bool]] = {
+    ".png": _starts_with(b"\x89PNG\r\n\x1a\n"),
+    ".jpg": _starts_with(b"\xff\xd8\xff"),
+    ".jpeg": _starts_with(b"\xff\xd8\xff"),
+    ".gif": _starts_with_any(b"GIF87a", b"GIF89a"),
+    ".webp": _webp,
+    ".bmp": _starts_with(b"BM"),
+    ".pdf": _starts_with(b"%PDF"),
+}
+
+
+def file_extension(filename: str) -> str:
+    return os.path.splitext(filename or "")[1].lower()
+
+
+def validate_upload(filename: Optional[str], content: bytes) -> None:
+    if not filename:
+        raise AppHTTPError(400, "请选择文件")
+    ext = file_extension(filename)
+    if ext not in ALLOWED_EXTENSIONS:
+        raise AppHTTPError(400, "不支持的文件类型")
+    if len(content) > MAX_UPLOAD_BYTES:
+        raise AppHTTPError(400, "文件过大,单文件不超过 20MB")
+    magic_check = _MAGIC_CHECKS.get(ext)
+    if magic_check is not None and not magic_check(content):
+        raise AppHTTPError(400, "文件内容与扩展名不符")
+
+
+def validate_image_upload(filename: Optional[str], content: bytes) -> None:
+    if not filename:
+        raise AppHTTPError(400, "请选择文件")
+    ext = file_extension(filename)
+    if ext not in IMAGE_EXTENSIONS:
+        raise AppHTTPError(400, "不支持的文件类型")
+    validate_upload(filename, content)

+ 39 - 0
backend/tests/test_upload_validate.py

@@ -0,0 +1,39 @@
+import pytest
+from app.core.response import AppHTTPError
+from app.core.upload_validate import validate_upload, validate_image_upload
+
+PNG = b"\x89PNG\r\n\x1a\n" + b"\x00" * 16
+JPEG = b"\xff\xd8\xff" + b"\x00" * 16
+PDF = b"%PDF-1.4" + b"\x00" * 16
+EXE_MZ = b"MZ" + b"\x00" * 64
+
+
+def test_rejects_exe_extension():
+    with pytest.raises(AppHTTPError) as ei:
+        validate_upload("malware.exe", EXE_MZ)
+    assert ei.value.status_code == 400
+
+
+def test_rejects_png_extension_with_exe_magic():
+    with pytest.raises(AppHTTPError) as ei:
+        validate_upload("fake.png", EXE_MZ)
+    assert ei.value.status_code == 400
+    assert "类型" in ei.value.message or "格式" in ei.value.message or "内容" in ei.value.message
+
+
+def test_accepts_real_png():
+    validate_upload("a.png", PNG)
+
+
+def test_accepts_jpeg_and_pdf():
+    validate_upload("a.jpg", JPEG)
+    validate_upload("a.pdf", PDF)
+
+
+def test_image_upload_rejects_pdf():
+    with pytest.raises(AppHTTPError):
+        validate_image_upload("x.pdf", PDF)
+
+
+def test_image_upload_accepts_png():
+    validate_image_upload("x.png", PNG)