Merge upstream PR pgjones#332 from pgjones/hypercorn: Fix WSGI Violations

Original PR: pgjones#332

Author: Vicente Ruiz

src/hypercorn/app_wrappers.py

@@ -84,16 +84,31 @@ async def handle_http(
         await send({"type": "http.response.body", "body": b"", "more_body": False})
 
     def run_app(self, environ: dict, send: Callable) -> None:
-        headers: list[tuple[bytes, bytes]]
+        headers: list[tuple[bytes, bytes]] = []
         response_started = False
+        headers_sent = False
         status_code: int | None = None
 
         def start_response(
             status: str,
             response_headers: list[tuple[str, str]],
             exc_info: Exception | None = None,
         ) -> None:
-            nonlocal headers, response_started, status_code
+            nonlocal headers, response_started, status_code, headers_sent
+
+            if response_started and exc_info is None:
+                raise RuntimeError(
+                    "start_response cannot be called again without the exc_info parameter"
+                )
+            elif exc_info is not None:
+                try:
+                    if headers_sent:
+                        # The headers have already been sent and we can no longer change
+                        # the status_code and headers. reraise this exception in accordance
+                        # with the WSGI specification.
+                        raise exc_info[1].with_traceback(exc_info[2])
+                finally:
+                    exc_info = None  # Delete reference to exc_info to avoid circular references
 
             raw, _ = status.split(" ", 1)
             status_code = int(raw)
@@ -106,16 +121,35 @@ def start_response(
         response_body = self.app(environ, start_response)
 
         try:
-            first_chunk = True
             for output in response_body:
-                if first_chunk:
+                # Per the WSGI specification in PEP-3333, the start_response callable
+                # must not actually transmit the response headers. Instead, it must
+                # store them for the server to transmit only after the first iteration
+                # of the application return value that yields a non-empty bytestring.
+                #
+                # We therefore delay sending the http.response.start event until after
+                # we receive a non-empty byte string from the application return value.
+                if output and not headers_sent:
                     if not response_started:
                         raise RuntimeError("WSGI app did not call start_response")
 
+                    # Send the http.response.start event with the status and headers, flagging
+                    # that this was completed so they aren't sent twice.
                     send({"type": "http.response.start", "status": status_code, "headers": headers})
-                    first_chunk = False
+                    headers_sent = True
 
                 send({"type": "http.response.body", "body": output, "more_body": True})
+
+            # If we still haven't sent the headers by this point, then we received no
+            # non-empty byte strings from the application return value. This can happen when
+            # handling certain HTTP methods that don't include a response body like HEAD.
+            # In those cases we still need to send the http.response.start event with the
+            # status code and headers, but we need to ensure they haven't been sent previously.
+            if not headers_sent:
+                if not response_started:
+                    raise RuntimeError("WSGI app did not call start_response")
+
+                send({"type": "http.response.start", "status": status_code, "headers": headers})
         finally:
             if hasattr(response_body, "close"):
                 response_body.close()

tests/test_app_wrappers.py

@@ -10,22 +10,22 @@
 
 from hypercorn.app_wrappers import _build_environ, InvalidPathError, WSGIWrapper
 from hypercorn.typing import ASGIReceiveEvent, ASGISendEvent, ConnectionState, HTTPScope
-
-
-def echo_body(environ: dict, start_response: Callable) -> list[bytes]:
-    status = "200 OK"
-    output = environ["wsgi.input"].read()
-    headers = [
-        ("Content-Type", "text/plain; charset=utf-8"),
-        ("Content-Length", str(len(output))),
-    ]
-    start_response(status, headers)
-    return [output]
+from .wsgi_applications import (
+    wsgi_app_echo_body,
+    wsgi_app_generator,
+    wsgi_app_generator_delayed_start_response,
+    wsgi_app_generator_multiple_start_response_after_body,
+    wsgi_app_generator_no_body,
+    wsgi_app_multiple_start_response_no_exc_info,
+    wsgi_app_no_body,
+    wsgi_app_no_start_response,
+    wsgi_app_simple,
+)
 
 
 @pytest.mark.trio
 async def test_wsgi_trio() -> None:
-    app = WSGIWrapper(echo_body, 2**16)
+    app = WSGIWrapper(wsgi_app_echo_body, 2**16)
     scope: HTTPScope = {
         "http_version": "1.1",
         "asgi": {},
@@ -52,12 +52,12 @@ async def _send(message: ASGISendEvent) -> None:
 
     await app(scope, receive_channel.receive, _send, trio.to_thread.run_sync, trio.from_thread.run)
     assert messages == [
+        {"body": bytearray(b""), "type": "http.response.body", "more_body": True},
         {
             "headers": [(b"content-type", b"text/plain; charset=utf-8"), (b"content-length", b"0")],
             "status": 200,
             "type": "http.response.start",
         },
-        {"body": bytearray(b""), "type": "http.response.body", "more_body": True},
         {"body": bytearray(b""), "type": "http.response.body", "more_body": False},
     ]
 
@@ -83,7 +83,7 @@ def _call_soon(func: Callable, *args: Any) -> Any:
 
 @pytest.mark.asyncio
 async def test_wsgi_asyncio() -> None:
-    app = WSGIWrapper(echo_body, 2**16)
+    app = WSGIWrapper(wsgi_app_echo_body, 2**16)
     scope: HTTPScope = {
         "http_version": "1.1",
         "asgi": {},
@@ -100,21 +100,24 @@ async def test_wsgi_asyncio() -> None:
         "extensions": {},
         "state": ConnectionState({}),
     }
-    messages = await _run_app(app, scope)
+    messages = await _run_app(app, scope, b"Hello, world!")
     assert messages == [
         {
-            "headers": [(b"content-type", b"text/plain; charset=utf-8"), (b"content-length", b"0")],
+            "headers": [
+                (b"content-type", b"text/plain; charset=utf-8"),
+                (b"content-length", b"13"),
+            ],
             "status": 200,
             "type": "http.response.start",
         },
-        {"body": bytearray(b""), "type": "http.response.body", "more_body": True},
-        {"body": bytearray(b""), "type": "http.response.body", "more_body": False},
+        {"body": b"Hello, world!", "type": "http.response.body", "more_body": True},
+        {"body": b"", "type": "http.response.body", "more_body": False},
     ]
 
 
 @pytest.mark.asyncio
 async def test_max_body_size() -> None:
-    app = WSGIWrapper(echo_body, 4)
+    app = WSGIWrapper(wsgi_app_echo_body, 4)
     scope: HTTPScope = {
         "http_version": "1.1",
         "asgi": {},
@@ -138,13 +141,9 @@ async def test_max_body_size() -> None:
     ]
 
 
-def no_start_response(environ: dict, start_response: Callable) -> list[bytes]:
-    return [b"result"]
-
-
 @pytest.mark.asyncio
 async def test_no_start_response() -> None:
-    app = WSGIWrapper(no_start_response, 2**16)
+    app = WSGIWrapper(wsgi_app_no_start_response, 2**16)
     scope: HTTPScope = {
         "http_version": "1.1",
         "asgi": {},
@@ -206,3 +205,151 @@ def test_build_environ_root_path() -> None:
     }
     with pytest.raises(InvalidPathError):
         _build_environ(scope, b"")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("wsgi_app", [wsgi_app_simple, wsgi_app_generator])
+async def test_wsgi_protocol(wsgi_app: Callable) -> None:
+    app = WSGIWrapper(wsgi_app, 2**16)
+    scope: HTTPScope = {
+        "http_version": "1.1",
+        "asgi": {},
+        "method": "GET",
+        "headers": [],
+        "path": "/",
+        "root_path": "/",
+        "query_string": b"a=b",
+        "raw_path": b"/",
+        "scheme": "http",
+        "type": "http",
+        "client": ("localhost", 80),
+        "server": None,
+        "extensions": {},
+        "state": ConnectionState({}),
+    }
+
+    messages = await _run_app(app, scope)
+    assert messages == [
+        {
+            "headers": [(b"x-test-header", b"Test-Value")],
+            "status": 200,
+            "type": "http.response.start",
+        },
+        {"body": b"Hello, ", "type": "http.response.body", "more_body": True},
+        {"body": b"world!", "type": "http.response.body", "more_body": True},
+        {"body": b"", "type": "http.response.body", "more_body": False},
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("wsgi_app", [wsgi_app_no_body, wsgi_app_generator_no_body])
+async def test_wsgi_protocol_no_body(wsgi_app: Callable) -> None:
+    app = WSGIWrapper(wsgi_app, 2**16)
+    scope: HTTPScope = {
+        "http_version": "1.1",
+        "asgi": {},
+        "method": "GET",
+        "headers": [],
+        "path": "/",
+        "root_path": "/",
+        "query_string": b"a=b",
+        "raw_path": b"/",
+        "scheme": "http",
+        "type": "http",
+        "client": ("localhost", 80),
+        "server": None,
+        "extensions": {},
+        "state": ConnectionState({}),
+    }
+
+    messages = await _run_app(app, scope)
+    assert messages == [
+        {
+            "headers": [(b"x-test-header", b"Test-Value")],
+            "status": 200,
+            "type": "http.response.start",
+        },
+        {"body": b"", "type": "http.response.body", "more_body": False},
+    ]
+
+
+@pytest.mark.asyncio
+async def test_wsgi_protocol_overwrite_start_response() -> None:
+    app = WSGIWrapper(wsgi_app_generator_delayed_start_response, 2**16)
+    scope: HTTPScope = {
+        "http_version": "1.1",
+        "asgi": {},
+        "method": "GET",
+        "headers": [],
+        "path": "/",
+        "root_path": "/",
+        "query_string": b"a=b",
+        "raw_path": b"/",
+        "scheme": "http",
+        "type": "http",
+        "client": ("localhost", 80),
+        "server": None,
+        "extensions": {},
+        "state": ConnectionState({}),
+    }
+
+    messages = await _run_app(app, scope)
+    assert messages == [
+        {"body": b"", "type": "http.response.body", "more_body": True},
+        {
+            "headers": [(b"x-test-header", b"New-Value")],
+            "status": 500,
+            "type": "http.response.start",
+        },
+        {"body": b"Hello, ", "type": "http.response.body", "more_body": True},
+        {"body": b"world!", "type": "http.response.body", "more_body": True},
+        {"body": b"", "type": "http.response.body", "more_body": False},
+    ]
+
+
+@pytest.mark.asyncio
+async def test_wsgi_protocol_multiple_start_response_no_exc_info() -> None:
+    app = WSGIWrapper(wsgi_app_multiple_start_response_no_exc_info, 2**16)
+    scope: HTTPScope = {
+        "http_version": "1.1",
+        "asgi": {},
+        "method": "GET",
+        "headers": [],
+        "path": "/",
+        "root_path": "/",
+        "query_string": b"a=b",
+        "raw_path": b"/",
+        "scheme": "http",
+        "type": "http",
+        "client": ("localhost", 80),
+        "server": None,
+        "extensions": {},
+        "state": ConnectionState({}),
+    }
+
+    with pytest.raises(RuntimeError):
+        await _run_app(app, scope)
+
+
+@pytest.mark.asyncio
+async def test_wsgi_protocol_multiple_start_response_after_body() -> None:
+    app = WSGIWrapper(wsgi_app_generator_multiple_start_response_after_body, 2**16)
+    scope: HTTPScope = {
+        "http_version": "1.1",
+        "asgi": {},
+        "method": "GET",
+        "headers": [],
+        "path": "/",
+        "root_path": "/",
+        "query_string": b"a=b",
+        "raw_path": b"/",
+        "scheme": "http",
+        "type": "http",
+        "client": ("localhost", 80),
+        "server": None,
+        "extensions": {},
+        "state": ConnectionState({}),
+    }
+
+    with pytest.raises(ValueError):
+        await _run_app(app, scope)