From 3b23283d0b8efe35be4508ad565575bf50ae21cd Mon Sep 17 00:00:00 2001 From: "alex.oleshkevich" Date: Tue, 6 Jun 2023 12:07:55 +0200 Subject: [PATCH 1/5] add request_max_size option --- starlette/applications.py | 3 +++ starlette/routing.py | 24 ++++++++++++++++++++++-- tests/test_applications.py | 20 ++++++++++++++++++++ tests/test_routing.py | 37 +++++++++++++++++++++++++++++++++++++ 4 files changed, 82 insertions(+), 2 deletions(-) diff --git a/starlette/applications.py b/starlette/applications.py index 013364be3..232bfa77e 100644 --- a/starlette/applications.py +++ b/starlette/applications.py @@ -61,6 +61,7 @@ def __init__( on_startup: typing.Optional[typing.Sequence[typing.Callable]] = None, on_shutdown: typing.Optional[typing.Sequence[typing.Callable]] = None, lifespan: typing.Optional[Lifespan["AppType"]] = None, + request_max_size: int = 2621440, # 2.5mb ) -> None: # The lifespan context function is a newer style that replaces # on_startup / on_shutdown handlers. Use one or the other, not both. @@ -78,6 +79,7 @@ def __init__( ) self.user_middleware = [] if middleware is None else list(middleware) self.middleware_stack: typing.Optional[ASGIApp] = None + self.request_max_size = request_max_size def build_middleware_stack(self) -> ASGIApp: debug = self.debug @@ -117,6 +119,7 @@ def url_path_for(self, __name: str, **path_params: typing.Any) -> URLPath: async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: scope["app"] = self + scope["request_max_size"] = self.request_max_size if self.middleware_stack is None: self.middleware_stack = self.build_middleware_stack() await self.middleware_stack(scope, receive, send) diff --git a/starlette/routing.py b/starlette/routing.py index 52cf174e1..f9c056795 100644 --- a/starlette/routing.py +++ b/starlette/routing.py @@ -17,7 +17,7 @@ from starlette.middleware import Middleware from starlette.requests import Request from starlette.responses import PlainTextResponse, RedirectResponse -from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send +from starlette.types import ASGIApp, Lifespan, Message, Receive, Scope, Send from starlette.websockets import WebSocket, WebSocketClose @@ -205,12 +205,14 @@ def __init__( methods: typing.Optional[typing.List[str]] = None, name: typing.Optional[str] = None, include_in_schema: bool = True, + request_max_size: typing.Optional[int] = None, ) -> None: assert path.startswith("/"), "Routed paths must start with '/'" self.path = path self.endpoint = endpoint self.name = get_name(endpoint) if name is None else name self.include_in_schema = include_in_schema + self.request_max_size = request_max_size endpoint_handler = endpoint while isinstance(endpoint_handler, functools.partial): @@ -273,7 +275,25 @@ async def handle(self, scope: Scope, receive: Receive, send: Send) -> None: ) await response(scope, receive, send) else: - await self.app(scope, receive, send) + bytes_read = 0 + max_body_size = self.request_max_size or scope.get("request_max_size") + + async def receive_wrapper() -> Message: + nonlocal bytes_read + message = await receive() + if message["type"] != "http.request" or max_body_size is None: + return message + + body = message.get("body", b"") + bytes_read += len(body) + if bytes_read > max_body_size: + raise HTTPException( + status_code=413, detail="Request Entity Too Large" + ) + + return message + + await self.app(scope, receive_wrapper, send) def __eq__(self, other: typing.Any) -> bool: return ( diff --git a/tests/test_applications.py b/tests/test_applications.py index e30ec9295..8fad636d2 100644 --- a/tests/test_applications.py +++ b/tests/test_applications.py @@ -207,6 +207,26 @@ def test_500(test_client_factory): assert response.json() == {"detail": "Server Error"} +def test_413(test_client_factory): + async def read_view(request): + content = await request.body() + return JSONResponse(content.decode()) + + app = Starlette( + request_max_size=10, # 10 bytes + routes=[Route("/", endpoint=read_view, methods=["POST"])], + ) + + client = test_client_factory(app, raise_server_exceptions=True) + response = client.post("/", data=b"youshallnotpass") + assert response.status_code == 413 + assert response.text == "Request Entity Too Large" + + response = client.post("/", data=b"ok") + assert response.status_code == 200 + assert response.text == '"ok"' + + def test_websocket_raise_websocket_exception(client): with client.websocket_connect("/ws-raise-websocket") as session: response = session.receive() diff --git a/tests/test_routing.py b/tests/test_routing.py index 298745d40..f38ebac09 100644 --- a/tests/test_routing.py +++ b/tests/test_routing.py @@ -1125,3 +1125,40 @@ async def startup() -> None: ... # pragma: nocover router.on_event("startup")(startup) + + +def test_request_body_size_limit(test_client_factory): + async def read_view(request): + content = await request.body() + return JSONResponse(content.decode()) + + app = Starlette( + routes=[Route("/", endpoint=read_view, methods=["POST"], request_max_size=10)], + ) + + client = test_client_factory(app, raise_server_exceptions=True) + response = client.post("/", data=b"youshallnotpass") + assert response.status_code == 413 + assert response.text == "Request Entity Too Large" + + response = client.post("/", data=b"ok") + assert response.status_code == 200 + assert response.text == '"ok"' + + +def test_request_body_size_limit_route_has_higher_precedense(test_client_factory): + async def read_view(request): + content = await request.body() + return JSONResponse(content.decode()) + + app = Starlette( + request_max_size=5, + routes=[Route("/", endpoint=read_view, methods=["POST"], request_max_size=24)], + ) + + # app caps at 5 bytes, route has 24 byte limit + # payload of size 12 bytes should pass + client = test_client_factory(app, raise_server_exceptions=True) + response = client.post("/", data=b"youshallpass") # 12 bytes + assert response.status_code == 200 + assert response.text == '"youshallpass"' From 1882aa1b828fe1e97f934ae53eca62b389fa4f62 Mon Sep 17 00:00:00 2001 From: "alex.oleshkevich" Date: Thu, 22 Jun 2023 17:24:29 +0200 Subject: [PATCH 2/5] add prefix to scope attribute --- starlette/applications.py | 2 +- starlette/routing.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/starlette/applications.py b/starlette/applications.py index 232bfa77e..b3771f175 100644 --- a/starlette/applications.py +++ b/starlette/applications.py @@ -119,7 +119,7 @@ def url_path_for(self, __name: str, **path_params: typing.Any) -> URLPath: async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: scope["app"] = self - scope["request_max_size"] = self.request_max_size + scope["starlette.request_max_size"] = self.request_max_size if self.middleware_stack is None: self.middleware_stack = self.build_middleware_stack() await self.middleware_stack(scope, receive, send) diff --git a/starlette/routing.py b/starlette/routing.py index f9c056795..ba9da9aed 100644 --- a/starlette/routing.py +++ b/starlette/routing.py @@ -276,7 +276,9 @@ async def handle(self, scope: Scope, receive: Receive, send: Send) -> None: await response(scope, receive, send) else: bytes_read = 0 - max_body_size = self.request_max_size or scope.get("request_max_size") + max_body_size = self.request_max_size or scope.get( + "starlette.request_max_size" + ) async def receive_wrapper() -> Message: nonlocal bytes_read From cae3aaaf758c02102b2e77ac2cfeb0e5bf328b0b Mon Sep 17 00:00:00 2001 From: "alex.oleshkevich" Date: Wed, 28 Jun 2023 21:00:01 +0200 Subject: [PATCH 3/5] document `request_max_size` application parameter --- starlette/applications.py | 1 + 1 file changed, 1 insertion(+) diff --git a/starlette/applications.py b/starlette/applications.py index 109804f32..d983fcf62 100644 --- a/starlette/applications.py +++ b/starlette/applications.py @@ -42,6 +42,7 @@ class Starlette: * **lifespan** - A lifespan context function, which can be used to perform startup and shutdown tasks. This is a newer style that replaces the `on_startup` and `on_shutdown` handlers. Use one or the other, not both. + * **request_max_size** - Integer (in bytes) that sets the maximum request body size. """ def __init__( From c49e492aa36b3af1cf0e13832331e4db50821760 Mon Sep 17 00:00:00 2001 From: "alex.oleshkevich" Date: Wed, 28 Jun 2023 21:02:38 +0200 Subject: [PATCH 4/5] document request_max_size in routing --- docs/routing.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/routing.md b/docs/routing.md index fd1558793..7ce81773d 100644 --- a/docs/routing.md +++ b/docs/routing.md @@ -331,3 +331,12 @@ The `endpoint` argument can be one of: * An async function, which accepts a single `websocket` argument. * A class that implements the ASGI interface, such as Starlette's [WebSocketEndpoint](endpoints.md#websocketendpoint). + + +## Request body size limit + +Pass `request_max_size` to override application-wide request body limit for a particulal route: + +```python +Route('/', homepage, request_max_size=1024**8), # 8mb +``` From d269fd5f8ddec7e1bacf6d5a525bd8dc328fe7e9 Mon Sep 17 00:00:00 2001 From: "alex.oleshkevich" Date: Wed, 28 Jun 2023 21:04:00 +0200 Subject: [PATCH 5/5] fix formatting --- starlette/applications.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/starlette/applications.py b/starlette/applications.py index d983fcf62..8338ac14f 100644 --- a/starlette/applications.py +++ b/starlette/applications.py @@ -42,7 +42,7 @@ class Starlette: * **lifespan** - A lifespan context function, which can be used to perform startup and shutdown tasks. This is a newer style that replaces the `on_startup` and `on_shutdown` handlers. Use one or the other, not both. - * **request_max_size** - Integer (in bytes) that sets the maximum request body size. + * **request_max_size** - Integer (in bytes) that sets the maximum request body size. """ def __init__(