shanxia2.0/backend/app/core/doclib_upload_route.py

"""在 FastAPI 解析 multipart 前限制文档请求流，保护临时磁盘。

认证仍由 get_current_user 执行；这里只对缺少 Bearer 的请求提前拒绝。
MultiPartException 让 Starlette 关闭已创建的临时文件，再统一转为 413。
"""
from fastapi import HTTPException
from fastapi.routing import APIRoute
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.formparsers import MultiPartException
from starlette.requests import Request

from app.core.config import settings


class DoclibRoute(APIRoute):
    def get_route_handler(self):
        handler = super().get_route_handler()
        if not self.path.endswith('/web/doclib/upload'):
            return handler

        async def limited_upload(request: Request):
            authorization = request.headers.get('authorization', '')
            if not authorization.lower().startswith('bearer ') or not authorization[7:].strip():
                raise HTTPException(401, '请先登录', headers={'WWW-Authenticate': 'Bearer'})
            if 'multipart/form-data' not in request.headers.get('content-type', '').lower():
                raise HTTPException(415, '文档上传仅支持 multipart/form-data')
            # 总请求体允许 64 KiB 的 multipart 边界/字段开销；文件本体仍精确校验。
            maximum = settings.DOCLIB_MAX_UPLOAD_BYTES + 64 * 1024
            length = request.headers.get('content-length')
            if length:
                try:
                    declared = int(length)
                except ValueError:
                    raise HTTPException(400, '无效 Content-Length')
                if declared < 0:
                    raise HTTPException(400, '无效 Content-Length')
                if declared > maximum:
                    raise HTTPException(413, '上传内容超过文档大小限制')
            received = 0
            exceeded = False

            async def receive():
                nonlocal received, exceeded
                message = await request.receive()
                if message['type'] == 'http.request':
                    received += len(message.get('body', b''))
                    if received > maximum:
                        exceeded = True
                        raise MultiPartException('上传内容超过文档大小限制')
                return message

            bounded = Request(request.scope, receive)
            try:
                return await handler(bounded)
            except StarletteHTTPException:
                if exceeded:
                    raise HTTPException(413, '上传内容超过文档大小限制')
                raise

        return limited_upload