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/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/sync/client.rst b/docs/reference/sync/client.rst index fdc772ed..77e95e58 100644 --- a/docs/reference/sync/client.rst +++ b/docs/reference/sync/client.rst @@ -10,6 +10,15 @@ Opening a connection .. autofunction:: unix_connect +Reconnecting automatically +--------------------------- + +.. autofunction:: reconnect + +.. autofunction:: unix_reconnect + +.. autofunction:: process_exception + Using a connection ------------------ diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index 6c6c1149..f49af72d 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -177,10 +177,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` 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 +188,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: diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 509eb9df..04212971 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -1,16 +1,28 @@ 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 collections.abc import Generator, Iterator, Sequence +from types import TracebackType from typing import Any, Callable, Literal, TypeVar, cast -from ..client import ClientProtocol -from ..datastructures import HeadersLike -from ..exceptions import InvalidProxyMessage, InvalidProxyStatus, ProxyError +from ..asyncio.client import process_exception +from ..client import ClientProtocol, backoff +from ..datastructures import Headers, HeadersLike +from ..exceptions import ( + InvalidProxyMessage, + InvalidProxyStatus, + InvalidStatus, + ProxyError, + SecurityError, +) from ..extensions.base import ClientExtensionFactory from ..extensions.permessage_deflate import enable_client_permessage_deflate from ..headers import validate_subprotocols @@ -24,7 +36,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 +79,7 @@ def __init__( ) -> None: self.protocol: ClientProtocol self.response_rcvd = threading.Event() + self.pending_legacy_warning = True super().__init__( sock, protocol, @@ -74,6 +89,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 +158,500 @@ def recv_events(self) -> None: self.response_rcvd.set() +# This is spelled in lower case because it's exposed as a callable in the API. +class reconnect: + """ + Connect to the WebSocket server at ``uri``, with support for reconnecting. + + This class returns a :class:`ClientConnection` instance, which you can + use to send and receive messages. + + :func:`reconnect` may be used as a context manager for a single + connection attempt:: + + from websockets.sync.client import reconnect + + with reconnect(...) as websocket: + ... + + The connection is closed automatically when exiting the context. + + :func:`reconnect` may be used 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. + + For a single connection attempt without a context manager, see + :func:`connect`. + + Args: + uri: URI of the WebSocket server. + sock: Preexisting TCP socket. ``sock`` overrides the host and port + from ``uri``. You may call :func:`socket.create_connection` to + create a suitable TCP socket. + ssl: Configuration for enabling TLS on the connection. + server_hostname: Host name for the TLS handshake. ``server_hostname`` + overrides the host name from ``uri``. + origin: Value of the ``Origin`` header, for servers that require it. + extensions: List of supported extensions, in order in which they + should be negotiated and run. + subprotocols: List of supported subprotocols, in order of decreasing + preference. + compression: The "permessage-deflate" extension is enabled by default. + Set ``compression`` to :obj:`None` to disable it. See the + :doc:`compression guide <../../topics/compression>` for details. + additional_headers: Arbitrary HTTP headers to add to the handshake + request. + user_agent_header: Value of the ``User-Agent`` request header. + It defaults to ``"Python/x.y.z websockets/X.Y"``. + Setting it to :obj:`None` removes the header. + proxy: If a proxy is configured, it is used by default. Set ``proxy`` + to :obj:`None` to disable the proxy or to the address of a proxy + to override the system configuration. See the :doc:`proxy docs + <../../topics/proxies>` for details. + proxy_ssl: Configuration for enabling TLS on the proxy connection. + proxy_server_hostname: Host name for the TLS handshake with the proxy. + ``proxy_server_hostname`` overrides the host name from ``proxy``. + process_exception: When reconnecting automatically, tell whether an + error is transient or fatal. The default behavior is defined by + :func:`process_exception`. Refer to its documentation for details. + open_timeout: Timeout for opening the connection in seconds. + :obj:`None` disables the timeout. + ping_interval: Interval between keepalive pings in seconds. + :obj:`None` disables keepalive. + ping_timeout: Timeout for keepalive pings in seconds. + :obj:`None` disables timeouts. + close_timeout: Timeout for closing the connection in seconds. + :obj:`None` disables the timeout. + reconnect_delays: Delays in seconds between reconnection attempts. + Default is exponential backoff with 5s jitter, capped at 60s. + max_size: Maximum size of incoming messages in bytes. + :obj:`None` disables the limit. You may pass a ``(max_message_size, + max_fragment_size)`` tuple to set different limits for messages and + fragments when you expect long messages sent in short fragments. + max_queue: High-water mark of the buffer where frames are received. + It defaults to 16 frames. The low-water mark defaults to ``max_queue + // 4``. You may pass a ``(high, low)`` tuple to set the high-water + and low-water marks. If you want to disable flow control entirely, + you may set it to ``None``, although that's a bad idea. + logger: Logger for this client. + It defaults to ``logging.getLogger("websockets.client")``. + See the :doc:`logging guide <../../topics/logging>` for details. + create_connection: Factory for the :class:`ClientConnection` managing + the connection. Set it to a wrapper or a subclass to customize + connection handling. + + Any other keyword arguments are passed to :func:`~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. + InvalidProxy: If ``proxy`` isn't a valid proxy. + OSError: If the TCP connection fails. + InvalidHandshake: If the opening handshake fails. + TimeoutError: If the opening handshake times out. + + """ + + def __init__( + self, + uri: str, + *, + # 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 + + # 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") + kwargs["unix"] = True + kwargs["path"] = path + + 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() + + proxy = self.proxy + if kwargs.get("unix", False): + proxy = None + if proxy is True: + proxy = get_proxy(self.ws_uri) + + if kwargs.pop("unix", False): + 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.""" + # TCP connection is already established. + if self.sock is not None: + sock = self.sock + else: + sock = self.open_socket(deadline) + + try: + # Disable Nagle algorithm + + if not self.open_socket_kwargs.get("unix", False): + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True) + + # Initialize TLS wrapper and perform TLS handshake + + if self.ws_uri.secure: + sock = self.enable_tls(sock, deadline) + + # Initialize WebSocket protocol + + protocol = ClientProtocol( + self.ws_uri, + **self.protocol_kwargs, # type: ignore + ) + + # Initialize WebSocket connection + + # 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 or port" + ) + + # Strip credentials to avoid leaking them to a different origin. + if self.additional_headers is not None: + self.additional_headers = Headers( + ( + (key, value) + for key, value in Headers(self.additional_headers).raw_items() + if key.lower() + not in ["authorization", "cookie", "proxy-authorization"] + ) + ) + + return new_uri + + def connect(self) -> ClientConnection: + """Connect to a WebSocket server, following redirects.""" + 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 reconnect(...) as ...: ... + + def __enter__(self) -> ClientConnection: + if hasattr(self, "connection"): + raise RuntimeError("reconnect() 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 + + def connect( uri: str, *, @@ -158,16 +682,16 @@ def connect( logger: LoggerLike | None = None, # Escape hatch for advanced customization create_connection: type[ClientConnection] | None = None, + # Deprecation + legacy: bool = False, # Other keyword arguments are passed to socket.create_connection **kwargs: Any, ) -> ClientConnection: """ Connect to the WebSocket server at ``uri``. - This function returns a :class:`ClientConnection` instance, which you can - use to send and receive messages. - - :func:`connect` may be used as a context manager:: + :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 +700,20 @@ 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 = await connect(..., legacy=True) + + In that case, you're responsible for closing the connection with + :meth:`ClientConnection.close` wh`en no longer needed. + + .. When the ``legacy`` flag is enabled, :func:`connect` returns directly a + .. :class:`ClientConnection` and iterating that connection yields messages. + .. This is different from the default behavior of returning an object that + .. can be iterated to reconnect automatically. + Args: uri: URI of the WebSocket server. sock: Preexisting TCP socket. ``sock`` overrides the host and port @@ -227,8 +765,14 @@ def connect( create_connection: Factory for the :class:`ClientConnection` managing the connection. Set it to a wrapper or a subclass to customize connection handling. + legacy: Set this to :obj:`True` to use :func:`connect` outside a + context manager without triggering a deprecation warning. 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,179 +782,80 @@ def connect( TimeoutError: If the opening handshake times out. """ + # Backwards compatibility: connect can return a ClientConnection. + legacy = kwargs.pop("legacy", False) + + connection = 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, + ).connect() + 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) - - # Initialize TLS wrapper and perform TLS handshake - - 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) - # Initialize WebSocket protocol +def unix_connect( + path: str | None = None, + uri: str | None = None, + *, + legacy: bool = False, + **kwargs: Any, +) -> ClientConnection: + """ + Connect to a WebSocket server listening on a Unix socket. - protocol = ClientProtocol( - ws_uri, - origin=origin, - extensions=extensions, - subprotocols=subprotocols, - max_size=max_size, - logger=logger, - ) + This function accepts the same keyword arguments as :func:`connect`. - # Initialize WebSocket connection + It's only available on Unix. - # 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 + It's mainly useful for debugging servers listening on Unix sockets. - try: - connection.handshake( - additional_headers, - user_agent_header, - deadline.timeout(), - ) - except Exception: - connection.close_socket() - connection.recv_events_thread.join() - raise + Args: + path: File system path to the Unix socket. + uri: URI of the WebSocket server. ``uri`` defaults to + ``ws://localhost/`` or, when a ``ssl`` is provided, to + ``wss://localhost/``. - connection.start_keepalive() - return connection + """ + 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, legacy=legacy, **kwargs) -def unix_connect( +def unix_reconnect( path: str | None = None, uri: str | None = None, **kwargs: Any, -) -> ClientConnection: +) -> reconnect: """ - Connect to a WebSocket server listening on a Unix socket. + Connect to a WebSocket server listening on a Unix socket, with support + for reconnecting. - This function accepts the same keyword arguments as :func:`connect`. + This function accepts the same keyword arguments as :func:`reconnect`. It's only available on Unix. @@ -429,7 +874,7 @@ def unix_connect( uri = "ws://localhost/" else: uri = "wss://localhost/" - return connect(uri=uri, unix=True, path=path, **kwargs) + return reconnect(uri=uri, unix=True, path=path, **kwargs) try: diff --git a/src/websockets/sync/connection.py b/src/websockets/sync/connection.py index 74ac997a..1cefc9c4 100644 --- a/src/websockets/sync/connection.py +++ b/src/websockets/sync/connection.py @@ -317,6 +317,11 @@ def recv(self, timeout: float | None = None, decode: bool | None = None) -> Data :meth:`recv_streaming` concurrently. """ +<<<<<<< HEAD + self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) try: return self.recv_messages.get(timeout, decode) except EOFError: @@ -387,6 +392,11 @@ def recv_streaming(self, decode: bool | None = None) -> Iterator[Data]: :meth:`recv_streaming` concurrently. """ +<<<<<<< HEAD + self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) try: yield from self.recv_messages.get_iter(decode) return @@ -466,6 +476,11 @@ def send( TypeError: If ``message`` doesn't have a supported type. """ +<<<<<<< HEAD + self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) # Unfragmented message — this case must be handled first because # strings and bytes-like objects are iterable. @@ -591,6 +606,11 @@ def close( reason: WebSocket close reason. """ +<<<<<<< HEAD + self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) 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 +667,11 @@ def ping( the corresponding pong wasn't received yet. """ +<<<<<<< HEAD + self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -684,6 +709,11 @@ def pong(self, data: DataLike = b"") -> None: ConnectionClosed: When the connection is closed. """ +<<<<<<< HEAD + self.maybe_raise_legacy_warning() +======= + self.raise_legacy_warning() +>>>>>>> e97cd2f (wip) if isinstance(data, BytesLike): data = bytes(data) elif isinstance(data, str): @@ -696,6 +726,16 @@ def pong(self, data: DataLike = b"") -> None: # Private methods +<<<<<<< HEAD + def maybe_raise_legacy_warning(self) -> None: + pass # see override in ClientConnection +======= + def raise_legacy_warning(self) -> None: + # No-op here. websockets.sync.client.ClientConnection overrides this + # method to warn when connect() is used outside a context manager. + pass +>>>>>>> e97cd2f (wip) + def process_event(self, event: Event) -> None: """ Process one incoming event. diff --git a/src/websockets/trio/client.py b/src/websockets/trio/client.py index 1bdde69c..cd79f3d6 100644 --- a/src/websockets/trio/client.py +++ b/src/websockets/trio/client.py @@ -135,8 +135,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` is designed to be called as an asynchronous context manager + yielding a :class:`ClientConnection`, which you can then use to receive and + send messages:: :func:`connect` may be used as an asynchronous context manager:: @@ -147,8 +148,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 +163,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 @@ -220,6 +225,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. @@ -532,7 +541,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 ...: ... diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index 6f6070c6..6f0a0c92 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.""" 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..a79e9bcc 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 @@ -8,8 +9,10 @@ import threading import time import unittest +import warnings from unittest.mock import patch +from websockets.client import backoff from websockets.exceptions import ( InvalidHandshake, InvalidMessage, @@ -18,6 +21,7 @@ InvalidStatus, InvalidURI, ProxyError, + SecurityError, ) from websockets.extensions.permessage_deflate import PerMessageDeflate from websockets.sync.client import * @@ -30,20 +34,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.socket) + 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 +168,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.socket) + 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.socket) + 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 or 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.socket)) 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): @@ -346,6 +629,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 +757,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") @@ -684,25 +1001,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 +1022,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 +1060,19 @@ def test_unsupported_compression(self): "unsupported compression: False", ) + def test_reentrancy(self): + """reconnect() 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), + "reconnect() isn't reentrant", + ) + class BackwardsCompatibilityTests(DeprecationTestCase): def test_ssl_context_argument(self): @@ -750,3 +1081,76 @@ 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_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.ping() + # Later calls don't trigger a warning + client.pong() + +LEGACY_WARNING = ( + "connect() must be used as a context manager: " + "with connect(...) as websocket: ...; alternatively, use " + "websocket = connect(..., legacy=True) to connect directly" +) + + +class LegacyConnectTests(DeprecationTestCase): + def test_legacy_warns_on_first_use(self): + """connect() without a context manager warns on first use.""" + with run_server() as server: + client = connect(get_uri(server)) + with self.assertDeprecationWarning(LEGACY_WARNING): + client.send("1 + 1") + self.assertEqual(client.recv(), "2") + client.close() + + def test_legacy_warns_on_message_iteration(self): + """connect() without a context manager warns when iterating messages.""" + with run_server() as server: + client = connect(get_uri(server) + "/no-op") + with self.assertDeprecationWarning(LEGACY_WARNING): + messages = list(client) + self.assertEqual(messages, []) + + def test_legacy_warns_once(self): + """connect() without a context manager warns only once.""" + with run_server() as server: + client = connect(get_uri(server)) + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + client.send("1 + 1") + self.assertEqual(client.recv(), "2") + client.close() + self.assertEqual(len(recorded_warnings), 1) + + def test_legacy_true_doesnt_warn(self): + """connect(legacy=True) doesn't warn.""" + with run_server() as server: + client = connect(get_uri(server), legacy=True) + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + client.send("1 + 1") + self.assertEqual(client.recv(), "2") + client.close() + self.assertEqual(recorded_warnings, []) + + def test_context_manager_doesnt_warn(self): + """connect() doesn't warn when used as a context manager.""" + with run_server() as server: + with warnings.catch_warnings(record=True) as recorded_warnings: + warnings.simplefilter("always") + with connect(get_uri(server)) as client: + client.send("1 + 1") + self.assertEqual(client.recv(), "2") + self.assertEqual(recorded_warnings, []) 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..dc7c68c9 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,7 +342,7 @@ 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. async with connect("ws://invalid/redirect", stream=stream): @@ -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") @@ -537,7 +542,7 @@ async def test_connection(self): 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 +726,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") 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()