From 36068f49fad0d4e3d8c228db021a68ef389229a8 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Fri, 24 Nov 2023 14:20:26 +0100 Subject: [PATCH 1/8] Add `LimitRequestMiddleware` --- docs/middleware.md | 30 +++++++++-- starlette/middleware/limits.py | 66 ++++++++++++++++++++++++ tests/middleware/test_limits.py | 90 +++++++++++++++++++++++++++++++++ 3 files changed, 183 insertions(+), 3 deletions(-) create mode 100644 starlette/middleware/limits.py create mode 100644 tests/middleware/test_limits.py diff --git a/docs/middleware.md b/docs/middleware.md index 92ac5886a..4ca1971ad 100644 --- a/docs/middleware.md +++ b/docs/middleware.md @@ -106,8 +106,7 @@ The following arguments are supported: * `max_age` - Session expiry time in seconds. Defaults to 2 weeks. If set to `None` then the cookie will last as long as the browser session. * `same_site` - SameSite flag prevents the browser from sending session cookie along with cross-site requests. Defaults to `'lax'`. * `https_only` - Indicate that Secure flag should be set (can be used with HTTPS only). Defaults to `False`. -* `domain` - Domain of the cookie used to share cookie between subdomains or cross-domains. The browser defaults the domain to the same host that set the cookie, excluding subdomains [refrence](https://developer.mozilla.org/en-US/docs/Web/HTTP/Cookies#domain_attribute). - +* `domain` - Domain of the cookie used to share cookie between subdomains or cross-domains. The browser defaults the domain to the same host that set the cookie, excluding subdomains [refrence](https://developer.mozilla.org/en-US/docs/Web/HTTP/Cookies#domain_attribute). ## HTTPSRedirectMiddleware @@ -184,6 +183,31 @@ The following arguments are supported: The middleware won't GZip responses that already have a `Content-Encoding` set, to prevent them from being encoded twice. +## LimitRequestMiddleware + +Limits the body size of incoming requests. If the incoming request has a body +larger than the limit, then a `413 Content Too Large` response will be sent. + +```python +from starlette.applications import Starlette +from starlette.middleware import Middleware +from starlette.middleware.limits import LimitRequestMiddleware + + +routes = ... + +middleware = [ + Middleware(LimitRequestMiddleware, max_body_size=1024) +] + +app = Starlette(routes=routes, middleware=middleware) +``` + +The following arguments are supported: + +* `max_body_size` - Send a "413 - Content Too Large" on requests that surpass the maximum allowed body size. + Defaults to `2.5 * 1024 * 1024` bytes (2.5MB). + ## BaseHTTPMiddleware An abstract class that allows you to write ASGI middleware against a request/response @@ -573,7 +597,7 @@ import time class MonitoringMiddleware: def __init__(self, app): self.app = app - + async def __call__(self, scope, receive, send): start = time.time() try: diff --git a/starlette/middleware/limits.py b/starlette/middleware/limits.py new file mode 100644 index 000000000..919a51fae --- /dev/null +++ b/starlette/middleware/limits.py @@ -0,0 +1,66 @@ +"""Middleware that limits the body size of incoming requests.""" +from starlette.datastructures import Headers +from starlette.responses import PlainTextResponse +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +DEFAULT_MAX_BODY_SIZE = 2_621_440 # 2.5MB + + +class ContentTooLarge(Exception): + def __init__(self, max_body_size: int) -> None: + self.max_body_size = max_body_size + + +class LimitRequestMiddleware: + def __init__( + self, app: ASGIApp, max_body_size: int = DEFAULT_MAX_BODY_SIZE + ) -> None: + self.app = app + self.max_body_size = max_body_size + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": # pragma: no cover + return await self.app(scope, receive, send) + + headers = Headers(scope=scope) + content_length = headers.get("content-length") + if content_length is not None: + if int(content_length) > self.max_body_size: + return await _content_too_large_app(scope)(scope, receive, send) + + # NOTE: The server makes sure that the content-length header sent by the + # client is the same as the length of the body. + # Ref.: https://github.com/django/asgiref/issues/422 + return await self.app(scope, receive, send) + + total_size = 0 + response_started = False + + async def wrap_send(message: Message) -> None: + nonlocal response_started + if message["type"] == "http.response.start": + response_started = True + await send(message) + + async def wrap_receive() -> Message: + nonlocal total_size + message = await receive() + if message["type"] == "http.request": + chunk_size = len(message["body"]) + total_size += chunk_size + if total_size > self.max_body_size: + raise ContentTooLarge(self.max_body_size) + return message + + try: + await self.app(scope, wrap_receive, wrap_send) + except ContentTooLarge as exc: + # NOTE: If response has already started, we can't return a 413, because the + # headers have already been sent. + if not response_started: + return await _content_too_large_app(scope)(scope, receive, send) + raise exc + + +def _content_too_large_app(scope: Scope) -> PlainTextResponse: + return PlainTextResponse("Content Too Large", status_code=413) diff --git a/tests/middleware/test_limits.py b/tests/middleware/test_limits.py new file mode 100644 index 000000000..cadb27bb0 --- /dev/null +++ b/tests/middleware/test_limits.py @@ -0,0 +1,90 @@ +from typing import AsyncGenerator, Callable + +import pytest + +from starlette.applications import Starlette +from starlette.middleware import Middleware +from starlette.middleware.limits import ContentTooLarge, LimitRequestMiddleware +from starlette.requests import Request +from starlette.routing import Route +from starlette.testclient import TestClient +from starlette.types import Message, Receive, Scope, Send + + +async def echo_app(scope: Scope, receive: Receive, send: Send) -> None: + while True: + message = await receive() + more_body = message.get("more_body", False) + body = message.get("body", b"") + + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": body, "more_body": more_body}) + + if not more_body: + break + + +app = LimitRequestMiddleware(echo_app, max_body_size=1024) + + +def test_no_op(test_client_factory: Callable[..., TestClient]) -> None: + client = test_client_factory(app) + + response = client.post("/", content="Small payload") + assert response.status_code == 200 + assert response.text == "Small payload" + + +def test_content_too_large(test_client_factory: Callable[..., TestClient]) -> None: + client = test_client_factory(app) + + response = client.post("/", content="X" * 1025) + assert response.status_code == 413 + assert response.text == "Content Too Large" + + +def test_content_too_large_on_streaming_body( + test_client_factory: Callable[..., TestClient] +) -> None: + client = test_client_factory(app) + + response = client.post("/", content=[b"X" * 1025]) + assert response.status_code == 413 + assert response.text == "Content Too Large" + + +@pytest.mark.anyio +async def test_content_too_large_on_started_response() -> None: + scope: Scope = {"type": "http", "method": "POST", "path": "/", "headers": []} + + async def receive() -> AsyncGenerator[Message, None]: + yield {"type": "http.request", "body": b"X" * 1024, "more_body": True} + yield {"type": "http.request", "body": b"X", "more_body": False} + + async def send(message: Message) -> None: + ... + + rcv = receive() + + with pytest.raises(ContentTooLarge) as ctx: + await app(scope, rcv.__anext__, send) + assert ctx.value.max_body_size == 1024 + + await rcv.aclose() + + +def test_content_too_large_on_starlette( + test_client_factory: Callable[..., TestClient] +) -> None: + async def read_body_endpoint(request: Request) -> None: + await request.body() + + app = Starlette( + routes=[Route("/", read_body_endpoint, methods=["POST"])], + middleware=[Middleware(LimitRequestMiddleware, max_body_size=1024)], + ) + client = test_client_factory(app) + + response = client.post("/", content=[b"X" * 1024, b"X"]) + assert response.status_code == 413 + assert response.text == "Content Too Large" From 4d0415d354e13d3a5cb78017149a0739d159719f Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Mon, 27 Nov 2023 15:29:43 +0100 Subject: [PATCH 2/8] Apply comments --- starlette/middleware/limits.py | 23 +++++++++-------------- tests/middleware/test_limits.py | 16 +++++++++++++--- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/starlette/middleware/limits.py b/starlette/middleware/limits.py index 919a51fae..42f4bc8a6 100644 --- a/starlette/middleware/limits.py +++ b/starlette/middleware/limits.py @@ -11,7 +11,7 @@ def __init__(self, max_body_size: int) -> None: self.max_body_size = max_body_size -class LimitRequestMiddleware: +class LimitRequestSizeMiddleware: def __init__( self, app: ASGIApp, max_body_size: int = DEFAULT_MAX_BODY_SIZE ) -> None: @@ -22,17 +22,6 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": # pragma: no cover return await self.app(scope, receive, send) - headers = Headers(scope=scope) - content_length = headers.get("content-length") - if content_length is not None: - if int(content_length) > self.max_body_size: - return await _content_too_large_app(scope)(scope, receive, send) - - # NOTE: The server makes sure that the content-length header sent by the - # client is the same as the length of the body. - # Ref.: https://github.com/django/asgiref/issues/422 - return await self.app(scope, receive, send) - total_size = 0 response_started = False @@ -52,15 +41,21 @@ async def wrap_receive() -> Message: raise ContentTooLarge(self.max_body_size) return message + headers = Headers(scope=scope) + content_length = headers.get("content-length") + if content_length is not None: + if int(content_length) > self.max_body_size: + return await _content_too_large_app()(scope, receive, send) + try: await self.app(scope, wrap_receive, wrap_send) except ContentTooLarge as exc: # NOTE: If response has already started, we can't return a 413, because the # headers have already been sent. if not response_started: - return await _content_too_large_app(scope)(scope, receive, send) + return await _content_too_large_app()(scope, receive, send) raise exc -def _content_too_large_app(scope: Scope) -> PlainTextResponse: +def _content_too_large_app() -> PlainTextResponse: return PlainTextResponse("Content Too Large", status_code=413) diff --git a/tests/middleware/test_limits.py b/tests/middleware/test_limits.py index cadb27bb0..b8eb6214f 100644 --- a/tests/middleware/test_limits.py +++ b/tests/middleware/test_limits.py @@ -4,7 +4,7 @@ from starlette.applications import Starlette from starlette.middleware import Middleware -from starlette.middleware.limits import ContentTooLarge, LimitRequestMiddleware +from starlette.middleware.limits import ContentTooLarge, LimitRequestSizeMiddleware from starlette.requests import Request from starlette.routing import Route from starlette.testclient import TestClient @@ -24,7 +24,7 @@ async def echo_app(scope: Scope, receive: Receive, send: Send) -> None: break -app = LimitRequestMiddleware(echo_app, max_body_size=1024) +app = LimitRequestSizeMiddleware(echo_app, max_body_size=1024) def test_no_op(test_client_factory: Callable[..., TestClient]) -> None: @@ -81,10 +81,20 @@ async def read_body_endpoint(request: Request) -> None: app = Starlette( routes=[Route("/", read_body_endpoint, methods=["POST"])], - middleware=[Middleware(LimitRequestMiddleware, max_body_size=1024)], + middleware=[Middleware(LimitRequestSizeMiddleware, max_body_size=1024)], ) client = test_client_factory(app) response = client.post("/", content=[b"X" * 1024, b"X"]) assert response.status_code == 413 assert response.text == "Content Too Large" + + +def test_content_too_large_and_content_length_mismatch( + test_client_factory: Callable[..., TestClient] +) -> None: + client = test_client_factory(app) + + response = client.post("/", content="X" * 1025, headers={"Content-Length": "1024"}) + assert response.status_code == 413 + assert response.text == "Content Too Large" From de2a25b89407050c1aefef8959a575885f993729 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Mon, 27 Nov 2023 15:40:25 +0100 Subject: [PATCH 3/8] Use LimitRequestMiddleware again --- starlette/middleware/limits.py | 2 +- tests/middleware/test_limits.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/starlette/middleware/limits.py b/starlette/middleware/limits.py index 42f4bc8a6..44176c8b1 100644 --- a/starlette/middleware/limits.py +++ b/starlette/middleware/limits.py @@ -11,7 +11,7 @@ def __init__(self, max_body_size: int) -> None: self.max_body_size = max_body_size -class LimitRequestSizeMiddleware: +class LimitRequestMiddleware: def __init__( self, app: ASGIApp, max_body_size: int = DEFAULT_MAX_BODY_SIZE ) -> None: diff --git a/tests/middleware/test_limits.py b/tests/middleware/test_limits.py index b8eb6214f..5d48409e1 100644 --- a/tests/middleware/test_limits.py +++ b/tests/middleware/test_limits.py @@ -4,7 +4,7 @@ from starlette.applications import Starlette from starlette.middleware import Middleware -from starlette.middleware.limits import ContentTooLarge, LimitRequestSizeMiddleware +from starlette.middleware.limits import ContentTooLarge, LimitRequestMiddleware from starlette.requests import Request from starlette.routing import Route from starlette.testclient import TestClient @@ -24,7 +24,7 @@ async def echo_app(scope: Scope, receive: Receive, send: Send) -> None: break -app = LimitRequestSizeMiddleware(echo_app, max_body_size=1024) +app = LimitRequestMiddleware(echo_app, max_body_size=1024) def test_no_op(test_client_factory: Callable[..., TestClient]) -> None: @@ -81,7 +81,7 @@ async def read_body_endpoint(request: Request) -> None: app = Starlette( routes=[Route("/", read_body_endpoint, methods=["POST"])], - middleware=[Middleware(LimitRequestSizeMiddleware, max_body_size=1024)], + middleware=[Middleware(LimitRequestMiddleware, max_body_size=1024)], ) client = test_client_factory(app) From c4e7aad4c5beae118890348af0fcb9d434ef7368 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Tue, 28 Nov 2023 14:32:43 +0100 Subject: [PATCH 4/8] Apply all comments and make the middleware "multilayer" --- docs/middleware.md | 6 +-- starlette/middleware/limits.py | 37 ++++++++++------- tests/middleware/test_limits.py | 74 +++++++++++++++++++++++++++++---- 3 files changed, 92 insertions(+), 25 deletions(-) diff --git a/docs/middleware.md b/docs/middleware.md index 4ca1971ad..23768d7d1 100644 --- a/docs/middleware.md +++ b/docs/middleware.md @@ -183,7 +183,7 @@ The following arguments are supported: The middleware won't GZip responses that already have a `Content-Encoding` set, to prevent them from being encoded twice. -## LimitRequestMiddleware +## LimitBodySizeMiddleware Limits the body size of incoming requests. If the incoming request has a body larger than the limit, then a `413 Content Too Large` response will be sent. @@ -191,13 +191,13 @@ larger than the limit, then a `413 Content Too Large` response will be sent. ```python from starlette.applications import Starlette from starlette.middleware import Middleware -from starlette.middleware.limits import LimitRequestMiddleware +from starlette.middleware.limits import LimitBodySizeMiddleware routes = ... middleware = [ - Middleware(LimitRequestMiddleware, max_body_size=1024) + Middleware(LimitBodySizeMiddleware, max_body_size=1024) ] app = Starlette(routes=routes, middleware=middleware) diff --git a/starlette/middleware/limits.py b/starlette/middleware/limits.py index 44176c8b1..bf4aa661c 100644 --- a/starlette/middleware/limits.py +++ b/starlette/middleware/limits.py @@ -1,19 +1,24 @@ """Middleware that limits the body size of incoming requests.""" from starlette.datastructures import Headers +from starlette.exceptions import HTTPException from starlette.responses import PlainTextResponse from starlette.types import ASGIApp, Message, Receive, Scope, Send DEFAULT_MAX_BODY_SIZE = 2_621_440 # 2.5MB +MAX_BODY_SIZE_KEY = "starlette.max_body_size" +_content_too_large_app = PlainTextResponse("Content Too Large", status_code=413) -class ContentTooLarge(Exception): + +class ContentTooLarge(HTTPException): def __init__(self, max_body_size: int) -> None: self.max_body_size = max_body_size + super().__init__(detail="Content Too Large", status_code=413) -class LimitRequestMiddleware: +class LimitBodySizeMiddleware: def __init__( - self, app: ASGIApp, max_body_size: int = DEFAULT_MAX_BODY_SIZE + self, app: ASGIApp, *, max_body_size: int = DEFAULT_MAX_BODY_SIZE ) -> None: self.app = app self.max_body_size = max_body_size @@ -22,8 +27,12 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": # pragma: no cover return await self.app(scope, receive, send) + scope[MAX_BODY_SIZE_KEY] = self.max_body_size + total_size = 0 response_started = False + headers = Headers(scope=scope) + content_length = headers.get("content-length") async def wrap_send(message: Message) -> None: nonlocal response_started @@ -33,19 +42,20 @@ async def wrap_send(message: Message) -> None: async def wrap_receive() -> Message: nonlocal total_size + + if content_length is not None: + if int(content_length) > scope[MAX_BODY_SIZE_KEY]: + raise ContentTooLarge(self.max_body_size) + message = await receive() + if message["type"] == "http.request": chunk_size = len(message["body"]) total_size += chunk_size - if total_size > self.max_body_size: + if total_size > scope[MAX_BODY_SIZE_KEY]: raise ContentTooLarge(self.max_body_size) - return message - headers = Headers(scope=scope) - content_length = headers.get("content-length") - if content_length is not None: - if int(content_length) > self.max_body_size: - return await _content_too_large_app()(scope, receive, send) + return message try: await self.app(scope, wrap_receive, wrap_send) @@ -53,9 +63,6 @@ async def wrap_receive() -> Message: # NOTE: If response has already started, we can't return a 413, because the # headers have already been sent. if not response_started: - return await _content_too_large_app()(scope, receive, send) + if "app" not in scope: + return await _content_too_large_app(scope, receive, send) raise exc - - -def _content_too_large_app() -> PlainTextResponse: - return PlainTextResponse("Content Too Large", status_code=413) diff --git a/tests/middleware/test_limits.py b/tests/middleware/test_limits.py index 5d48409e1..260ba338e 100644 --- a/tests/middleware/test_limits.py +++ b/tests/middleware/test_limits.py @@ -4,9 +4,10 @@ from starlette.applications import Starlette from starlette.middleware import Middleware -from starlette.middleware.limits import ContentTooLarge, LimitRequestMiddleware +from starlette.middleware.limits import ContentTooLarge, LimitBodySizeMiddleware from starlette.requests import Request -from starlette.routing import Route +from starlette.responses import Response +from starlette.routing import Mount, Route from starlette.testclient import TestClient from starlette.types import Message, Receive, Scope, Send @@ -24,7 +25,7 @@ async def echo_app(scope: Scope, receive: Receive, send: Send) -> None: break -app = LimitRequestMiddleware(echo_app, max_body_size=1024) +app = LimitBodySizeMiddleware(echo_app, max_body_size=1024) def test_no_op(test_client_factory: Callable[..., TestClient]) -> None: @@ -73,18 +74,26 @@ async def send(message: Message) -> None: await rcv.aclose() +async def read_body_endpoint(request: Request) -> Response: + body = b"" + async for chunk in request.stream(): + body += chunk + return Response(content=body) + + def test_content_too_large_on_starlette( test_client_factory: Callable[..., TestClient] ) -> None: - async def read_body_endpoint(request: Request) -> None: - await request.body() - app = Starlette( routes=[Route("/", read_body_endpoint, methods=["POST"])], - middleware=[Middleware(LimitRequestMiddleware, max_body_size=1024)], + middleware=[Middleware(LimitBodySizeMiddleware, max_body_size=1024)], ) client = test_client_factory(app) + response = client.post("/", content=b"X" * 1024) + assert response.status_code == 200 + assert response.text == "X" * 1024 + response = client.post("/", content=[b"X" * 1024, b"X"]) assert response.status_code == 413 assert response.text == "Content Too Large" @@ -98,3 +107,54 @@ def test_content_too_large_and_content_length_mismatch( response = client.post("/", content="X" * 1025, headers={"Content-Length": "1024"}) assert response.status_code == 413 assert response.text == "Content Too Large" + + +def test_inner_middleware_overrides_outer_middleware( + test_client_factory: Callable[..., TestClient] +) -> None: + outer_app = LimitBodySizeMiddleware( + LimitBodySizeMiddleware( + echo_app, + max_body_size=2048, + ), + max_body_size=1024, + ) + + client = test_client_factory(outer_app) + + response = client.post("/", content="X" * 2049) + assert response.status_code == 413 + assert response.text == "Content Too Large" + + response = client.post("/", content="X" * 2048) + assert response.status_code == 200 + assert response.text == "X" * 2048 + + +def test_multiple_middleware_on_starlette( + test_client_factory: Callable[..., TestClient] +) -> None: + app = Starlette( + routes=[ + Route("/outer", read_body_endpoint, methods=["POST"]), + Mount( + "/inner", + routes=[Route("/", read_body_endpoint, methods=["POST"])], + middleware=[Middleware(LimitBodySizeMiddleware, max_body_size=2048)], + ), + ], + middleware=[Middleware(LimitBodySizeMiddleware, max_body_size=1024)], + ) + client = test_client_factory(app) + + # response = client.post("/outer", content="X" * 1025) + # assert response.status_code == 413 + # assert response.text == "Content Too Large" + + # response = client.post("/outer", content="X" * 1025) + # assert response.status_code == 413 + # assert response.text == "Content Too Large" + + response = client.post("/inner", content="X" * 1025) + assert response.status_code == 200 + assert response.text == "X" * 1025 From 054bdcd2d8b8bbc45120251071bba76f00e41553 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Tue, 28 Nov 2023 14:36:03 +0100 Subject: [PATCH 5/8] Remove the `HTTPException` --- starlette/middleware/limits.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/starlette/middleware/limits.py b/starlette/middleware/limits.py index bf4aa661c..2531c24b2 100644 --- a/starlette/middleware/limits.py +++ b/starlette/middleware/limits.py @@ -1,6 +1,5 @@ """Middleware that limits the body size of incoming requests.""" from starlette.datastructures import Headers -from starlette.exceptions import HTTPException from starlette.responses import PlainTextResponse from starlette.types import ASGIApp, Message, Receive, Scope, Send @@ -10,10 +9,9 @@ _content_too_large_app = PlainTextResponse("Content Too Large", status_code=413) -class ContentTooLarge(HTTPException): +class ContentTooLarge(Exception): def __init__(self, max_body_size: int) -> None: self.max_body_size = max_body_size - super().__init__(detail="Content Too Large", status_code=413) class LimitBodySizeMiddleware: @@ -63,6 +61,5 @@ async def wrap_receive() -> Message: # NOTE: If response has already started, we can't return a 413, because the # headers have already been sent. if not response_started: - if "app" not in scope: - return await _content_too_large_app(scope, receive, send) + return await _content_too_large_app(scope, receive, send) raise exc From f1515864c7647f79c563d7b4f167c29cfb2a79f8 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Tue, 28 Nov 2023 14:39:58 +0100 Subject: [PATCH 6/8] Uncomment tests --- tests/middleware/test_limits.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/middleware/test_limits.py b/tests/middleware/test_limits.py index 260ba338e..77f31cecf 100644 --- a/tests/middleware/test_limits.py +++ b/tests/middleware/test_limits.py @@ -147,14 +147,18 @@ def test_multiple_middleware_on_starlette( ) client = test_client_factory(app) - # response = client.post("/outer", content="X" * 1025) - # assert response.status_code == 413 - # assert response.text == "Content Too Large" + response = client.post("/outer", content="X" * 1024) + assert response.status_code == 200 + assert response.text == "X" * 1024 - # response = client.post("/outer", content="X" * 1025) - # assert response.status_code == 413 - # assert response.text == "Content Too Large" + response = client.post("/outer", content="X" * 1025) + assert response.status_code == 413 + assert response.text == "Content Too Large" - response = client.post("/inner", content="X" * 1025) + response = client.post("/inner", content="X" * 2048) assert response.status_code == 200 - assert response.text == "X" * 1025 + assert response.text == "X" * 2048 + + response = client.post("/inner", content="X" * 2049) + assert response.status_code == 413 + assert response.text == "Content Too Large" From e14c8205512e14dd0d2fa8214b7f42b401e9ddec Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Tue, 28 Nov 2023 14:41:46 +0100 Subject: [PATCH 7/8] Remove the `_content_too_large_app` --- starlette/middleware/limits.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/starlette/middleware/limits.py b/starlette/middleware/limits.py index 2531c24b2..784505581 100644 --- a/starlette/middleware/limits.py +++ b/starlette/middleware/limits.py @@ -6,8 +6,6 @@ DEFAULT_MAX_BODY_SIZE = 2_621_440 # 2.5MB MAX_BODY_SIZE_KEY = "starlette.max_body_size" -_content_too_large_app = PlainTextResponse("Content Too Large", status_code=413) - class ContentTooLarge(Exception): def __init__(self, max_body_size: int) -> None: @@ -61,5 +59,6 @@ async def wrap_receive() -> Message: # NOTE: If response has already started, we can't return a 413, because the # headers have already been sent. if not response_started: - return await _content_too_large_app(scope, receive, send) + response = PlainTextResponse("Content Too Large", status_code=413) + return await response(scope, receive, send) raise exc From 50ab7a6ef6cb6bfb74f397b491232e9d6664529f Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 29 Nov 2023 10:52:26 -0600 Subject: [PATCH 8/8] wip on splitting into two middlewares --- starlette/middleware/limits.py | 26 ++++++++++++++++++-------- tests/middleware/test_limits.py | 26 ++++++++++++++++++-------- 2 files changed, 36 insertions(+), 16 deletions(-) diff --git a/starlette/middleware/limits.py b/starlette/middleware/limits.py index 784505581..5e625f611 100644 --- a/starlette/middleware/limits.py +++ b/starlette/middleware/limits.py @@ -12,24 +12,34 @@ def __init__(self, max_body_size: int) -> None: self.max_body_size = max_body_size +class SetBodySizeLimit: + def __init__(self, app: ASGIApp, max_body_size: int) -> None: + self.app = app + self.max_body_size = max_body_size + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + scope[MAX_BODY_SIZE_KEY] = self.max_body_size + await self.app(scope, receive, send) + class LimitBodySizeMiddleware: def __init__( - self, app: ASGIApp, *, max_body_size: int = DEFAULT_MAX_BODY_SIZE + self, app: ASGIApp ) -> None: self.app = app - self.max_body_size = max_body_size async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] != "http": # pragma: no cover return await self.app(scope, receive, send) - scope[MAX_BODY_SIZE_KEY] = self.max_body_size - total_size = 0 response_started = False headers = Headers(scope=scope) content_length = headers.get("content-length") + max_body_size = scope.get(MAX_BODY_SIZE_KEY, None) + if not max_body_size: + return await self.app(scope, receive, send) + async def wrap_send(message: Message) -> None: nonlocal response_started if message["type"] == "http.response.start": @@ -40,16 +50,16 @@ async def wrap_receive() -> Message: nonlocal total_size if content_length is not None: - if int(content_length) > scope[MAX_BODY_SIZE_KEY]: - raise ContentTooLarge(self.max_body_size) + if int(content_length) > max_body_size: + raise ContentTooLarge(max_body_size) message = await receive() if message["type"] == "http.request": chunk_size = len(message["body"]) total_size += chunk_size - if total_size > scope[MAX_BODY_SIZE_KEY]: - raise ContentTooLarge(self.max_body_size) + if total_size > max_body_size: + raise ContentTooLarge(max_body_size) return message diff --git a/tests/middleware/test_limits.py b/tests/middleware/test_limits.py index 77f31cecf..204b1c78e 100644 --- a/tests/middleware/test_limits.py +++ b/tests/middleware/test_limits.py @@ -4,12 +4,12 @@ from starlette.applications import Starlette from starlette.middleware import Middleware -from starlette.middleware.limits import ContentTooLarge, LimitBodySizeMiddleware +from starlette.middleware.limits import ContentTooLarge, LimitBodySizeMiddleware, SetBodySizeLimit from starlette.requests import Request from starlette.responses import Response from starlette.routing import Mount, Route from starlette.testclient import TestClient -from starlette.types import Message, Receive, Scope, Send +from starlette.types import ASGIApp, Message, Receive, Scope, Send async def echo_app(scope: Scope, receive: Receive, send: Send) -> None: @@ -25,7 +25,7 @@ async def echo_app(scope: Scope, receive: Receive, send: Send) -> None: break -app = LimitBodySizeMiddleware(echo_app, max_body_size=1024) +app = SetBodySizeLimit(LimitBodySizeMiddleware(echo_app), max_body_size=1024) def test_no_op(test_client_factory: Callable[..., TestClient]) -> None: @@ -85,7 +85,7 @@ def test_content_too_large_on_starlette( test_client_factory: Callable[..., TestClient] ) -> None: app = Starlette( - routes=[Route("/", read_body_endpoint, methods=["POST"])], + routes=[Mount("/", routes=[Route("/", read_body_endpoint, methods=["POST"])], middleware=[Middleware(LimitBodySizeMiddleware)])], middleware=[Middleware(LimitBodySizeMiddleware, max_body_size=1024)], ) client = test_client_factory(app) @@ -112,10 +112,20 @@ def test_content_too_large_and_content_length_mismatch( def test_inner_middleware_overrides_outer_middleware( test_client_factory: Callable[..., TestClient] ) -> None: - outer_app = LimitBodySizeMiddleware( - LimitBodySizeMiddleware( - echo_app, - max_body_size=2048, + class CopyScopeMiddleware: + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + scope = dict(scope) + await self.app(scope, receive, send) + + outer_app = SetBodySizeLimit( + CopyScopeMiddleware( + SetBodySizeLimit( + LimitBodySizeMiddleware(echo_app), + max_body_size=2048, + ) ), max_body_size=1024, )