后台更新 服务优化

backend/ali_oss.py

@@ -7,7 +7,9 @@
 import base64
 import io
 import mimetypes
+import os
 import re
+import tempfile
 import threading
 import time
 
@@ -27,6 +29,8 @@ from config import (
 
 _client = None
 _client_lock = threading.Lock()
+_REMOTE_IMAGE_MAX_BYTES = int(os.getenv("OSS_UPLOAD_IMAGE_MAX_BYTES", str(10 * 1024 * 1024)))
+_REMOTE_IMAGE_TIMEOUT = (5, 30)
 
 
 def get_client():
@@ -65,19 +69,72 @@ def get_presigned_url(key: str, expires: int = 7 * 24 * 3600) -> str:
 
 
 
-def upload_file(file_content: bytes, key: str):
-    """上传字节内容到 MinIO，返回可访问链接"""
+class _LimitedReader:
+    """Keep a streamed upload from accepting an unexpectedly huge object."""
+
+    def __init__(self, source, max_bytes: int):
+        self._source = source
+        self._max_bytes = max_bytes
+        self._read_bytes = 0
+
+    def read(self, size=-1):
+        remaining = self._max_bytes - self._read_bytes
+        if remaining < 0:
+            raise ValueError("upload exceeds configured size limit")
+        read_size = remaining + 1 if size is None or size < 0 else min(size, remaining + 1)
+        data = self._source.read(read_size)
+        if not data:
+            return data
+        self._read_bytes += len(data)
+        if self._read_bytes > self._max_bytes:
+            raise ValueError("upload exceeds configured size limit")
+        return data
+
+    def seek(self, offset, whence=0):
+        position = self._source.seek(offset, whence)
+        self._read_bytes = max(0, position)
+        return position
+
+    def tell(self):
+        return self._source.tell()
+
+
+def upload_fileobj(file_obj, key: str, max_bytes: int = 0, content_type: str = ""):
+    """Stream a file-like object to S3/MinIO without materializing it."""
     key = key.lstrip("/")
+    if isinstance(file_obj, (bytes, bytearray)):
+        file_obj = io.BytesIO(file_obj)
+    if not hasattr(file_obj, "read"):
+        raise TypeError("file_obj must be bytes or a readable file-like object")
+    stream_size = None
+    try:
+        file_obj.seek(0)
+        file_obj.seek(0, os.SEEK_END)
+        stream_size = file_obj.tell()
+        file_obj.seek(0)
+    except (AttributeError, OSError, TypeError, ValueError):
+        try:
+            file_obj.seek(0)
+        except (AttributeError, OSError, TypeError, ValueError):
+            pass
+    if max_bytes and max_bytes > 0 and stream_size is not None and stream_size > max_bytes:
+        raise ValueError("upload exceeds configured size limit")
     client = get_client()
-    client.put_object(
-        Bucket=bucket,          # 存储桶名称
-        Key=key,                # 对象名称
-        Body=io.BytesIO(file_content) if isinstance(file_content, (bytes, bytearray)) else file_content,
-        ContentType=_guess_content_type(key),
+    body = _LimitedReader(file_obj, max_bytes) if max_bytes and max_bytes > 0 else file_obj
+    client.upload_fileobj(
+        body,
+        bucket,
+        key,
+        ExtraArgs={"ContentType": content_type or _guess_content_type(key)},
     )
     return file_url_pre + key
 
 
+def upload_file(file_content: bytes, key: str):
+    """上传字节内容到 MinIO，返回可访问链接"""
+    return upload_fileobj(file_content, key)
+
+
 def upload_data_url(data_url: str, prefix: str = "history", key_hint: str = "") -> str:
     """
     将 base64 data URL 上传到对象存储，返回图片链接
@@ -89,14 +146,19 @@ def upload_data_url(data_url: str, prefix: str = "history", key_hint: str = "")
     if not match:
         raise ValueError('无效的 data URL 格式')
     ext = 'png' if match.group(1).lower() in ('png', 'webp') else 'jpg'
-    file_content = base64.b64decode(match.group(2))
+    encoded_payload = match.group(2)
+    if len(encoded_payload) > ((max(_REMOTE_IMAGE_MAX_BYTES, 0) + 2) // 3) * 4:
+        raise ValueError("image exceeds configured upload size limit")
+    file_content = base64.b64decode(encoded_payload, validate=True)
+    if len(file_content) > _REMOTE_IMAGE_MAX_BYTES:
+        raise ValueError("image exceeds configured upload size limit")
     ts = int(time.time() * 1000)
     key = f"{bucket_path}{prefix}/{ts}{key_hint}.{ext}"
-    return upload_file(file_content, key)
+    return upload_fileobj(io.BytesIO(file_content), key, _REMOTE_IMAGE_MAX_BYTES)
 
 
 def upload_data_urls(data_urls: list, prefix: str = "history") -> list:
-    """批量上传 base64 图片到对象存储，返回图片链接列表"""
+    """批量上传图片；每张图���完成上传后立即释放其缓冲区。"""
     urls = []
     ts = int(time.time() * 1000)
     for i, data_url in enumerate(data_urls or []):
@@ -107,12 +169,37 @@ def upload_data_urls(data_urls: list, prefix: str = "history") -> list:
             if not match:
                 continue
             ext = 'png' if match.group(1).lower() in ('png', 'webp') else 'jpg'
-            file_content = base64.b64decode(match.group(2))
-        else:
-            file_content = requests.get(data_url).content
-            ext = "png"
-        key = f"{bucket_path}{prefix}/{ts}_{i}.{ext}"
-        urls.append(upload_file(file_content, key))
+            encoded_payload = match.group(2)
+            if len(encoded_payload) > ((max(_REMOTE_IMAGE_MAX_BYTES, 0) + 2) // 3) * 4:
+                raise ValueError("image exceeds configured upload size limit")
+            file_content = base64.b64decode(encoded_payload, validate=True)
+            if len(file_content) > _REMOTE_IMAGE_MAX_BYTES:
+                raise ValueError("image exceeds configured upload size limit")
+            key = f"{bucket_path}{prefix}/{ts}_{i}.{ext}"
+            urls.append(upload_fileobj(io.BytesIO(file_content), key, _REMOTE_IMAGE_MAX_BYTES))
+            del file_content
+            continue
+
+        with requests.get(data_url, stream=True, timeout=_REMOTE_IMAGE_TIMEOUT) as response:
+            response.raise_for_status()
+            content_length = response.headers.get("Content-Length")
+            if content_length and int(content_length) > _REMOTE_IMAGE_MAX_BYTES:
+                raise ValueError("image exceeds configured upload size limit")
+            content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
+            extension = mimetypes.guess_extension(content_type) or ".png"
+            ext = extension.lstrip(".") or "png"
+            key = f"{bucket_path}{prefix}/{ts}_{i}.{ext}"
+            with tempfile.SpooledTemporaryFile(max_size=2 * 1024 * 1024, mode="w+b") as image_file:
+                total = 0
+                for chunk in response.iter_content(chunk_size=64 * 1024):
+                    if not chunk:
+                        continue
+                    total += len(chunk)
+                    if total > _REMOTE_IMAGE_MAX_BYTES:
+                        raise ValueError("image exceeds configured upload size limit")
+                    image_file.write(chunk)
+                image_file.seek(0)
+                urls.append(upload_fileobj(image_file, key, _REMOTE_IMAGE_MAX_BYTES, content_type))
     return urls