33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
from fastapi import Request, HTTPException
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
from app.core.config import get_settings
|
|
|
|
settings = get_settings()
|
|
|
|
|
|
class SecurityMiddleware(BaseHTTPMiddleware):
|
|
async def dispatch(self, request: Request, call_next):
|
|
if request.method in ("POST", "PUT", "PATCH"):
|
|
content_length = request.headers.get("content-length")
|
|
if content_length and int(content_length) > settings.MAX_UPLOAD_SIZE:
|
|
raise HTTPException(
|
|
status_code=413,
|
|
detail=f"文件大小不能超过 {settings.MAX_UPLOAD_SIZE // 1024 // 1024}MB",
|
|
)
|
|
response = await call_next(request)
|
|
return response
|
|
|
|
|
|
ALLOWED_EXTENSIONS = {".docx", ".txt", ".pdf"}
|
|
|
|
|
|
def validate_file_extension(filename: str) -> str:
|
|
import os
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
if ext not in ALLOWED_EXTENSIONS:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail=f"不支持的文件类型: {ext}，允许的类型: {', '.join(ALLOWED_EXTENSIONS)}",
|
|
)
|
|
return ext
|