Refine seller store implementation

2026-08-31 14:37:28 +08:00
parent 957289f717
commit 2a667b4f3b
27 changed files with 1650 additions and 624 deletions
@@ -1,12 +1,20 @@
 """图片/文件存储抽象：本地文件系统（开发兜底）+ 七牛（生产）。"""
 from __future__ import annotations

+import asyncio
+import io
+import ipaddress
 import mimetypes
+import socket
 import shutil
 import uuid
+import warnings
+from dataclasses import dataclass
 from pathlib import Path
+from urllib.parse import urljoin, urlparse

 import httpx
+from PIL import Image, ImageOps, UnidentifiedImageError

 from config import get_settings

@@ -26,6 +34,157 @@ class ResourceSizeLimitError(ValueError):
        super().__init__(detail)


+class UnsafeRemoteUrlError(ValueError):
+    """远程资源 URL 指向回环、内网或其他保留网络，不能由服务端代请求。"""
+
+
+class ResourceDownloadTimeoutError(TimeoutError):
+    """远程资源下载超过连接、读取或总耗时限制。"""
+
+
+class ImageValidationError(ValueError):
+    """图片字节无法解码，或不符合格式与像素限制。"""
+
+
+@dataclass(frozen=True)
+class NormalizedImage:
+    """经真实格式校验后的图片字节及可落库元数据。"""
+
+    data: bytes
+    content_type: str
+    width: int
+    height: int
+    image_format: str
+    orientation_normalized: bool
+
+
+_IMAGE_CONTENT_TYPES = {
+    "JPEG": "image/jpeg",
+    "PNG": "image/png",
+    "WEBP": "image/webp",
+    "GIF": "image/gif",
+    "BMP": "image/bmp",
+}
+
+
+def normalize_image_bytes(
+    data: bytes,
+    max_bytes: int | None = None,
+    max_pixels: int | None = None,
+) -> NormalizedImage:
+    """按真实图片内容校验、纠正 EXIF 方向，并返回可信 MIME 与宽高。
+
+    不盲从 URL 后缀或客户端 Content-Type。为避免无意义的质量损耗，仅当 EXIF
+    方向不是默认值时才重编码；动画图不重编码，以免丢失帧。
+    """
+    settings = get_settings()
+    byte_limit = max_bytes if max_bytes is not None else settings.image_upload_max_bytes
+    pixel_limit = max_pixels if max_pixels is not None else settings.image_max_pixels
+    if byte_limit <= 0 or pixel_limit <= 0:
+        raise ValueError("图片大小和像素上限必须大于 0")
+    if not data:
+        raise ImageValidationError("图片内容为空")
+    if len(data) > byte_limit:
+        raise ResourceSizeLimitError(byte_limit, len(data))
+
+    try:
+        # Pillow 的默认 bomb 保护不能替代项目自定义像素上限；把 warning 升级为错误。
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", Image.DecompressionBombWarning)
+            with Image.open(io.BytesIO(data)) as image:
+                image_format = (image.format or "").upper()
+                if image_format not in _IMAGE_CONTENT_TYPES:
+                    raise ImageValidationError(f"不支持的图片格式：{image_format or '未知'}")
+
+                width, height = image.size
+                if width <= 0 or height <= 0 or width * height > pixel_limit:
+                    raise ImageValidationError(
+                        f"图片像素超过上限（最多 {pixel_limit:,}，当前 {width:,} × {height:,}）"
+                    )
+
+                # 强制实际解码，拒绝截断/伪造图片而不是只相信文件头。
+                image.load()
+                frame_count = getattr(image, "n_frames", 1)
+                orientation = image.getexif().get(274, 1) if frame_count == 1 else 1
+                normalized = orientation not in (None, 1)
+                output = data
+                if normalized:
+                    corrected = ImageOps.exif_transpose(image)
+                    width, height = corrected.size
+                    if image_format == "JPEG" and corrected.mode not in {"RGB", "L"}:
+                        corrected = corrected.convert("RGB")
+                    buffer = io.BytesIO()
+                    corrected.save(buffer, format=image_format)
+                    output = buffer.getvalue()
+    except (UnidentifiedImageError, OSError, Image.DecompressionBombError) as exc:
+        raise ImageValidationError("文件不是可用的图片，或图片数据已损坏") from exc
+    except Image.DecompressionBombWarning as exc:
+        raise ImageValidationError("图片像素过大，已拒绝解码") from exc
+
+    if len(output) > byte_limit:
+        raise ResourceSizeLimitError(byte_limit, len(output))
+    return NormalizedImage(
+        data=output,
+        content_type=_IMAGE_CONTENT_TYPES[image_format],
+        width=width,
+        height=height,
+        image_format=image_format,
+        orientation_normalized=normalized,
+    )
+
+
+def _is_public_ip(address: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
+    """仅允许可公开路由的地址；IPv4-mapped IPv6 需按其 IPv4 实际地址判断。"""
+    if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped:
+        return _is_public_ip(address.ipv4_mapped)
+    return address.is_global
+
+
+async def validate_remote_url(url: str) -> None:
+    """校验服务端允许代请求的公网 HTTP(S) URL。
+
+    该校验用于代理、采集转存、AI 参考图和导出。项目自己的 `/media/` 文件
+    在调用方会先直读磁盘；其他回环或私网 URL 一律不能通过网络访问。
+    """
+    try:
+        parsed = urlparse(url)
+        port = parsed.port
+    except ValueError as exc:
+        raise UnsafeRemoteUrlError("图片地址端口无效") from exc
+
+    if parsed.scheme.lower() not in {"http", "https"}:
+        raise UnsafeRemoteUrlError("仅允许 http/https 图片地址")
+    if not parsed.hostname:
+        raise UnsafeRemoteUrlError("图片地址缺少主机名")
+    if parsed.username is not None or parsed.password is not None:
+        raise UnsafeRemoteUrlError("图片地址不能包含用户名或密码")
+
+    host = parsed.hostname.rstrip(".")
+    if host.lower() in {"localhost", "localhost.localdomain"}:
+        raise UnsafeRemoteUrlError("不允许请求 localhost 或本机网络地址")
+
+    try:
+        addresses = {ipaddress.ip_address(host)}
+    except ValueError:
+        try:
+            infos = await asyncio.get_running_loop().getaddrinfo(
+                host,
+                port or (443 if parsed.scheme.lower() == "https" else 80),
+                type=socket.SOCK_STREAM,
+            )
+        except socket.gaierror as exc:
+            raise UnsafeRemoteUrlError("图片地址域名无法解析") from exc
+        addresses = {ipaddress.ip_address(info[4][0]) for info in infos}
+
+    if not addresses:
+        raise UnsafeRemoteUrlError("图片地址域名未解析到有效地址")
+    private_addresses = [str(address) for address in addresses if not _is_public_ip(address)]
+    if private_addresses:
+        raise UnsafeRemoteUrlError(
+            f"不允许请求非公网地址（{', '.join(sorted(private_addresses))}）"
+        )
+
+
 def _ext_from_content_type(content_type: str) -> str:
    ctype = (content_type or "").split(";")[0].strip().lower()
    mapping = {
@@ -48,41 +207,78 @@ def image_extension(data: bytes, content_type: str = "") -> str:
        return ".webp"
    if data[:3] == b"GIF":
        return ".gif"
+    if data[:2] == b"\xff\xd8":
+        return ".jpg"
+    if data[:2] == b"BM":
+        return ".bmp"
    return _ext_from_content_type(content_type)


 async def download_bytes(
    url: str,
    referer: str | None = None,
-    timeout: float = 60.0,
+    timeout: float | None = None,
    max_bytes: int | None = None,
 ) -> tuple[bytes, str]:
-    """流式下载远程图片并限制体积。返回 (bytes, content_type)。"""
-    limit = max_bytes or get_settings().image_download_max_bytes
+    """安全地流式下载公网资源，并限制地址、重定向、体积和耗时。"""
+    settings = get_settings()
+    limit = max_bytes or settings.image_download_max_bytes
    if limit <= 0:
        raise ValueError("图片下载大小上限必须大于 0")
-    headers = {"Referer": referer} if referer else {}
-    async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
-        async with client.stream("GET", url, headers=headers) as resp:
-            resp.raise_for_status()
-            length = resp.headers.get("content-length")
-            if length:
-                try:
-                    declared_size = int(length)
-                except ValueError:
-                    declared_size = None
-                if declared_size is not None and declared_size > limit:
-                    raise ResourceSizeLimitError(limit, declared_size)
+    total_timeout = timeout if timeout is not None else settings.image_download_total_timeout
+    if total_timeout <= 0:
+        raise ValueError("图片下载总超时必须大于 0")
+    if settings.image_download_connect_timeout <= 0 or settings.image_download_read_timeout <= 0:
+        raise ValueError("图片下载连接和读取超时必须大于 0")
+    if settings.image_download_max_redirects < 0:
+        raise ValueError("图片下载重定向上限不能小于 0")

-            chunks: list[bytes] = []
-            total = 0
-            async for chunk in resp.aiter_bytes():
-                total += len(chunk)
-                if total > limit:
-                    raise ResourceSizeLimitError(limit, total)
-                chunks.append(chunk)
-            ctype = (resp.headers.get("content-type") or "application/octet-stream").split(";")[0].strip()
-            return b"".join(chunks), ctype
+    headers = {"Referer": referer} if referer else {}
+    http_timeout = httpx.Timeout(
+        connect=settings.image_download_connect_timeout,
+        read=settings.image_download_read_timeout,
+        write=settings.image_download_read_timeout,
+        pool=settings.image_download_connect_timeout,
+    )
+    current_url = url
+    try:
+        async with asyncio.timeout(total_timeout):
+            async with httpx.AsyncClient(timeout=http_timeout, follow_redirects=False) as client:
+                for redirect_count in range(settings.image_download_max_redirects + 1):
+                    await validate_remote_url(current_url)
+                    async with client.stream("GET", current_url, headers=headers) as resp:
+                        if resp.is_redirect:
+                            location = resp.headers.get("location")
+                            if not location:
+                                raise httpx.RemoteProtocolError("重定向响应缺少 Location")
+                            if redirect_count >= settings.image_download_max_redirects:
+                                raise httpx.TooManyRedirects("图片重定向次数超过上限")
+                            current_url = urljoin(str(resp.url), location)
+                            continue
+
+                        resp.raise_for_status()
+                        length = resp.headers.get("content-length")
+                        if length:
+                            try:
+                                declared_size = int(length)
+                            except ValueError:
+                                declared_size = None
+                            if declared_size is not None and declared_size > limit:
+                                raise ResourceSizeLimitError(limit, declared_size)
+
+                        chunks: list[bytes] = []
+                        total = 0
+                        async for chunk in resp.aiter_bytes():
+                            total += len(chunk)
+                            if total > limit:
+                                raise ResourceSizeLimitError(limit, total)
+                            chunks.append(chunk)
+                        ctype = (resp.headers.get("content-type") or "application/octet-stream").split(";")[0].strip()
+                        return b"".join(chunks), ctype
+    except (TimeoutError, httpx.TimeoutException) as exc:
+        raise ResourceDownloadTimeoutError(f"图片下载超时（最多 {total_timeout:g} 秒）") from exc
+
+    raise RuntimeError("图片下载未返回有效响应")


 def read_local_bytes(path: Path, max_bytes: int | None = None) -> bytes: