14641300_體育班的肉便器男孩 13

admin

from __future__ import annotations

import functools
import http
import logging
import re
import ssl as ssl_module
from collections.abc import Awaitable, Mapping, Sequence
from types import TracebackType
from typing import Any, Callable, Coroutine, Self

import trio
import trio.abc

from ..asyncio.server import basic_auth
from ..extensions.base import ServerExtensionFactory
from ..extensions.permessage_deflate import enable_server_permessage_deflate
from ..frames import CloseCode
from ..headers import validate_subprotocols
from ..http11 import SERVER, Request, Response
from ..protocol import CONNECTING, OPEN, Event
from ..server import ServerProtocol
from ..typing import LoggerLike, Origin, StatusLike, Subprotocol
from ..utils import get_socket_name
from .connection import Connection, broadcast
from .utils import race_events


__all__ = [
    "broadcast",
    "serve",
    "ServerConnection",
    "Server",
    "basic_auth",
]


class ServerConnection(Connection):
    """
    :mod:`trio` implementation of a WebSocket server connection.

    :class:`ServerConnection` provides :meth:`recv` and :meth:`send` methods for
    receiving and sending messages.

    It supports asynchronous iteration to receive messages::

        async for message in websocket:
            await process(message)

    The iterator exits normally when the connection is closed with close code
    1000 (OK) or 1001 (going away) or without a close code. It raises a
    :exc:`~websockets.exceptions.ConnectionClosedError` when the connection is
    closed with any other code.

    The ``ping_interval``, ``ping_timeout``, ``close_timeout``, and
    ``max_queue`` arguments have the same meaning as in :func:`serve`.

    Args:
        nursery: Trio nursery.
        stream: Trio stream connected to a WebSocket client.
        protocol: Sans-I/O connection.
        server: Server that manages this connection.

    """

    def __init__(
        self,
        nursery: trio.Nursery,
        stream: trio.abc.Stream,
        protocol: ServerProtocol,
        server: Server,
        *,
        ping_interval: float | None = 20,
        ping_timeout: float | None = 20,
        close_timeout: float | None = 10,
        max_queue: int | None | tuple[int | None, int | None] = 16,
    ) -> None:
        self.protocol: ServerProtocol
        super().__init__(
            nursery,
            stream,
            protocol,
            ping_interval=ping_interval,
            ping_timeout=ping_timeout,
            close_timeout=close_timeout,
            max_queue=max_queue,
        )
        self.server = server
        self.request_rcvd: trio.Event = trio.Event()
        self.username: str  # see basic_auth()
        self.handler: Callable[[ServerConnection], Awaitable[None]]  # see route()
        self.handler_kwargs: Mapping[str, Any]  # see route()

    def respond(self, status: StatusLike, text: str) -> Response:
        """
        Create a plain text HTTP response.

        ``process_request`` and ``process_response`` may call this method to
        return an HTTP response instead of performing the WebSocket opening
        handshake.

        You can modify the response before returning it, for example by changing
        HTTP headers.

        Args:
            status: HTTP status code.
            text: HTTP response body; it will be encoded to UTF-8.

        Returns:
            HTTP response to send to the client.

        """
        return self.protocol.reject(status, text)

    async def handshake(
        self,
        process_request: (
            Callable[
                [ServerConnection, Request],
                Awaitable[Response | None] | Response | None,
            ]
            | None
        ) = None,
        process_response: (
            Callable[
                [ServerConnection, Request, Response],
                Awaitable[Response | None] | Response | None,
            ]
            | None
        ) = None,
        server_header: str | None = SERVER,
    ) -> None:
        """
        Perform the opening handshake.

        """
        await race_events(self.request_rcvd, self.stream_closed)

        if self.request is not None:
            response = None

            if process_request is not None:
                try:
                    response = process_request(self, self.request)
                    if isinstance(response, Awaitable):
                        response = await response
                except Exception as exc:
                    self.protocol.handshake_exc = exc
                    self.logger.error("process_request failed", exc_info=True)
                    response = self.protocol.reject(
                        http.HTTPStatus.INTERNAL_SERVER_ERROR,
                        (
                            "Failed to open a WebSocket connection.\n"
                            "See server log for more information.\n"
                        ),
                    )

            if response is None:
                self.response = self.protocol.accept(self.request)
            else:
                assert isinstance(response, Response)  # help mypy
                self.response = response

            if server_header is not None:
                self.response.headers["Server"] = server_header

            response = None

            if process_response is not None:
                try:
                    response = process_response(self, self.request, self.response)
                    if isinstance(response, Awaitable):
                        response = await response
                except Exception as exc:
                    self.protocol.handshake_exc = exc
                    self.logger.error("process_response failed", exc_info=True)
                    response = self.protocol.reject(
                        http.HTTPStatus.INTERNAL_SERVER_ERROR,
                        (
                            "Failed to open a WebSocket connection.\n"
                            "See server log for more information.\n"
                        ),
                    )

            if response is not None:
                assert isinstance(response, Response)  # help mypy
                self.response = response

            # Reject the connection if the server started closing during the
            # opening handshake. Don't yield before send_response() to avoid
            # a race condition after checking if the server is closing.
            if (
                self.response.status_code == http.HTTPStatus.SWITCHING_PROTOCOLS
                and self.server.closing
            ):
                self.response = self.protocol.reject(
                    http.HTTPStatus.SERVICE_UNAVAILABLE,
                    "Server is shutting down.\n",
                )

            # Don't respond if the connection was closed during the handshake.
            if self.state is CONNECTING:
                async with self.send_context(expected_state=CONNECTING):
                    self.protocol.send_response(self.response)

    def process_event(self, event: Event) -> None:
        """
        Process one incoming event.

        """
        # First event - handshake request.
        if self.request is None:
            assert isinstance(event, Request)
            self.request = event
            self.request_rcvd.set()
        # Later events - frames.
        else:
            super().process_event(event)


class Server(trio.abc.AsyncResource):
    """
    WebSocket server returned by :func:`serve`.

    Args:
        listeners: List of Trio listeners accepting new connections.
        handler: Handler for one connection. It receives a Trio stream.
        logger: Logger for this server.
            It defaults to ``logging.getLogger("websockets.server")``.
            See the :doc:`logging guide <../../topics/logging>` for details.

    """

    def __init__(
        self,
        listeners: list[trio.SocketListener],
        handler: Callable[[trio.abc.Stream], Coroutine[Any, Any, None]],
        logger: LoggerLike | None = None,
    ) -> None:
        self.listeners = listeners
        self.handler = handler
        if logger is None:
            logger = logging.getLogger("websockets.server")
        self.logger = logger

        # Keep track of active connections.
        # Trio keeps track of connection handler tasks.
        self.all_connections: set[ServerConnection] = set()

        # Completed when all handlers are done.
        self.handlers_waiter = trio.Event()

        self.closing = False

    @property
    def connections(self) -> set[ServerConnection]:
        """
        Set of active connections.

        This property contains all connections that completed the opening
        handshake successfully and didn't start the closing handshake yet.
        It can be useful in combination with :func:`~broadcast`.

        """
        return {
            connection
            for connection in self.all_connections
            if connection.protocol.state is OPEN
        }

    async def serve_forever(
        self,
        task_status: trio.TaskStatus[Server] = trio.TASK_STATUS_IGNORED,
    ) -> None:
        # Running handlers in a dedicated nursery makes it possible to close
        # listeners while handlers finish running. The nursery for listeners
        # is created in trio.serve_listeners().
        async with trio.open_nursery() as self.handler_nursery:
            # Wrap trio.serve_listeners() in another nursery to return the
            # Server object in task_status instead of a list of listeners.
            async with trio.open_nursery() as self.serve_nursery:
                await self.serve_nursery.start(
                    functools.partial(
                        trio.serve_listeners,
                        self.handler,
                        self.listeners,
                        handler_nursery=self.handler_nursery,
                    )
                )
                for listener in self.listeners:
                    self.logger.info(
                        "server listening on %s",
                        # listener.socket is a Trio socket, not a socket.socket,
                        # but it offers the same APIs used by get_socket_name().
                        get_socket_name(listener.socket),  # type: ignore
                    )
                task_status.started(self)
        # When the nursery for handlers has exited, all handlers have returned.
        self.handlers_waiter.set()

    # Shutting down the server cleanly when serve_forever() is canceled would be
    # the most idiomatic in Trio. However, that would require shielding too many
    # asynchronous operations, including the TLS & WebSocket opening handshakes.

    async def aclose(
        self,
        close_connections: bool = True,
        code: CloseCode | int = CloseCode.GOING_AWAY,
        reason: str = "",
    ) -> None:
        """
        Close the server.

        * Close the TCP listeners.
        * When ``close_connections`` is :obj:`True`, which is the default,
          close existing connections. Specifically:

          * Reject opening WebSocket connections with an HTTP 503 (service
            unavailable) error. This happens when the server accepted the TCP
            connection but didn't complete the opening handshake before closing.
          * Close open WebSocket connections with close code 1001 (going away).
            ``code`` and ``reason`` can be customized, for example to use code
            1012 (service restart).

        * Wait until all connection handlers have returned.

        :meth:`aclose` is idempotent.

        """
        self.logger.info("server closing")

        # Stop accepting new connections.
        self.serve_nursery.cancel_scope.cancel()

        # Reject OPENING connections with HTTP 503 — see handshake().
        self.closing = True

        # Close OPEN connections.
        if close_connections:
            for connection in self.all_connections:
                if connection.protocol.state is OPEN:  # pragma: no branch
                    self.handler_nursery.start_soon(connection.aclose, code, reason)

        # Wait until all connection handlers have returned.
        await self.handlers_waiter.wait()

        self.logger.info("server closed")

    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        await self.aclose()


async def serve(
    handler: Callable[[ServerConnection], Awaitable[None]],
    port: int | None = None,
    *,
    # TCP/TLS
    host: str | bytes | None = None,
    backlog: int | None = None,
    listeners: list[trio.SocketListener] | None = None,
    ssl: ssl_module.SSLContext | None = None,
    # WebSocket
    origins: Sequence[Origin | re.Pattern[str] | None] | None = None,
    extensions: Sequence[ServerExtensionFactory] | None = None,
    subprotocols: Sequence[Subprotocol] | None = None,
    select_subprotocol: (
        Callable[
            [ServerConnection, Sequence[Subprotocol]],
            Subprotocol | None,
        ]
        | None
    ) = None,
    compression: str | None = "deflate",
    # HTTP
    process_request: (
        Callable[
            [ServerConnection, Request],
            Awaitable[Response | None] | Response | None,
        ]
        | None
    ) = None,
    process_response: (
        Callable[
            [ServerConnection, Request, Response],
            Awaitable[Response | None] | Response | None,
        ]
        | None
    ) = None,
    server_header: str | None = SERVER,
    # Timeouts
    open_timeout: float | None = 10,
    ping_interval: float | None = 20,
    ping_timeout: float | None = 20,
    close_timeout: float | None = 10,
    # Limits
    max_size: int | None | tuple[int | None, int | None] = 2**20,
    max_queue: int | None | tuple[int | None, int | None] = 16,
    # Logging
    logger: LoggerLike | None = None,
    # Escape hatch for advanced customization