admin/sourcehub
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first commit

Browse Source
This commit is contained in:
linkitapps
2026-09-14 23:22:10 +09:00
commit 9f9c8fb69b
4207 changed files with 1333187 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.venv/Lib/site-packages/websockets/trio/__init__.py
View File
.venv/Lib/site-packages/websockets/trio/client.py
+800
View File
@@ -0,0 +1,800 @@
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 ..client import ClientProtocol, backoff, process_exception
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, PathLike, 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``.
:func:`connect` should be treated as an asynchronous context manager
yielding a :class:`ClientConnection`, which can then receive and send
messages::
from websockets.trio.client import connect
async with connect(...) as websocket:
...
The connection is closed automatically when exiting the context.
:func:`connect` can also be treated 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.
:func:`connect` cannot be awaited directly. This is because it runs a task
to manage the connection and Trio doesn't support spawning tasks without a
context that ensures completion.
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:`~websockets.client.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`.
For example, you can set ``host`` and ``port`` to connect to a different
host and port from those found in ``uri``. This only changes the destination
of the TCP connection. The host name from ``uri`` is still used in the TLS
handshake for secure connections and in the ``Host`` header.
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()
unix = kwargs.pop("unix", False)
proxy = self.proxy
if unix:
proxy = None
if proxy is True:
proxy = get_proxy(self.ws_uri)
if unix:
return await trio.open_unix_socket(kwargs.pop("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."""
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,
)
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
return connection
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: PathLike | 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/``.
"""
stream = kwargs.get("stream")
if path is None and stream is None:
raise ValueError("missing path argument")
elif path is not None and stream is not None:
raise ValueError("path is incompatible with stream")
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
.venv/Lib/site-packages/websockets/trio/connection.py
+1255
View File
File diff suppressed because it is too large Load Diff
.venv/Lib/site-packages/websockets/trio/messages.py
+291
View File
@@ -0,0 +1,291 @@
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()
# Break reference cycle to allow immediate garbage collection:
# Connection -> Assembler -> bound methods -> Connection.
self.pause = lambda: None
self.resume = lambda: None
.venv/Lib/site-packages/websockets/trio/router.py
+207
View File
@@ -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)
.venv/Lib/site-packages/websockets/trio/server.py
+658
View File
@@ -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)
.venv/Lib/site-packages/websockets/trio/utils.py
+36
View File
@@ -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