diff --git a/recognization/code/transport.py b/recognization/code/transport.py index 292a95e..66e94a1 100644 --- a/recognization/code/transport.py +++ b/recognization/code/transport.py @@ -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"