From a694ee4a09380b26d681bcdf91bdc3258804e107 Mon Sep 17 00:00:00 2001 From: Vytautas Liuolia Date: Wed, 16 Feb 2022 13:34:22 +0100 Subject: [PATCH] Fix WSGI middleware not to explode quadratically in the case of a larger body (#1329) * Fix WSGI middleware not to explode quadratically in the case of a larger body. * Clean up more_body and body stream handling, get rid of while-else. * Update body type in `build_environ`'s tests (although they pass regardless...) * Address typing issues. --- tests/middleware/test_wsgi.py | 18 ++++++++++++++++-- uvicorn/middleware/wsgi.py | 21 ++++++++++++++------- 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/tests/middleware/test_wsgi.py b/tests/middleware/test_wsgi.py index a9673c3f6..8b35fd1ad 100644 --- a/tests/middleware/test_wsgi.py +++ b/tests/middleware/test_wsgi.py @@ -1,5 +1,6 @@ +import io import sys -from typing import List +from typing import AsyncGenerator, List import httpx import pytest @@ -67,6 +68,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() -> AsyncGenerator[bytes, None]: + 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. @@ -120,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" diff --git a/uvicorn/middleware/wsgi.py b/uvicorn/middleware/wsgi.py index 74bdfada4..f58236acf 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,17 @@ 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) - while more_body: - body_message: HTTPRequestEvent = await receive() # type: ignore[assignment] - body += body_message.get("body", b"") - more_body = 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) + body.seek(0) environ = build_environ(self.scope, message, body) self.loop = asyncio.get_event_loop() wsgi = self.loop.run_in_executor(