websockets_impl.py

from __future__ import annotations

import asyncio
import http
import logging
from typing import (
    Any,
    List,
    Literal,
    Optional,
    Sequence,
    Tuple,
    Union,
    cast,
)
from urllib.parse import unquote

import websockets
from websockets.datastructures import Headers
from websockets.exceptions import ConnectionClosed
from websockets.extensions.permessage_deflate import ServerPerMessageDeflateFactory
from websockets.legacy.server import HTTPResponse
from websockets.server import WebSocketServerProtocol
from websockets.typing import Subprotocol

from uvicorn._types import (
    ASGISendEvent,
    WebSocketAcceptEvent,
    WebSocketCloseEvent,
    WebSocketConnectEvent,
    WebSocketDisconnectEvent,
    WebSocketReceiveEvent,
    WebSocketResponseBodyEvent,
    WebSocketResponseStartEvent,
    WebSocketScope,
    WebSocketSendEvent,
)
from uvicorn.config import Config
from uvicorn.logging import TRACE_LOG_LEVEL
from uvicorn.protocols.utils import (
    ClientDisconnected,
    get_local_addr,
    get_path_with_query_string,
    get_remote_addr,
    is_ssl,
)
from uvicorn.server import ServerState


class Server:
    closing = False

    def register(self, ws: WebSocketServerProtocol) -> None:
        pass

    def unregister(self, ws: WebSocketServerProtocol) -> None:
        pass

    def is_serving(self) -> bool:
        return not self.closing


class WebSocketProtocol(WebSocketServerProtocol):
    extra_headers: List[Tuple[str, str]]

    def __init__(
        self,
        config: Config,
        server_state: ServerState,
        app_state: dict[str, Any],
        _loop: asyncio.AbstractEventLoop | None = None,
    ):
        if not config.loaded:
            config.load()

        self.config = config
        self.app = config.loaded_app
        self.loop = _loop or asyncio.get_event_loop()
        self.root_path = config.root_path
        self.app_state = app_state

        # Shared server state
        self.connections = server_state.connections
        self.tasks = server_state.tasks

        # Connection state
        self.transport: asyncio.Transport = None  # type: ignore[assignment]
        self.server: tuple[str, int] | None = None
        self.client: tuple[str, int] | None = None
        self.scheme: Literal["wss", "ws"] = None  # type: ignore[assignment]

        # Connection events
        self.scope: WebSocketScope
        self.handshake_started_event = asyncio.Event()
        self.handshake_completed_event = asyncio.Event()
        self.closed_event = asyncio.Event()
        self.initial_response: HTTPResponse | None = None
        self.connect_sent = False
        self.lost_connection_before_handshake = False
        self.accepted_subprotocol: Subprotocol | None = None

        self.ws_server: Server = Server()  # type: ignore[assignment]

        extensions = []
        if self.config.ws_per_message_deflate:
            extensions.append(ServerPerMessageDeflateFactory())

        super().__init__(
            ws_handler=self.ws_handler,
            ws_server=self.ws_server,  # type: ignore[arg-type]
            max_size=self.config.ws_max_size,
            max_queue=self.config.ws_max_queue,
            ping_interval=self.config.ws_ping_interval,
            ping_timeout=self.config.ws_ping_timeout,
            extensions=extensions,
            logger=logging.getLogger("uvicorn.error"),
        )
        self.server_header = None
        self.extra_headers = [
            (name.decode("latin-1"), value.decode("latin-1"))
            for name, value in server_state.default_headers
        ]

    def connection_made(  # type: ignore[override]
        self, transport: asyncio.Transport
    ) -> None:
        self.connections.add(self)
        self.transport = transport
        self.server = get_local_addr(transport)
        self.client = get_remote_addr(transport)
        self.scheme = "wss" if is_ssl(transport) else "ws"

        if self.logger.isEnabledFor(TRACE_LOG_LEVEL):
            prefix = "%s:%d - " % self.client if self.client else ""
            self.logger.log(TRACE_LOG_LEVEL, "%sWebSocket connection made", prefix)

        super().connection_made(transport)

    def connection_lost(self, exc: Optional[Exception]) -> None:
        self.connections.remove(self)

        if self.logger.isEnabledFor(TRACE_LOG_LEVEL):
            prefix = "%s:%d - " % self.client if self.client else ""
            self.logger.log(TRACE_LOG_LEVEL, "%sWebSocket connection lost", prefix)

        self.lost_connection_before_handshake = (
            not self.handshake_completed_event.is_set()
        )
        self.handshake_completed_event.set()
        super().connection_lost(exc)
        if exc is None:
            self.transport.close()

    def shutdown(self) -> None:
        self.ws_server.closing = True
        if self.handshake_completed_event.is_set():
            self.fail_connection(1012)
        else:
            self.send_500_response()
        self.transport.close()

    def on_task_complete(self, task: asyncio.Task) -> None:
        self.tasks.discard(task)

    async def process_request(
        self, path: str, headers: Headers
    ) -> Optional[HTTPResponse]:
        """
        This hook is called to determine if the websocket should return
        an HTTP response and close.

        Our behavior here is to start the ASGI application, and then wait
        for either `accept` or `close` in order to determine if we should
        close the connection.
        """
        path_portion, _, query_string = path.partition("?")

        websockets.legacy.handshake.check_request(headers)

        subprotocols = []
        for header in headers.get_all("Sec-WebSocket-Protocol"):
            subprotocols.extend([token.strip() for token in header.split(",")])

        asgi_headers = [
            (name.encode("ascii"), value.encode("ascii", errors="surrogateescape"))
            for name, value in headers.raw_items()
        ]
        path = unquote(path_portion)
        full_path = self.root_path + path
        full_raw_path = self.root_path.encode("ascii") + path_portion.encode("ascii")

        self.scope = {
            "type": "websocket",
            "asgi": {"version": self.config.asgi_version, "spec_version": "2.3"},
            "http_version": "1.1",
            "scheme": self.scheme,
            "server": self.server,
            "client": self.client,
            "root_path": self.root_path,
            "path": full_path,
            "raw_path": full_raw_path,
            "query_string": query_string.encode("ascii"),
            "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)
        self.tasks.add(task)
        await self.handshake_started_event.wait()
        return self.initial_response

    def process_subprotocol(
        self, headers: Headers, available_subprotocols: Optional[Sequence[Subprotocol]]
    ) -> Optional[Subprotocol]:
        """
        We override the standard 'process_subprotocol' behavior here so that
        we return whatever subprotocol is sent in the 'accept' message.
        """
        return self.accepted_subprotocol

    def send_500_response(self) -> None:
        msg = b"Internal Server Error"
        content = [
            b"HTTP/1.1 500 Internal Server Error\r\n"
            b"content-type: text/plain; charset=utf-8\r\n",
            b"content-length: " + str(len(msg)).encode("ascii") + b"\r\n",
            b"connection: close\r\n",
            b"\r\n",
            msg,
        ]
        self.transport.write(b"".join(content))
        # Allow handler task to terminate cleanly, as websockets doesn't cancel it by
        # itself (see https://github.com/encode/uvicorn/issues/920)
        self.handshake_started_event.set()

    async def ws_handler(  # type: ignore[override]
        self, protocol: WebSocketServerProtocol, path: str
    ) -> Any:
        """
        This is the main handler function for the 'websockets' implementation
        to call into. We just wait for close then return, and instead allow
        'send' and 'receive' events to drive the flow.
        """
        self.handshake_completed_event.set()
        await self.wait_closed()

    async def run_asgi(self) -> None:
        """
        Wrapper around the ASGI callable, handling exceptions and unexpected
        termination states.
        """
        try:
            result = await self.app(self.scope, self.asgi_receive, self.asgi_send)
        except ClientDisconnected:
            self.closed_event.set()
            self.transport.close()
        except BaseException as exc:
            self.closed_event.set()
            msg = "Exception in ASGI application\n"
            self.logger.error(msg, exc_info=exc)
            if not self.handshake_started_event.is_set():
                self.send_500_response()
            else:
                await self.handshake_completed_event.wait()
            self.transport.close()
        else:
            self.closed_event.set()
            if not self.handshake_started_event.is_set():
                msg = "ASGI callable returned without sending handshake."
                self.logger.error(msg)
                self.send_500_response()
                self.transport.close()
            elif result is not None:
                msg = "ASGI callable should return None, but returned '%s'."
                self.logger.error(msg, result)
                await self.handshake_completed_event.wait()
                self.transport.close()

    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
                        (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" %d',
                    self.scope["client"],
                    get_path_with_query_string(self.scope),
                    message["status"],
                )
                # websockets requires the status to be an enum. look it up.
                status = http.HTTPStatus(message["status"])
                headers = [
                    (name.decode("latin-1"), value.decode("latin-1"))
                    for name, value in message.get("headers", [])
                ]
                self.initial_response = (status, headers, b"")
                self.handshake_started_event.set()

            else:
                msg = (
                    "Expected ASGI message 'websocket.accept', 'websocket.close', "
                    "or 'websocket.http.response.start' but got '%s'."
                )
                raise RuntimeError(msg % message_type)

        elif not self.closed_event.is_set() and self.initial_response is None:
            await self.handshake_completed_event.wait()

            try:
                if message_type == "websocket.send":
                    message = cast("WebSocketSendEvent", message)
                    bytes_data = message.get("bytes")
                    text_data = message.get("text")
                    data = text_data if bytes_data is None else bytes_data
                    await self.send(data)  # type: ignore[arg-type]

                elif message_type == "websocket.close":
                    message = cast("WebSocketCloseEvent", message)
                    code = message.get("code", 1000)
                    reason = message.get("reason", "") or ""
                    await self.close(code, reason)
                    self.closed_event.set()

                else:
                    msg = (
                        "Expected ASGI message 'websocket.send' or 'websocket.close',"
                        " but got '%s'."
                    )
                    raise RuntimeError(msg % message_type)
            except ConnectionClosed as exc:
                raise ClientDisconnected from exc

        elif self.initial_response is not None:
            if message_type == "websocket.http.response.body":
                message = cast("WebSocketResponseBodyEvent", message)
                body = self.initial_response[2] + message["body"]
                self.initial_response = self.initial_response[:2] + (body,)
                if not message.get("more_body", False):
                    self.closed_event.set()
            else:
                msg = (
                    "Expected ASGI message 'websocket.http.response.body' "
                    "but got '%s'."
                )
                raise RuntimeError(msg % message_type)

        else:
            msg = (
                "Unexpected ASGI message '%s', after sending 'websocket.close' "
                "or response already completed."
            )
            raise RuntimeError(msg % message_type)

    async def asgi_receive(
        self,
    ) -> Union[
        "WebSocketDisconnectEvent", "WebSocketConnectEvent", "WebSocketReceiveEvent"
    ]:
        if not self.connect_sent:
            self.connect_sent = True
            return {"type": "websocket.connect"}

        await self.handshake_completed_event.wait()

        if self.lost_connection_before_handshake:
            # If the handshake failed or the app closed before handshake completion,
            # use 1006 Abnormal Closure.
            return {"type": "websocket.disconnect", "code": 1006}

        if self.closed_event.is_set():
            return {"type": "websocket.disconnect", "code": 1005}

        try:
            data = await self.recv()
        except ConnectionClosed as exc:
            self.closed_event.set()
            if self.ws_server.closing:
                return {"type": "websocket.disconnect", "code": 1012}
            return {"type": "websocket.disconnect", "code": exc.code}

        if isinstance(data, str):
            return {"type": "websocket.receive", "text": data}
        return {"type": "websocket.receive", "bytes": data}