feat: 支持S3预签名上传

2025-12-31 03:25:09 +08:00
parent 287d9aeedb
commit 8e57a497f7
25 changed files with 825 additions and 336 deletions
@@ -0,0 +1,26 @@
+from tortoise import connections
+
+
+async def create_presign_upload_session_table():
+    conn = connections.get("default")
+    await conn.execute_script(
+        """
+        CREATE TABLE IF NOT EXISTS presignuploadsession (
+            id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
+            upload_id VARCHAR(36) NOT NULL UNIQUE,
+            file_name VARCHAR(255) NOT NULL,
+            file_size BIGINT NOT NULL,
+            save_path VARCHAR(512) NOT NULL,
+            mode VARCHAR(10) NOT NULL,
+            expire_value INT NOT NULL DEFAULT 1,
+            expire_style VARCHAR(20) NOT NULL DEFAULT 'day',
+            created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            expires_at TIMESTAMP NOT NULL
+        );
+        CREATE INDEX IF NOT EXISTS idx_presignuploadsession_upload_id ON presignuploadsession (upload_id);
+        """
+    )
+
+
+async def migrate():
+    await create_presign_upload_session_table()
@@ -64,6 +64,25 @@ class KeyValue(Model):
    )


+class PresignUploadSession(models.Model):
+    """预签名上传会话模型"""
+    id = fields.IntField(pk=True)
+    upload_id = fields.CharField(max_length=36, unique=True, index=True)
+    file_name = fields.CharField(max_length=255)
+    file_size = fields.BigIntField()
+    save_path = fields.CharField(max_length=512)
+    mode = fields.CharField(max_length=10)  # "direct" 或 "proxy"
+    expire_value = fields.IntField(default=1)
+    expire_style = fields.CharField(max_length=20, default="day")
+    created_at = fields.DatetimeField(auto_now_add=True)
+    expires_at = fields.DatetimeField()  # 会话过期时间
+
+    async def is_expired(self):
+        """检查会话是否已过期"""
+        return self.expires_at < await get_now()
+
+
 file_codes_pydantic = pydantic_model_creator(FileCodes, name="FileCodes")
 upload_chunk_pydantic = pydantic_model_creator(UploadChunk, name="UploadChunk")
 key_value_pydantic = pydantic_model_creator(KeyValue, name="KeyValue")
+presign_upload_session_pydantic = pydantic_model_creator(PresignUploadSession, name="PresignUploadSession")
@@ -1,4 +1,5 @@
 from pydantic import BaseModel
+from typing import Optional


 class SelectFileModel(BaseModel):
@@ -15,3 +16,21 @@ class InitChunkUploadModel(BaseModel):
 class CompleteUploadModel(BaseModel):
    expire_value: int
    expire_style: str
+
+
+# 预签名上传相关模型
+class PresignUploadInitRequest(BaseModel):
+    """预签名上传初始化请求"""
+    file_name: str
+    file_size: int
+    expire_value: int = 1
+    expire_style: str = "day"
+
+
+class PresignUploadInitResponse(BaseModel):
+    """预签名上传初始化响应"""
+    upload_id: str
+    upload_url: str
+    mode: str  # "direct" 或 "proxy"
+    save_path: str
+    expires_in: int  # URL过期时间（秒）
@@ -1,21 +1,70 @@
+import datetime
 import hashlib
+import os
 import uuid
+from datetime import timedelta
+from urllib.parse import unquote

 from fastapi import APIRouter, Form, UploadFile, File, Depends, HTTPException
 from starlette import status

 from apps.admin.dependencies import share_required_login
-from apps.base.models import FileCodes, UploadChunk
-from apps.base.schemas import SelectFileModel, InitChunkUploadModel, CompleteUploadModel
+from apps.base.models import FileCodes, UploadChunk, PresignUploadSession
+from apps.base.schemas import SelectFileModel, InitChunkUploadModel, CompleteUploadModel, PresignUploadInitRequest
 from apps.base.utils import get_expire_info, get_file_path_name, ip_limit, get_chunk_file_path_name
 from core.response import APIResponse
 from core.settings import settings
 from core.storage import storages, FileStorageInterface
-from core.utils import get_select_token
+from core.utils import get_select_token, get_now, sanitize_filename

 share_api = APIRouter(prefix="/share", tags=["分享"])


+# ============ 公共服务层 ============
+class FileUploadService:
+    """统一的文件上传服务"""
+
+    @staticmethod
+    async def generate_file_path(file_name: str, upload_id: str = None) -> tuple[str, str, str, str, str]:
+        """统一的路径生成"""
+        today = datetime.datetime.now()
+        storage_path = settings.storage_path.strip("/")
+        file_uuid = upload_id or uuid.uuid4().hex
+        filename = await sanitize_filename(unquote(file_name))
+        base_path = f"share/data/{today.strftime('%Y/%m/%d')}/{file_uuid}"
+        path = f"{storage_path}/{base_path}" if storage_path else base_path
+        prefix, suffix = os.path.splitext(filename)
+        save_path = f"{path}/{filename}"
+        return path, suffix, prefix, filename, save_path
+
+    @staticmethod
+    async def create_file_record(
+        file_name: str,
+        file_size: int,
+        file_path: str,
+        expire_value: int,
+        expire_style: str,
+        **extra_fields
+    ) -> str:
+        """统一创建FileCodes记录，返回code"""
+        expired_at, expired_count, used_count, code = await get_expire_info(expire_value, expire_style)
+        prefix, suffix = os.path.splitext(file_name)
+
+        await FileCodes.create(
+            code=code,
+            prefix=prefix,
+            suffix=suffix,
+            uuid_file_name=file_name,
+            file_path=file_path,
+            size=file_size,
+            expired_at=expired_at,
+            expired_count=expired_count,
+            used_count=used_count,
+            **extra_fields
+        )
+        return code
+
+
 async def validate_file_size(file: UploadFile, max_size: int) -> int:
    size = file.size
    if size is None:
@@ -425,3 +474,144 @@ async def complete_upload(upload_id: str, data: CompleteUploadModel, ip: str = D
        except Exception:
            pass
        raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"文件合并失败: {str(e)}")
+
+
+# ============ 预签名上传API ============
+presign_api = APIRouter(prefix="/presign", tags=["预签名上传"])
+
+PRESIGN_SESSION_EXPIRES = 900  # 15分钟
+
+
+async def _get_valid_session(upload_id: str, expected_mode: str = None) -> PresignUploadSession:
+    """获取并验证会话"""
+    session = await PresignUploadSession.filter(upload_id=upload_id).first()
+    if not session:
+        raise HTTPException(404, "上传会话不存在")
+    if await session.is_expired():
+        await session.delete()
+        raise HTTPException(404, "上传会话已过期")
+    if expected_mode and session.mode != expected_mode:
+        raise HTTPException(400, f"此会话不支持{expected_mode}模式")
+    return session
+
+
+@presign_api.post("/upload/init", dependencies=[Depends(share_required_login)])
+async def presign_upload_init(data: PresignUploadInitRequest, ip: str = Depends(ip_limit["upload"])):
+    """初始化预签名上传，S3返回直传URL，其他存储返回代理URL"""
+    if data.file_size > settings.uploadSize:
+        raise HTTPException(403, f"文件大小超过限制，最大为 {settings.uploadSize / (1024*1024):.2f} MB")
+    if data.expire_style not in settings.expireStyle:
+        raise HTTPException(400, "过期时间类型错误")
+
+    upload_id = uuid.uuid4().hex
+    path, _, _, filename, save_path = await FileUploadService.generate_file_path(data.file_name, upload_id)
+
+    storage: FileStorageInterface = storages[settings.file_storage]()
+    presigned_url = await storage.generate_presigned_upload_url(save_path, PRESIGN_SESSION_EXPIRES)
+
+    mode = "direct" if presigned_url else "proxy"
+    upload_url = presigned_url or f"/api/presign/upload/proxy/{upload_id}"
+
+    await PresignUploadSession.create(
+        upload_id=upload_id,
+        file_name=filename,
+        file_size=data.file_size,
+        save_path=save_path,
+        mode=mode,
+        expire_value=data.expire_value,
+        expire_style=data.expire_style,
+        expires_at=await get_now() + timedelta(seconds=PRESIGN_SESSION_EXPIRES),
+    )
+
+    ip_limit["upload"].add_ip(ip)
+    return APIResponse(detail={
+        "upload_id": upload_id,
+        "upload_url": upload_url,
+        "mode": mode,
+        "expires_in": PRESIGN_SESSION_EXPIRES,
+    })
+
+
+@presign_api.put("/upload/proxy/{upload_id}", dependencies=[Depends(share_required_login)])
+async def presign_upload_proxy(upload_id: str, file: UploadFile = File(...), ip: str = Depends(ip_limit["upload"])):
+    """代理模式上传，服务器转存到存储后端"""
+    session = await _get_valid_session(upload_id, expected_mode="proxy")
+
+    file_size = await validate_file_size(file, settings.uploadSize)
+    if abs(file_size - session.file_size) > 1024:
+        raise HTTPException(400, "文件大小与声明不符")
+
+    storage: FileStorageInterface = storages[settings.file_storage]()
+    try:
+        await storage.save_file(file, session.save_path)
+    except Exception as e:
+        raise HTTPException(500, f"文件保存失败: {str(e)}")
+
+    code = await FileUploadService.create_file_record(
+        session.file_name, file_size, os.path.dirname(session.save_path),
+        session.expire_value, session.expire_style
+    )
+
+    await session.delete()
+    ip_limit["upload"].add_ip(ip)
+    return APIResponse(detail={"code": code, "name": session.file_name})
+
+
+@presign_api.post("/upload/confirm/{upload_id}", dependencies=[Depends(share_required_login)])
+async def presign_upload_confirm(upload_id: str, ip: str = Depends(ip_limit["upload"])):
+    """直传确认，客户端完成S3直传后调用获取分享码"""
+    session = await _get_valid_session(upload_id, expected_mode="direct")
+
+    storage: FileStorageInterface = storages[settings.file_storage]()
+    if not await storage.file_exists(session.save_path):
+        raise HTTPException(404, "文件未上传或上传失败")
+
+    code = await FileUploadService.create_file_record(
+        session.file_name, session.file_size, os.path.dirname(session.save_path),
+        session.expire_value, session.expire_style
+    )
+
+    await session.delete()
+    ip_limit["upload"].add_ip(ip)
+    return APIResponse(detail={"code": code, "name": session.file_name})
+
+
+@presign_api.get("/upload/status/{upload_id}", dependencies=[Depends(share_required_login)])
+async def presign_upload_status(upload_id: str):
+    """查询上传会话状态"""
+    session = await PresignUploadSession.filter(upload_id=upload_id).first()
+    if not session:
+        raise HTTPException(404, "上传会话不存在")
+
+    return APIResponse(detail={
+        "upload_id": session.upload_id,
+        "file_name": session.file_name,
+        "file_size": session.file_size,
+        "mode": session.mode,
+        "created_at": session.created_at.isoformat(),
+        "expires_at": session.expires_at.isoformat(),
+        "is_expired": await session.is_expired(),
+    })
+
+
+@presign_api.delete("/upload/{upload_id}", dependencies=[Depends(share_required_login)])
+async def presign_upload_cancel(upload_id: str):
+    """取消上传会话"""
+    session = await PresignUploadSession.filter(upload_id=upload_id).first()
+    if not session:
+        raise HTTPException(404, "上传会话不存在")
+
+    if session.mode == "direct":
+        storage: FileStorageInterface = storages[settings.file_storage]()
+        try:
+            if await storage.file_exists(session.save_path):
+                temp_file_code = FileCodes(
+                    file_path=os.path.dirname(session.save_path),
+                    uuid_file_name=os.path.basename(session.save_path),
+                )
+                await storage.delete_file(temp_file_code)
+        except Exception:
+            pass
+
+    await session.delete()
+    return APIResponse(detail={"message": "上传会话已取消"})