tngo/product-recognition-demo
1
0
Fork 0
forked from vnguyen/product-recognition-demo
Code Pull requests Activity

fix: reject oversized bodies before app response

Browse source
This commit is contained in:
kiyo9w 2026-09-28 12:40:59 +09:00
parent 1794640f9d
commit f2345e482c
1 changed files with 19 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

42
recognization/code/transport.py
View file

@ -116,41 +116,37 @@ class BodyLimit:
if scope.get("type") != "http":
await self.app(scope, receive, send)
return
seen = 0
blocked = False
class TooLarge(Exception):
pass
seen = 0
async def limited() -> Dict[str, Any]:
nonlocal seen, blocked
if blocked:
return {"type": "http.request", "body": b"", "more_body": False}
nonlocal seen
message = await receive()
if message.get("type") != "http.request":
return message
body = message.get("body") or b""
seen += len(body)
if seen > self.max_bytes:
blocked = True
await send(
{
"type": "http.response.start",
"status": 413,
"headers": [(b"content-type", b"application/json")],
}
)
await send(
{
"type": "http.response.body",
"body": b'{"error":{"code":"PAYLOAD_TOO_LARGE","message":"request body too large"}}',
}
)
return {"type": "http.request", "body": b"", "more_body": False}
raise TooLarge
return message
try:
await self.app(scope, limited, send)
except RuntimeError:
if not blocked:
raise
except TooLarge:
await send(
{
"type": "http.response.start",
"status": 413,
"headers": [(b"content-type", b"application/json")],
}
)
await send(
{
"type": "http.response.body",
"body": b'{"error":{"code":"PAYLOAD_TOO_LARGE","message":"request body too large"}}',
}
)
SERVICE_TYPE = "_dorayaki._tcp"