diff --git a/docs/howto/upgrade.rst b/docs/howto/upgrade.rst index 8cfd7b4b..c6daf7d3 100644 --- a/docs/howto/upgrade.rst +++ b/docs/howto/upgrade.rst @@ -161,7 +161,7 @@ network errors and server errors (HTTP 500, 502, 503, or 504) are considered retryable. You can customize this behavior with the ``process_exception`` argument of :func:`~asyncio.client.connect`. -See :func:`~asyncio.client.process_exception` for more information. +See :func:`~client.process_exception` for more information. Here's how to revert to the behavior of the original implementation:: diff --git a/docs/project/changelog.rst b/docs/project/changelog.rst index fa2540c1..a7e73eba 100644 --- a/docs/project/changelog.rst +++ b/docs/project/changelog.rst @@ -32,6 +32,19 @@ notice. *In development* +New features +............ + +* Added support for reconnecting automatically by using + :func:`~sync.client.reconnect` as an iterator to the :mod:`threading` + implementation. + +* :func:`~sync.client.connect` now follows redirects in the :mod:`threading` + implementation. + +* :func:`~sync.client.connect` can connect to another host and port than those + specified in the URI in the :mod:`threading` implementation. + .. _17.0.1: 17.0.1 diff --git a/docs/reference/asyncio/client.rst b/docs/reference/asyncio/client.rst index 72c7dce3..324f9bf9 100644 --- a/docs/reference/asyncio/client.rst +++ b/docs/reference/asyncio/client.rst @@ -12,8 +12,6 @@ Opening a connection .. autofunction:: unix_connect :async: -.. autofunction:: process_exception - Using a connection ------------------ diff --git a/docs/reference/features.rst b/docs/reference/features.rst index 2bc505b5..01909bdb 100644 --- a/docs/reference/features.rst +++ b/docs/reference/features.rst @@ -149,7 +149,7 @@ Client +------------------------------------+--------+--------+--------+--------+--------+ | Close connection on context exit | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ - | Reconnect automatically | ✅ | ❌ | ✅ | — | ✅ | + | Reconnect automatically | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ | Configure ``Origin`` header | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ @@ -161,7 +161,7 @@ Client +------------------------------------+--------+--------+--------+--------+--------+ | Connect to non-ASCII IRIs | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ - | Follow HTTP redirects | ✅ | ❌ | ✅ | — | ✅ | + | Follow HTTP redirects | ✅ | ✅ | ✅ | — | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ | Perform HTTP Basic Authentication | ✅ | ✅ | ✅ | ✅ | ✅ | +------------------------------------+--------+--------+--------+--------+--------+ diff --git a/docs/reference/sansio/client.rst b/docs/reference/sansio/client.rst index 12f88b8e..b5ab8b1c 100644 --- a/docs/reference/sansio/client.rst +++ b/docs/reference/sansio/client.rst @@ -56,3 +56,5 @@ Client (`Sans-I/O`_) .. autoproperty:: close_reason .. autoproperty:: close_exc + +.. autofunction:: process_exception diff --git a/docs/reference/sync/client.rst b/docs/reference/sync/client.rst index fdc772ed..414c27ab 100644 --- a/docs/reference/sync/client.rst +++ b/docs/reference/sync/client.rst @@ -8,8 +8,12 @@ Opening a connection .. autofunction:: connect +.. autofunction:: reconnect + .. autofunction:: unix_connect +.. autofunction:: unix_reconnect + Using a connection ------------------ diff --git a/docs/reference/trio/client.rst b/docs/reference/trio/client.rst index 23fc70e3..a4a2b0b6 100644 --- a/docs/reference/trio/client.rst +++ b/docs/reference/trio/client.rst @@ -18,8 +18,6 @@ Opening a connection .. autofunction:: unix_connect :async: -.. autofunction:: process_exception - Using a connection ------------------ diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index 6c6c1149..941e94e9 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -11,10 +11,9 @@ from types import TracebackType from typing import Any, Callable, Literal, cast -from ..client import ClientProtocol, backoff +from ..client import ClientProtocol, backoff, process_exception from ..datastructures import Headers, HeadersLike from ..exceptions import ( - InvalidMessage, InvalidProxyMessage, InvalidProxyStatus, InvalidStatus, @@ -128,59 +127,14 @@ def process_event(self, event: Event) -> None: 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:: + :func:`connect` should be treated as an asynchronous context manager + yielding a :class:`ClientConnection`, which can then receive and send + messages:: from websockets.asyncio.client import connect @@ -189,8 +143,8 @@ class connect: The connection is closed automatically when exiting the context. - :func:`connect` can be used as an infinite asynchronous iterator to - reconnect automatically on errors:: + :func:`connect` can also be treated as an infinite asynchronous iterator + to reconnect automatically on errors:: async for websocket in connect(...): try: @@ -225,7 +179,8 @@ class connect: <../../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. + :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. @@ -555,13 +510,7 @@ def process_redirect(self, exc: Exception) -> Exception | str: 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: + async def connect(self) -> ClientConnection: try: async with asyncio.timeout(self.open_timeout): for _ in range(MAX_REDIRECTS): @@ -606,6 +555,12 @@ async def __await_impl__(self) -> ClientConnection: # Re-raise exception with an informative error message. raise TimeoutError("timed out during opening handshake") from exc + # ... = await connect(...) + + def __await__(self) -> Generator[Any, None, ClientConnection]: + # Create a suitable iterator by calling __await__ on a coroutine. + return self.connect().__await__() + # async with connect(...) as ...: ... async def __aenter__(self) -> ClientConnection: diff --git a/src/websockets/client.py b/src/websockets/client.py index be85e26f..1525a21c 100644 --- a/src/websockets/client.py +++ b/src/websockets/client.py @@ -357,6 +357,53 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) +def process_exception(exc: Exception) -> Exception | None: + """ + Determine whether a connection error is retryable or fatal. + + When reconnecting automatically with ``async for ... in connect(...)`` in + the :mod:`asyncio` and :mod:`trio` implementations or with ``for ... in + reconnect(...)`` in the :mod:`threading` implementation, 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:`OSError` and :exc:`asyncio.TimeoutError`: network errors; + * :exc:`~websockets.exceptions.InvalidMessage` when it stems from an + :exc:`EOFError`: also 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 + + BACKOFF_INITIAL_DELAY = float(os.environ.get("WEBSOCKETS_BACKOFF_INITIAL_DELAY", "5")) BACKOFF_MIN_DELAY = float(os.environ.get("WEBSOCKETS_BACKOFF_MIN_DELAY", "3.1")) BACKOFF_MAX_DELAY = float(os.environ.get("WEBSOCKETS_BACKOFF_MAX_DELAY", "90.0")) diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 509eb9df..ec43a737 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -1,16 +1,27 @@ from __future__ import annotations import logging +import os import socket import ssl as ssl_module import threading +import time +import traceback +import urllib.parse import warnings -from collections.abc import Sequence -from typing import Any, Callable, Literal, TypeVar, cast - -from ..client import ClientProtocol -from ..datastructures import HeadersLike -from ..exceptions import InvalidProxyMessage, InvalidProxyStatus, ProxyError +from collections.abc import Generator, Iterator, Sequence +from types import TracebackType +from typing import Any, Callable, Literal, TypeVar, cast, overload + +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 @@ -24,7 +35,9 @@ from .utils import Deadline -__all__ = ["connect", "unix_connect", "ClientConnection"] +__all__ = ["connect", "unix_connect", "reconnect", "unix_reconnect", "ClientConnection"] + +MAX_REDIRECTS = int(os.environ.get("WEBSOCKETS_MAX_REDIRECTS", "10")) class ClientConnection(Connection): @@ -65,6 +78,7 @@ def __init__( ) -> None: self.protocol: ClientProtocol self.response_rcvd = threading.Event() + self.pending_legacy_warning = True super().__init__( sock, protocol, @@ -74,6 +88,21 @@ def __init__( max_queue=max_queue, ) + def __enter__(self) -> ClientConnection: + self.pending_legacy_warning = False + return super().__enter__() + + def maybe_raise_legacy_warning(self) -> None: + if self.pending_legacy_warning: + self.pending_legacy_warning = False + warnings.warn( # deprecated in 17.1 + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly", + DeprecationWarning, + stacklevel=3, + ) + def handshake( self, additional_headers: HeadersLike | None = None, @@ -128,6 +157,497 @@ def recv_events(self) -> None: self.response_rcvd.set() +class reconnect: + """ + Similar to :func:`connect`, with support for automatic reconnection. + + :func:`reconnect` can also be treated as an infinite iterator to reconnect + automatically on errors:: + + for websocket in reconnect(...): + 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:`reconnect` accepts the same arguments as :func:`connect`, minus the + ``legacy`` flag, plus those listed below. It raises the same exceptions. + + Args: + 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. + reconnect_delays: Delays in seconds between reconnection attempts. + Default is exponential backoff with 5s jitter, capped at 60s. + + .. admonition:: Why is :func:`reconnect` a separate API from :func:`connect`? + :class: tip + + A new API was necessary to maintain backwards compatibility with this + historical behavior of :func:`connect`:: + + websocket = connect(...) + for message in websocket: + ... + + Once the deprecation period elapses, :func:`connect` will be changed to + behave like :func:`reconnect` by default. + + """ + + def __init__( + self, + 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, + 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 socket.create_connection + **kwargs: Any, + ) -> None: + # 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, + ) + + 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.sock = sock + 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_socket_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, + ) + + def open_socket(self, deadline: Deadline) -> socket.socket: + """Open a TCP or Unix connection to the server, possibly through a proxy.""" + kwargs = self.open_socket_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: + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + sock.settimeout(deadline.timeout()) + sock.connect(kwargs.pop("path")) + except Exception: + sock.close() + raise + + elif proxy is not None: + proxy_parsed = parse_proxy(proxy) + + if proxy_parsed.scheme[:5] == "socks": + sock = connect_socks_proxy( + proxy_parsed, + self.ws_uri, + deadline, + # websockets is consistent with the socket module while + # python_socks is consistent across implementations. + local_addr=kwargs.pop("source_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" + ) + sock = connect_http_proxy( + proxy_parsed, + self.ws_uri, + deadline, + 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("address", (self.ws_uri.host, self.ws_uri.port)) + kwargs.setdefault("timeout", deadline.timeout()) + sock = socket.create_connection(**kwargs) + + sock.settimeout(None) + return sock + + def enable_tls(self, sock: socket.socket, deadline: Deadline) -> socket.socket: + """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 + sock.settimeout(deadline.timeout()) + if self.proxy_ssl is None: + sock = ssl.wrap_socket(sock, server_hostname=server_hostname) + else: + sock_2 = SSLSSLSocket(sock, ssl, server_hostname=server_hostname) + # Let's pretend that sock is a socket, even though it isn't. + sock = cast(socket.socket, sock_2) + sock.settimeout(None) + return sock + + def open_connection(self, deadline: Deadline) -> ClientConnection: + """Create a WebSocket connection.""" + if self.sock is None: + sock = self.open_socket(deadline) + else: + sock = self.sock + + try: + if sock.family in {socket.AF_INET, socket.AF_INET6}: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) + + if self.ws_uri.secure: + sock = self.enable_tls(sock, deadline) + + protocol = ClientProtocol( + self.ws_uri, + **self.protocol_kwargs, # type: ignore + ) + + # self.create_connection defaults to ClientConnection. + connection = self.create_connection( + sock, + protocol, + **self.connection_kwargs, # type: ignore + ) + except Exception: + sock.close() + raise + + try: + connection.handshake( + self.additional_headers, + self.user_agent_header, + deadline.timeout(), + ) + except Exception: + connection.close_socket() + connection.recv_events_thread.join() + 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 socket, it is closed and cannot be reused. + if self.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.open_socket_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_socket_kwargs.get("address") is not None: + return ValueError( + f"cannot follow cross-origin redirect to {new_uri} " + f"with an explicit host and 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 + + def connect(self) -> ClientConnection: + """Connect to a WebSocket server, following redirects.""" + # 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(self.open_timeout) + + for _ in range(MAX_REDIRECTS): + try: + connection = self.open_connection(deadline) + 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") + + # with connect(...) as ...: ... + + def __enter__(self) -> ClientConnection: + if hasattr(self, "connection"): + raise RuntimeError("connect() isn't reentrant") + self.connection = self.connect() + return self.connection + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + exc_traceback: TracebackType | None, + ) -> None: + try: + self.connection.close() + finally: + del self.connection + + # for ... in reconnect(...): ... + + def __iter__(self) -> Iterator[ClientConnection]: + delays: Generator[float] | None = None + while True: + try: + 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(), + ) + time.sleep(delay) + + else: + # The connection succeeded. Reset backoff. + delays = None + + +@overload +def connect( + uri: str, + *, + # TCP/TLS + sock: socket.socket | None = ..., + ssl: ssl_module.SSLContext | None = ..., + server_hostname: str | None = ..., + # WebSocket + origin: Origin | None = ..., + extensions: Sequence[ClientExtensionFactory] | None = ..., + subprotocols: Sequence[Subprotocol] | None = ..., + compression: str | None = ..., + # HTTP + additional_headers: HeadersLike | None = ..., + user_agent_header: str | None = ..., + proxy: str | Literal[True] | None = ..., + proxy_ssl: ssl_module.SSLContext | None = ..., + proxy_server_hostname: str | None = ..., + # Timeouts + open_timeout: float | None = ..., + ping_interval: float | None = ..., + ping_timeout: float | None = ..., + close_timeout: float | None = ..., + # Limits + max_size: int | None | tuple[int | None, int | None] = ..., + max_queue: int | None | tuple[int | None, int | None] = ..., + # Logging + logger: LoggerLike | None = ..., + # Escape hatch for advanced customization + create_connection: type[ClientConnection] | None = ..., + # Backwards and forwards compatibility + legacy: Literal[True] | None = ..., + # Other keyword arguments are passed to socket.create_connection + **kwargs: Any, +) -> ClientConnection: ... + + +@overload +def connect( + uri: str, + *, + # TCP/TLS + sock: socket.socket | None = ..., + ssl: ssl_module.SSLContext | None = ..., + server_hostname: str | None = ..., + # WebSocket + origin: Origin | None = ..., + extensions: Sequence[ClientExtensionFactory] | None = ..., + subprotocols: Sequence[Subprotocol] | None = ..., + compression: str | None = ..., + # HTTP + additional_headers: HeadersLike | None = ..., + user_agent_header: str | None = ..., + proxy: str | Literal[True] | None = ..., + proxy_ssl: ssl_module.SSLContext | None = ..., + proxy_server_hostname: str | None = ..., + # Timeouts + open_timeout: float | None = ..., + ping_interval: float | None = ..., + ping_timeout: float | None = ..., + close_timeout: float | None = ..., + # Limits + max_size: int | None | tuple[int | None, int | None] = ..., + max_queue: int | None | tuple[int | None, int | None] = ..., + # Logging + logger: LoggerLike | None = ..., + # Escape hatch for advanced customization + create_connection: type[ClientConnection] | None = ..., + # Backwards and forwards compatibility + legacy: Literal[False], + # Other keyword arguments are passed to socket.create_connection + **kwargs: Any, +) -> reconnect: ... + + def connect( uri: str, *, @@ -158,16 +678,16 @@ def connect( logger: LoggerLike | None = None, # Escape hatch for advanced customization create_connection: type[ClientConnection] | None = None, + # Backwards and forwards compatibility + legacy: bool | None = None, # Other keyword arguments are passed to socket.create_connection **kwargs: Any, -) -> ClientConnection: +) -> ClientConnection | reconnect: """ 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:: + :func:`connect` should be treated as a context manager yielding a + :class:`ClientConnection`, which can then receive and send messages:: from websockets.sync.client import connect @@ -176,6 +696,23 @@ def connect( The connection is closed automatically when exiting the context. + Use :func:`reconnect` to reconnect automatically on errors. + + For backwards compatibility, :func:`connect` may be called directly:: + + websocket = connect(..., legacy=True) + + In that case, you're responsible for closing the connection with + :meth:`ClientConnection.close` when no longer needed. + + When the ``legacy`` flag is enabled, :func:`connect` returns directly a + :class:`ClientConnection` and iterating that connection yields messages. + Currently, this is the default behavior when ``legacy`` isn't specified. + + When the ``legacy`` flag is explicitly disabled, :func:`connect` behaves + like :func:`reconnect`: using it as an iterator returns a new connection + at each iteration, making it easy to reconnect automatically on errors. + Args: uri: URI of the WebSocket server. sock: Preexisting TCP socket. ``sock`` overrides the host and port @@ -224,11 +761,17 @@ def connect( logger: Logger for this client. It defaults to ``logging.getLogger("websockets.client")``. See the :doc:`logging guide <../../topics/logging>` for details. + legacy: Set to :obj:`True` to opt into the historical behavior of + returning a :class:`ClientConnection`, without deprecation warning. 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`. + For example, you can set ``address`` to a ``(host, port)`` tuple 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. @@ -238,175 +781,96 @@ def connect( TimeoutError: If the opening handshake times out. """ + connecter = reconnect( + uri, + sock=sock, + ssl=ssl, + server_hostname=server_hostname, + origin=origin, + extensions=extensions, + subprotocols=subprotocols, + compression=compression, + additional_headers=additional_headers, + user_agent_header=user_agent_header, + proxy=proxy, + proxy_ssl=proxy_ssl, + proxy_server_hostname=proxy_server_hostname, + open_timeout=open_timeout, + ping_interval=ping_interval, + ping_timeout=ping_timeout, + close_timeout=close_timeout, + max_size=max_size, + max_queue=max_queue, + logger=logger, + create_connection=create_connection, + **kwargs, + ) + # For backwards compatibility, connect defaults to the historical behavior. + # For forwards compatibility, the future behavior can be chosen explicitly. + if legacy is False: + return connecter + connection = connecter.connect() + # Users can opt in to the historical behavior to remain unaffected when the + # future behavior becomes the default. + if legacy: + connection.pending_legacy_warning = False + return connection - # 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, - ) - - ws_uri = parse_uri(uri) - if not 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 - - # 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 is incompatible with sock") - - if unix: - proxy = None - if sock is not None: - proxy = None - if proxy is True: - proxy = get_proxy(ws_uri) - - # 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) - - 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) - - elif proxy is not None: - proxy_parsed = parse_proxy(proxy) - - if proxy_parsed.scheme[:5] == "socks": - sock = connect_socks_proxy( - proxy_parsed, - ws_uri, - deadline, - # websockets is consistent with the socket module while - # python_socks is consistent across implementations. - local_addr=kwargs.pop("source_address", None), - ) - - elif proxy_parsed.scheme[:4] == "http": - if proxy_parsed.scheme != "https" and proxy_ssl is not None: - raise ValueError( - "proxy_ssl argument is incompatible with an http:// proxy" - ) - sock = connect_http_proxy( - proxy_parsed, - ws_uri, - deadline, - user_agent_header=user_agent_header, - ssl=proxy_ssl, - server_hostname=proxy_server_hostname, - **kwargs, - ) - - else: - raise AssertionError("parse_proxy returned unsupported proxy") - - else: # proxy is None - kwargs.setdefault("timeout", deadline.timeout()) - sock = socket.create_connection( - (ws_uri.host, ws_uri.port), - **kwargs, - ) - - sock.settimeout(None) - - # Disable Nagle algorithm - if not unix: - sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) +def unix_reconnect( + path: str | None = None, + uri: str | None = None, + **kwargs: Any, +) -> reconnect: + """ + Similar to :func:`unix_connect`, with support for automatic reconnection. - # Initialize TLS wrapper and perform TLS handshake + Refer to the documentation of :func:`reconnect` for details on its behavior. - if ws_uri.secure: - if ssl is None: - ssl = ssl_module.create_default_context() - if server_hostname is None: - server_hostname = ws_uri.host - sock.settimeout(deadline.timeout()) - if proxy_ssl is None: - sock = ssl.wrap_socket(sock, server_hostname=server_hostname) - else: - sock_2 = SSLSSLSocket(sock, ssl, server_hostname=server_hostname) - # Let's pretend that sock is a socket, even though it isn't. - sock = cast(socket.socket, sock_2) - sock.settimeout(None) + """ + sock = kwargs.get("sock") + 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 is incompatible with sock") - # Initialize WebSocket protocol + 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/" - protocol = ClientProtocol( - ws_uri, - origin=origin, - extensions=extensions, - subprotocols=subprotocols, - max_size=max_size, - logger=logger, - ) + return reconnect(uri=uri, unix=True, path=path, **kwargs) - # Initialize WebSocket connection - # create_connection defaults to ClientConnection. - connection = create_connection( - sock, - protocol, - ping_interval=ping_interval, - ping_timeout=ping_timeout, - close_timeout=close_timeout, - max_queue=max_queue, - ) - except Exception: - if sock is not None: - sock.close() - raise +@overload +def unix_connect( + path: str | None = ..., + uri: str | None = ..., + *, + legacy: Literal[True] | None = ..., + **kwargs: Any, +) -> ClientConnection: ... - try: - connection.handshake( - additional_headers, - user_agent_header, - deadline.timeout(), - ) - except Exception: - connection.close_socket() - connection.recv_events_thread.join() - raise - connection.start_keepalive() - return connection +@overload +def unix_connect( + path: str | None = ..., + uri: str | None = ..., + *, + legacy: Literal[False], + **kwargs: Any, +) -> reconnect: ... def unix_connect( path: str | None = None, uri: str | None = None, + *, + legacy: bool | None = None, **kwargs: Any, -) -> ClientConnection: +) -> ClientConnection | reconnect: """ Connect to a WebSocket server listening on a Unix socket. @@ -423,13 +887,17 @@ def unix_connect( ``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) + connecter = unix_reconnect(path, uri, **kwargs) + # For backwards compatibility, connect defaults to the historical behavior. + # For forwards compatibility, the future behavior can be chosen explicitly. + if legacy is False: + return connecter + connection = connecter.connect() + # Users can opt in to the historical behavior to remain unaffected when the + # future behavior becomes the default. + if legacy: + connection.pending_legacy_warning = False + return connection try: diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index 74ac997a..b2b7415c 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -317,6 +317,7 @@ def recv(self, timeout: float | None = None, decode: bool | None = None) -> Data :meth:`recv_streaming` concurrently. """ + self.maybe_raise_legacy_warning() try: return self.recv_messages.get(timeout, decode) except EOFError: @@ -387,6 +388,7 @@ def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]: :meth:`recv_streaming` concurrently. """ + self.maybe_raise_legacy_warning() try: yield from self.recv_messages.get_iter(decode) return @@ -466,6 +468,7 @@ def send( TypeError: If ``message`` doesn't have a supported type. """ + self.maybe_raise_legacy_warning() # Unfragmented message — this case must be handled first because # strings and bytes-like objects are iterable. @@ -591,6 +594,7 @@ def close( reason: WebSocket close reason. """ + self.maybe_raise_legacy_warning() try: # The context manager takes care of waiting for the TCP connection # to terminate after calling a method that sends a close frame. @@ -647,6 +651,7 @@ def ping( the corresponding pong wasn't received yet. """ + self.maybe_raise_legacy_warning() if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -684,6 +689,7 @@ def pong(self, data: DataLike = b"") -> None: ConnectionClosed: When the connection is closed. """ + self.maybe_raise_legacy_warning() if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -696,6 +702,9 @@ def pong(self, data: DataLike = b"") -> None: # Private methods + def maybe_raise_legacy_warning(self) -> None: + pass # see override in ClientConnection + def process_event(self, event: Event) -> None: """ Process one incoming event. diff --git a/src/websockets/sync/server.py b/src/websockets/sync/server.py index 0d1c19fb..97a8d9d4 100644 --- a/src/websockets/sync/server.py +++ b/src/websockets/sync/server.py @@ -686,7 +686,7 @@ def sock_handler(sock: socket.socket, addr: Any) -> None: try: # Disable Nagle algorithm - if not unix: + if sock.family in {socket.AF_INET, socket.AF_INET6}: sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) # Perform TLS handshake diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 1bdde69c..63dc4ac1 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -11,8 +11,7 @@ import trio -from ..asyncio.client import process_exception -from ..client import ClientProtocol, backoff +from ..client import ClientProtocol, backoff, process_exception from ..datastructures import Headers, HeadersLike from ..exceptions import ( InvalidProxyMessage, @@ -135,10 +134,9 @@ 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:: + :func:`connect` is designed to be called as an asynchronous context manager + yielding a :class:`ClientConnection`, which you can then use to receive and + send messages:: from websockets.trio.client import connect @@ -147,8 +145,8 @@ class connect: The connection is closed automatically when exiting the context. - :func:`connect` can be used as an infinite asynchronous iterator to - reconnect automatically on errors:: + :func:`connect` can also be treated as an infinite asynchronous iterator + to reconnect automatically on errors:: async for websocket in connect(...): try: @@ -162,6 +160,10 @@ class connect: 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 @@ -192,7 +194,8 @@ class connect: ``proxy_server_hostname`` overrides the host name from ``proxy``. process_exception: When reconnecting automatically, tell whether an error is transient or fatal. The default behavior is defined by - :func:`process_exception`. Refer to its documentation for details. + :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. @@ -220,6 +223,10 @@ class connect: 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. @@ -319,15 +326,16 @@ def __init__( 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 kwargs.get("unix", False): + if unix: proxy = None if proxy is True: proxy = get_proxy(self.ws_uri) - if kwargs.pop("unix", False): - return await trio.open_unix_socket(kwargs["path"]) + if unix: + return await trio.open_unix_socket(kwargs.pop("path")) elif proxy is not None: proxy_parsed = parse_proxy(proxy) @@ -384,7 +392,6 @@ async def enable_tls(self, stream: trio.abc.Stream) -> trio.abc.Stream: async def open_connection(self, nursery: trio.Nursery) -> ClientConnection: """Create a WebSocket connection.""" - # TCP connection is already established. if self.stream is None: stream = await self.open_tcp_stream() else: @@ -412,8 +419,6 @@ async def open_connection(self, nursery: trio.Nursery) -> ClientConnection: self.user_agent_header, ) - return connection - except trio.Cancelled: await trio.aclose_forcefully(stream) # The nursery running this coroutine was canceled. @@ -427,6 +432,8 @@ async def open_connection(self, nursery: trio.Nursery) -> ClientConnection: 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. @@ -532,7 +539,7 @@ async def connect(self, nursery: trio.Nursery) -> ClientConnection: # 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, ...) + # Do not define __await__ for ... = await nursery.start(connect, ...) # because it doesn't look idiomatic in Trio. # async with connect(...) as ...: ... @@ -653,11 +660,18 @@ def unix_connect( ``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) diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index 6f6070c6..66fa7f81 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -52,11 +52,19 @@ async def few_redirects(): class ClientTests(unittest.IsolatedAsyncioTestCase): - async def test_connection(self): - """Client connects to server.""" + async def test_context_manager(self): + """Client connects to server and disconnects automatically.""" async with serve(*args) as server: async with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_direct_connection(self): + """Client connects to server directly.""" + async with serve(*args) as server: + client = await connect(get_uri(server)) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" @@ -519,13 +527,23 @@ async def junk(reader, writer): class SecureClientTests(unittest.IsolatedAsyncioTestCase): - async def test_connection(self): - """Client connects to server securely.""" + async def test_context_manager(self): + """Client connects to server securely and disconnects automatically.""" async with serve(*args, ssl=SERVER_CONTEXT) as server: async with connect(get_uri(server), ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") ssl_object = client.transport.get_extra_info("ssl_object") self.assertEqual(ssl_object.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_direct_connection(self): + """Client connects to server securely and directly.""" + async with serve(*args, ssl=SERVER_CONTEXT) as server: + client = await connect(get_uri(server), ssl=CLIENT_CONTEXT) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + ssl_object = client.transport.get_extra_info("ssl_object") + self.assertEqual(ssl_object.version()[:3], "TLS") async def test_set_server_hostname_implicitly(self): """Client sets server_hostname to the host in the WebSocket URI.""" @@ -899,12 +917,21 @@ async def test_https_proxy_invalid_server_certificate(self): @unittest.skipUnless(hasattr(socket, "AF_UNIX"), "this test requires Unix sockets") class UnixClientTests(unittest.IsolatedAsyncioTestCase): - async def test_connection(self): - """Client connects to server over a Unix socket.""" + async def test_context_manager(self): + """Client connects to Unix server and disconnects automatically.""" with temp_unix_socket_path() as path: async with unix_serve(handler, path): async with unix_connect(path) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_direct_connection(self): + """Client connects to Unix server directly.""" + with temp_unix_socket_path() as path: + async with unix_serve(handler, path): + client = await unix_connect(path) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") async def test_set_host_header(self): """Client sets the Host header to the host in the WebSocket URI.""" @@ -933,14 +960,25 @@ def redirect(connection, request): "cannot follow cross-origin redirect to ws://other/ with a Unix socket", ) - async def test_secure_connection(self): - """Client connects to server securely over a Unix socket.""" + async def test_secure_context_manager(self): + """Client connects to Unix server securely and disconnects automatically.""" with temp_unix_socket_path() as path: async with unix_serve(handler, path, ssl=SERVER_CONTEXT): async with unix_connect(path, ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") ssl_object = client.transport.get_extra_info("ssl_object") self.assertEqual(ssl_object.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") + + async def test_secure_direct_connection(self): + """Client connects to Unix server securely and directly.""" + with temp_unix_socket_path() as path: + async with unix_serve(handler, path, ssl=SERVER_CONTEXT): + client = await unix_connect(path, ssl=CLIENT_CONTEXT) + self.addAsyncCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + ssl_object = client.transport.get_extra_info("ssl_object") + self.assertEqual(ssl_object.version()[:3], "TLS") async def test_set_server_hostname(self): """Client sets server_hostname to the host in the WebSocket URI.""" @@ -955,6 +993,13 @@ async def test_set_server_hostname(self): ssl_object = client.transport.get_extra_info("ssl_object") self.assertEqual(ssl_object.server_hostname, "overridden") + async def test_non_existing_path(self): + """Client attempts to connect to a non-existing Unix socket path.""" + with temp_unix_socket_path() as path: + with self.assertRaises(FileNotFoundError): + async with unix_connect(path + ".doesnotexist"): + self.fail("did not raise") + class ClientUsageErrorsTests(unittest.IsolatedAsyncioTestCase): async def test_ssl_without_secure_uri(self): @@ -988,7 +1033,7 @@ async def test_proxy_ssl_without_https_proxy(self): "proxy_ssl argument is incompatible with an http:// proxy", ) - async def test_https_proxy_without_ssl(self): + async def test_https_proxy_without_proxy_ssl(self): """Client rejects proxy_ssl=None when proxy is HTTPS.""" with self.assertRaises(ValueError) as raised: await connect( diff --git a/tests/asyncio/test_router.py b/tests/asyncio/test_router.py index ea8e14bf..c9651e9c 100644 --- a/tests/asyncio/test_router.py +++ b/tests/asyncio/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) async def echo(websocket, count): @@ -48,37 +55,28 @@ async def test_router_matches_paths_and_extracts_parameters(self): messages = await alist(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - async def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: async with connect(get_uri(server) + "/?a=b") as client: await self.assertEval(client, "ws.request.path", "/?a=b") async def test_redirect(self): """Router redirects connections according to redirect_to.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: async with connect(get_uri(server) + "/r") as client: await self.assertEval(client, "ws.request.path", "/") async def test_secure_redirect(self): """Router redirects connections according to redirect_to when TLS is enabled.""" - async with route(self.url_map, "localhost", 0, ssl=SERVER_CONTEXT) as server: + async with route(url_map, "localhost", 0, ssl=SERVER_CONTEXT) as server: async with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: await self.assertEval(client, "ws.request.path", "/") @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - async with route(self.url_map, "localhost", 0, ssl=True) as server: + async with route(url_map, "localhost", 0, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): @@ -91,7 +89,7 @@ async def test_force_secure_redirect(self): @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - async with route(self.url_map, "localhost", 0, server_name="other") as server: + async with route(url_map, "localhost", 0, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -102,7 +100,7 @@ async def test_force_redirect_server_name(self): async def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - async with route(self.url_map, "localhost", 0) as server: + async with route(url_map, "localhost", 0) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -118,7 +116,7 @@ def process_request(ws, request): ws.process_request_ran = True async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -130,7 +128,7 @@ async def process_request(ws, request): ws.process_request_ran = True async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -142,7 +140,7 @@ def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): @@ -159,7 +157,7 @@ async def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") async with route( - self.url_map, "localhost", 0, process_request=process_request + url_map, "localhost", 0, process_request=process_request ) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): @@ -177,9 +175,7 @@ async def handler(self, connection): connection.my_router_ran = True return await super().handler(connection) - async with route( - self.url_map, "localhost", 0, create_router=MyRouter - ) as server: + async with route(url_map, "localhost", 0, create_router=MyRouter) as server: async with connect(get_uri(server)) as client: await self.assertEval(client, "ws.my_router_ran", "True") diff --git a/tests/sync/server.py b/tests/sync/server.py index 78d1c674..194829fe 100644 --- a/tests/sync/server.py +++ b/tests/sync/server.py @@ -8,11 +8,15 @@ from websockets.sync.server import serve, unix_serve +def get_host_port(server): + return server.socket.getsockname() + + def get_uri(server, secure=None): if secure is None: secure = isinstance(server.socket, ssl.SSLSocket) # hack protocol = "wss" if secure else "ws" - host, port = server.socket.getsockname() + host, port = get_host_port(server) return f"{protocol}://{host}:{port}" @@ -69,9 +73,9 @@ def run_router(url_map, **kwargs): @contextlib.contextmanager def run_unix_server_or_router( - path, unix_serve_or_route, handler_or_url_map, + path, **kwargs, ): with unix_serve_or_route(handler_or_url_map, path, **kwargs) as server: @@ -85,8 +89,8 @@ def run_unix_server_or_router( def run_unix_server(path, handler=handler, **kwargs): - return run_unix_server_or_router(path, unix_serve, handler, **kwargs) + return run_unix_server_or_router(unix_serve, handler, path, **kwargs) def run_unix_router(path, url_map, **kwargs): - return run_unix_server_or_router(path, unix_route, url_map, **kwargs) + return run_unix_server_or_router(unix_route, url_map, path, **kwargs) diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index d4d42c31..88d54ab6 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -1,3 +1,4 @@ +import contextlib import http import logging import os @@ -10,6 +11,7 @@ import unittest from unittest.mock import patch +from websockets.client import backoff from websockets.exceptions import ( InvalidHandshake, InvalidMessage, @@ -18,6 +20,7 @@ InvalidStatus, InvalidURI, ProxyError, + SecurityError, ) from websockets.extensions.permessage_deflate import PerMessageDeflate from websockets.sync.client import * @@ -30,20 +33,57 @@ DeprecationTestCase, temp_unix_socket_path, ) -from .server import get_uri, run_server, run_unix_server +from .server import get_host_port, get_uri, run_server, run_unix_server + + +def short_backoff(): + defaults = backoff.__defaults__ + yield from backoff( + defaults[0] * MS, + defaults[1] * MS, + defaults[2] * MS, + defaults[3], + ) + + +@contextlib.contextmanager +def few_redirects(): + from websockets.sync import client + + max_redirects = client.MAX_REDIRECTS + client.MAX_REDIRECTS = 2 + try: + yield + finally: + client.MAX_REDIRECTS = max_redirects class ClientTests(unittest.TestCase): - def test_connection(self): - """Client connects to server and the handshake succeeds.""" + def test_context_manager(self): + """Client connects to server and disconnects automatically.""" with run_server() as server: with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_direct_connection(self): + """Client connects to server directly.""" + with run_server() as server: + client = connect(get_uri(server), legacy=True) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + + def test_explicit_host_port(self): + """Client connects using an explicit host / port.""" + with run_server() as server: + address = get_host_port(server) + with connect("ws://overridden/", address=address) as client: + self.assertEqual(client.protocol.state.name, "OPEN") def test_existing_socket(self): """Client connects using a pre-existing socket.""" with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Use a non-existing domain to ensure we connect via sock. with connect("ws://invalid/", sock=sock) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -127,6 +167,248 @@ def create_connection(*args, **kwargs): ) as client: self.assertTrue(client.create_connection_ran) + def test_reconnect(self): + """Client reconnects to server.""" + iterations = 0 + successful = 0 + + def process_request(connection, request): + nonlocal iterations + iterations += 1 + # Retriable errors + if iterations == 1: + time.sleep(3 * MS) + elif iterations == 2: + connection.socket.close() + elif iterations == 3: + return connection.respond(http.HTTPStatus.SERVICE_UNAVAILABLE, "🚒") + # Fatal error + elif iterations == 6: + return connection.respond(http.HTTPStatus.PAYMENT_REQUIRED, "💸") + + with run_server(process_request=process_request) as server: + with self.assertRaises(InvalidStatus) as raised: + for client in reconnect( + get_uri(server), + open_timeout=3 * MS, + reconnect_delays=short_backoff, + ): + self.assertEqual(client.protocol.state.name, "OPEN") + successful += 1 + + self.assertEqual( + str(raised.exception), + "server rejected WebSocket connection: HTTP 402", + ) + self.assertEqual(iterations, 6) + self.assertEqual(successful, 2) + + def test_reconnect_with_custom_process_exception(self): + """Client runs process_exception to tell if errors are retryable or fatal.""" + iteration = 0 + + def process_request(connection, request): + nonlocal iteration + iteration += 1 + if iteration == 1: + return connection.respond(http.HTTPStatus.SERVICE_UNAVAILABLE, "🚒") + return connection.respond(http.HTTPStatus.IM_A_TEAPOT, "🫖") + + def process_exception(exc): + if isinstance(exc, InvalidStatus): + if 500 <= exc.response.status_code < 600: + return None + if exc.response.status_code == 418: + return Exception("🫖 💔 ☕️") + self.fail("unexpected exception") + + with run_server(process_request=process_request) as server: + with self.assertRaises(Exception) as raised: + for _ in reconnect( + get_uri(server), + process_exception=process_exception, + reconnect_delays=short_backoff, + ): + self.fail("did not raise") + + self.assertEqual(iteration, 2) + self.assertEqual( + str(raised.exception), + "🫖 💔 ☕️", + ) + + def test_reconnect_with_custom_process_exception_raising_exception(self): + """Client supports raising an exception in process_exception.""" + + def process_request(connection, request): + return connection.respond(http.HTTPStatus.IM_A_TEAPOT, "🫖") + + def process_exception(exc): + if isinstance(exc, InvalidStatus) and exc.response.status_code == 418: + raise Exception("🫖 💔 ☕️") + self.fail("unexpected exception") + + with run_server(process_request=process_request) as server: + with self.assertRaises(Exception) as raised: + for _ in reconnect( + get_uri(server), + process_exception=process_exception, + reconnect_delays=short_backoff, + ): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "🫖 💔 ☕️", + ) + + def test_redirect(self): + """Client follows redirect.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with connect(get_uri(server) + "/redirect") as client: + self.assertEqual(client.protocol.uri.path, "/") + + def test_cross_origin_redirect(self): + """Client follows redirect to a secure URI on a different origin.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(process_request=redirect) as server: + with run_server() as other_server: + with connect(get_uri(server)): + self.assertFalse(server.connections) + self.assertTrue(other_server.connections) + + @few_redirects() + def test_redirect_limit(self): + """Client stops following redirects after limit is reached.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = request.path + return response + + with run_server(process_request=redirect) as server: + with self.assertRaises(SecurityError) as raised: + with connect(get_uri(server)): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "more than 2 redirects", + ) + + def test_redirect_with_explicit_host_port(self): + """Client follows redirect with an explicit host / port.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + address = get_host_port(server) + with connect("ws://overridden/redirect", address=address) as client: + self.assertEqual(client.protocol.uri.path, "/") + + def test_cross_origin_redirect_with_explicit_host_port(self): + """Client doesn't follow cross-origin redirect with an explicit host / port.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "ws://other/" + return response + + with run_server(process_request=redirect) as server: + address = get_host_port(server) + with self.assertRaises(ValueError) as raised: + with connect("ws://overridden/", address=address): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "cannot follow cross-origin redirect to ws://other/ " + "with an explicit host and port", + ) + + def test_redirect_with_existing_socket(self): + """Client doesn't follow redirect when using a pre-existing socket.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with socket.create_connection(get_host_port(server)) as sock: + with self.assertRaises(ValueError) as raised: + # Use a non-existing domain to ensure we connect via sock. + with connect("ws://invalid/redirect", sock=sock): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "cannot follow redirect to ws://invalid/ with a preexisting socket", + ) + + def test_cross_origin_redirect_strips_credentials(self): + """Client strips credentials when following a cross-origin redirect.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(process_request=redirect) as server: + with run_server() as other_server: + with connect( + get_uri(server), + additional_headers={ + "Authorization": "Bearer secret", + "Cookie": "session=secret", + "Proxy-Authorization": "Basic secret", + "X-Custom": "keep", + }, + ) as client: + self.assertNotIn("Authorization", client.request.headers) + self.assertNotIn("Cookie", client.request.headers) + self.assertNotIn("Proxy-Authorization", client.request.headers) + self.assertIn("X-Custom", client.request.headers) + + def test_same_origin_redirect_preserves_credentials(self): + """Client preserves credentials when following a same-origin redirect.""" + + def redirect(connection, request): + if request.path == "/redirect": + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "/" + return response + + with run_server(process_request=redirect) as server: + with connect( + get_uri(server) + "/redirect", + additional_headers={ + "Authorization": "Bearer secret", + "Cookie": "session=secret", + "Proxy-Authorization": "Basic secret", + "X-Custom": "keep", + }, + ) as client: + self.assertIn("Authorization", client.request.headers) + self.assertIn("Cookie", client.request.headers) + self.assertIn("Proxy-Authorization", client.request.headers) + def test_invalid_uri(self): """Client receives an invalid URI.""" with self.assertRaises(InvalidURI): @@ -295,12 +577,21 @@ def handle(self): class SecureClientTests(unittest.TestCase): - def test_connection(self): - """Client connects to server securely.""" + def test_context_manager(self): + """Client connects to server securely and disconnects automatically.""" with run_server(ssl=SERVER_CONTEXT) as server: with connect(get_uri(server), ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.socket.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_direct_connection(self): + """Client connects to server directly.""" + with run_server(ssl=SERVER_CONTEXT) as server: + client = connect(get_uri(server), ssl=CLIENT_CONTEXT, legacy=True) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.socket.version()[:3], "TLS") def test_set_server_hostname_implicitly(self): """Client sets server_hostname to the host in the WebSocket URI.""" @@ -346,6 +637,40 @@ def test_reject_invalid_server_hostname(self): str(raised.exception), ) + def test_cross_origin_redirect(self): + """Client follows redirect to a secure URI on a different origin.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = get_uri(other_server) + return response + + with run_server(ssl=SERVER_CONTEXT, process_request=redirect) as server: + with run_server(ssl=SERVER_CONTEXT) as other_server: + with connect(get_uri(server), ssl=CLIENT_CONTEXT): + self.assertFalse(server.connections) + self.assertTrue(other_server.connections) + + def test_redirect_to_insecure_uri(self): + """Client doesn't follow redirect from secure URI to non-secure URI.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = insecure_uri + return response + + with run_server(ssl=SERVER_CONTEXT, process_request=redirect) as server: + with self.assertRaises(SecurityError) as raised: + secure_uri = get_uri(server) + insecure_uri = secure_uri.replace("wss://", "ws://") + with connect(secure_uri, ssl=CLIENT_CONTEXT): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + f"cannot follow redirect to non-secure URI {insecure_uri}", + ) + @unittest.skipUnless("mitmproxy" in sys.modules, "mitmproxy not installed") class SocksProxyClientTests(ProxyMixin, unittest.TestCase): @@ -440,7 +765,7 @@ def test_explicit_socks_proxy(self): def test_ignore_proxy_with_existing_socket(self): """Client connects using a pre-existing socket.""" with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Use a non-existing domain to ensure we connect via sock. with connect("ws://invalid/", sock=sock) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -637,12 +962,21 @@ def test_https_proxy_invalid_server_certificate(self): @unittest.skipUnless(hasattr(socket, "AF_UNIX"), "this test requires Unix sockets") class UnixClientTests(unittest.TestCase): - def test_connection(self): - """Client connects to server over a Unix socket.""" + def test_context_manager(self): + """Client connects to Unix server and disconnects automatically.""" with temp_unix_socket_path() as path: with run_unix_server(path): with unix_connect(path) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_direct_connection(self): + """Client connects to Unix server directly.""" + with temp_unix_socket_path() as path: + with run_unix_server(path): + client = unix_connect(path, legacy=True) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") def test_set_host_header(self): """Client sets the Host header to the host in the WebSocket URI.""" @@ -652,13 +986,42 @@ def test_set_host_header(self): with unix_connect(path, uri="ws://overridden/") as client: self.assertEqual(client.request.headers["Host"], "overridden") - def test_secure_connection(self): - """Client connects to server securely over a Unix socket.""" + def test_cross_origin_redirect(self): + """Client doesn't follows redirect to a URI on a different origin.""" + + def redirect(connection, request): + response = connection.respond(http.HTTPStatus.FOUND, "") + response.headers["Location"] = "ws://other/" + return response + + with temp_unix_socket_path() as path: + with run_unix_server(path, process_request=redirect): + with self.assertRaises(ValueError) as raised: + with unix_connect(path): + self.fail("did not raise") + + self.assertEqual( + str(raised.exception), + "cannot follow cross-origin redirect to ws://other/ with a Unix socket", + ) + + def test_secure_context_manager(self): + """Client connects to Unix server securely and disconnects automatically.""" with temp_unix_socket_path() as path: with run_unix_server(path, ssl=SERVER_CONTEXT): with unix_connect(path, ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(client.socket.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") + + def test_secure_direct_connection(self): + """Client connects to Unix server securely and directly.""" + with temp_unix_socket_path() as path: + with run_unix_server(path, ssl=SERVER_CONTEXT): + client = unix_connect(path, ssl=CLIENT_CONTEXT, legacy=True) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.socket.version()[:3], "TLS") def test_set_server_hostname(self): """Client sets server_hostname to the host in the WebSocket URI.""" @@ -670,6 +1033,13 @@ def test_set_server_hostname(self): ) as client: self.assertEqual(client.socket.server_hostname, "overridden") + def test_non_existing_path(self): + """Client attempts to connect to a non-existing Unix socket path.""" + with temp_unix_socket_path() as path: + with self.assertRaises(FileNotFoundError): + with unix_connect(path + ".doesnotexist"): + self.fail("did not raise") + class ClientUsageErrorsTests(unittest.TestCase): def test_ssl_without_secure_uri(self): @@ -684,25 +1054,17 @@ def test_ssl_without_secure_uri(self): def test_proxy_ssl_without_https_proxy(self): """Client rejects proxy_ssl when proxy isn't HTTPS.""" with self.assertRaises(ValueError) as raised: - connect( + with connect( "ws://localhost/", proxy="http://localhost:8080", proxy_ssl=CLIENT_CONTEXT, - ) + ): + self.fail("did not raise") self.assertEqual( str(raised.exception), "proxy_ssl argument is incompatible with an http:// proxy", ) - def test_unix_without_path_or_sock(self): - """Unix client requires path when sock isn't provided.""" - with self.assertRaises(ValueError) as raised: - unix_connect() - self.assertEqual( - str(raised.exception), - "missing path argument", - ) - def test_unsupported_proxy(self): """Client rejects unsupported proxy.""" with self.assertRaises(InvalidProxy) as raised: @@ -713,6 +1075,15 @@ def test_unsupported_proxy(self): "other://localhost:58080 isn't a valid proxy: scheme other isn't supported", ) + def test_unix_without_path_or_sock(self): + """Unix client requires path when sock isn't provided.""" + with self.assertRaises(ValueError) as raised: + unix_connect() + self.assertEqual( + str(raised.exception), + "missing path argument", + ) + def test_unix_with_path_and_sock(self): """Unix client rejects path when sock is provided.""" sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) @@ -742,6 +1113,19 @@ def test_unsupported_compression(self): "unsupported compression: False", ) + def test_reentrancy(self): + """Client isn't reentrant.""" + with run_server() as server: + connecter = reconnect(get_uri(server)) + with connecter: + with self.assertRaises(RuntimeError) as raised: + with connecter: + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "connect() isn't reentrant", + ) + class BackwardsCompatibilityTests(DeprecationTestCase): def test_ssl_context_argument(self): @@ -750,3 +1134,64 @@ def test_ssl_context_argument(self): with self.assertDeprecationWarning("ssl_context was renamed to ssl"): with connect(get_uri(server), ssl_context=CLIENT_CONTEXT): pass + + def test_set_legacy_flag_explicitly(self): + """Client connects to server with legacy=True.""" + with run_server() as server: + client = connect(get_uri(server), legacy=True) + self.addCleanup(client.close) + self.assertIsInstance(client, ClientConnection) + + def test_unset_legacy_flag_explicitly(self): + """Client connects to server with legacy=False.""" + with run_server() as server: + client = connect(get_uri(server), legacy=False) + self.assertIsInstance(client, reconnect) + + def test_unix_set_legacy_flag_explicitly(self): + """Client connects to server with legacy=True.""" + with temp_unix_socket_path() as path: + with run_unix_server(path): + client = unix_connect(path, legacy=True) + self.addCleanup(client.close) + self.assertIsInstance(client, ClientConnection) + + def test_unix_unset_legacy_flag_explicitly(self): + """Client connects to server with legacy=False.""" + with temp_unix_socket_path() as path: + with run_unix_server(path): + client = unix_connect(path, legacy=False) + self.assertIsInstance(client, reconnect) + + def test_direct_connection_without_legacy_flag(self): + """Client connects to server without legacy=True.""" + with run_server() as server: + client = connect(get_uri(server)) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + # First call of a public API triggers a warning + with self.assertDeprecationWarning( + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly" + ): + client.send("2 + 2") + # Later calls don't trigger a warning + self.assertEqual(client.recv(), "4") + + def test_direct_unix_connection_without_legacy_flag(self): + """Client connects to Unix server without legacy=True.""" + with temp_unix_socket_path() as path: + with run_unix_server(path): + client = unix_connect(path) + self.addCleanup(client.close) + self.assertEqual(client.protocol.state.name, "OPEN") + # First call of a public API triggers a warning + with self.assertDeprecationWarning( + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly" + ): + client.send("2 + 2") + # Later calls don't trigger a warning + self.assertEqual(client.recv(), "4") diff --git a/tests/sync/test_router.py b/tests/sync/test_router.py index cf04a868..4d2b2ad9 100644 --- a/tests/sync/test_router.py +++ b/tests/sync/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) def echo(websocket, count): @@ -48,49 +55,28 @@ def test_router_matches_paths_and_extracts_parameters(self): messages = list(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - with run_router(self.url_map) as server: + with run_router(url_map) as server: with connect(get_uri(server) + "/?a=b") as client: self.assertEval(client, "ws.request.path", "/?a=b") def test_redirect(self): """Router redirects connections according to redirect_to.""" - with run_router(self.url_map, server_name="localhost") as server: - with self.assertRaises(InvalidStatus) as raised: - with connect(get_uri(server) + "/r"): - self.fail("did not raise") - self.assertEqual( - raised.exception.response.headers["Location"], - "ws://localhost/", - ) + with run_router(url_map) as server: + with connect(get_uri(server) + "/r") as client: + self.assertEval(client, "ws.request.path", "/") def test_secure_redirect(self): - """Router redirects connections to a wss:// URI when TLS is enabled.""" - with run_router( - self.url_map, server_name="localhost", ssl=SERVER_CONTEXT - ) as server: - with self.assertRaises(InvalidStatus) as raised: - with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT): - self.fail("did not raise") - self.assertEqual( - raised.exception.response.headers["Location"], - "wss://localhost/", - ) + """Router redirects connections according to redirect_to when TLS is enabled.""" + with run_router(url_map, ssl=SERVER_CONTEXT) as server: + with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: + self.assertEval(client, "ws.request.path", "/") - @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) + @patch("websockets.sync.client.reconnect.process_redirect", lambda _, exc: exc) def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - with run_router(self.url_map, ssl=True) as server: + with run_router(url_map, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r"): @@ -100,10 +86,10 @@ def test_force_secure_redirect(self): redirect_uri + "/", ) - @patch("websockets.asyncio.client.connect.process_redirect", lambda _, exc: exc) + @patch("websockets.sync.client.reconnect.process_redirect", lambda _, exc: exc) def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - with run_router(self.url_map, server_name="other") as server: + with run_router(url_map, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -114,7 +100,7 @@ def test_force_redirect_server_name(self): def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - with run_router(self.url_map) as server: + with run_router(url_map) as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -129,7 +115,7 @@ def test_process_request_returning_none(self): def process_request(ws, request): ws.process_request_ran = True - with run_router(self.url_map, process_request=process_request) as server: + with run_router(url_map, process_request=process_request) as server: with connect(get_uri(server) + "/") as client: self.assertEval(client, "ws.process_request_ran", "True") @@ -139,7 +125,7 @@ def test_process_request_returning_response(self): def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - with run_router(self.url_map, process_request=process_request) as server: + with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -156,7 +142,7 @@ def handler(self, connection): connection.my_router_ran = True return super().handler(connection) - with run_router(self.url_map, create_router=MyRouter) as server: + with run_router(url_map, create_router=MyRouter) as server: with connect(get_uri(server)) as client: self.assertEval(client, "ws.my_router_ran", "True") diff --git a/tests/sync/test_server.py b/tests/sync/test_server.py index ee01a5ee..92280ac3 100644 --- a/tests/sync/test_server.py +++ b/tests/sync/test_server.py @@ -29,6 +29,7 @@ ) from .server import ( EvalShellMixin, + get_host_port, get_uri, handler, run_server, @@ -320,7 +321,7 @@ def test_timeout_before_handshake_request(self): """Server times out before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: with run_server(open_timeout=MS) as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") @@ -334,7 +335,7 @@ def test_connection_closed_before_handshake_request(self): """Server reads EOF before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: with run_server() as server: - with socket.create_connection(server.socket.getsockname()): + with socket.create_connection(get_host_port(server)): # Wait for the server to receive the connection, then close it. time.sleep(MS) @@ -360,7 +361,7 @@ def test_junk_handshake_request(self): """Server closes the connection when receiving non-HTTP request from client.""" with self.assertLogs("websockets.server", logging.DEBUG) as logs: with run_server() as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: sock.send(b"HELO relay.invalid\r\n") # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") @@ -502,14 +503,14 @@ def test_connection(self): def test_timeout_during_tls_handshake(self): """Server times out before receiving TLS handshake request from client.""" with run_server(ssl=SERVER_CONTEXT, open_timeout=MS) as server: - with socket.create_connection(server.socket.getsockname()) as sock: + with socket.create_connection(get_host_port(server)) as sock: # Wait for the server to close the connection. self.assertEqual(sock.recv(4096), b"") def test_connection_closed_during_tls_handshake(self): """Server reads EOF before receiving TLS handshake request from client.""" with run_server(ssl=SERVER_CONTEXT) as server: - with socket.create_connection(server.socket.getsockname()): + with socket.create_connection(get_host_port(server)): # Wait for the server to receive the connection, then close it. time.sleep(MS) diff --git a/tests/trio/server.py b/tests/trio/server.py index 6e9ef417..7b69e491 100644 --- a/tests/trio/server.py +++ b/tests/trio/server.py @@ -10,8 +10,8 @@ from websockets.trio.server import serve -def get_host_port(listeners): - for listener in listeners: +def get_host_port(server): + for listener in server.listeners: if listener.socket.family == socket.AF_INET: # pragma: no branch return listener.socket.getsockname() raise AssertionError("expected at least one IPv4 socket") @@ -24,7 +24,7 @@ def get_uri(server, secure=None): for cell in server.handler.__closure__ ) # l33t hack protocol = "wss" if secure else "ws" - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) return f"{protocol}://{host}:{port}" @@ -53,19 +53,23 @@ async def assertEval(self, client, expr, value): self.assertEqual(await client.recv(), value) -kwargs = {"port": 0, "host": "localhost"} - - @contextlib.asynccontextmanager async def run_server_or_route( serve_or_route, handler_or_url_map, - **overrides, + port=0, + host="localhost", + **kwargs, ): - merged_kwargs = {**kwargs, **overrides} async with trio.open_nursery() as nursery: server = await nursery.start( - functools.partial(serve_or_route, handler_or_url_map, **merged_kwargs) + functools.partial( + serve_or_route, + handler_or_url_map, + port, + host=host, + **kwargs, + ) ) try: yield server @@ -76,9 +80,9 @@ async def run_server_or_route( nursery.cancel_scope.cancel() -def run_server(handler=handler, **overrides): - return run_server_or_route(serve, handler, **overrides) +def run_server(handler=handler, **kwargs): + return run_server_or_route(serve, handler, **kwargs) -def run_router(url_map, **overrides): - return run_server_or_route(route, url_map, **overrides) +def run_router(url_map, **kwargs): + return run_server_or_route(route, url_map, **kwargs) diff --git a/tests/trio/test_client.py b/tests/trio/test_client.py index 6afd35f8..89269e04 100644 --- a/tests/trio/test_client.py +++ b/tests/trio/test_client.py @@ -54,23 +54,24 @@ async def few_redirects(): class ClientTests(IsolatedTrioTestCase): - async def test_connection(self): - """Client connects to server.""" + async def test_context_manager(self): + """Client connects to server and disconnects automatically.""" async with run_server() as server: async with connect(get_uri(server)) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_explicit_host_port(self): """Client connects using an explicit host / port.""" async with run_server() as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect("ws://overridden/", host=host, port=port) as client: self.assertEqual(client.protocol.state.name, "OPEN") async def test_existing_stream(self): """Client connects using a pre-existing stream.""" async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) # Use a non-existing domain to ensure we connect via stream. async with connect("ws://invalid/", stream=stream) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -306,7 +307,7 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect( "ws://overridden/redirect", host=host, port=port ) as client: @@ -321,7 +322,7 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) with self.assertRaises(ValueError) as raised: async with connect("ws://overridden/", host=host, port=port): self.fail("did not raise") @@ -341,9 +342,9 @@ def redirect(connection, request): return response async with run_server(process_request=redirect) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) with self.assertRaises(ValueError) as raised: - # Use a non-existing domain to ensure we connect via sock. + # Use a non-existing domain to ensure we connect via stream. async with connect("ws://invalid/redirect", stream=stream): self.fail("did not raise") @@ -505,7 +506,11 @@ async def junk(stream): async with trio.open_nursery() as nursery: try: listeners = await nursery.start(trio.serve_tcp, junk, 0) - host, port = get_host_port(listeners) + host, port = next( + listener + for listener in listeners + if listener.socket.family == socket.AF_INET + ).socket.getsockname() with self.assertRaises(InvalidMessage) as raised: async with connect(f"ws://{host}:{port}"): self.fail("did not raise") @@ -524,8 +529,8 @@ async def junk(stream): class SecureClientTests(IsolatedTrioTestCase): - async def test_connection(self): - """Client connects to server securely.""" + async def test_context_manager(self): + """Client connects to server securely and disconnects automatically.""" async with run_server(ssl=SERVER_CONTEXT) as server: async with connect( get_uri(server, secure=True), ssl=CLIENT_CONTEXT @@ -533,11 +538,12 @@ async def test_connection(self): self.assertEqual(client.protocol.state.name, "OPEN") ssl_object = client.stream._ssl_object self.assertEqual(ssl_object.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_set_server_hostname_implicitly(self): """Client sets server_hostname to the host in the WebSocket URI.""" async with run_server(ssl=SERVER_CONTEXT) as server: - host, port = get_host_port(server.listeners) + host, port = get_host_port(server) async with connect( "wss://overridden/", host=host, port=port, ssl=CLIENT_CONTEXT ) as client: @@ -721,7 +727,7 @@ async def test_explicit_socks_proxy(self): async def test_ignore_proxy_with_existing_stream(self): """Cli ent connects using a pre-existing stream.""" async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) # Use a non-existing domain to ensure we connect via stream. async with connect("ws://invalid/", stream=stream) as client: self.assertEqual(client.protocol.state.name, "OPEN") @@ -914,12 +920,13 @@ async def test_https_proxy_invalid_server_certificate(self): @unittest.skipUnless(hasattr(socket, "AF_UNIX"), "this test requires Unix sockets") class UnixClientTests(IsolatedTrioTestCase): - async def test_connection(self): - """Client connects to server over a Unix socket.""" + async def test_context_manager(self): + """Client connects to Unix server and disconnects automatically.""" with temp_unix_socket_path() as path: with run_unix_server(path): async with unix_connect(path) as client: self.assertEqual(client.protocol.state.name, "OPEN") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_set_host_header(self): """Client sets the Host header to the host in the WebSocket URI.""" @@ -948,14 +955,15 @@ def redirect(connection, request): "cannot follow cross-origin redirect to ws://other/ with a Unix socket", ) - async def test_secure_connection(self): - """Client connects to server securely over a Unix socket.""" + async def test_secure_context_manager(self): + """Client connects to Unix server securely.""" with temp_unix_socket_path() as path: with run_unix_server(path, ssl=SERVER_CONTEXT): async with unix_connect(path, ssl=CLIENT_CONTEXT) as client: self.assertEqual(client.protocol.state.name, "OPEN") ssl_object = client.stream._ssl_object self.assertEqual(ssl_object.version()[:3], "TLS") + self.assertEqual(client.protocol.state.name, "CLOSED") async def test_set_server_hostname(self): """Client sets server_hostname to the host in the WebSocket URI.""" @@ -970,6 +978,13 @@ async def test_set_server_hostname(self): ssl_object = client.stream._ssl_object self.assertEqual(ssl_object.server_hostname, "overridden") + async def test_non_existing_path(self): + """Client attempts to connect to a non-existing Unix socket path.""" + with temp_unix_socket_path() as path: + with self.assertRaises(FileNotFoundError): + async with unix_connect(path + ".doesnotexist"): + self.fail("did not raise") + class ClientUsageErrorsTests(IsolatedTrioTestCase): async def test_ssl_without_secure_uri(self): @@ -1006,6 +1021,27 @@ async def test_unsupported_proxy(self): "other://localhost:51080 isn't a valid proxy: scheme other isn't supported", ) + async def test_unix_without_path_or_sock(self): + """Unix client requires path when sock isn't provided.""" + with self.assertRaises(ValueError) as raised: + async with unix_connect(): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "missing path argument", + ) + + async def test_unix_with_path_and_stream(self): + """Unix client rejects path when stream is provided.""" + stream, _ = trio.testing.memory_stream_pair() + with self.assertRaises(ValueError) as raised: + async with unix_connect(path="/", stream=stream): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "path is incompatible with stream", + ) + async def test_invalid_subprotocol(self): """Client rejects single value of subprotocols.""" with self.assertRaises(TypeError) as raised: diff --git a/tests/trio/test_router.py b/tests/trio/test_router.py index a85f5bf0..e303a3bd 100644 --- a/tests/trio/test_router.py +++ b/tests/trio/test_router.py @@ -16,6 +16,13 @@ from werkzeug.routing import Map, Rule except ImportError: pass +else: + url_map = Map( + [ + Rule("/", endpoint=handler), + Rule("/r", redirect_to="/"), + ] + ) async def echo(websocket, count): @@ -48,37 +55,28 @@ async def test_router_matches_paths_and_extracts_parameters(self): messages = await alist(client) self.assertEqual(messages, ["hello", "hello", "hello"]) - @property # avoids an import-time dependency on werkzeug - def url_map(self): - return Map( - [ - Rule("/", endpoint=handler), - Rule("/r", redirect_to="/"), - ] - ) - async def test_route_with_query_string(self): """Router ignores query strings when matching paths.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: async with connect(get_uri(server) + "/?a=b") as client: await self.assertEval(client, "ws.request.path", "/?a=b") async def test_redirect(self): """Router redirects connections according to redirect_to.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: async with connect(get_uri(server) + "/r") as client: await self.assertEval(client, "ws.request.path", "/") async def test_secure_redirect(self): """Router redirects connections according to redirect_to when TLS is enabled.""" - async with run_router(self.url_map, ssl=SERVER_CONTEXT) as server: + async with run_router(url_map, ssl=SERVER_CONTEXT) as server: async with connect(get_uri(server) + "/r", ssl=CLIENT_CONTEXT) as client: await self.assertEval(client, "ws.request.path", "/") @patch("websockets.trio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_secure_redirect(self): """Router redirects ws:// connections to a wss:// URI when ssl=True.""" - async with run_router(self.url_map, ssl=True) as server: + async with run_router(url_map, ssl=True) as server: redirect_uri = get_uri(server, secure=True) with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): @@ -91,7 +89,7 @@ async def test_force_secure_redirect(self): @patch("websockets.trio.client.connect.process_redirect", lambda _, exc: exc) async def test_force_redirect_server_name(self): """Router redirects connections to the host declared in server_name.""" - async with run_router(self.url_map, server_name="other") as server: + async with run_router(url_map, server_name="other") as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/r"): self.fail("did not raise") @@ -102,7 +100,7 @@ async def test_force_redirect_server_name(self): async def test_not_found(self): """Router rejects requests to unknown paths with an HTTP 404 error.""" - async with run_router(self.url_map) as server: + async with run_router(url_map) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/n"): self.fail("did not raise") @@ -117,7 +115,7 @@ async def test_process_request_function_returning_none(self): def process_request(ws, request): ws.process_request_ran = True - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -127,7 +125,7 @@ async def test_process_request_coroutine_returning_none(self): async def process_request(ws, request): ws.process_request_ran = True - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: async with connect(get_uri(server) + "/") as client: await self.assertEval(client, "ws.process_request_ran", "True") @@ -137,7 +135,7 @@ async def test_process_request_function_returning_response(self): def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -152,7 +150,7 @@ async def test_process_request_coroutine_returning_response(self): async def process_request(ws, request): return ws.respond(http.HTTPStatus.FORBIDDEN, "Forbidden") - async with run_router(self.url_map, process_request=process_request) as server: + async with run_router(url_map, process_request=process_request) as server: with self.assertRaises(InvalidStatus) as raised: async with connect(get_uri(server) + "/"): self.fail("did not raise") @@ -169,6 +167,6 @@ async def handler(self, connection): connection.my_router_ran = True return await super().handler(connection) - async with run_router(self.url_map, create_router=MyRouter) as server: + async with run_router(url_map, create_router=MyRouter) as server: async with connect(get_uri(server)) as client: await self.assertEval(client, "ws.my_router_ran", "True") diff --git a/tests/trio/test_server.py b/tests/trio/test_server.py index 8018175e..ffd7337c 100644 --- a/tests/trio/test_server.py +++ b/tests/trio/test_server.py @@ -2,6 +2,7 @@ import hmac import http import logging +import socket import trio @@ -64,7 +65,11 @@ async def test_connection_handler_raises_exception(self): async def test_existing_listeners(self): """Server receives connection using pre-existing listeners.""" listeners = await trio.open_tcp_listeners(0, host="localhost") - host, port = get_host_port(listeners) + host, port = next( + listener + for listener in listeners + if listener.socket.family == socket.AF_INET + ).socket.getsockname() # Unset the default values of port and host set by run_server. async with run_server(port=None, host=None, listeners=listeners): async with connect(f"ws://{host}:{port}/") as client: # type: ignore @@ -416,7 +421,7 @@ async def test_timeout_before_handshake_request(self): """Server times out before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server(open_timeout=MS) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) try: # Wait for the server to close the connection. self.assertEqual(await stream.receive_some(4096), b"") @@ -433,7 +438,7 @@ async def test_connection_closed_before_handshake_request(self): """Server reads EOF before receiving handshake request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.aclose() self.assertExceptionLogged( @@ -459,7 +464,7 @@ async def test_junk_handshake_request(self): """Server closes the connection when receiving non-HTTP request from client.""" with self.assertLogs("websockets", logging.DEBUG) as logs: async with run_server() as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.send_all(b"HELO relay.invalid\r\n") try: # Wait for the server to close the connection. @@ -612,7 +617,7 @@ async def test_connection(self): async def test_timeout_during_tls_handshake(self): """Server times out before receiving TLS handshake request from client.""" async with run_server(ssl=SERVER_CONTEXT, open_timeout=MS) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) try: # Wait for the server to close the connection. self.assertEqual(await stream.receive_some(4096), b"") @@ -622,7 +627,7 @@ async def test_timeout_during_tls_handshake(self): async def test_connection_closed_during_tls_handshake(self): """Server reads EOF before receiving TLS handshake request from client.""" async with run_server(ssl=SERVER_CONTEXT) as server: - stream = await trio.open_tcp_stream(*get_host_port(server.listeners)) + stream = await trio.open_tcp_stream(*get_host_port(server)) await stream.aclose()