Initial commit: AuroraMusic monorepo
This commit is contained in:
Kris
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,785 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import ssl as ssl_module
|
||||
import traceback
|
||||
import urllib.parse
|
||||
from collections.abc import AsyncIterator, Generator, Sequence
|
||||
from types import TracebackType
|
||||
from typing import Any, Callable, Literal
|
||||
|
||||
import trio
|
||||
|
||||
from ..asyncio.client import process_exception
|
||||
from ..client import ClientProtocol, backoff
|
||||
from ..datastructures import Headers, HeadersLike
|
||||
from ..exceptions import (
|
||||
InvalidProxyMessage,
|
||||
InvalidProxyStatus,
|
||||
InvalidStatus,
|
||||
ProxyError,
|
||||
SecurityError,
|
||||
)
|
||||
from ..extensions.base import ClientExtensionFactory
|
||||
from ..extensions.permessage_deflate import enable_client_permessage_deflate
|
||||
from ..headers import validate_subprotocols
|
||||
from ..http11 import USER_AGENT, Response
|
||||
from ..protocol import CONNECTING, Event
|
||||
from ..proxy import Proxy, get_proxy, parse_proxy, prepare_connect_request
|
||||
from ..streams import StreamReader
|
||||
from ..typing import LoggerLike, Origin, Subprotocol
|
||||
from ..uri import WebSocketURI, parse_uri
|
||||
from .connection import Connection
|
||||
from .utils import race_events
|
||||
|
||||
|
||||
__all__ = ["connect", "unix_connect", "ClientConnection"]
|
||||
|
||||
MAX_REDIRECTS = int(os.environ.get("WEBSOCKETS_MAX_REDIRECTS", "10"))
|
||||
|
||||
|
||||
class ClientConnection(Connection):
|
||||
"""
|
||||
:mod:`trio` implementation of a WebSocket client connection.
|
||||
|
||||
:class:`ClientConnection` provides :meth:`recv` and :meth:`send` coroutines
|
||||
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:`connect`.
|
||||
|
||||
Args:
|
||||
nursery: Trio nursery.
|
||||
stream: Trio stream connected to a WebSocket server.
|
||||
protocol: Sans-I/O connection.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
nursery: trio.Nursery,
|
||||
stream: trio.abc.Stream,
|
||||
protocol: ClientProtocol,
|
||||
*,
|
||||
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: ClientProtocol
|
||||
super().__init__(
|
||||
nursery,
|
||||
stream,
|
||||
protocol,
|
||||
ping_interval=ping_interval,
|
||||
ping_timeout=ping_timeout,
|
||||
close_timeout=close_timeout,
|
||||
max_queue=max_queue,
|
||||
)
|
||||
self.response_rcvd = trio.Event()
|
||||
|
||||
async def handshake(
|
||||
self,
|
||||
additional_headers: HeadersLike | None = None,
|
||||
user_agent_header: str | None = USER_AGENT,
|
||||
) -> None:
|
||||
"""
|
||||
Perform the opening handshake.
|
||||
|
||||
"""
|
||||
self.request = self.protocol.connect()
|
||||
if additional_headers is not None:
|
||||
self.request.headers.update(additional_headers)
|
||||
if user_agent_header is not None:
|
||||
self.request.headers.setdefault("User-Agent", user_agent_header)
|
||||
async with self.send_context(expected_state=CONNECTING):
|
||||
self.protocol.send_request(self.request)
|
||||
|
||||
await race_events(self.response_rcvd, self.stream_closed)
|
||||
|
||||
# self.protocol.handshake_exc is set when the connection is lost before
|
||||
# receiving a response, when the response cannot be parsed, or when the
|
||||
# response fails the handshake.
|
||||
|
||||
if self.protocol.handshake_exc is not None:
|
||||
raise self.protocol.handshake_exc
|
||||
|
||||
def process_event(self, event: Event) -> None:
|
||||
"""
|
||||
Process one incoming event.
|
||||
|
||||
"""
|
||||
# First event - handshake response.
|
||||
if self.response is None:
|
||||
assert isinstance(event, Response)
|
||||
self.response = event
|
||||
self.response_rcvd.set()
|
||||
# Later events - frames.
|
||||
else:
|
||||
super().process_event(event)
|
||||
|
||||
|
||||
# This is spelled in lower case because it's exposed as a callable in the API.
|
||||
class connect:
|
||||
"""
|
||||
Connect to the WebSocket server at ``uri``.
|
||||
|
||||
This coroutine returns a :class:`ClientConnection` instance, which you can
|
||||
use to send and receive messages.
|
||||
|
||||
:func:`connect` may be used as an asynchronous context manager::
|
||||
|
||||
from websockets.trio.client import connect
|
||||
|
||||
async with connect(...) as websocket:
|
||||
...
|
||||
|
||||
The connection is closed automatically when exiting the context.
|
||||
|
||||
:func:`connect` can be used as an infinite asynchronous iterator to
|
||||
reconnect automatically on errors::
|
||||
|
||||
async for websocket in connect(...):
|
||||
try:
|
||||
...
|
||||
except websockets.exceptions.ConnectionClosed:
|
||||
continue
|
||||
|
||||
If the connection fails with a transient error, it is retried with
|
||||
exponential backoff. If it fails with a fatal error, the exception is
|
||||
raised, breaking out of the loop.
|
||||
|
||||
The connection is closed automatically after each iteration of the loop.
|
||||
|
||||
Args:
|
||||
uri: URI of the WebSocket server.
|
||||
stream: Preexisting TCP stream. ``stream`` overrides the host and port
|
||||
from ``uri``. You may call :func:`~trio.open_tcp_stream` to create a
|
||||
suitable TCP stream.
|
||||
ssl: Configuration for enabling TLS on the connection.
|
||||
server_hostname: Host name for the TLS handshake. ``server_hostname``
|
||||
overrides the host name from ``uri``.
|
||||
origin: Value of the ``Origin`` header, for servers that require it.
|
||||
extensions: List of supported extensions, in order in which they
|
||||
should be negotiated and run.
|
||||
subprotocols: List of supported subprotocols, in order of decreasing
|
||||
preference.
|
||||
compression: The "permessage-deflate" extension is enabled by default.
|
||||
Set ``compression`` to :obj:`None` to disable it. See the
|
||||
:doc:`compression guide <../../topics/compression>` for details.
|
||||
additional_headers: Arbitrary HTTP headers to add to the handshake
|
||||
request.
|
||||
user_agent_header: Value of the ``User-Agent`` request header.
|
||||
It defaults to ``"Python/x.y.z websockets/X.Y"``.
|
||||
Setting it to :obj:`None` removes the header.
|
||||
proxy: If a proxy is configured, it is used by default. Set ``proxy``
|
||||
to :obj:`None` to disable the proxy or to the address of a proxy
|
||||
to override the system configuration. See the :doc:`proxy docs
|
||||
<../../topics/proxies>` for details.
|
||||
proxy_ssl: Configuration for enabling TLS on the proxy connection.
|
||||
proxy_server_hostname: Host name for the TLS handshake with the proxy.
|
||||
``proxy_server_hostname`` overrides the host name from ``proxy``.
|
||||
process_exception: When reconnecting automatically, tell whether an
|
||||
error is transient or fatal. The default behavior is defined by
|
||||
:func:`process_exception`. Refer to its documentation for details.
|
||||
open_timeout: Timeout for opening the connection in seconds.
|
||||
:obj:`None` disables the timeout.
|
||||
ping_interval: Interval between keepalive pings in seconds.
|
||||
:obj:`None` disables keepalive.
|
||||
ping_timeout: Timeout for keepalive pings in seconds.
|
||||
:obj:`None` disables timeouts.
|
||||
close_timeout: Timeout for closing the connection in seconds.
|
||||
:obj:`None` disables the timeout.
|
||||
reconnect_delays: Delays in seconds between reconnection attempts.
|
||||
Default is exponential backoff with 5s jitter, capped at 60s.
|
||||
max_size: Maximum size of incoming messages in bytes.
|
||||
:obj:`None` disables the limit. You may pass a ``(max_message_size,
|
||||
max_fragment_size)`` tuple to set different limits for messages and
|
||||
fragments when you expect long messages sent in short fragments.
|
||||
max_queue: High-water mark of the buffer where frames are received.
|
||||
It defaults to 16 frames. The low-water mark defaults to ``max_queue
|
||||
// 4``. You may pass a ``(high, low)`` tuple to set the high-water
|
||||
and low-water marks. If you want to disable flow control entirely,
|
||||
you may set it to ``None``, although that's a bad idea.
|
||||
logger: Logger for this client.
|
||||
It defaults to ``logging.getLogger("websockets.client")``.
|
||||
See the :doc:`logging guide <../../topics/logging>` for details.
|
||||
create_connection: Factory for the :class:`ClientConnection` managing
|
||||
the connection. Set it to a wrapper or a subclass to customize
|
||||
connection handling.
|
||||
|
||||
Any other keyword arguments are passed to :func:`~trio.open_tcp_stream`.
|
||||
|
||||
Raises:
|
||||
InvalidURI: If ``uri`` isn't a valid WebSocket URI.
|
||||
InvalidProxy: If ``proxy`` isn't a valid proxy.
|
||||
OSError: If the TCP connection fails.
|
||||
InvalidHandshake: If the opening handshake fails.
|
||||
TimeoutError: If the opening handshake times out.
|
||||
|
||||
"""
|
||||
|
||||
# Arguments of type SSLContext don't render correctly in the documentation
|
||||
# because of https://github.com/sphinx-doc/sphinx/issues/13838.
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
uri: str,
|
||||
*,
|
||||
# TCP/TLS
|
||||
stream: trio.abc.Stream | None = None,
|
||||
ssl: ssl_module.SSLContext | None = None,
|
||||
server_hostname: str | None = None,
|
||||
# WebSocket
|
||||
origin: Origin | None = None,
|
||||
extensions: Sequence[ClientExtensionFactory] | None = None,
|
||||
subprotocols: Sequence[Subprotocol] | None = None,
|
||||
compression: str | None = "deflate",
|
||||
# HTTP
|
||||
additional_headers: HeadersLike | None = None,
|
||||
user_agent_header: str | None = USER_AGENT,
|
||||
proxy: str | Literal[True] | None = True,
|
||||
proxy_ssl: ssl_module.SSLContext | None = None,
|
||||
proxy_server_hostname: str | None = None,
|
||||
process_exception: Callable[[Exception], Exception | None] = process_exception,
|
||||
# Timeouts
|
||||
open_timeout: float | None = 10,
|
||||
ping_interval: float | None = 20,
|
||||
ping_timeout: float | None = 20,
|
||||
close_timeout: float | None = 10,
|
||||
reconnect_delays: Callable[[], Generator[float]] = backoff,
|
||||
# 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
|
||||
create_connection: type[ClientConnection] | None = None,
|
||||
# Other keyword arguments are passed to trio.open_tcp_stream
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self.uri = uri
|
||||
self.ws_uri = parse_uri(uri)
|
||||
if not self.ws_uri.secure and ssl is not None:
|
||||
raise ValueError("ssl argument is incompatible with a ws:// URI")
|
||||
|
||||
if subprotocols is not None:
|
||||
validate_subprotocols(subprotocols)
|
||||
|
||||
if compression == "deflate":
|
||||
extensions = enable_client_permessage_deflate(extensions)
|
||||
elif compression is not None:
|
||||
raise ValueError(f"unsupported compression: {compression}")
|
||||
|
||||
if logger is None:
|
||||
logger = logging.getLogger("websockets.client")
|
||||
|
||||
if create_connection is None:
|
||||
create_connection = ClientConnection
|
||||
|
||||
self.stream = stream
|
||||
self.ssl = ssl
|
||||
self.server_hostname = server_hostname
|
||||
self.additional_headers = additional_headers
|
||||
self.user_agent_header = user_agent_header
|
||||
self.proxy = proxy
|
||||
self.proxy_ssl = proxy_ssl
|
||||
self.proxy_server_hostname = proxy_server_hostname
|
||||
self.process_exception = process_exception
|
||||
self.open_timeout = open_timeout
|
||||
self.reconnect_delays = reconnect_delays
|
||||
self.logger = logger
|
||||
self.create_connection = create_connection
|
||||
self.open_tcp_stream_kwargs = kwargs
|
||||
self.protocol_kwargs = dict(
|
||||
origin=origin,
|
||||
extensions=extensions,
|
||||
subprotocols=subprotocols,
|
||||
max_size=max_size,
|
||||
logger=logger,
|
||||
)
|
||||
self.connection_kwargs = dict(
|
||||
ping_interval=ping_interval,
|
||||
ping_timeout=ping_timeout,
|
||||
close_timeout=close_timeout,
|
||||
max_queue=max_queue,
|
||||
)
|
||||
|
||||
async def open_tcp_stream(self) -> trio.abc.Stream:
|
||||
"""Open a TCP or Unix connection to the server, possibly through a proxy."""
|
||||
kwargs = self.open_tcp_stream_kwargs.copy()
|
||||
|
||||
proxy = self.proxy
|
||||
if kwargs.get("unix", False):
|
||||
proxy = None
|
||||
if proxy is True:
|
||||
proxy = get_proxy(self.ws_uri)
|
||||
|
||||
if kwargs.pop("unix", False):
|
||||
return await trio.open_unix_socket(kwargs["path"])
|
||||
|
||||
elif proxy is not None:
|
||||
proxy_parsed = parse_proxy(proxy)
|
||||
|
||||
if proxy_parsed.scheme[:5] == "socks":
|
||||
return await connect_socks_proxy(
|
||||
proxy_parsed,
|
||||
self.ws_uri,
|
||||
# websockets is consistent with trio while python_socks is
|
||||
# consistent across implementations.
|
||||
local_addr=kwargs.pop("local_address", None),
|
||||
)
|
||||
|
||||
elif proxy_parsed.scheme[:4] == "http":
|
||||
if proxy_parsed.scheme != "https" and self.proxy_ssl is not None:
|
||||
raise ValueError(
|
||||
"proxy_ssl argument is incompatible with an http:// proxy"
|
||||
)
|
||||
return await connect_http_proxy(
|
||||
proxy_parsed,
|
||||
self.ws_uri,
|
||||
user_agent_header=self.user_agent_header,
|
||||
ssl=self.proxy_ssl,
|
||||
server_hostname=self.proxy_server_hostname,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError("parse_proxy returned unsupported proxy")
|
||||
|
||||
else: # proxy is None
|
||||
kwargs.setdefault("host", self.ws_uri.host)
|
||||
kwargs.setdefault("port", self.ws_uri.port)
|
||||
return await trio.open_tcp_stream(**kwargs)
|
||||
|
||||
async def enable_tls(self, stream: trio.abc.Stream) -> trio.abc.Stream:
|
||||
"""Enable TLS on the connection."""
|
||||
if self.ssl is None:
|
||||
ssl = ssl_module.create_default_context()
|
||||
else:
|
||||
ssl = self.ssl
|
||||
if self.server_hostname is None:
|
||||
server_hostname = self.ws_uri.host
|
||||
else:
|
||||
server_hostname = self.server_hostname
|
||||
ssl_stream = trio.SSLStream(
|
||||
stream,
|
||||
ssl,
|
||||
server_hostname=server_hostname,
|
||||
https_compatible=True,
|
||||
)
|
||||
await ssl_stream.do_handshake()
|
||||
return ssl_stream
|
||||
|
||||
async def open_connection(self, nursery: trio.Nursery) -> ClientConnection:
|
||||
"""Create a WebSocket connection."""
|
||||
# TCP connection is already established.
|
||||
if self.stream is None:
|
||||
stream = await self.open_tcp_stream()
|
||||
else:
|
||||
stream = self.stream
|
||||
|
||||
try:
|
||||
if self.ws_uri.secure:
|
||||
stream = await self.enable_tls(stream)
|
||||
|
||||
protocol = ClientProtocol(
|
||||
self.ws_uri,
|
||||
**self.protocol_kwargs, # type: ignore
|
||||
)
|
||||
|
||||
# self.create_connection defaults to ClientConnection.
|
||||
connection = self.create_connection(
|
||||
nursery,
|
||||
stream,
|
||||
protocol,
|
||||
**self.connection_kwargs, # type: ignore
|
||||
)
|
||||
|
||||
await connection.handshake(
|
||||
self.additional_headers,
|
||||
self.user_agent_header,
|
||||
)
|
||||
|
||||
return connection
|
||||
|
||||
except trio.Cancelled:
|
||||
await trio.aclose_forcefully(stream)
|
||||
# The nursery running this coroutine was canceled.
|
||||
# The next checkpoint raises trio.Cancelled.
|
||||
# aclose_forcefully() never returns.
|
||||
raise AssertionError("nursery should be canceled")
|
||||
except Exception:
|
||||
# Always close the connection even though keep-alive is the default
|
||||
# in HTTP/1.1 because the current implementation ties opening the
|
||||
# TCP/TLS connection with initializing the WebSocket protocol.
|
||||
await trio.aclose_forcefully(stream)
|
||||
raise
|
||||
|
||||
def process_redirect(self, exc: Exception) -> Exception | str:
|
||||
"""
|
||||
Determine whether a connection error is a redirect that can be followed.
|
||||
|
||||
Return the new URI if it's a valid redirect. Else, return an exception.
|
||||
|
||||
"""
|
||||
if not (
|
||||
isinstance(exc, InvalidStatus)
|
||||
and exc.response.status_code
|
||||
in [
|
||||
300, # Multiple Choices
|
||||
301, # Moved Permanently
|
||||
302, # Found
|
||||
303, # See Other
|
||||
307, # Temporary Redirect
|
||||
308, # Permanent Redirect
|
||||
]
|
||||
and "Location" in exc.response.headers
|
||||
):
|
||||
return exc
|
||||
|
||||
old_ws_uri = self.ws_uri
|
||||
new_uri = urllib.parse.urljoin(self.uri, exc.response.headers["Location"])
|
||||
new_ws_uri = parse_uri(new_uri)
|
||||
|
||||
# If connect() received a stream, it is closed and cannot be reused.
|
||||
if self.stream is not None:
|
||||
return ValueError(
|
||||
f"cannot follow redirect to {new_uri} with a preexisting stream"
|
||||
)
|
||||
|
||||
# TLS downgrade is forbidden.
|
||||
if old_ws_uri.secure and not new_ws_uri.secure:
|
||||
return SecurityError(f"cannot follow redirect to non-secure URI {new_uri}")
|
||||
|
||||
# Apply restrictions to cross-origin redirects.
|
||||
if (
|
||||
old_ws_uri.secure != new_ws_uri.secure
|
||||
or old_ws_uri.host != new_ws_uri.host
|
||||
or old_ws_uri.port != new_ws_uri.port
|
||||
):
|
||||
# Cross-origin redirects on Unix sockets don't quite make sense.
|
||||
if self.open_tcp_stream_kwargs.get("unix", False):
|
||||
return ValueError(
|
||||
f"cannot follow cross-origin redirect to {new_uri} "
|
||||
f"with a Unix socket"
|
||||
)
|
||||
# Cross-origin redirects when host and port are overridden are ill-defined.
|
||||
if (
|
||||
self.open_tcp_stream_kwargs.get("host") is not None
|
||||
or self.open_tcp_stream_kwargs.get("port") is not None
|
||||
):
|
||||
return ValueError(
|
||||
f"cannot follow cross-origin redirect to {new_uri} "
|
||||
f"with an explicit host or port"
|
||||
)
|
||||
|
||||
# Strip credentials to avoid leaking them to a different origin.
|
||||
if self.additional_headers is not None:
|
||||
self.additional_headers = Headers(
|
||||
(
|
||||
(key, value)
|
||||
for key, value in Headers(self.additional_headers).raw_items()
|
||||
if key.lower()
|
||||
not in ["authorization", "cookie", "proxy-authorization"]
|
||||
)
|
||||
)
|
||||
|
||||
return new_uri
|
||||
|
||||
async def connect(self, nursery: trio.Nursery) -> ClientConnection:
|
||||
try:
|
||||
with (
|
||||
trio.CancelScope()
|
||||
if self.open_timeout is None
|
||||
else trio.fail_after(self.open_timeout)
|
||||
):
|
||||
for _ in range(MAX_REDIRECTS):
|
||||
try:
|
||||
connection = await self.open_connection(nursery)
|
||||
except Exception as exc:
|
||||
exc_or_uri = self.process_redirect(exc)
|
||||
if isinstance(exc_or_uri, Exception):
|
||||
# Response isn't a valid redirect; raise the exception.
|
||||
if exc_or_uri is exc:
|
||||
raise
|
||||
else:
|
||||
raise exc_or_uri from exc
|
||||
else:
|
||||
# Response is a valid redirect; follow it.
|
||||
self.uri = exc_or_uri
|
||||
self.ws_uri = parse_uri(exc_or_uri)
|
||||
continue
|
||||
|
||||
else:
|
||||
connection.start_keepalive()
|
||||
return connection
|
||||
else:
|
||||
raise SecurityError(f"more than {MAX_REDIRECTS} redirects")
|
||||
|
||||
except trio.TooSlowError as exc:
|
||||
# Re-raise exception with an informative error message.
|
||||
raise TimeoutError("timed out during opening handshake") from exc
|
||||
|
||||
# Do not define __await__ for... = await nursery.start(connect, ...)
|
||||
# because it doesn't look idiomatic in Trio.
|
||||
|
||||
# async with connect(...) as ...: ...
|
||||
|
||||
async def __aenter__(self) -> ClientConnection:
|
||||
await self.__aenter_nursery__()
|
||||
try:
|
||||
self.connection = await self.connect(self.nursery)
|
||||
return self.connection
|
||||
except BaseException as exc:
|
||||
await self.__aexit_nursery__(type(exc), exc, exc.__traceback__)
|
||||
raise AssertionError("expected __aexit_nursery__ to re-raise the exception")
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
try:
|
||||
try:
|
||||
await self.connection.aclose()
|
||||
finally:
|
||||
del self.connection
|
||||
finally:
|
||||
await self.__aexit_nursery__(exc_type, exc_value, traceback)
|
||||
|
||||
async def __aenter_nursery__(self) -> None:
|
||||
if hasattr(self, "nursery_manager"):
|
||||
raise RuntimeError("connect() isn't reentrant")
|
||||
self.nursery_manager = trio.open_nursery()
|
||||
self.nursery = await self.nursery_manager.__aenter__()
|
||||
|
||||
async def __aexit_nursery__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> None:
|
||||
# We need a nursery to start the recv_events and keepalive coroutines.
|
||||
# They aren't expected to raise exceptions; instead they catch and log
|
||||
# all unexpected errors. To keep the nursery an implementation detail,
|
||||
# unwrap exceptions raised by user code — per the second option here:
|
||||
# https://trio.readthedocs.io/en/stable/reference-core.html#designing-for-multiple-errors
|
||||
try:
|
||||
await self.nursery_manager.__aexit__(exc_type, exc_value, traceback)
|
||||
except BaseException as exc:
|
||||
assert isinstance(exc, BaseExceptionGroup)
|
||||
try:
|
||||
trio._util.raise_single_exception_from_group(exc)
|
||||
except trio._util.MultipleExceptionError:
|
||||
raise AssertionError(
|
||||
"unexpected multiple exceptions; please file a bug report"
|
||||
) from exc
|
||||
finally:
|
||||
del self.nursery_manager
|
||||
|
||||
# async for ... in connect(...):
|
||||
|
||||
async def __aiter__(self) -> AsyncIterator[ClientConnection]:
|
||||
delays: Generator[float] | None = None
|
||||
while True:
|
||||
try:
|
||||
async with self as connection:
|
||||
yield connection
|
||||
except Exception as exc:
|
||||
# Determine whether the exception is retryable or fatal.
|
||||
# The API of process_exception is "return an exception or None";
|
||||
# "raise an exception" is also supported because it's a frequent
|
||||
# mistake. It isn't documented in order to keep the API simple.
|
||||
try:
|
||||
new_exc = self.process_exception(exc)
|
||||
except Exception as raised_exc:
|
||||
new_exc = raised_exc
|
||||
|
||||
# The connection failed with a fatal error.
|
||||
# Raise the exception and exit the loop.
|
||||
if new_exc is exc:
|
||||
raise
|
||||
if new_exc is not None:
|
||||
raise new_exc from exc
|
||||
|
||||
# The connection failed with a retryable error.
|
||||
# Start or continue backoff and reconnect.
|
||||
if delays is None:
|
||||
delays = self.reconnect_delays()
|
||||
delay = next(delays)
|
||||
self.logger.info(
|
||||
"connect failed; reconnecting in %.1f seconds: %s",
|
||||
delay,
|
||||
traceback.format_exception_only(exc)[0].strip(),
|
||||
)
|
||||
await trio.sleep(delay)
|
||||
|
||||
else:
|
||||
# The connection succeeded. Reset backoff.
|
||||
delays = None
|
||||
|
||||
|
||||
def unix_connect(
|
||||
path: str | None = None,
|
||||
uri: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> connect:
|
||||
"""
|
||||
Connect to a WebSocket server listening on a Unix socket.
|
||||
|
||||
This function accepts the same keyword arguments as :func:`connect`.
|
||||
|
||||
It's only available on Unix.
|
||||
|
||||
It's mainly useful for debugging servers listening on Unix sockets.
|
||||
|
||||
Args:
|
||||
path: File system path to the Unix socket.
|
||||
uri: URI of the WebSocket server. ``uri`` defaults to
|
||||
``ws://localhost/`` or, when a ``ssl`` argument is provided, to
|
||||
``wss://localhost/``.
|
||||
|
||||
"""
|
||||
if uri is None:
|
||||
if kwargs.get("ssl") is None:
|
||||
uri = "ws://localhost/"
|
||||
else:
|
||||
uri = "wss://localhost/"
|
||||
return connect(uri=uri, unix=True, path=path, **kwargs)
|
||||
|
||||
|
||||
try:
|
||||
from python_socks import ProxyType
|
||||
from python_socks.async_.trio import Proxy as SocksProxy
|
||||
|
||||
except ImportError:
|
||||
|
||||
async def connect_socks_proxy(
|
||||
proxy: Proxy,
|
||||
ws_uri: WebSocketURI,
|
||||
**kwargs: Any,
|
||||
) -> trio.abc.Stream:
|
||||
raise ImportError("connecting through a SOCKS proxy requires python-socks")
|
||||
|
||||
else:
|
||||
SOCKS_PROXY_TYPES = {
|
||||
"socks5h": ProxyType.SOCKS5,
|
||||
"socks5": ProxyType.SOCKS5,
|
||||
"socks4a": ProxyType.SOCKS4,
|
||||
"socks4": ProxyType.SOCKS4,
|
||||
}
|
||||
|
||||
SOCKS_PROXY_RDNS = {
|
||||
"socks5h": True,
|
||||
"socks5": False,
|
||||
"socks4a": True,
|
||||
"socks4": False,
|
||||
}
|
||||
|
||||
async def connect_socks_proxy(
|
||||
proxy: Proxy,
|
||||
ws_uri: WebSocketURI,
|
||||
**kwargs: Any,
|
||||
) -> trio.abc.Stream:
|
||||
"""Connect via a SOCKS proxy and return the socket."""
|
||||
socks_proxy = SocksProxy(
|
||||
SOCKS_PROXY_TYPES[proxy.scheme],
|
||||
proxy.host,
|
||||
proxy.port,
|
||||
proxy.username,
|
||||
proxy.password,
|
||||
SOCKS_PROXY_RDNS[proxy.scheme],
|
||||
)
|
||||
# connect() is documented to raise OSError.
|
||||
# socks_proxy.connect() re-raises trio.TooSlowError as ProxyTimeoutError.
|
||||
# Wrap other exceptions in ProxyError, a subclass of InvalidHandshake.
|
||||
try:
|
||||
return trio.SocketStream(
|
||||
await socks_proxy.connect(ws_uri.host, ws_uri.port, **kwargs)
|
||||
)
|
||||
except OSError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise ProxyError("failed to connect to SOCKS proxy") from exc
|
||||
|
||||
|
||||
async def read_connect_response(stream: trio.abc.Stream) -> Response:
|
||||
reader = StreamReader()
|
||||
parser = Response.parse(
|
||||
reader.read_line,
|
||||
reader.read_exact,
|
||||
reader.read_to_eof,
|
||||
proxy=True,
|
||||
)
|
||||
try:
|
||||
while True:
|
||||
data = await stream.receive_some(4096)
|
||||
if data:
|
||||
reader.feed_data(data)
|
||||
else:
|
||||
reader.feed_eof()
|
||||
next(parser)
|
||||
except StopIteration as exc:
|
||||
assert isinstance(exc.value, Response) # help mypy
|
||||
response = exc.value
|
||||
if 200 <= response.status_code < 300:
|
||||
return response
|
||||
else:
|
||||
raise InvalidProxyStatus(response)
|
||||
except Exception as exc:
|
||||
raise InvalidProxyMessage(
|
||||
"did not receive a valid HTTP response from proxy"
|
||||
) from exc
|
||||
|
||||
|
||||
async def connect_http_proxy(
|
||||
proxy: Proxy,
|
||||
ws_uri: WebSocketURI,
|
||||
*,
|
||||
user_agent_header: str | None = None,
|
||||
ssl: ssl_module.SSLContext | None = None,
|
||||
server_hostname: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> trio.abc.Stream:
|
||||
stream: trio.abc.Stream
|
||||
stream = await trio.open_tcp_stream(proxy.host, proxy.port, **kwargs)
|
||||
|
||||
try:
|
||||
# Initialize TLS wrapper and perform TLS handshake
|
||||
if proxy.scheme == "https":
|
||||
if ssl is None:
|
||||
ssl = ssl_module.create_default_context()
|
||||
if server_hostname is None:
|
||||
server_hostname = proxy.host
|
||||
ssl_stream = trio.SSLStream(
|
||||
stream,
|
||||
ssl,
|
||||
server_hostname=server_hostname,
|
||||
https_compatible=True,
|
||||
)
|
||||
await ssl_stream.do_handshake()
|
||||
stream = ssl_stream
|
||||
|
||||
# Send CONNECT request to the proxy and read response.
|
||||
request = prepare_connect_request(proxy, ws_uri, user_agent_header)
|
||||
await stream.send_all(request)
|
||||
await read_connect_response(stream)
|
||||
|
||||
except (trio.Cancelled, Exception):
|
||||
await trio.aclose_forcefully(stream)
|
||||
raise
|
||||
|
||||
return stream
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,286 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import codecs
|
||||
import math
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, Callable, Literal, overload
|
||||
|
||||
import trio
|
||||
|
||||
from ..exceptions import ConcurrencyError
|
||||
from ..frames import BINARY, CONT, TEXT, Frame
|
||||
from ..typing import Data
|
||||
|
||||
|
||||
__all__ = ["Assembler"]
|
||||
|
||||
UTF8Decoder = codecs.getincrementaldecoder("utf-8")
|
||||
|
||||
|
||||
class Assembler:
|
||||
"""
|
||||
Assemble messages from frames.
|
||||
|
||||
:class:`Assembler` expects only data frames. The stream of frames must
|
||||
respect the protocol; if it doesn't, the behavior is undefined.
|
||||
|
||||
Args:
|
||||
pause: Called when the buffer of frames goes above the high water mark;
|
||||
should pause reading from the network.
|
||||
resume: Called when the buffer of frames goes below the low water mark;
|
||||
should resume reading from the network.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
high: int | None = None,
|
||||
low: int | None = None,
|
||||
pause: Callable[[], Any] = lambda: None,
|
||||
resume: Callable[[], Any] = lambda: None,
|
||||
) -> None:
|
||||
# Queue of incoming frames.
|
||||
self.send_frames: trio.MemorySendChannel[Frame]
|
||||
self.recv_frames: trio.MemoryReceiveChannel[Frame]
|
||||
self.send_frames, self.recv_frames = trio.open_memory_channel(math.inf)
|
||||
|
||||
# We cannot put a hard limit on the size of the queue because a single
|
||||
# call to Protocol.data_received() could produce thousands of frames,
|
||||
# which must be buffered. Instead, we pause reading when the buffer goes
|
||||
# above the high limit and we resume when it goes under the low limit.
|
||||
if high is not None and low is None:
|
||||
low = high // 4
|
||||
if high is None and low is not None:
|
||||
high = low * 4
|
||||
if high is not None and low is not None:
|
||||
if low < 0:
|
||||
raise ValueError("low must be positive or equal to zero")
|
||||
if high < low:
|
||||
raise ValueError("high must be greater than or equal to low")
|
||||
self.high, self.low = high, low
|
||||
self.pause = pause
|
||||
self.resume = resume
|
||||
self.paused = False
|
||||
|
||||
# This flag prevents concurrent calls to get() by user code.
|
||||
self.get_in_progress = False
|
||||
|
||||
# This flag marks the end of the connection.
|
||||
self.closed = False
|
||||
|
||||
@overload
|
||||
async def get(self, decode: Literal[True]) -> str: ...
|
||||
|
||||
@overload
|
||||
async def get(self, decode: Literal[False]) -> bytes: ...
|
||||
|
||||
@overload
|
||||
async def get(self, decode: bool | None = None) -> Data: ...
|
||||
|
||||
async def get(self, decode: bool | None = None) -> Data:
|
||||
"""
|
||||
Read the next message.
|
||||
|
||||
:meth:`get` returns a single :class:`str` or :class:`bytes`.
|
||||
|
||||
If the message is fragmented, :meth:`get` waits until the last frame is
|
||||
received, then it reassembles the message and returns it. To receive
|
||||
messages frame by frame, use :meth:`get_iter` instead.
|
||||
|
||||
Args:
|
||||
decode: :obj:`False` disables UTF-8 decoding of text frames and
|
||||
returns :class:`bytes`. :obj:`True` forces UTF-8 decoding of
|
||||
binary frames and returns :class:`str`.
|
||||
|
||||
Raises:
|
||||
EOFError: If the stream of frames has ended.
|
||||
UnicodeDecodeError: If a text frame contains invalid UTF-8.
|
||||
ConcurrencyError: If two coroutines run :meth:`get` or
|
||||
:meth:`get_iter` concurrently.
|
||||
|
||||
"""
|
||||
if self.get_in_progress:
|
||||
raise ConcurrencyError("get() or get_iter() is already running")
|
||||
self.get_in_progress = True
|
||||
|
||||
# Locking with get_in_progress prevents concurrent execution
|
||||
# until get() fetches a complete message or is canceled.
|
||||
|
||||
try:
|
||||
# Fetch the first frame.
|
||||
try:
|
||||
frame = await self.recv_frames.receive()
|
||||
except trio.EndOfChannel:
|
||||
raise EOFError("stream of frames ended")
|
||||
self.maybe_resume()
|
||||
assert frame.opcode is TEXT or frame.opcode is BINARY
|
||||
if decode is None:
|
||||
decode = frame.opcode is TEXT
|
||||
frames = [frame]
|
||||
|
||||
# Fetch subsequent frames for fragmented messages.
|
||||
while not frame.fin:
|
||||
try:
|
||||
frame = await self.recv_frames.receive()
|
||||
except trio.Cancelled:
|
||||
# Put frames already received back into the queue
|
||||
# so that future calls to get() can return them.
|
||||
# Bypass the statistics() method for performance.
|
||||
state = self.send_frames._state
|
||||
assert not state.receive_tasks, "no task should receive"
|
||||
assert not state.data, "queue should be empty"
|
||||
for frame in frames:
|
||||
self.send_frames.send_nowait(frame)
|
||||
raise
|
||||
except trio.EndOfChannel:
|
||||
raise EOFError("stream of frames ended")
|
||||
self.maybe_resume()
|
||||
assert frame.opcode is CONT
|
||||
frames.append(frame)
|
||||
|
||||
finally:
|
||||
self.get_in_progress = False
|
||||
|
||||
# This converts frame.data to bytes when it's a bytearray.
|
||||
data = b"".join(frame.data for frame in frames)
|
||||
if decode:
|
||||
return data.decode()
|
||||
else:
|
||||
return data
|
||||
|
||||
@overload
|
||||
def get_iter(self, decode: Literal[True]) -> AsyncIterator[str]: ...
|
||||
|
||||
@overload
|
||||
def get_iter(self, decode: Literal[False]) -> AsyncIterator[bytes]: ...
|
||||
|
||||
@overload
|
||||
def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]: ...
|
||||
|
||||
async def get_iter(self, decode: bool | None = None) -> AsyncIterator[Data]:
|
||||
"""
|
||||
Stream the next message.
|
||||
|
||||
Iterating the return value of :meth:`get_iter` asynchronously yields a
|
||||
:class:`str` or :class:`bytes` for each frame in the message.
|
||||
|
||||
The iterator must be fully consumed before calling :meth:`get_iter` or
|
||||
:meth:`get` again. Else, :exc:`ConcurrencyError` is raised.
|
||||
|
||||
This method only makes sense for fragmented messages. If messages aren't
|
||||
fragmented, use :meth:`get` instead.
|
||||
|
||||
Args:
|
||||
decode: :obj:`False` disables UTF-8 decoding of text frames and
|
||||
returns :class:`bytes`. :obj:`True` forces UTF-8 decoding of
|
||||
binary frames and returns :class:`str`.
|
||||
|
||||
Raises:
|
||||
EOFError: If the stream of frames has ended.
|
||||
UnicodeDecodeError: If a text frame contains invalid UTF-8.
|
||||
ConcurrencyError: If two coroutines run :meth:`get` or
|
||||
:meth:`get_iter` concurrently.
|
||||
|
||||
"""
|
||||
if self.get_in_progress:
|
||||
raise ConcurrencyError("get() or get_iter() is already running")
|
||||
self.get_in_progress = True
|
||||
|
||||
# Locking with get_in_progress prevents concurrent execution
|
||||
# until get_iter() fetches a complete message or is canceled.
|
||||
|
||||
# If get_iter() raises an exception e.g. in decoder.decode(),
|
||||
# get_in_progress remains set and the connection becomes unusable.
|
||||
|
||||
# Yield the first frame.
|
||||
try:
|
||||
frame = await self.recv_frames.receive()
|
||||
except trio.Cancelled:
|
||||
self.get_in_progress = False
|
||||
raise
|
||||
except trio.EndOfChannel:
|
||||
raise EOFError("stream of frames ended")
|
||||
self.maybe_resume()
|
||||
assert frame.opcode is TEXT or frame.opcode is BINARY
|
||||
if decode is None:
|
||||
decode = frame.opcode is TEXT
|
||||
if decode:
|
||||
decoder = UTF8Decoder()
|
||||
yield decoder.decode(frame.data, frame.fin)
|
||||
else:
|
||||
# Convert to bytes when frame.data is a bytearray.
|
||||
yield bytes(frame.data)
|
||||
|
||||
# Yield subsequent frames for fragmented messages.
|
||||
while not frame.fin:
|
||||
# We cannot handle trio.Cancelled because we don't buffer
|
||||
# previous fragments — we're streaming them. Canceling get_iter()
|
||||
# here will leave the assembler in a stuck state. Future calls to
|
||||
# get() or get_iter() will raise ConcurrencyError.
|
||||
try:
|
||||
frame = await self.recv_frames.receive()
|
||||
except trio.EndOfChannel:
|
||||
raise EOFError("stream of frames ended")
|
||||
self.maybe_resume()
|
||||
assert frame.opcode is CONT
|
||||
if decode:
|
||||
yield decoder.decode(frame.data, frame.fin)
|
||||
else:
|
||||
# Convert to bytes when frame.data is a bytearray.
|
||||
yield bytes(frame.data)
|
||||
|
||||
self.get_in_progress = False
|
||||
|
||||
def put(self, frame: Frame) -> None:
|
||||
"""
|
||||
Add ``frame`` to the next message.
|
||||
|
||||
Raises:
|
||||
EOFError: If the stream of frames has ended.
|
||||
|
||||
"""
|
||||
if self.closed:
|
||||
raise EOFError("stream of frames ended")
|
||||
|
||||
self.send_frames.send_nowait(frame)
|
||||
self.maybe_pause()
|
||||
|
||||
def maybe_pause(self) -> None:
|
||||
"""Pause the writer if queue is above the high water mark."""
|
||||
# Skip if flow control is disabled.
|
||||
if self.high is None:
|
||||
return
|
||||
|
||||
# Bypass the statistics() method for performance.
|
||||
# Check for "> high" to support high = 0.
|
||||
if len(self.send_frames._state.data) > self.high and not self.paused:
|
||||
self.paused = True
|
||||
self.pause()
|
||||
|
||||
def maybe_resume(self) -> None:
|
||||
"""Resume the writer if queue is below the low water mark."""
|
||||
# Skip if flow control is disabled.
|
||||
if self.low is None:
|
||||
return
|
||||
|
||||
# Bypass the statistics() method for performance.
|
||||
# Check for "<= low" to support low = 0.
|
||||
if len(self.send_frames._state.data) <= self.low and self.paused:
|
||||
self.paused = False
|
||||
self.resume()
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
End the stream of frames.
|
||||
|
||||
Calling :meth:`close` concurrently with :meth:`get`, :meth:`get_iter`,
|
||||
or :meth:`put` is safe. They will raise :exc:`EOFError`.
|
||||
|
||||
"""
|
||||
if self.closed:
|
||||
return
|
||||
|
||||
self.closed = True
|
||||
|
||||
# Unblock get() or get_iter().
|
||||
self.send_frames.close()
|
||||
@@ -0,0 +1,207 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import http
|
||||
import ssl as ssl_module
|
||||
import urllib.parse
|
||||
from typing import Any, Awaitable, Callable, Literal
|
||||
|
||||
import trio
|
||||
|
||||
from ..http11 import Request, Response
|
||||
from .server import Server, ServerConnection, serve
|
||||
|
||||
|
||||
__all__ = ["route", "Router"]
|
||||
|
||||
|
||||
try:
|
||||
from werkzeug.exceptions import NotFound
|
||||
from werkzeug.routing import Map, RequestRedirect
|
||||
|
||||
except ImportError:
|
||||
|
||||
async def route(
|
||||
url_map: Map,
|
||||
*args: Any,
|
||||
server_name: str | None = None,
|
||||
ssl: ssl_module.SSLContext | Literal[True] | None = None,
|
||||
create_router: type[Router] | None = None,
|
||||
task_status: trio.TaskStatus[Server] = trio.TASK_STATUS_IGNORED,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
raise ImportError("route() requires werkzeug")
|
||||
|
||||
else:
|
||||
|
||||
async def route(
|
||||
url_map: Map,
|
||||
*args: Any,
|
||||
server_name: str | None = None,
|
||||
ssl: ssl_module.SSLContext | Literal[True] | None = None,
|
||||
create_router: type[Router] | None = None,
|
||||
task_status: trio.TaskStatus[Server] = trio.TASK_STATUS_IGNORED,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""
|
||||
Create a WebSocket server dispatching connections to different handlers.
|
||||
|
||||
This feature requires the third-party library `werkzeug`_:
|
||||
|
||||
.. code-block:: console
|
||||
|
||||
$ pip install werkzeug
|
||||
|
||||
.. _werkzeug: https://werkzeug.palletsprojects.com/
|
||||
|
||||
:func:`route` accepts the same arguments as
|
||||
:func:`~websockets.trio.server.serve`, except as described below.
|
||||
|
||||
The first argument is a :class:`werkzeug.routing.Map` that maps URL patterns
|
||||
to connection handlers. In addition to the connection, handlers receive
|
||||
parameters captured in the URL as keyword arguments.
|
||||
|
||||
Here's an example::
|
||||
|
||||
from websockets.trio.router import route
|
||||
from werkzeug.routing import Map, Rule
|
||||
|
||||
async def channel_handler(websocket, channel_id):
|
||||
...
|
||||
|
||||
url_map = Map([
|
||||
Rule("/channel/<uuid:channel_id>", endpoint=channel_handler),
|
||||
...
|
||||
])
|
||||
|
||||
# set this event to exit the server
|
||||
stop = trio.Event()
|
||||
|
||||
with trio.open_nursery() as nursery:
|
||||
server = await nursery.start(route, url_map, ...)
|
||||
async with server:
|
||||
await stop.wait()
|
||||
|
||||
Refer to the documentation of :mod:`werkzeug.routing` for details.
|
||||
|
||||
If you define redirects with ``Rule(..., redirect_to=...)`` in the URL map,
|
||||
when the server runs behind a reverse proxy that modifies the ``Host``
|
||||
header or terminates TLS, you need additional configuration:
|
||||
|
||||
* Set ``server_name`` to the name of the server as seen by clients. When
|
||||
not provided, websockets uses the value of the ``Host`` header.
|
||||
|
||||
* Set ``ssl=True`` to generate ``wss://`` URIs without enabling TLS.
|
||||
Under the hood, this bind the URL map with a ``url_scheme`` of
|
||||
``wss://`` instead of ``ws://``.
|
||||
|
||||
There is no need to specify ``websocket=True`` in each rule. It is added
|
||||
automatically.
|
||||
|
||||
Args:
|
||||
url_map: Mapping of URL patterns to connection handlers.
|
||||
server_name: Name of the server as seen by clients. If :obj:`None`,
|
||||
websockets uses the value of the ``Host`` header.
|
||||
ssl: Configuration for enabling TLS on the connection. Set it to
|
||||
:obj:`True` if a reverse proxy terminates TLS connections.
|
||||
create_router: Factory for the :class:`Router` dispatching requests to
|
||||
handlers. Set it to a wrapper or a subclass to customize routing.
|
||||
task_status: For compatibility with :meth:`nursery.start
|
||||
<trio.Nursery.start>`.
|
||||
|
||||
"""
|
||||
url_scheme = "ws" if ssl is None else "wss"
|
||||
if ssl is not True and ssl is not None:
|
||||
kwargs["ssl"] = ssl
|
||||
|
||||
if create_router is None:
|
||||
create_router = Router
|
||||
|
||||
router = create_router(url_map, server_name, url_scheme)
|
||||
|
||||
_process_request: (
|
||||
Callable[
|
||||
[ServerConnection, Request],
|
||||
Awaitable[Response | None] | Response | None,
|
||||
]
|
||||
| None
|
||||
) = kwargs.pop("process_request", None)
|
||||
if _process_request is None:
|
||||
process_request: Callable[
|
||||
[ServerConnection, Request],
|
||||
Awaitable[Response | None] | Response | None,
|
||||
] = router.route_request
|
||||
else:
|
||||
|
||||
async def process_request(
|
||||
connection: ServerConnection,
|
||||
request: Request,
|
||||
) -> Response | None:
|
||||
response = _process_request(connection, request)
|
||||
if isinstance(response, Awaitable):
|
||||
response = await response
|
||||
if response is not None:
|
||||
return response
|
||||
return router.route_request(connection, request)
|
||||
|
||||
return await serve(
|
||||
router.handler,
|
||||
*args,
|
||||
process_request=process_request,
|
||||
task_status=task_status,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class Router:
|
||||
"""WebSocket router supporting :func:`route`."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
url_map: Map,
|
||||
server_name: str | None = None,
|
||||
url_scheme: str = "ws",
|
||||
) -> None:
|
||||
self.url_map = url_map
|
||||
self.server_name = server_name
|
||||
self.url_scheme = url_scheme
|
||||
for rule in self.url_map.iter_rules():
|
||||
rule.websocket = True
|
||||
|
||||
def get_server_name(self, connection: ServerConnection, request: Request) -> str:
|
||||
if self.server_name is None:
|
||||
return request.headers["Host"]
|
||||
else:
|
||||
return self.server_name
|
||||
|
||||
def redirect(self, connection: ServerConnection, url: str) -> Response:
|
||||
response = connection.respond(http.HTTPStatus.FOUND, f"Found at {url}")
|
||||
response.headers["Location"] = url
|
||||
return response
|
||||
|
||||
def not_found(self, connection: ServerConnection) -> Response:
|
||||
return connection.respond(http.HTTPStatus.NOT_FOUND, "Not Found")
|
||||
|
||||
def route_request(
|
||||
self, connection: ServerConnection, request: Request
|
||||
) -> Response | None:
|
||||
"""Route incoming request."""
|
||||
url_map_adapter = self.url_map.bind(
|
||||
server_name=self.get_server_name(connection, request),
|
||||
url_scheme=self.url_scheme,
|
||||
)
|
||||
try:
|
||||
parsed = urllib.parse.urlparse(request.path)
|
||||
handler, kwargs = url_map_adapter.match(
|
||||
path_info=parsed.path,
|
||||
query_args=parsed.query,
|
||||
)
|
||||
except RequestRedirect as redirect:
|
||||
return self.redirect(connection, redirect.new_url)
|
||||
except NotFound:
|
||||
return self.not_found(connection)
|
||||
connection.handler, connection.handler_kwargs = handler, kwargs
|
||||
return None
|
||||
|
||||
async def handler(self, connection: ServerConnection) -> None:
|
||||
"""Handle a connection."""
|
||||
return await connection.handler(connection, **connection.handler_kwargs)
|
||||
@@ -0,0 +1,658 @@
|
||||
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
|
||||
create_connection: type[ServerConnection] | None = None,
|
||||
# Compatibility with trio.Nursery.start()
|
||||
task_status: trio.TaskStatus[Server] = trio.TASK_STATUS_IGNORED,
|
||||
) -> None:
|
||||
"""
|
||||
Create a WebSocket server listening on ``port``.
|
||||
|
||||
Whenever a client connects, the server creates a :class:`ServerConnection`,
|
||||
performs the opening handshake, and delegates to the ``handler`` coroutine.
|
||||
|
||||
The handler receives the :class:`ServerConnection` instance, which you can
|
||||
use to send and receive messages.
|
||||
|
||||
Once the handler completes, either normally or with an exception, the server
|
||||
performs the closing handshake and closes the connection.
|
||||
|
||||
When using :func:`serve` with :meth:`nursery.start <trio.Nursery.start>`,
|
||||
you get back a :class:`Server` object. Treat it as an asynchronous context
|
||||
manager to ensure that the server will be closed gracefully::
|
||||
|
||||
from websockets.trio.server import serve
|
||||
|
||||
async def handler(websocket):
|
||||
...
|
||||
|
||||
# set this event to exit the server
|
||||
stop = trio.Event()
|
||||
|
||||
with trio.open_nursery() as nursery:
|
||||
server = await nursery.start(serve, handler, port)
|
||||
async with server:
|
||||
await stop.wait()
|
||||
|
||||
Alternatively, to stop the server gracefully, call its
|
||||
:meth:`~Server.aclose` method::
|
||||
|
||||
with trio.open_nursery() as nursery:
|
||||
server = await nursery.start(serve, handler, port)
|
||||
try:
|
||||
await stop.wait()
|
||||
finally:
|
||||
await server.aclose()
|
||||
|
||||
Args:
|
||||
handler: Connection handler. It receives the WebSocket connection,
|
||||
which is a :class:`ServerConnection`, in argument.
|
||||
port: TCP port the server listens on.
|
||||
See :func:`~trio.open_tcp_listeners` for details.
|
||||
host: Network interfaces the server binds to.
|
||||
See :func:`~trio.open_tcp_listeners` for details.
|
||||
backlog: Listen backlog. See :func:`~trio.open_tcp_listeners` for
|
||||
details.
|
||||
listeners: Preexisting TCP listeners. ``listeners`` replaces ``port``,
|
||||
``host``, and ``backlog``. See :func:`trio.serve_listeners` for
|
||||
details.
|
||||
ssl: Configuration for enabling TLS on the connection.
|
||||
origins: Acceptable values of the ``Origin`` header, for defending
|
||||
against Cross-Site WebSocket Hijacking attacks. Values can be
|
||||
:class:`str` to test for an exact match or regular expressions
|
||||
compiled by :func:`re.compile` to test against a pattern. Include
|
||||
:obj:`None` in the list if the lack of an origin is acceptable.
|
||||
extensions: List of supported extensions, in order in which they
|
||||
should be negotiated and run.
|
||||
subprotocols: List of supported subprotocols, in order of decreasing
|
||||
preference.
|
||||
select_subprotocol: Callback for selecting a subprotocol among
|
||||
those supported by the client and the server. It receives a
|
||||
:class:`ServerConnection` (not a
|
||||
:class:`~websockets.server.ServerProtocol`!) instance and a list of
|
||||
subprotocols offered by the client. Other than the first argument,
|
||||
it has the same behavior as the
|
||||
:meth:`ServerProtocol.select_subprotocol
|
||||
<websockets.server.ServerProtocol.select_subprotocol>` method.
|
||||
compression: The "permessage-deflate" extension is enabled by default.
|
||||
Set ``compression`` to :obj:`None` to disable it. See the
|
||||
:doc:`compression guide <../../topics/compression>` for details.
|
||||
process_request: Intercept the request during the opening handshake.
|
||||
Return an HTTP response to force the response or :obj:`None` to
|
||||
continue normally. When you force an HTTP 101 Continue response, the
|
||||
handshake is successful. Else, the connection is aborted.
|
||||
``process_request`` may be a function or a coroutine.
|
||||
process_response: Intercept the response during the opening handshake.
|
||||
Return an HTTP response to force the response or :obj:`None` to
|
||||
continue normally. When you force an HTTP 101 Continue response, the
|
||||
handshake is successful. Else, the connection is aborted.
|
||||
``process_response`` may be a function or a coroutine.
|
||||
server_header: Value of the ``Server`` response header.
|
||||
It defaults to ``"Python/x.y.z websockets/X.Y"``. Setting it to
|
||||
:obj:`None` removes the header.
|
||||
open_timeout: Timeout for opening connections in seconds.
|
||||
:obj:`None` disables the timeout.
|
||||
ping_interval: Interval between keepalive pings in seconds.
|
||||
:obj:`None` disables keepalive.
|
||||
ping_timeout: Timeout for keepalive pings in seconds.
|
||||
:obj:`None` disables timeouts.
|
||||
close_timeout: Timeout for closing connections in seconds.
|
||||
:obj:`None` disables the timeout.
|
||||
max_size: Maximum size of incoming messages in bytes.
|
||||
:obj:`None` disables the limit. You may pass a ``(max_message_size,
|
||||
max_fragment_size)`` tuple to set different limits for messages and
|
||||
fragments when you expect long messages sent in short fragments.
|
||||
max_queue: High-water mark of the buffer where frames are received.
|
||||
It defaults to 16 frames. The low-water mark defaults to ``max_queue
|
||||
// 4``. You may pass a ``(high, low)`` tuple to set the high-water
|
||||
and low-water marks. If you want to disable flow control entirely,
|
||||
you may set it to ``None``, although that's a bad idea.
|
||||
logger: Logger for this server.
|
||||
It defaults to ``logging.getLogger("websockets.server")``.
|
||||
See the :doc:`logging guide <../../topics/logging>` for details.
|
||||
create_connection: Factory for the :class:`ServerConnection` managing
|
||||
the connection. Set it to a wrapper or a subclass to customize
|
||||
connection handling.
|
||||
task_status: For compatibility with :meth:`nursery.start
|
||||
<trio.Nursery.start>`.
|
||||
|
||||
"""
|
||||
|
||||
# Process parameters
|
||||
|
||||
if subprotocols is not None:
|
||||
validate_subprotocols(subprotocols)
|
||||
|
||||
if compression == "deflate":
|
||||
extensions = enable_server_permessage_deflate(extensions)
|
||||
elif compression is not None:
|
||||
raise ValueError(f"unsupported compression: {compression}")
|
||||
|
||||
if create_connection is None:
|
||||
create_connection = ServerConnection
|
||||
|
||||
# Create listeners
|
||||
|
||||
if listeners is None:
|
||||
if port is None:
|
||||
raise ValueError("port is required when listeners is not provided")
|
||||
listeners = await trio.open_tcp_listeners(port, host=host, backlog=backlog)
|
||||
else:
|
||||
if port is not None:
|
||||
raise ValueError("port is incompatible with listeners")
|
||||
if host is not None:
|
||||
raise ValueError("host is incompatible with listeners")
|
||||
if backlog is not None:
|
||||
raise ValueError("backlog is incompatible with listeners")
|
||||
|
||||
async def stream_handler(stream: trio.abc.Stream) -> None:
|
||||
"""
|
||||
Handle the lifecycle of a WebSocket connection.
|
||||
|
||||
Since this coroutine doesn't have a caller that can handle
|
||||
exceptions, it attempts to log relevant ones.
|
||||
|
||||
It guarantees that the TCP connection is closed before exiting.
|
||||
|
||||
"""
|
||||
async with trio.open_nursery() as nursery:
|
||||
try:
|
||||
# Apply open_timeout to the TLS and WebSocket handshake.
|
||||
with (
|
||||
trio.CancelScope()
|
||||
if open_timeout is None
|
||||
else trio.fail_after(open_timeout)
|
||||
):
|
||||
# Enable TLS.
|
||||
if ssl is not None:
|
||||
# Wrap with SSLStream here rather than with TLSListener
|
||||
# in order to include the TLS handshake within open_timeout.
|
||||
stream = trio.SSLStream(
|
||||
stream,
|
||||
ssl,
|
||||
server_side=True,
|
||||
https_compatible=True,
|
||||
)
|
||||
assert isinstance(stream, trio.SSLStream) # help mypy
|
||||
try:
|
||||
await stream.do_handshake()
|
||||
except trio.BrokenResourceError:
|
||||
return
|
||||
|
||||
# Create a closure to give select_subprotocol access to connection.
|
||||
protocol_select_subprotocol: (
|
||||
Callable[
|
||||
[ServerProtocol, Sequence[Subprotocol]],
|
||||
Subprotocol | None,
|
||||
]
|
||||
| None
|
||||
) = None
|
||||
if select_subprotocol is not None:
|
||||
|
||||
def protocol_select_subprotocol(
|
||||
protocol: ServerProtocol,
|
||||
subprotocols: Sequence[Subprotocol],
|
||||
) -> Subprotocol | None:
|
||||
# mypy doesn't know that select_subprotocol is immutable.
|
||||
assert select_subprotocol is not None
|
||||
# Ensure this function is only used in the intended context.
|
||||
assert protocol is connection.protocol
|
||||
return select_subprotocol(connection, subprotocols)
|
||||
|
||||
# Initialize WebSocket protocol.
|
||||
protocol = ServerProtocol(
|
||||
origins=origins,
|
||||
extensions=extensions,
|
||||
subprotocols=subprotocols,
|
||||
select_subprotocol=protocol_select_subprotocol,
|
||||
max_size=max_size,
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
# Initialize WebSocket connection.
|
||||
connection = create_connection(
|
||||
nursery,
|
||||
stream,
|
||||
protocol,
|
||||
server,
|
||||
ping_interval=ping_interval,
|
||||
ping_timeout=ping_timeout,
|
||||
close_timeout=close_timeout,
|
||||
max_queue=max_queue,
|
||||
)
|
||||
|
||||
await connection.handshake(
|
||||
process_request,
|
||||
process_response,
|
||||
server_header,
|
||||
)
|
||||
|
||||
if connection.protocol.state is not OPEN:
|
||||
await connection.close_stream()
|
||||
return
|
||||
|
||||
server.all_connections.add(connection)
|
||||
connection.start_keepalive()
|
||||
try:
|
||||
await handler(connection)
|
||||
except Exception:
|
||||
connection.logger.error("connection handler failed", exc_info=True)
|
||||
await connection.aclose(CloseCode.INTERNAL_ERROR)
|
||||
else:
|
||||
await connection.aclose()
|
||||
finally:
|
||||
server.all_connections.discard(connection)
|
||||
|
||||
except Exception:
|
||||
# Don't leak connections when the opening handshake times out or
|
||||
# an unexpected error occurs.
|
||||
await trio.aclose_forcefully(stream)
|
||||
|
||||
# The server variable is captured by the closure of conn_handler().
|
||||
server = Server(listeners, stream_handler, logger)
|
||||
await server.serve_forever(task_status=task_status)
|
||||
@@ -0,0 +1,36 @@
|
||||
import trio
|
||||
|
||||
|
||||
__all__ = ["race_events"]
|
||||
|
||||
|
||||
# Based on https://trio.readthedocs.io/en/stable/reference-core.html#custom-supervisors
|
||||
|
||||
|
||||
async def jockey(event: trio.Event, cancel_scope: trio.CancelScope) -> None:
|
||||
await event.wait()
|
||||
cancel_scope.cancel()
|
||||
|
||||
|
||||
async def race_events(*events: trio.Event) -> None:
|
||||
"""
|
||||
Wait for any of the given events to be set.
|
||||
|
||||
Args:
|
||||
*events: The events to wait for.
|
||||
|
||||
"""
|
||||
if not events:
|
||||
raise ValueError("no events provided")
|
||||
|
||||
try:
|
||||
async with trio.open_nursery() as nursery:
|
||||
for event in events:
|
||||
nursery.start_soon(jockey, event, nursery.cancel_scope)
|
||||
except BaseExceptionGroup as exc:
|
||||
try:
|
||||
trio._util.raise_single_exception_from_group(exc)
|
||||
except trio._util.MultipleExceptionError:
|
||||
raise AssertionError(
|
||||
"race_events should be canceled; please file a bug report"
|
||||
) from exc
|
||||
Reference in New Issue
Block a user