72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
"""限制请求体大小：声明的 Content-Length 超限直接 413，没声明（分块）或声明得不实的，边读边数
|
|
|
|
FastAPI 会先把整个请求体收完再调用接口函数，超大的请求体已经占满了带宽和内存
|
|
和数据文件上传的 UploadSizeLimit 同一出发点；限制随路径而不同，所以用规则表，第一条命中的生效
|
|
"""
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
|
|
from fastapi.responses import JSONResponse
|
|
from starlette.exceptions import HTTPException
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
from app.common.response import error_content
|
|
|
|
|
|
class BodyTooLarge(HTTPException):
|
|
"""边读边数时超限；message 由全局 HTTP 异常处理放进信封，和声明超限的响应文案一致"""
|
|
|
|
def __init__(self, message: str):
|
|
super().__init__(413)
|
|
self.message = message
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BodyLimitRule:
|
|
methods: frozenset[str]
|
|
path: re.Pattern[str]
|
|
max_bytes: int
|
|
message: str
|
|
|
|
|
|
class BodyLimit:
|
|
def __init__(self, app: ASGIApp, rules: list[BodyLimitRule]):
|
|
self.app = app
|
|
self.rules = rules
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
rule = self._match(scope) if scope["type"] == "http" else None
|
|
if rule is None:
|
|
await self.app(scope, receive, send)
|
|
return
|
|
length = dict(scope["headers"]).get(b"content-length", b"")
|
|
if length.isdigit() and int(length) > rule.max_bytes:
|
|
response = JSONResponse(error_content(413, rule.message), 413)
|
|
await response(scope, receive, send)
|
|
return
|
|
received = 0
|
|
|
|
async def counted() -> Message:
|
|
nonlocal received
|
|
message = await receive()
|
|
if message["type"] == "http.request":
|
|
received += len(message.get("body", b""))
|
|
if received > rule.max_bytes:
|
|
# 抛出的异常由应用里统一的 HTTP 异常处理转成 413 信封
|
|
raise BodyTooLarge(rule.message)
|
|
return message
|
|
|
|
await self.app(scope, counted, send)
|
|
|
|
def _match(self, scope: Scope) -> BodyLimitRule | None:
|
|
return next(
|
|
(
|
|
rule
|
|
for rule in self.rules
|
|
if scope["method"] in rule.methods
|
|
and rule.path.fullmatch(scope["path"])
|
|
),
|
|
None,
|
|
)
|