From 921841bc30244f8da8be6cf562af0b2dd917fee2 Mon Sep 17 00:00:00 2001 From: Vytautas Liuolia Date: Sun, 16 Jan 2022 18:40:23 +0100 Subject: [PATCH 1/4] Fix WSGI middleware not to explode quadratically in the case of a larger body. --- tests/middleware/test_wsgi.py | 13 +++++++++++++ uvicorn/middleware/wsgi.py | 14 ++++++++++---- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/tests/middleware/test_wsgi.py b/tests/middleware/test_wsgi.py index a9673c3f6..a555fc5be 100644 --- a/tests/middleware/test_wsgi.py +++ b/tests/middleware/test_wsgi.py @@ -67,6 +67,19 @@ async def test_wsgi_post() -> None: assert response.text == '{"example": 123}' +@pytest.mark.asyncio +async def test_wsgi_put_more_body() -> None: + async def generate_body(): + for _ in range(1024): + yield b"123456789abcdef\n" * 64 + + app = WSGIMiddleware(echo_body) + async with httpx.AsyncClient(app=app, base_url="http://testserver") as client: + response = await client.put("/", content=generate_body()) + assert response.status_code == 200 + assert response.text == "123456789abcdef\n" * 64 * 1024 + + @pytest.mark.asyncio async def test_wsgi_exception() -> None: # Note that we're testing the WSGI app directly here. diff --git a/uvicorn/middleware/wsgi.py b/uvicorn/middleware/wsgi.py index 74bdfada4..e29857bf6 100644 --- a/uvicorn/middleware/wsgi.py +++ b/uvicorn/middleware/wsgi.py @@ -19,7 +19,9 @@ from uvicorn._types import Environ, ExcInfo, StartResponse, WSGIApp -def build_environ(scope: HTTPScope, message: ASGIReceiveEvent, body: bytes) -> Environ: +def build_environ( + scope: HTTPScope, message: ASGIReceiveEvent, body: io.BytesIO +) -> Environ: """ Builds a scope and request message into a WSGI environ object. """ @@ -31,7 +33,7 @@ def build_environ(scope: HTTPScope, message: ASGIReceiveEvent, body: bytes) -> E "SERVER_PROTOCOL": "HTTP/%s" % scope["http_version"], "wsgi.version": (1, 0), "wsgi.url_scheme": scope.get("scheme", "http"), - "wsgi.input": io.BytesIO(body), + "wsgi.input": body, "wsgi.errors": sys.stdout, "wsgi.multithread": True, "wsgi.multiprocess": True, @@ -105,12 +107,16 @@ async def __call__( self, receive: ASGIReceiveCallable, send: ASGISendCallable ) -> None: message: HTTPRequestEvent = await receive() # type: ignore[assignment] - body = message.get("body", b"") + body = io.BytesIO(message.get("body", b"")) more_body = message.get("more_body", False) + if more_body: + body.seek(0, io.SEEK_END) while more_body: body_message: HTTPRequestEvent = await receive() # type: ignore[assignment] - body += body_message.get("body", b"") + body.write(body_message.get("body", b"")) more_body = body_message.get("more_body", False) + else: + body.seek(0) environ = build_environ(self.scope, message, body) self.loop = asyncio.get_event_loop() wsgi = self.loop.run_in_executor( From d0490f66226cce95549b784ea74005f4a2f3163b Mon Sep 17 00:00:00 2001 From: Vytautas Liuolia Date: Sun, 16 Jan 2022 19:30:04 +0100 Subject: [PATCH 2/4] Clean up more_body and body stream handling, get rid of while-else. --- uvicorn/middleware/wsgi.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/uvicorn/middleware/wsgi.py b/uvicorn/middleware/wsgi.py index e29857bf6..abb0d7b27 100644 --- a/uvicorn/middleware/wsgi.py +++ b/uvicorn/middleware/wsgi.py @@ -111,11 +111,12 @@ async def __call__( more_body = message.get("more_body", False) if more_body: body.seek(0, io.SEEK_END) - while more_body: - body_message: HTTPRequestEvent = await receive() # type: ignore[assignment] - body.write(body_message.get("body", b"")) - more_body = body_message.get("more_body", False) - else: + while more_body: + body_message: HTTPRequestEvent = ( + await receive() + ) # type: ignore[assignment] + body.write(body_message.get("body", b"")) + more_body = body_message.get("more_body", False) body.seek(0) environ = build_environ(self.scope, message, body) self.loop = asyncio.get_event_loop() From ef8374503fbf45892e90e904aefea1cca079e7b3 Mon Sep 17 00:00:00 2001 From: Vytautas Liuolia Date: Sun, 16 Jan 2022 19:35:44 +0100 Subject: [PATCH 3/4] Update body type in `build_environ`'s tests (although they pass regardless...) --- tests/middleware/test_wsgi.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/middleware/test_wsgi.py b/tests/middleware/test_wsgi.py index a555fc5be..b58c75c76 100644 --- a/tests/middleware/test_wsgi.py +++ b/tests/middleware/test_wsgi.py @@ -1,3 +1,4 @@ +import io import sys from typing import List @@ -133,6 +134,6 @@ def test_build_environ_encoding() -> None: "body": b"", "more_body": False, } - environ = build_environ(scope, message, b"") + environ = build_environ(scope, message, io.BytesIO(b"")) assert environ["PATH_INFO"] == "/文".encode("utf8").decode("latin-1") assert environ["HTTP_KEY"] == "value1,value2" From 54c3fba1a140a887a0bf3b658aa8b43bb68e0e6f Mon Sep 17 00:00:00 2001 From: Vytautas Liuolia Date: Fri, 11 Feb 2022 19:45:46 +0100 Subject: [PATCH 4/4] Address typing issues. --- tests/middleware/test_wsgi.py | 4 ++-- uvicorn/middleware/wsgi.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/middleware/test_wsgi.py b/tests/middleware/test_wsgi.py index b58c75c76..8b35fd1ad 100644 --- a/tests/middleware/test_wsgi.py +++ b/tests/middleware/test_wsgi.py @@ -1,6 +1,6 @@ import io import sys -from typing import List +from typing import AsyncGenerator, List import httpx import pytest @@ -70,7 +70,7 @@ async def test_wsgi_post() -> None: @pytest.mark.asyncio async def test_wsgi_put_more_body() -> None: - async def generate_body(): + async def generate_body() -> AsyncGenerator[bytes, None]: for _ in range(1024): yield b"123456789abcdef\n" * 64 diff --git a/uvicorn/middleware/wsgi.py b/uvicorn/middleware/wsgi.py index abb0d7b27..f58236acf 100644 --- a/uvicorn/middleware/wsgi.py +++ b/uvicorn/middleware/wsgi.py @@ -113,8 +113,8 @@ async def __call__( body.seek(0, io.SEEK_END) while more_body: body_message: HTTPRequestEvent = ( - await receive() - ) # type: ignore[assignment] + await receive() # type: ignore[assignment] + ) body.write(body_message.get("body", b"")) more_body = body_message.get("more_body", False) body.seek(0)