349 lines
12 KiB
Python
349 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import socket
|
|
import ssl as ssl_module
|
|
import threading
|
|
import warnings
|
|
from collections.abc import Sequence
|
|
from typing import Any
|
|
|
|
from ..client import ClientProtocol
|
|
from ..datastructures import HeadersLike
|
|
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 ..typing import LoggerLike, Origin, Subprotocol
|
|
from ..uri import parse_uri
|
|
from .connection import Connection
|
|
from .utils import Deadline
|
|
|
|
|
|
__all__ = ["connect", "unix_connect", "ClientConnection"]
|
|
|
|
|
|
class ClientConnection(Connection):
|
|
"""
|
|
:mod:`threading` implementation of a WebSocket client connection.
|
|
|
|
:class:`ClientConnection` provides :meth:`recv` and :meth:`send` methods for
|
|
receiving and sending messages.
|
|
|
|
It supports iteration to receive messages::
|
|
|
|
for message in websocket:
|
|
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 ``close_timeout`` and ``max_queue`` arguments have the same meaning as
|
|
in :func:`connect`.
|
|
|
|
Args:
|
|
socket: Socket connected to a WebSocket server.
|
|
protocol: Sans-I/O connection.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
socket: socket.socket,
|
|
protocol: ClientProtocol,
|
|
*,
|
|
close_timeout: float | None = 10,
|
|
max_queue: int | None | tuple[int | None, int | None] = 16,
|
|
) -> None:
|
|
self.protocol: ClientProtocol
|
|
self.response_rcvd = threading.Event()
|
|
super().__init__(
|
|
socket,
|
|
protocol,
|
|
close_timeout=close_timeout,
|
|
max_queue=max_queue,
|
|
)
|
|
|
|
def handshake(
|
|
self,
|
|
additional_headers: HeadersLike | None = None,
|
|
user_agent_header: str | None = USER_AGENT,
|
|
timeout: float | None = None,
|
|
) -> None:
|
|
"""
|
|
Perform the opening handshake.
|
|
|
|
"""
|
|
with self.send_context(expected_state=CONNECTING):
|
|
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["User-Agent"] = user_agent_header
|
|
self.protocol.send_request(self.request)
|
|
|
|
if not self.response_rcvd.wait(timeout):
|
|
raise TimeoutError("timed out during handshake")
|
|
|
|
# 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)
|
|
|
|
def recv_events(self) -> None:
|
|
"""
|
|
Read incoming data from the socket and process events.
|
|
|
|
"""
|
|
try:
|
|
super().recv_events()
|
|
finally:
|
|
# If the connection is closed during the handshake, unblock it.
|
|
self.response_rcvd.set()
|
|
|
|
|
|
def connect(
|
|
uri: str,
|
|
*,
|
|
# TCP/TLS
|
|
sock: socket.socket | 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,
|
|
additional_headers: HeadersLike | None = None,
|
|
user_agent_header: str | None = USER_AGENT,
|
|
compression: str | None = "deflate",
|
|
# Timeouts
|
|
open_timeout: float | None = 10,
|
|
close_timeout: float | None = 10,
|
|
# Limits
|
|
max_size: 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,
|
|
**kwargs: Any,
|
|
) -> ClientConnection:
|
|
"""
|
|
Connect to the WebSocket server at ``uri``.
|
|
|
|
This function returns a :class:`ClientConnection` instance, which you can
|
|
use to send and receive messages.
|
|
|
|
:func:`connect` may be used as a context manager::
|
|
|
|
from websockets.sync.client import connect
|
|
|
|
with connect(...) as websocket:
|
|
...
|
|
|
|
The connection is closed automatically when exiting the context.
|
|
|
|
Args:
|
|
uri: URI of the WebSocket server.
|
|
sock: Preexisting TCP socket. ``sock`` overrides the host and port
|
|
from ``uri``. You may call :func:`socket.create_connection` to
|
|
create a suitable TCP socket.
|
|
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.
|
|
additional_headers (HeadersLike | None): 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.
|
|
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.
|
|
open_timeout: Timeout for opening the connection in seconds.
|
|
:obj:`None` disables the timeout.
|
|
close_timeout: Timeout for closing the connection in seconds.
|
|
:obj:`None` disables the timeout.
|
|
max_size: Maximum size of incoming messages in bytes.
|
|
:obj:`None` disables the limit.
|
|
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:`~socket.create_connection`.
|
|
|
|
Raises:
|
|
InvalidURI: If ``uri`` isn't a valid WebSocket URI.
|
|
OSError: If the TCP connection fails.
|
|
InvalidHandshake: If the opening handshake fails.
|
|
TimeoutError: If the opening handshake times out.
|
|
|
|
"""
|
|
|
|
# Process parameters
|
|
|
|
# Backwards compatibility: ssl used to be called ssl_context.
|
|
if ssl is None and "ssl_context" in kwargs:
|
|
ssl = kwargs.pop("ssl_context")
|
|
warnings.warn( # deprecated in 13.0 - 2024-08-20
|
|
"ssl_context was renamed to ssl",
|
|
DeprecationWarning,
|
|
)
|
|
|
|
wsuri = parse_uri(uri)
|
|
if not wsuri.secure and ssl is not None:
|
|
raise ValueError("ssl argument is incompatible with a ws:// URI")
|
|
|
|
# Private APIs for unix_connect()
|
|
unix: bool = kwargs.pop("unix", False)
|
|
path: str | None = kwargs.pop("path", None)
|
|
|
|
if unix:
|
|
if path is None and sock is None:
|
|
raise ValueError("missing path argument")
|
|
elif path is not None and sock is not None:
|
|
raise ValueError("path and sock arguments are incompatible")
|
|
|
|
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}")
|
|
|
|
# Calculate timeouts on the TCP, TLS, and WebSocket handshakes.
|
|
# The TCP and TLS timeouts must be set on the socket, then removed
|
|
# to avoid conflicting with the WebSocket timeout in handshake().
|
|
deadline = Deadline(open_timeout)
|
|
|
|
if create_connection is None:
|
|
create_connection = ClientConnection
|
|
|
|
try:
|
|
# Connect socket
|
|
|
|
if sock is None:
|
|
if unix:
|
|
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
|
sock.settimeout(deadline.timeout())
|
|
assert path is not None # mypy cannot figure this out
|
|
sock.connect(path)
|
|
else:
|
|
kwargs.setdefault("timeout", deadline.timeout())
|
|
sock = socket.create_connection((wsuri.host, wsuri.port), **kwargs)
|
|
sock.settimeout(None)
|
|
|
|
# Disable Nagle algorithm
|
|
|
|
if not unix:
|
|
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True)
|
|
|
|
# Initialize TLS wrapper and perform TLS handshake
|
|
|
|
if wsuri.secure:
|
|
if ssl is None:
|
|
ssl = ssl_module.create_default_context()
|
|
if server_hostname is None:
|
|
server_hostname = wsuri.host
|
|
sock.settimeout(deadline.timeout())
|
|
sock = ssl.wrap_socket(sock, server_hostname=server_hostname)
|
|
sock.settimeout(None)
|
|
|
|
# Initialize WebSocket protocol
|
|
|
|
protocol = ClientProtocol(
|
|
wsuri,
|
|
origin=origin,
|
|
extensions=extensions,
|
|
subprotocols=subprotocols,
|
|
max_size=max_size,
|
|
logger=logger,
|
|
)
|
|
|
|
# Initialize WebSocket connection
|
|
|
|
connection = create_connection(
|
|
sock,
|
|
protocol,
|
|
close_timeout=close_timeout,
|
|
max_queue=max_queue,
|
|
)
|
|
except Exception:
|
|
if sock is not None:
|
|
sock.close()
|
|
raise
|
|
|
|
try:
|
|
connection.handshake(
|
|
additional_headers,
|
|
user_agent_header,
|
|
deadline.timeout(),
|
|
)
|
|
except Exception:
|
|
connection.close_socket()
|
|
connection.recv_events_thread.join()
|
|
raise
|
|
|
|
return connection
|
|
|
|
|
|
def unix_connect(
|
|
path: str | None = None,
|
|
uri: str | None = None,
|
|
**kwargs: Any,
|
|
) -> ClientConnection:
|
|
"""
|
|
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`` is provided, to
|
|
``wss://localhost/``.
|
|
|
|
"""
|
|
if uri is None:
|
|
# Backwards compatibility: ssl used to be called ssl_context.
|
|
if kwargs.get("ssl") is None and kwargs.get("ssl_context") is None:
|
|
uri = "ws://localhost/"
|
|
else:
|
|
uri = "wss://localhost/"
|
|
return connect(uri=uri, unix=True, path=path, **kwargs)
|