R3ap3r/auroramusic
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial commit: AuroraMusic monorepo

Browse Source
This commit is contained in:
Kris
2026-08-29 21:21:51 -05:00
commit 18f124a3d6
8881 changed files with 2283044 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/__init__.py
View File
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/__pycache__/__init__.cpython-312.pyc
BIN
View File
Binary file not shown.
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/__pycache__/client.cpython-312.pyc
BIN
View File
Binary file not shown.
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/__pycache__/connection.cpython-312.pyc
BIN
View File
Binary file not shown.
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/__pycache__/messages.cpython-312.pyc
BIN
View File
Binary file not shown.
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/__pycache__/router.cpython-312.pyc
BIN
View File
Binary file not shown.
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/__pycache__/server.cpython-312.pyc
BIN
View File
Binary file not shown.
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/client.py
+831
View File
@@ -0,0 +1,831 @@
from __future__ import annotations
import asyncio
import logging
import os
import socket
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, cast
from ..client import ClientProtocol, backoff
from ..datastructures import Headers, HeadersLike
from ..exceptions import (
InvalidMessage,
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
__all__ = ["connect", "unix_connect", "ClientConnection"]
MAX_REDIRECTS = int(os.environ.get("WEBSOCKETS_MAX_REDIRECTS", "10"))
class ClientConnection(Connection):
"""
:mod:`asyncio` 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 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``, ``max_queue``,
and ``write_limit`` arguments have the same meaning as in :func:`connect`.
Args:
protocol: Sans-I/O connection.
"""
def __init__(
self,
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,
write_limit: int | tuple[int, int | None] = 2**15,
) -> None:
self.protocol: ClientProtocol
super().__init__(
protocol,
ping_interval=ping_interval,
ping_timeout=ping_timeout,
close_timeout=close_timeout,
max_queue=max_queue,
write_limit=write_limit,
)
self.response_rcvd: asyncio.Future[None] = self.loop.create_future()
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 asyncio.wait(
[self.response_rcvd, self.connection_lost_waiter],
return_when=asyncio.FIRST_COMPLETED,
)
# 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_result(None)
# Later events - frames.
else:
super().process_event(event)
def process_exception(exc: Exception) -> Exception | None:
"""
Determine whether a connection error is retryable or fatal.
When reconnecting automatically with ``async for ... in connect(...)``, if a
connection attempt fails, :func:`process_exception` is called to determine
whether to retry connecting or to raise the exception.
This function defines the default behavior, which is to retry on:
* :exc:`EOFError`, :exc:`OSError`, :exc:`asyncio.TimeoutError`: network
errors;
* :exc:`~websockets.exceptions.InvalidStatus` when the status code is 500,
502, 503, or 504: server or proxy errors.
All other exceptions are considered fatal.
You can change this behavior with the ``process_exception`` argument of
:func:`connect`.
Return :obj:`None` if the exception is retryable i.e. when the error could
be transient and trying to reconnect with the same parameters could succeed.
The exception will be logged at the ``INFO`` level.
Return an exception, either ``exc`` or a new exception, if the exception is
fatal i.e. when trying to reconnect will most likely produce the same error.
That exception will be raised, breaking out of the retry loop.
"""
# This catches python-socks' ProxyConnectionError and ProxyTimeoutError.
if isinstance(exc, (OSError, TimeoutError)):
return None
if isinstance(exc, InvalidMessage) and isinstance(exc.__cause__, EOFError):
return None
if isinstance(exc, InvalidStatus) and exc.response.status_code in [
500, # Internal Server Error
502, # Bad Gateway
503, # Service Unavailable
504, # Gateway Timeout
]:
return None
return exc
# 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.asyncio.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.
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.
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.
write_limit: High-water mark of write buffer in bytes. It is passed to
:meth:`~asyncio.WriteTransport.set_write_buffer_limits`. It defaults
to 32 KiB. You may pass a ``(high, low)`` tuple to set the
high-water and low-water marks.
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 the event loop's
:meth:`~asyncio.loop.create_connection` method.
For example:
* You can set ``sock`` to provide a preexisting TCP socket. You may call
:func:`socket.create_connection` (not to be confused with the event loop's
:meth:`~asyncio.loop.create_connection` method) to create a suitable
client socket and customize it.
* You can set ``ssl`` to a :class:`~ssl.SSLContext` to enforce TLS settings.
When connecting to a ``wss://`` URI, if ``ssl`` isn't provided, a TLS
context is created with :func:`~ssl.create_default_context`.
* You can set ``server_hostname`` to override the host name from ``uri`` in
the TLS handshake.
* You can configure ``ssl_handshake_timeout`` and ``ssl_shutdown_timeout``.
* 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.
When using a proxy:
* Prefix keyword arguments with ``proxy_`` for configuring TLS between the
client and an HTTPS proxy: ``proxy_ssl``, ``proxy_server_hostname``,
``proxy_ssl_handshake_timeout``, and ``proxy_ssl_shutdown_timeout``.
* Use the standard keyword arguments for configuring TLS between the proxy
and the WebSocket server: ``ssl``, ``server_hostname``,
``ssl_handshake_timeout``, and ``ssl_shutdown_timeout``.
* Other keyword arguments are used only for connecting to the proxy.
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.
"""
def __init__(
self,
uri: str,
*,
# 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,
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,
write_limit: int | tuple[int, int | None] = 2**15,
# Logging
logger: LoggerLike | None = None,
# Escape hatches for advanced customization
create_connection: type[ClientConnection] | None = None,
# Other keyword arguments are passed to loop.create_connection
**kwargs: Any,
) -> None:
self.uri = uri
self.ws_uri = parse_uri(uri)
if not self.ws_uri.secure and kwargs.get("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.additional_headers = additional_headers
self.user_agent_header = user_agent_header
self.proxy = proxy
self.process_exception = process_exception
self.open_timeout = open_timeout
self.reconnect_delays = reconnect_delays
self.logger = logger
self.create_connection_kwargs = kwargs
def factory(uri: WebSocketURI) -> ClientConnection:
# This is a protocol in the Sans-I/O implementation of websockets.
protocol = ClientProtocol(
uri,
origin=origin,
extensions=extensions,
subprotocols=subprotocols,
max_size=max_size,
logger=logger,
)
# This is a connection in websockets and a protocol in asyncio.
# create_connection defaults to ClientConnection.
connection = create_connection(
protocol,
ping_interval=ping_interval,
ping_timeout=ping_timeout,
close_timeout=close_timeout,
max_queue=max_queue,
write_limit=write_limit,
)
return connection
self.factory = factory
async def open_tcp_connection(self) -> ClientConnection:
"""Open a TCP or Unix connection to the server, possibly through a proxy."""
loop = asyncio.get_running_loop()
kwargs = self.create_connection_kwargs.copy()
proxy = self.proxy
if kwargs.get("unix", False):
proxy = None
if kwargs.get("sock") is not None:
proxy = None
if proxy is True:
proxy = get_proxy(self.ws_uri)
def factory() -> ClientConnection:
return self.factory(self.ws_uri)
if self.ws_uri.secure:
kwargs.setdefault("ssl", True)
if kwargs.get("ssl") is None:
raise ValueError("ssl=None is incompatible with a wss:// URI")
kwargs.setdefault("server_hostname", self.ws_uri.host)
if kwargs.pop("unix", False):
_, connection = await loop.create_unix_connection(factory, **kwargs)
return connection
if proxy is not None:
proxy_parsed = parse_proxy(proxy)
if proxy_parsed.scheme[:5] == "socks":
# Connect to the server through the SOCKS proxy.
sock = await connect_socks_proxy(
proxy_parsed,
self.ws_uri,
# python_socks uses local_addr like asyncio.
local_addr=kwargs.pop("local_addr", None),
)
# Initialize WebSocket connection via the proxy.
_, connection = await loop.create_connection(
factory,
sock=sock,
**kwargs,
)
return connection
elif proxy_parsed.scheme[:4] == "http":
# Split keyword arguments between the proxy and the server.
all_kwargs, proxy_kwargs, kwargs = kwargs, {}, {}
for key, value in all_kwargs.items():
if key.startswith("ssl") or key == "server_hostname":
kwargs[key] = value
elif key.startswith("proxy_"):
proxy_kwargs[key[6:]] = value
else:
proxy_kwargs[key] = value
# Validate the proxy_ssl argument.
if proxy_parsed.scheme == "https":
proxy_kwargs.setdefault("ssl", True)
if proxy_kwargs.get("ssl") is None:
raise ValueError(
"proxy_ssl=None is incompatible with an https:// proxy"
)
else:
if proxy_kwargs.get("ssl") is not None:
raise ValueError(
"proxy_ssl argument is incompatible with an http:// proxy"
)
# Connect to the server through the HTTP proxy.
transport = await connect_http_proxy(
proxy_parsed,
self.ws_uri,
user_agent_header=self.user_agent_header,
**proxy_kwargs,
)
# Initialize WebSocket connection via the proxy.
connection = factory()
transport.set_protocol(connection)
ssl = kwargs.pop("ssl", None)
if ssl is True:
ssl = ssl_module.create_default_context()
if ssl is not None:
new_transport = await loop.start_tls(
transport, connection, ssl, **kwargs
)
assert new_transport is not None # help mypy
transport = new_transport
connection.connection_made(transport)
return connection
else:
raise AssertionError("parse_proxy returned unsupported proxy")
else: # proxy is None
# Connect to the server directly.
if kwargs.get("sock") is None:
kwargs.setdefault("host", self.ws_uri.host)
kwargs.setdefault("port", self.ws_uri.port)
# Initialize WebSocket connection.
_, connection = await loop.create_connection(factory, **kwargs)
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 socket, it is closed and cannot be reused.
if self.create_connection_kwargs.get("sock") is not None:
return ValueError(
f"cannot follow redirect to {new_uri} with a preexisting socket"
)
# 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.create_connection_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.create_connection_kwargs.get("host") is not None
or self.create_connection_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
# ... = await connect(...)
def __await__(self) -> Generator[Any, None, ClientConnection]:
# Create a suitable iterator by calling __await__ on a coroutine.
return self.__await_impl__().__await__()
async def __await_impl__(self) -> ClientConnection:
try:
async with asyncio.timeout(self.open_timeout):
for _ in range(MAX_REDIRECTS):
connection = await self.open_tcp_connection()
try:
await connection.handshake(
self.additional_headers,
self.user_agent_header,
)
except asyncio.CancelledError:
connection.transport.abort()
raise
except Exception as exc:
# Always close the connection even though keep-alive is
# the default in HTTP/1.1 because create_connection ties
# opening the network connection with initializing the
# protocol. In the current design of connect(), there is
# no easy way to reuse the network connection that works
# in every case nor to reinitialize the protocol.
connection.transport.abort()
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 TimeoutError as exc:
# Re-raise exception with an informative error message.
raise TimeoutError("timed out during opening handshake") from exc
# async with connect(...) as ...: ...
async def __aenter__(self) -> ClientConnection:
if hasattr(self, "connection"):
raise RuntimeError("connect() isn't reentrant")
self.connection = await self
return self.connection
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
try:
await self.connection.close()
finally:
del self.connection
# 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 asyncio.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_.asyncio import Proxy as SocksProxy
except ImportError:
async def connect_socks_proxy(
proxy: Proxy,
ws_uri: WebSocketURI,
**kwargs: Any,
) -> socket.socket:
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,
) -> socket.socket:
"""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() doesn't raise TimeoutError; it gets canceled.
# Wrap other exceptions in ProxyError, a subclass of InvalidHandshake.
try:
return 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
class HTTPProxyConnection(asyncio.Protocol):
def __init__(
self,
ws_uri: WebSocketURI,
proxy: Proxy,
user_agent_header: str | None = None,
):
self.ws_uri = ws_uri
self.proxy = proxy
self.user_agent_header = user_agent_header
self.reader = StreamReader()
self.parser = Response.parse(
self.reader.read_line,
self.reader.read_exact,
self.reader.read_to_eof,
proxy=True,
)
loop = asyncio.get_running_loop()
self.response: asyncio.Future[Response] = loop.create_future()
def run_parser(self) -> None:
try:
next(self.parser)
except StopIteration as exc:
response = exc.value
if 200 <= response.status_code < 300:
self.response.set_result(response)
else:
self.response.set_exception(InvalidProxyStatus(response))
except Exception as exc:
proxy_exc = InvalidProxyMessage(
"did not receive a valid HTTP response from proxy"
)
proxy_exc.__cause__ = exc
self.response.set_exception(proxy_exc)
def connection_made(self, transport: asyncio.BaseTransport) -> None:
transport = cast(asyncio.Transport, transport)
self.transport = transport
self.transport.write(
prepare_connect_request(self.proxy, self.ws_uri, self.user_agent_header)
)
def data_received(self, data: bytes) -> None:
self.reader.feed_data(data)
self.run_parser()
def eof_received(self) -> None:
self.reader.feed_eof()
self.run_parser()
def connection_lost(self, exc: Exception | None) -> None:
self.reader.feed_eof()
self.run_parser()
async def connect_http_proxy(
proxy: Proxy,
ws_uri: WebSocketURI,
*,
user_agent_header: str | None = None,
**kwargs: Any,
) -> asyncio.Transport:
transport, protocol = await asyncio.get_running_loop().create_connection(
lambda: HTTPProxyConnection(ws_uri, proxy, user_agent_header),
proxy.host,
proxy.port,
**kwargs,
)
try:
# This raises exceptions if the connection to the proxy fails.
await protocol.response
except (asyncio.CancelledError, Exception):
transport.abort()
raise
return transport
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/connection.py
+1253
View File
File diff suppressed because it is too large Load Diff
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/messages.py
+316
View File
@@ -0,0 +1,316 @@
from __future__ import annotations
import asyncio
import codecs
import collections
from collections.abc import AsyncIterator, Iterable
from typing import Any, Callable, Generic, Literal, TypeVar, overload
from ..exceptions import ConcurrencyError
from ..frames import BINARY, CONT, TEXT, Frame
from ..typing import Data
__all__ = ["Assembler"]
UTF8Decoder = codecs.getincrementaldecoder("utf-8")
T = TypeVar("T")
class SimpleQueue(Generic[T]):
"""
Simplified version of :class:`asyncio.Queue`.
Provides only the subset of functionality needed by :class:`Assembler`.
"""
def __init__(self) -> None:
self.loop = asyncio.get_running_loop()
self.get_waiter: asyncio.Future[None] | None = None
self.queue: collections.deque[T] = collections.deque()
def __len__(self) -> int:
return len(self.queue)
def put(self, item: T) -> None:
"""Put an item into the queue."""
self.queue.append(item)
if self.get_waiter is not None and not self.get_waiter.done():
self.get_waiter.set_result(None)
async def get(self, block: bool = True) -> T:
"""Remove and return an item from the queue, waiting if necessary."""
if not self.queue:
if not block:
raise EOFError("stream of frames ended")
assert self.get_waiter is None, "cannot call get() concurrently"
self.get_waiter = self.loop.create_future()
try:
await self.get_waiter
finally:
self.get_waiter.cancel()
self.get_waiter = None
return self.queue.popleft()
def reset(self, items: Iterable[T]) -> None:
"""Put back items into an empty, idle queue."""
assert self.get_waiter is None, "cannot reset() while get() is running"
assert not self.queue, "cannot reset() while queue isn't empty"
self.queue.extend(items)
def abort(self) -> None:
"""Close the queue, raising EOFError in get() if necessary."""
if self.get_waiter is not None and not self.get_waiter.done():
self.get_waiter.set_exception(EOFError("stream of frames ended"))
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.frames: SimpleQueue[Frame] = SimpleQueue()
# 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.
frame = await self.frames.get(not self.closed)
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.frames.get(not self.closed)
except asyncio.CancelledError:
# Put frames already received back into the queue
# so that future calls to get() can return them.
self.frames.reset(frames)
raise
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.frames.get(not self.closed)
except asyncio.CancelledError:
self.get_in_progress = False
raise
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 asyncio.CancelledError 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.
frame = await self.frames.get(not self.closed)
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.frames.put(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
# Check for "> high" to support high = 0.
if len(self.frames) > 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
# Check for "<= low" to support low = 0.
if len(self.frames) <= 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.frames.abort()
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/router.py
+227
View File
@@ -0,0 +1,227 @@
from __future__ import annotations
import http
import ssl as ssl_module
import urllib.parse
from typing import Any, Awaitable, Callable, Literal
from ..http11 import Request, Response
from .server import Server, ServerConnection, serve
__all__ = ["route", "unix_route", "Router"]
try:
from werkzeug.exceptions import NotFound
from werkzeug.routing import Map, RequestRedirect
except ImportError:
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,
**kwargs: Any,
) -> Server:
raise ImportError("route() requires werkzeug")
def unix_route(
url_map: Map,
path: str | None = None,
**kwargs: Any,
) -> Server:
raise ImportError("unix_route() requires werkzeug")
else:
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,
**kwargs: Any,
) -> Server:
"""
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.sync.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.asyncio.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 = asyncio.Event()
async with route(url_map, ...) as 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.
Like :func:`~websockets.sync.server.serve`, :func:`route` returns a
:class:`~websockets.sync.server.Server` that you can also run with
:meth:`~websockets.sync.server.Server.serve_forever`.
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.
"""
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 serve(
router.handler,
*args,
process_request=process_request,
**kwargs,
)
def unix_route(
url_map: Map,
path: str | None = None,
**kwargs: Any,
) -> Server:
"""
Create a WebSocket Unix server dispatching connections to different handlers.
:func:`unix_route` combines the behaviors of :func:`route` and
:func:`~websockets.asyncio.server.unix_serve`.
Args:
url_map: Mapping of URL patterns to connection handlers.
path: File system path to the Unix socket.
"""
return route(url_map, unix=True, path=path, **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)
apps/api/.venv/lib/python3.12/site-packages/websockets/asyncio/server.py
+911
View File
@@ -0,0 +1,911 @@
from __future__ import annotations
import asyncio
import hmac
import http
import logging
import re
import socket
from collections.abc import Awaitable, Coroutine, Generator, Iterable, Sequence
from types import TracebackType
from typing import Any, Callable, Mapping, Self, cast
from ..exceptions import InvalidHeader
from ..extensions.base import ServerExtensionFactory
from ..extensions.permessage_deflate import enable_server_permessage_deflate
from ..frames import CloseCode
from ..headers import (
build_www_authenticate_basic,
parse_authorization_basic,
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
__all__ = [
"broadcast",
"serve",
"unix_serve",
"ServerConnection",
"Server",
"basic_auth",
]
class ServerConnection(Connection):
"""
:mod:`asyncio` 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 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``, ``max_queue``,
and ``write_limit`` arguments have the same meaning as in :func:`serve`.
Args:
protocol: Sans-I/O connection.
server: Server that manages this connection.
"""
def __init__(
self,
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,
write_limit: int | tuple[int, int | None] = 2**15,
) -> None:
self.protocol: ServerProtocol
super().__init__(
protocol,
ping_interval=ping_interval,
ping_timeout=ping_timeout,
close_timeout=close_timeout,
max_queue=max_queue,
write_limit=write_limit,
)
self.server = server
self.request_rcvd: asyncio.Future[None] = self.loop.create_future()
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 asyncio.wait(
[self.request_rcvd, self.connection_lost_waiter],
return_when=asyncio.FIRST_COMPLETED,
)
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 not self.server.is_serving()
):
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_result(None)
# Later events - frames.
else:
super().process_event(event)
def connection_made(self, transport: asyncio.BaseTransport) -> None:
super().connection_made(transport)
# The handler task must be registered in self.handler_tasks now. If it
# was registered inside the task, a race condition could happen when
# closing the server after scheduling the task but before it executes.
handler_task = self.loop.create_task(self.server.handler(self))
self.server.handler_tasks.add(handler_task)
class Server:
"""
WebSocket server returned by :func:`serve`.
This class mirrors most of the API of :class:`asyncio.Server`, with the
following differences:
* You can invoke :func:`serve` as ``async with serve(...) as server: ...``
in addition to ``server = await serve(...)`` to start the server.
* It doesn't provide ``close_clients`` or ``abort_clients``; by default,
:meth:`close` closes existing connections with code 1001 (going away).
Args:
handler: Handler for one connection. It receives an asyncio protocol.
logger: Logger for this server.
It defaults to ``logging.getLogger("websockets.server")``.
See the :doc:`logging guide <../../topics/logging>` for details.
"""
def __init__(
self,
create_server: Callable[[], Coroutine[Any, Any, asyncio.Server]],
handler: Callable[[ServerConnection], Coroutine[Any, Any, None]],
logger: LoggerLike | None = None,
) -> None:
self.create_server = create_server
self.handler = handler
if logger is None:
logger = logging.getLogger("websockets.server")
self.logger = logger
# Keep track of active connections and connection handler tasks.
self.all_connections: set[ServerConnection] = set()
self.handler_tasks: set[asyncio.Task[None]] = set()
# Task responsible for closing the server and terminating connections.
self.close_task: asyncio.Task[None] | None = None
# Completed when the server is closed and connections are terminated.
loop = asyncio.get_running_loop()
self.handlers_waiter: asyncio.Future[None] = loop.create_future()
@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
}
def close(
self,
close_connections: bool = True,
code: CloseCode | int = CloseCode.GOING_AWAY,
reason: str = "",
) -> None:
"""
Close the server.
* Close the underlying :class:`asyncio.Server`.
* 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 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:`close` is idempotent.
"""
if self.close_task is None:
self.close_task = self.get_loop().create_task(
self._close(close_connections, code, reason)
)
async def _close(
self,
close_connections: bool = True,
code: CloseCode | int = CloseCode.GOING_AWAY,
reason: str = "",
) -> None:
"""
Implementation of :meth:`close`.
This calls :meth:`~asyncio.Server.close` on the underlying
:class:`asyncio.Server` object to stop accepting new connections and
then closes open connections.
"""
self.logger.info("server closing")
# Stop accepting new connections.
# Also reject OPENING connections with HTTP 503 — see handshake().
self.server.close()
# Close OPEN connections.
if close_connections:
close_tasks = [
asyncio.create_task(connection.close(code, reason))
for connection in self.all_connections
if connection.protocol.state is OPEN
]
# asyncio.wait doesn't accept an empty first argument.
if close_tasks:
await asyncio.wait(close_tasks)
# Wait until all TCP connections are closed.
await self.server.wait_closed()
# Wait until all connection handlers have returned.
# asyncio.wait doesn't accept an empty first argument.
if self.handler_tasks:
await asyncio.wait(self.handler_tasks)
# Tell wait_closed() to return.
self.handlers_waiter.set_result(None)
self.logger.info("server closed")
def get_loop(self) -> asyncio.AbstractEventLoop:
"""
See :meth:`asyncio.Server.get_loop`.
"""
return self.server.get_loop()
async def start_serving(self) -> None:
"""
See :meth:`asyncio.Server.start_serving`.
Typical use::
server = await serve(..., start_serving=False)
# perform additional setup here...
# ... then start the server
await server.start_serving()
"""
await self.server.start_serving()
async def serve_forever(self) -> None:
"""
See :meth:`asyncio.Server.serve_forever`.
Typical use::
server = await serve(...)
# this coroutine doesn't return
# canceling it stops the server
await server.serve_forever()
This is an alternative to using :func:`serve` as an asynchronous context
manager. Shutdown is triggered by canceling :meth:`serve_forever` or by
calling :meth:`~Server.close` from another task.
"""
try:
if not self.is_serving():
await self.start_serving()
# If close() is called, wait_closed() will return, and we'll exit.
await self.wait_closed()
except asyncio.CancelledError:
try:
self.close()
await self.wait_closed()
finally:
raise
def is_serving(self) -> bool:
"""
See :meth:`asyncio.Server.is_serving`.
"""
return self.server.is_serving()
async def wait_closed(self) -> None:
"""
Wait until the server is closed.
When :meth:`wait_closed` returns, all TCP connections are closed and
all connection handlers have returned.
To ensure a fast shutdown, a connection handler should always be
awaiting at least one of:
* :meth:`~ServerConnection.recv`: when the connection is closed,
it raises :exc:`~websockets.exceptions.ConnectionClosedOK`;
* :meth:`~ServerConnection.wait_closed`: when the connection is
closed, it returns.
Then the connection handler is immediately notified of the shutdown;
it can clean up and exit.
"""
await asyncio.shield(self.handlers_waiter)
@property
def sockets(self) -> tuple[socket.socket, ...]:
"""
See :attr:`asyncio.Server.sockets`.
"""
return self.server.sockets
async def _await(self) -> Self:
if not hasattr(self, "server"):
self.server = await self.create_server()
if self.server.is_serving():
for sock in self.server.sockets:
self.logger.info("server listening on %s", get_socket_name(sock))
return self
def __await__(self) -> Generator[Any, None, Self]:
# Create a suitable iterator by calling __await__ on a coroutine.
return self._await().__await__()
async def __aenter__(self) -> Self:
return await self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
self.close()
await self.wait_closed()
# serve() is declared as a function rather than a coroutine in order to support
# async with serve(...) as server: ... in addition to server = await serve(...).
def serve(
handler: Callable[[ServerConnection], Awaitable[None]],
host: str | None = None,
port: int | 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,
write_limit: int | tuple[int, int | None] = 2**15,
# Logging
logger: LoggerLike | None = None,
# Escape hatch for advanced customization
create_connection: type[ServerConnection] | None = None,
# Other keyword arguments are passed to loop.create_server
**kwargs: Any,
) -> Server:
"""
Create a WebSocket server listening on ``host`` and ``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.
This function returns a :class:`Server` object whose API mirrors
:class:`asyncio.Server`. Treat it as an asynchronous context manager to
serve requests and ensure that the server will be closed gracefully::
from websockets.asyncio.server import serve
async def handler(websocket):
...
# set this event to exit the server
stop = asyncio.Event()
async with serve(handler, host, port):
await stop.wait()
Alternatively, await it and call :meth:`~Server.serve_forever` to serve
requests, then cancel it or call :meth:`~Server.close` to stop the server::
server = await serve(handler, host, port)
await server.serve_forever()
The following pattern is functional but redundant: by the time the context
manager exits, :meth:`~Server.serve_forever` has already closed the server::
async with serve(handler, host, port) as server:
await server.serve_forever()
Args:
handler: Connection handler. It receives the WebSocket connection,
which is a :class:`ServerConnection`, in argument.
host: Network interfaces the server binds to.
See :meth:`~asyncio.loop.create_server` for details.
port: TCP port the server listens on.
See :meth:`~asyncio.loop.create_server` for details.
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.
write_limit: High-water mark of write buffer in bytes. It is passed to
:meth:`~asyncio.WriteTransport.set_write_buffer_limits`. It defaults
to 32 KiB. You may pass a ``(high, low)`` tuple to set the
high-water and low-water marks.
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.
Any other keyword arguments are passed to the event loop's
:meth:`~asyncio.loop.create_server` method.
For example:
* You can set ``ssl`` to a :class:`~ssl.SSLContext` to enable TLS.
* You can set ``sock`` to provide a preexisting TCP socket. You may call
:func:`socket.create_server` (not to be confused with the event loop's
:meth:`~asyncio.loop.create_server` method) to create a suitable server
socket and customize it.
* You can set ``start_serving`` to ``False`` to start accepting connections
only after you call :meth:`~Server.start_serving()` or
:meth:`~Server.serve_forever()`.
"""
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
if kwargs.get("ssl") is not None:
kwargs.setdefault("ssl_handshake_timeout", open_timeout)
kwargs.setdefault("ssl_shutdown_timeout", close_timeout)
async def create_server() -> asyncio.Server:
loop = asyncio.get_running_loop()
if kwargs.pop("unix", False):
return await loop.create_unix_server(protocol_factory, **kwargs)
else:
# mypy cannot tell that kwargs must provide sock when port is None.
return await loop.create_server(protocol_factory, host, port, **kwargs) # type: ignore[arg-type]
def protocol_factory() -> ServerConnection:
"""
Create an asyncio protocol for managing a WebSocket connection.
"""
# 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)
# This is a protocol in the Sans-I/O implementation of websockets.
protocol = ServerProtocol(
origins=origins,
extensions=extensions,
subprotocols=subprotocols,
select_subprotocol=protocol_select_subprotocol,
max_size=max_size,
logger=logger,
)
# This is a connection in websockets and a protocol in asyncio.
connection = create_connection(
protocol,
server,
ping_interval=ping_interval,
ping_timeout=ping_timeout,
close_timeout=close_timeout,
max_queue=max_queue,
write_limit=write_limit,
)
return connection
async def protocol_handler(connection: ServerConnection) -> 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.
"""
try:
# Apply open_timeout to the WebSocket handshake.
# Use ssl_handshake_timeout for the TLS handshake.
async with asyncio.timeout(open_timeout):
await connection.handshake(
process_request,
process_response,
server_header,
)
if connection.protocol.state is not OPEN:
connection.transport.abort()
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.close(CloseCode.INTERNAL_ERROR)
else:
await connection.close()
finally:
server.all_connections.discard(connection)
except Exception:
# Don't leak connections when the opening handshake times out or
# an unexpected error occurs.
connection.transport.abort()
finally:
server.handler_tasks.discard(asyncio.current_task())
# The server variable is captured by the closure of conn_handler() and
# protocol_factory().
server = Server(create_server, protocol_handler, logger)
return server
def unix_serve(
handler: Callable[[ServerConnection], Awaitable[None]],
path: str | None = None,
**kwargs: Any,
) -> Server:
"""
Create a WebSocket server listening on a Unix socket.
This function is identical to :func:`serve`, except the ``host`` and
``port`` arguments are replaced by ``path``. It's only available on Unix.
It's useful for deploying a server behind a reverse proxy such as nginx.
Args:
handler: Connection handler. It receives the WebSocket connection,
which is a :class:`ServerConnection`, in argument.
path: File system path to the Unix socket.
"""
return serve(handler, unix=True, path=path, **kwargs)
def is_credentials(credentials: Any) -> bool:
try:
username, password = credentials
except (TypeError, ValueError):
return False
else:
return isinstance(username, str) and isinstance(password, str)
def basic_auth(
realm: str = "",
credentials: tuple[str, str] | Iterable[tuple[str, str]] | None = None,
check_credentials: Callable[[str, str], Awaitable[bool] | bool] | None = None,
) -> Callable[[ServerConnection, Request], Awaitable[Response | None]]:
"""
Factory for ``process_request`` to enforce HTTP Basic Authentication.
:func:`basic_auth` is designed to integrate with :func:`serve` as follows::
from websockets.asyncio.server import basic_auth, serve
async with serve(
...,
process_request=basic_auth(
realm="my dev server",
credentials=("hello", "iloveyou"),
),
):
If authentication succeeds, the connection's ``username`` attribute is set.
If it fails, the server responds with an HTTP 401 Unauthorized status.
One of ``credentials`` or ``check_credentials`` must be provided; not both.
Args:
realm: Scope of protection. It should contain only ASCII characters
because the encoding of non-ASCII characters is undefined. Refer to
section 2.2 of :rfc:`7235` for details.
credentials: Hard coded authorized credentials. It can be a
``(username, password)`` pair or a list of such pairs.
check_credentials: Function or coroutine that verifies credentials.
It receives ``username`` and ``password`` arguments and returns
whether they're valid.
Raises:
TypeError: If ``credentials`` or ``check_credentials`` is wrong.
ValueError: If ``credentials`` and ``check_credentials`` are both
provided or both not provided.
"""
if (credentials is None) == (check_credentials is None):
raise ValueError("provide either credentials or check_credentials")
if credentials is not None:
if is_credentials(credentials):
credentials_list = [cast(tuple[str, str], credentials)]
elif isinstance(credentials, Iterable):
credentials_list = list(cast(Iterable[tuple[str, str]], credentials))
if not all(is_credentials(item) for item in credentials_list):
raise TypeError(f"invalid credentials argument: {credentials}")
else:
raise TypeError(f"invalid credentials argument: {credentials}")
credentials_dict = dict(credentials_list)
def check_credentials(username: str, password: str) -> bool:
try:
expected_password = credentials_dict[username]
except KeyError:
return False
return hmac.compare_digest(expected_password, password)
assert check_credentials is not None # help mypy
async def process_request(
connection: ServerConnection,
request: Request,
) -> Response | None:
"""
Perform HTTP Basic Authentication.
If it succeeds, set the connection's ``username`` attribute and return
:obj:`None`. If it fails, return an HTTP 401 Unauthorized responss.
"""
try:
authorization = request.headers["Authorization"]
except KeyError:
response = connection.respond(
http.HTTPStatus.UNAUTHORIZED,
"Missing credentials\n",
)
response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
return response
try:
username, password = parse_authorization_basic(authorization)
except InvalidHeader:
response = connection.respond(
http.HTTPStatus.UNAUTHORIZED,
"Unsupported credentials\n",
)
response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
return response
valid_credentials = check_credentials(username, password)
if isinstance(valid_credentials, Awaitable):
valid_credentials = await valid_credentials
if not valid_credentials:
response = connection.respond(
http.HTTPStatus.UNAUTHORIZED,
"Invalid credentials\n",
)
response.headers["WWW-Authenticate"] = build_www_authenticate_basic(realm)
return response
connection.username = username
return None
return process_request