Support websocket.http.response ASGI extension

tests/protocols/test_websocket.py

@@ -6,6 +6,7 @@
 import pytest
 
 from tests.protocols.test_http import HTTP_PROTOCOLS
+from tests.response import Response
 from tests.utils import run_server
 from uvicorn.config import Config
 from uvicorn.protocols.websockets.wsproto_impl import WSProtocol
@@ -923,11 +924,99 @@ async def app(scope, receive, send):
         assert message["type"] == "websocket.disconnect"
 
     async def websocket_session(url):
-        try:
+        with pytest.raises(websockets.exceptions.InvalidStatusCode) as exc_info:
+            async with websockets.client.connect(url):
+                pass  # pragma: no cover
+        assert exc_info.value.status_code == 403
+
+    config = Config(
+        app=app,
+        ws=ws_protocol_cls,
+        http=http_protocol_cls,
+        lifespan="off",
+        port=unused_tcp_port,
+    )
+    async with run_server(config):
+        await websocket_session(f"ws://127.0.0.1:{unused_tcp_port}")
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_server_reject_connection_with_response(
+    ws_protocol_cls, http_protocol_cls, unused_tcp_port: int
+):
+    async def app(scope, receive, send):
+        assert scope["type"] == "websocket"
+        assert "websocket.http.response" in scope["extensions"]
+
+        # Pull up first recv message.
+        message = await receive()
+        assert message["type"] == "websocket.connect"
+
+        # Reject the connection with a response
+        response = Response(b"goodbye", status_code=400)
+        await response(scope, receive, send)
+        message = await receive()
+        assert message["type"] == "websocket.disconnect"
+
+    async def websocket_session(url):
+        with pytest.raises(websockets.exceptions.InvalidStatusCode) as exc_info:
+            async with websockets.client.connect(url):
+                pass  # pragma: no cover
+        assert exc_info.value.status_code == 400
+        # Websockets module currently does not read the response body from the socket.
+
+    config = Config(
+        app=app,
+        ws=ws_protocol_cls,
+        http=http_protocol_cls,
+        lifespan="off",
+        port=unused_tcp_port,
+    )
+    async with run_server(config):
+        await websocket_session(f"ws://127.0.0.1:{unused_tcp_port}")
+
+
+@pytest.mark.anyio
+@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS)
+@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS)
+async def test_server_reject_connection_with_multibody_response(
+    ws_protocol_cls, http_protocol_cls, unused_tcp_port: int
+):
+    async def app(scope, receive, send):
+        assert scope["type"] == "websocket"
+        assert "websocket.http.response" in scope["extensions"]
+
+        # Pull up first recv message.
+        message = await receive()
+        assert message["type"] == "websocket.connect"
+        message = {
+            "type": "websocket.http.response.start",
+            "status": 400,
+            "headers": [(b"Content-Length", b"20"), (b"Content-Type", b"text/plain")],
+        }
+        await send(message)
+        message = {
+            "type": "websocket.http.response.body",
+            "body": b"x" * 10,
+            "more_body": True,
+        }
+        await send(message)
+        message = {
+            "type": "websocket.http.response.body",
+            "body": b"y" * 10,
+        }
+        await send(message)
+        message = await receive()
+        assert message["type"] == "websocket.disconnect"
+
+    async def websocket_session(url):
+        with pytest.raises(websockets.exceptions.InvalidStatusCode) as exc_info:
             async with websockets.client.connect(url):
                 pass  # pragma: no cover
-        except Exception:
-            pass
+        assert exc_info.value.status_code == 400
+        # Websockets module currently does not read the response body from the socket.
 
     config = Config(
         app=app,

tests/response.py

@@ -10,17 +10,18 @@ def __init__(self, content, status_code=200, headers=None, media_type=None):
         self.set_content_length()
 
     async def __call__(self, scope, receive, send) -> None:
+        prefix = "websocket." if scope["type"] == "websocket" else ""
         await send(
             {
-                "type": "http.response.start",
+                "type": prefix + "http.response.start",
                 "status": self.status_code,
                 "headers": [
                     [key.encode(), value.encode()]
                     for key, value in self.headers.items()
                 ],
             }
         )
-        await send({"type": "http.response.body", "body": self.body})
+        await send({"type": prefix + "http.response.body", "body": self.body})
 
     def render(self, content) -> bytes:
         if isinstance(content, bytes):

uvicorn/protocols/websockets/websockets_impl.py

@@ -46,6 +46,8 @@
         WebSocketConnectEvent,
         WebSocketDisconnectEvent,
         WebSocketReceiveEvent,
+        WebSocketResponseBodyEvent,
+        WebSocketResponseStartEvent,
         WebSocketScope,
         WebSocketSendEvent,
     )
@@ -102,6 +104,7 @@ def __init__(
         self.connect_sent = False
         self.lost_connection_before_handshake = False
         self.accepted_subprotocol: Optional[Subprotocol] = None
+        self.response_body: Optional[List[bytes]] = None
 
         self.ws_server: Server = Server()  # type: ignore[assignment]
 
@@ -203,6 +206,7 @@ async def process_request(
             "headers": asgi_headers,
             "subprotocols": subprotocols,
             "state": self.app_state.copy(),
+            "extensions": {"websocket.http.response": {}},
         }
         task = self.loop.create_task(self.run_asgi())
         task.add_done_callback(self.on_task_complete)
@@ -278,43 +282,82 @@ async def asgi_send(self, message: "ASGISendEvent") -> None:
         message_type = message["type"]
 
         if not self.handshake_started_event.is_set():
-            if message_type == "websocket.accept":
-                message = cast("WebSocketAcceptEvent", message)
-                self.logger.info(
-                    '%s - "WebSocket %s" [accepted]',
-                    self.scope["client"],
-                    get_path_with_query_string(self.scope),
-                )
-                self.initial_response = None
-                self.accepted_subprotocol = cast(
-                    Optional[Subprotocol], message.get("subprotocol")
-                )
-                if "headers" in message:
-                    self.extra_headers.extend(
-                        # ASGI spec requires bytes
-                        # But for compatibility we need to convert it to strings
+            if self.response_body is None:
+                if message_type == "websocket.accept":
+                    message = cast("WebSocketAcceptEvent", message)
+                    self.logger.info(
+                        '%s - "WebSocket %s" [accepted]',
+                        self.scope["client"],
+                        get_path_with_query_string(self.scope),
+                    )
+                    self.initial_response = None
+                    self.accepted_subprotocol = cast(
+                        Optional[Subprotocol], message.get("subprotocol")
+                    )
+                    if "headers" in message:
+                        self.extra_headers.extend(
+                            # ASGI spec requires bytes
+                            # But for compatibility we need to convert it to strings
+                            (name.decode("latin-1"), value.decode("latin-1"))
+                            for name, value in message["headers"]
+                        )
+                    self.handshake_started_event.set()
+
+                elif message_type == "websocket.close":
+                    message = cast("WebSocketCloseEvent", message)
+                    self.logger.info(
+                        '%s - "WebSocket %s" 403',
+                        self.scope["client"],
+                        get_path_with_query_string(self.scope),
+                    )
+                    self.initial_response = (http.HTTPStatus.FORBIDDEN, [], b"")
+                    self.handshake_started_event.set()
+                    self.closed_event.set()
+
+                elif message_type == "websocket.http.response.start":
+                    message = cast("WebSocketResponseStartEvent", message)
+                    self.logger.info(
+                        '%s - "WebSocket %s" %s',
+                        self.scope["client"],
+                        get_path_with_query_string(self.scope),
+                        message["status"],
+                    )
+                    try:
+                        status = http.HTTPStatus(message["status"])
+                    except AttributeError:
+                        status = http.HTTPStatus.FORBIDDEN
+                    headers = [
                         (name.decode("latin-1"), value.decode("latin-1"))
-                        for name, value in message["headers"]
+                        for name, value in message.get("headers", [])
+                    ]
+                    self.initial_response = (status, headers, b"")
+                    self.response_body = []
+
+                else:
+                    msg = (
+                        "Expected ASGI message 'websocket.accept', 'websocket.close', "
+                        "or 'websocket.http.response.start' "
+                        "but got '%s'."
                     )
-                self.handshake_started_event.set()
-
-            elif message_type == "websocket.close":
-                message = cast("WebSocketCloseEvent", message)
-                self.logger.info(
-                    '%s - "WebSocket %s" 403',
-                    self.scope["client"],
-                    get_path_with_query_string(self.scope),
-                )
-                self.initial_response = (http.HTTPStatus.FORBIDDEN, [], b"")
-                self.handshake_started_event.set()
-                self.closed_event.set()
-
+                    raise RuntimeError(msg % message_type)
             else:
-                msg = (
-                    "Expected ASGI message 'websocket.accept' or 'websocket.close', "
-                    "but got '%s'."
-                )
-                raise RuntimeError(msg % message_type)
+                if message_type == "websocket.http.response.body":
+                    message = cast("WebSocketResponseBodyEvent", message)
+                    self.response_body.append(message["body"])
+
+                    if not message.get("more_body", False):
+                        self.initial_response = self.initial_response[:2] + (
+                            b"".join(self.response_body),
+                        )
+                        self.response_body = None
+                        self.handshake_started_event.set()
+                        self.closed_event.set()
+                else:
+                    msg = (
+                        "Expected ASGI message 'websocket.http.response.body' "
+                        "but got '%s'."
+                    )
+                    raise RuntimeError(msg % message_type)
 
         elif not self.closed_event.is_set():
             await self.handshake_completed_event.wait()