61 lines
2.6 KiB
Python
61 lines
2.6 KiB
Python
"""在 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
|