feat: s3适配切片上传

Author: vastsa

apps/base/models.py

@@ -51,6 +51,7 @@ class UploadChunk(models.Model):
     file_name = fields.CharField(max_length=255)
     created_at = fields.DatetimeField(auto_now_add=True)
     completed = fields.BooleanField(default=False)
+    actual_size = fields.IntField(null=True)  # 实际接收的分片大小
 
 
 class KeyValue(Model):

apps/base/views.py

@@ -164,6 +164,16 @@ async def download_file(key: str, code: str, ip: str = Depends(ip_limit["error"]
 
 @chunk_api.post("/upload/init/", dependencies=[Depends(share_required_login)])
 async def init_chunk_upload(data: InitChunkUploadModel):
+    # 服务端校验：根据 total_chunks * chunk_size 计算理论最大上传量
+    total_chunks = (data.file_size + data.chunk_size - 1) // data.chunk_size
+    max_possible_size = total_chunks * data.chunk_size
+    if max_possible_size > settings.uploadSize:
+        max_size_mb = settings.uploadSize / (1024 * 1024)
+        raise HTTPException(
+            status_code=403,
+            detail=f"文件大小超过限制，最大为 {max_size_mb:.2f} MB"
+        )
+
     # # 秒传检查
     # existing = await FileCodes.filter(file_hash=data.file_hash).first()
     # if existing:
@@ -179,9 +189,30 @@ async def init_chunk_upload(data: InitChunkUploadModel):
     #             "name": f'{existing.prefix}{existing.suffix}'
     #         })
 
-    # 创建上传会话
+    # 断点续传：检查是否存在相同文件的未完成上传会话
+    existing_session = await UploadChunk.filter(
+        chunk_hash=data.file_hash,
+        chunk_index=-1,
+        file_size=data.file_size,
+        file_name=data.file_name,
+    ).first()
+
+    if existing_session:
+        # 复用已有会话，获取已上传的分片列表
+        uploaded_chunks = await UploadChunk.filter(
+            upload_id=existing_session.upload_id,
+            completed=True
+        ).values_list('chunk_index', flat=True)
+        return APIResponse(detail={
+            "existed": False,
+            "upload_id": existing_session.upload_id,
+            "chunk_size": existing_session.chunk_size,
+            "total_chunks": existing_session.total_chunks,
+            "uploaded_chunks": list(uploaded_chunks)
+        })
+
+    # 创建新的上传会话
     upload_id = uuid.uuid4().hex
-    total_chunks = (data.file_size + data.chunk_size - 1) // data.chunk_size
     await UploadChunk.create(
         upload_id=upload_id,
         chunk_index=-1,
@@ -191,17 +222,12 @@ async def init_chunk_upload(data: InitChunkUploadModel):
         chunk_hash=data.file_hash,
         file_name=data.file_name,
     )
-    # 获取已上传的分片列表
-    uploaded_chunks = await UploadChunk.filter(
-        upload_id=upload_id,
-        completed=True
-    ).values_list('chunk_index', flat=True)
     return APIResponse(detail={
         "existed": False,
         "upload_id": upload_id,
         "chunk_size": data.chunk_size,
         "total_chunks": total_chunks,
-        "uploaded_chunks": uploaded_chunks
+        "uploaded_chunks": []
     })
 
 
@@ -220,11 +246,52 @@ async def upload_chunk(
     if chunk_index < 0 or chunk_index >= chunk_info.total_chunks:
         raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="无效的分片索引")
 
+    # 检查是否已上传（支持断点续传）
+    existing_chunk = await UploadChunk.filter(
+        upload_id=upload_id,
+        chunk_index=chunk_index,
+        completed=True
+    ).first()
+    if existing_chunk:
+        return APIResponse(detail={"chunk_hash": existing_chunk.chunk_hash, "skipped": True})
+
     # 读取分片数据并计算哈希
     chunk_data = await chunk.read()
+    chunk_size = len(chunk_data)
+
+    # 校验分片大小不超过声明的 chunk_size
+    if chunk_size > chunk_info.chunk_size:
+        raise HTTPException(
+            status.HTTP_400_BAD_REQUEST,
+            detail=f"分片大小超过声明值: 最大 {chunk_info.chunk_size}, 实际 {chunk_size}"
+        )
+
+    # 计算已上传的实际总大小（防止绕过前端限制）
+    uploaded_chunks_list = await UploadChunk.filter(
+        upload_id=upload_id,
+        completed=True
+    ).all()
+    uploaded_size = sum(c.actual_size for c in uploaded_chunks_list if c.actual_size)
+    if uploaded_size + chunk_size > settings.uploadSize:
+        max_size_mb = settings.uploadSize / (1024 * 1024)
+        raise HTTPException(
+            status_code=403,
+            detail=f"累计上传大小超过限制，最大为 {max_size_mb:.2f} MB"
+        )
+    
     chunk_hash = hashlib.sha256(chunk_data).hexdigest()
 
-    # 更新或创建分片记录
+    # 获取文件路径
+    _, _, _, _, save_path = await get_chunk_file_path_name(chunk_info.file_name, upload_id)
+    
+    # 保存分片到存储
+    storage = storages[settings.file_storage]()
+    try:
+        await storage.save_chunk(upload_id, chunk_index, chunk_data, chunk_hash, save_path)
+    except Exception as e:
+        raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"分片保存失败: {str(e)}")
+    
+    # 更新或创建分片记录（保存成功后再记录，包含实际大小）
     await UploadChunk.update_or_create(
         upload_id=upload_id,
         chunk_index=chunk_index,
@@ -234,15 +301,59 @@ async def upload_chunk(
             'file_size': chunk_info.file_size,
             'total_chunks': chunk_info.total_chunks,
             'chunk_size': chunk_info.chunk_size,
-            'file_name': chunk_info.file_name
+            'file_name': chunk_info.file_name,
+            'actual_size': chunk_size
         }
     )
+    return APIResponse(detail={"chunk_hash": chunk_hash})
+
+
+@chunk_api.delete("/upload/{upload_id}", dependencies=[Depends(share_required_login)])
+async def cancel_upload(upload_id: str):
+    """取消上传并清理临时文件"""
+    chunk_info = await UploadChunk.filter(upload_id=upload_id, chunk_index=-1).first()
+    if not chunk_info:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="上传会话不存在")
+    
     # 获取文件路径
     _, _, _, _, save_path = await get_chunk_file_path_name(chunk_info.file_name, upload_id)
-    # 保存分片到存储
+    
+    # 清理存储中的临时文件
     storage = storages[settings.file_storage]()
-    await storage.save_chunk(upload_id, chunk_index, chunk_data, chunk_hash, save_path)
-    return APIResponse(detail={"chunk_hash": chunk_hash})
+    try:
+        await storage.clean_chunks(upload_id, save_path)
+    except Exception as e:
+        # 记录错误但不阻止删除数据库记录
+        pass
+    
+    # 清理数据库记录
+    await UploadChunk.filter(upload_id=upload_id).delete()
+    
+    return APIResponse(detail={"message": "上传已取消"})
+
+
+@chunk_api.get("/upload/status/{upload_id}", dependencies=[Depends(share_required_login)])
+async def get_upload_status(upload_id: str):
+    """获取上传状态"""
+    chunk_info = await UploadChunk.filter(upload_id=upload_id, chunk_index=-1).first()
+    if not chunk_info:
+        raise HTTPException(status.HTTP_404_NOT_FOUND, detail="上传会话不存在")
+    
+    # 获取已上传的分片列表
+    uploaded_chunks = await UploadChunk.filter(
+        upload_id=upload_id,
+        completed=True
+    ).values_list('chunk_index', flat=True)
+    
+    return APIResponse(detail={
+        "upload_id": upload_id,
+        "file_name": chunk_info.file_name,
+        "file_size": chunk_info.file_size,
+        "chunk_size": chunk_info.chunk_size,
+        "total_chunks": chunk_info.total_chunks,
+        "uploaded_chunks": list(uploaded_chunks),
+        "progress": len(uploaded_chunks) / chunk_info.total_chunks * 100
+    })
 
 
 @chunk_api.post("/upload/complete/{upload_id}", dependencies=[Depends(share_required_login)])
@@ -254,32 +365,63 @@ async def complete_upload(upload_id: str, data: CompleteUploadModel, ip: str = D
 
     storage = storages[settings.file_storage]()
     # 验证所有分片
-    completed_chunks = await UploadChunk.filter(
+    completed_chunks_list = await UploadChunk.filter(
         upload_id=upload_id,
         completed=True
-    ).count()
-    if completed_chunks != chunk_info.total_chunks:
+    ).all()
+    if len(completed_chunks_list) != chunk_info.total_chunks:
         raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="分片不完整")
+
+    # 计算实际上传的总大小并校验
+    actual_total_size = sum(c.actual_size for c in completed_chunks_list if c.actual_size)
+    if actual_total_size > settings.uploadSize:
+        # 清理已上传的分片
+        _, _, _, _, save_path = await get_chunk_file_path_name(chunk_info.file_name, upload_id)
+        try:
+            await storage.clean_chunks(upload_id, save_path)
+        except Exception:
+            pass
+        await UploadChunk.filter(upload_id=upload_id).delete()
+        max_size_mb = settings.uploadSize / (1024 * 1024)
+        raise HTTPException(
+            status_code=403,
+            detail=f"实际上传大小超过限制，最大为 {max_size_mb:.2f} MB"
+        )
+
     # 获取文件路径
     path, suffix, prefix, _, save_path = await get_chunk_file_path_name(chunk_info.file_name, upload_id)
-    # 合并文件并计算哈希
-    await storage.merge_chunks(upload_id, chunk_info, save_path)
-    # 创建文件记录
-    expired_at, expired_count, used_count, code = await get_expire_info(data.expire_value, data.expire_style)
-    await FileCodes.create(
-        code=code,
-        file_hash=chunk_info.chunk_hash,
-        is_chunked=True,
-        upload_id=upload_id,
-        size=chunk_info.file_size,
-        expired_at=expired_at,
-        expired_count=expired_count,
-        used_count=used_count,
-        file_path=path,
-        uuid_file_name=f"{prefix}{suffix}",
-        prefix=prefix,
-        suffix=suffix
-    )
-    # 清理临时文件
-    await storage.clean_chunks(upload_id, save_path)
-    return APIResponse(detail={"code": code, "name": chunk_info.file_name})
+    
+    try:
+        # 合并文件并计算哈希
+        _, file_hash = await storage.merge_chunks(upload_id, chunk_info, save_path)
+        # 创建文件记录
+        expired_at, expired_count, used_count, code = await get_expire_info(data.expire_value, data.expire_style)
+        await FileCodes.create(
+            code=code,
+            file_hash=file_hash,  # 使用合并后计算的哈希
+            is_chunked=True,
+            upload_id=upload_id,
+            size=actual_total_size,  # 使用实际上传大小而非前端声明值
+            expired_at=expired_at,
+            expired_count=expired_count,
+            used_count=used_count,
+            file_path=path,
+            uuid_file_name=f"{prefix}{suffix}",
+            prefix=prefix,
+            suffix=suffix
+        )
+        # 清理临时文件
+        await storage.clean_chunks(upload_id, save_path)
+        # 清理数据库中的分片记录
+        await UploadChunk.filter(upload_id=upload_id).delete()
+        ip_limit["upload"].add_ip(ip)
+        return APIResponse(detail={"code": code, "name": chunk_info.file_name})
+    except ValueError as e:
+        raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=str(e))
+    except Exception as e:
+        # 合并失败时清理临时文件
+        try:
+            await storage.clean_chunks(upload_id, save_path)
+        except Exception:
+            pass
+        raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"文件合并失败: {str(e)}")