Repository: aaugustin/websockets
Pull request #632 (2019-06-29 11:17:37)

Add features to connect()


Patch:
diff --git a/docs/api.rst b/docs/api.rst
index ef567ed..56372eb 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -59,6 +59,9 @@ Client
     .. autofunction:: connect(uri, *, create_protocol=None, ping_interval=20, ping_timeout=20, close_timeout=10, max_size=2 ** 20, max_queue=2 ** 5, read_limit=2 ** 16, write_limit=2 ** 16, loop=None, compression='deflate', origin=None, extensions=None, subprotocols=None, extra_headers=None, **kwds)
         :async:
 
+    .. autofunction:: unix_connect(path, uri="ws://localhost/", *, create_protocol=None, ping_interval=20, ping_timeout=20, close_timeout=10, max_size=2 ** 20, max_queue=2 ** 5, read_limit=2 ** 16, write_limit=2 ** 16, loop=None, compression='deflate', origin=None, extensions=None, subprotocols=None, extra_headers=None, **kwds)
+        :async:
+
     .. autoclass:: WebSocketClientProtocol(*, host=None, port=None, secure=None, ping_interval=20, ping_timeout=20, close_timeout=10, max_size=2 ** 20, max_queue=2 ** 5, read_limit=2 ** 16, write_limit=2 ** 16, loop=None, origin=None, extensions=None, subprotocols=None, extra_headers=None)
 
         .. automethod:: handshake
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 92cbce5..7a02ec0 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -56,6 +56,10 @@ Also:
 * :func:`~client.connect` handles redirects from the server during the
   handshake.
 
+* :func:`~client.connect` supports overriding ``host`` and ``port``.
+
+* Added :func:`~client.unix_connect` for connecting to Unix sockets.
+
 * Improved support for sending fragmented messages by accepting asynchronous
   iterators in :meth:`~protocol.WebSocketCommonProtocol.send`.
 
diff --git a/example/unix_client.py b/example/unix_client.py
new file mode 100755
index 0000000..577135b
--- /dev/null
+++ b/example/unix_client.py
@@ -0,0 +1,19 @@
+#!/usr/bin/env python
+
+# WS client example connecting to a Unix socket
+
+import asyncio
+import os.path
+import websockets
+
+async def hello():
+    socket_path = os.path.join(os.path.dirname(__file__), "socket")
+    async with websockets.unix_connect(socket_path) as websocket:
+        name = input("What's your name? ")
+        await websocket.send(name)
+        print(f"> {name}")
+
+        greeting = await websocket.recv()
+        print(f"< {greeting}")
+
+asyncio.get_event_loop().run_until_complete(hello())
diff --git a/example/unix_server.py b/example/unix_server.py
new file mode 100755
index 0000000..a6ec016
--- /dev/null
+++ b/example/unix_server.py
@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+
+# WS server example listening on a Unix socket
+
+import asyncio
+import os.path
+import websockets
+
+async def hello(websocket, path):
+    name = await websocket.recv()
+    print(f"< {name}")
+
+    greeting = f"Hello {name}!"
+
+    await websocket.send(greeting)
+    print(f"> {greeting}")
+
+socket_path = os.path.join(os.path.dirname(__file__), "socket")
+start_server = websockets.unix_serve(hello, socket_path)
+
+asyncio.get_event_loop().run_until_complete(start_server)
+asyncio.get_event_loop().run_forever()
diff --git a/src/websockets/client.py b/src/websockets/client.py
index 110e61f..10435c1 100644
--- a/src/websockets/client.py
+++ b/src/websockets/client.py
@@ -5,6 +5,7 @@ The :mod:`websockets.client` module defines a simple WebSocket client API.
 
 import asyncio
 import collections.abc
+import functools
 import logging
 import warnings
 from types import TracebackType
@@ -17,6 +18,7 @@ from .exceptions import (
     InvalidStatusCode,
     NegotiationError,
     RedirectHandshake,
+    SecurityError,
 )
 from .extensions.base import ClientExtensionFactory, Extension
 from .extensions.permessage_deflate import ClientPerMessageDeflateFactory
@@ -35,7 +37,7 @@ from .typing import Origin, Subprotocol
 from .uri import WebSocketURI, parse_uri
 
 
-__all__ = ["connect", "WebSocketClientProtocol"]
+__all__ = ["connect", "unix_connect", "WebSocketClientProtocol"]
 
 logger = logging.getLogger(__name__)
 
@@ -332,6 +334,11 @@ class Connect:
     a ``wss://`` URI, if this argument isn't provided explicitly, it's set to
     ``True``, which means Python's default :class:`~ssl.SSLContext` is used.
 
+    You can connect to a different host and port from those found in ``uri``
+    by setting ``host`` and ``port`` keyword arguments. This only changes the
+    destination of the TCP connection; the hostname from ``uri`` is still used
+    in the TLS handshake for secure connections and in the ``Host`` header.
+
     The behavior of the ``ping_interval``, ``ping_timeout``, ``close_timeout``,
     ``max_size``, ``max_queue``, ``read_limit``, and ``write_limit`` optional
     arguments is described in the documentation of
@@ -370,6 +377,7 @@ class Connect:
         self,
         uri: str,
         *,
+        path: Optional[str] = None,
         create_protocol: Optional[Type[WebSocketClientProtocol]] = None,
         ping_interval: float = 20,
         ping_timeout: float = 20,
@@ -410,8 +418,8 @@ class Connect:
         if loop is None:
             loop = asyncio.get_event_loop()
 
-        self._wsuri = parse_uri(uri)
-        if self._wsuri.secure:
+        wsuri = parse_uri(uri)
+        if wsuri.secure:
             kwargs.setdefault("ssl", True)
         elif kwargs.get("ssl") is not None:
             raise ValueError(
@@ -432,69 +440,83 @@ class Connect:
         elif compression is not None:
             raise ValueError(f"unsupported compression: {compression}")
 
-        self._create_protocol = create_protocol
-        self._ping_interval = ping_interval
-        self._ping_timeout = ping_timeout
-        self._close_timeout = close_timeout
-        self._max_size = max_size
-        self._max_queue = max_queue
-        self._read_limit = read_limit
-        self._write_limit = write_limit
-        self._loop = loop
-        self._legacy_recv = legacy_recv
-        self._klass = klass
-        self._timeout = timeout
-        self._compression = compression
-        self._origin = origin
-        self._extensions = extensions
-        self._subprotocols = subprotocols
-        self._extra_headers = extra_headers
-        self._kwargs = kwargs
-
-    async def _creating_connection(
-        self
-    ) -> Tuple[asyncio.Transport, WebSocketClientProtocol]:
-        if self._wsuri.secure:
-            self._kwargs.setdefault("ssl", True)
-
-        factory = lambda: self._create_protocol(
-            host=self._wsuri.host,
-            port=self._wsuri.port,
-            secure=self._wsuri.secure,
-            ping_interval=self._ping_interval,
-            ping_timeout=self._ping_timeout,
-            close_timeout=self._close_timeout,
-            max_size=self._max_size,
-            max_queue=self._max_queue,
-            read_limit=self._read_limit,
-            write_limit=self._write_limit,
-            loop=self._loop,
-            legacy_recv=self._legacy_recv,
-            origin=self._origin,
-            extensions=self._extensions,
-            subprotocols=self._subprotocols,
-            extra_headers=self._extra_headers,
+        factory = functools.partial(
+            create_protocol,
+            host=wsuri.host,
+            port=wsuri.port,
+            secure=wsuri.secure,
+            ping_interval=ping_interval,
+            ping_timeout=ping_timeout,
+            close_timeout=close_timeout,
+            max_size=max_size,
+            max_queue=max_queue,
+            read_limit=read_limit,
+            write_limit=write_limit,
+            loop=loop,
+            legacy_recv=legacy_recv,
+            origin=origin,
+            extensions=extensions,
+            subprotocols=subprotocols,
+            extra_headers=extra_headers,
         )
 
-        host: Optional[str]
-        port: Optional[int]
-        if self._kwargs.get("sock") is None:
-            host, port = self._wsuri.host, self._wsuri.port
+        if path is None:
+            host: Optional[str]
+            port: Optional[int]
+            if kwargs.get("sock") is None:
+                host, port = wsuri.host, wsuri.port
+            else:
+                # If sock is given, host and port shouldn't be specified.
+                host, port = None, None
+            # If host and port are given, override values from the URI.
+            host = kwargs.pop("host", host)
+            port = kwargs.pop("port", port)
+            create_connection = functools.partial(
+                loop.create_connection, factory, host, port, **kwargs
+            )
         else:
-            # If sock is given, host and port mustn't be specified.
-            host, port = None, None
+            create_connection = functools.partial(
+                loop.create_unix_connection, factory, path, **kwargs
+            )
+
+        # This is a coroutine function.
+        self._create_connection = create_connection
+        self._wsuri = wsuri
+        self._origin = origin
 
-        self._wsuri = self._wsuri
-        self._origin = self._origin
+    def handle_redirect(self, uri: str) -> None:
+        # Update the state of this instance to connect to a new URI.
+        old_wsuri = self._wsuri
+        new_wsuri = parse_uri(uri)
 
-        # This is a coroutine object.
-        # https://github.com/python/typeshed/pull/2756
-        transport, protocol = await self._loop.create_connection(  # type: ignore
-            factory, host, port, **self._kwargs
+        # Forbid TLS downgrade.
+        if old_wsuri.secure and not new_wsuri.secure:
+            raise SecurityError("redirect from WSS to WS")
+
+        same_origin = (
+            old_wsuri.host == new_wsuri.host and old_wsuri.port == new_wsuri.port
         )
-        transport = cast(asyncio.Transport, transport)
-        protocol = cast(WebSocketClientProtocol, protocol)
-        return transport, protocol
+
+        # Rewrite the host and port arguments for cross-origin redirects.
+        # This preserves connection overrides with the host and port
+        # arguments if the redirect points to the same host and port.
+        if not same_origin:
+            # Replace the host and port argument passed to the protocol factory.
+            factory = self._create_connection.args[0]
+            factory = functools.partial(
+                factory.func,
+                *factory.args,
+                **dict(factory.keywords, host=new_wsuri.host, port=new_wsuri.port),
+            )
+            # Replace the host and port argument passed to create_connection.
+            self._create_connection = functools.partial(
+                self._create_connection.func,
+                *(factory, new_wsuri.host, new_wsuri.port),
+                **self._create_connection.keywords,
+            )
+
+        # Set the new WebSocket URI. This suffices for same-origin redirects.
+        self._wsuri = new_wsuri
 
     # async with connect(...)
 
@@ -517,7 +539,10 @@ class Connect:
 
     async def __await_impl__(self) -> WebSocketClientProtocol:
         for redirects in range(self.MAX_REDIRECTS_ALLOWED):
-            transport, protocol = await self._creating_connection()
+            transport, protocol = await self._create_connection()
+            # https://github.com/python/typeshed/pull/2756
+            transport = cast(asyncio.Transport, transport)
+            protocol = cast(WebSocketClientProtocol, protocol)
 
             try:
                 try:
@@ -528,22 +553,17 @@ class Connect:
                         available_subprotocols=protocol.available_subprotocols,
                         extra_headers=protocol.extra_headers,
                     )
-                    break  # redirection chain ended
                 except Exception:
                     protocol.fail_connection()
                     await protocol.wait_closed()
                     raise
-            except RedirectHandshake as e:
-                wsuri = parse_uri(e.uri)
-                if self._wsuri.secure and not wsuri.secure:
-                    raise InvalidHandshake("redirect dropped TLS")
-                self._wsuri = wsuri
-                continue  # redirection chain continues
+                else:
+                    self.ws_client = protocol
+                    return protocol
+            except RedirectHandshake as exc:
+                self.handle_redirect(exc.uri)
         else:
-            raise InvalidHandshake("maximum redirects exceeded")
-
-        self.ws_client = protocol
-        return protocol
+            raise SecurityError("too many redirects")
 
     # yield from connect(...)
 
@@ -551,3 +571,20 @@ class Connect:
 
 
 connect = Connect
+
+
+def unix_connect(path: str, uri: str = "ws://localhost/", **kwargs: Any) -> Connect:
+    """
+    Similar to :func:`connect`, but for connecting to a Unix socket.
+
+    ``path`` is the path to the Unix socket. ``uri`` is the WebSocket URI.
+
+    This function calls the event loop's
+    :meth:`~asyncio.AbstractEventLoop.create_unix_connection` method.
+
+    It is only available on Unix.
+
+    It's mainly useful for debugging servers listening on Unix sockets.
+
+    """
+    return connect(uri=uri, path=path, **kwargs)
diff --git a/src/websockets/exceptions.py b/src/websockets/exceptions.py
index 36a8ed4..ce2c1e6 100644
--- a/src/websockets/exceptions.py
+++ b/src/websockets/exceptions.py
@@ -25,6 +25,7 @@ __all__ = [
     "NegotiationError",
     "PayloadTooBig",
     "RedirectHandshake",
+    "SecurityError",
     "WebSocketProtocolError",
 ]
 
@@ -52,6 +53,13 @@ class AbortHandshake(InvalidHandshake):
         super().__init__(message)
 
 
+class SecurityError(InvalidHandshake):
+    """
+    Exception raised when a HTTP request or response breaks security rules.
+
+    """
+
+
 class RedirectHandshake(InvalidHandshake):
     """
     Exception raised when a handshake gets redirected.
diff --git a/src/websockets/http.py b/src/websockets/http.py
index 6fbe5eb..04424c6 100644
--- a/src/websockets/http.py
+++ b/src/websockets/http.py
@@ -39,13 +39,6 @@ MAX_LINE = 4096
 USER_AGENT = f"Python/{sys.version[:3]} websockets/{websockets_version}"
 
 
-class SecurityError(ValueError):
-    """
-    HTTP request or response exceeds security limits.
-
-    """
-
-
 def d(value: bytes) -> str:
     """
     Decode a bytestring for interpolating into an error message.
@@ -211,6 +204,8 @@ async def read_headers(stream: asyncio.StreamReader) -> "Headers":
         headers[name] = value
 
     else:
+        from .exceptions import SecurityError  # avoid circular import
+
         raise SecurityError("too many HTTP headers")
 
     return headers
@@ -229,6 +224,8 @@ async def read_line(stream: asyncio.StreamReader) -> bytes:
     line = await stream.readline()
     # Security: this guarantees header values are small (hard-coded = 4 KiB)
     if len(line) > MAX_LINE:
+        from .exceptions import SecurityError  # avoid circular import
+
         raise SecurityError("line too long")
     # Not mandatory but safe - https://tools.ietf.org/html/rfc7230#section-3.5
     if not line.endswith(b"\r\n"):
diff --git a/src/websockets/server.py b/src/websockets/server.py
index 7c268c2..c02b67e 100644
--- a/src/websockets/server.py
+++ b/src/websockets/server.py
@@ -6,6 +6,7 @@ The :mod:`websockets.server` module defines a simple WebSocket server API.
 import asyncio
 import collections.abc
 import email.utils
+import functools
 import http
 import logging
 import socket
@@ -890,7 +891,8 @@ class Serve:
         elif compression is not None:
             raise ValueError(f"unsupported compression: {compression}")
 
-        factory = lambda: create_protocol(
+        factory = functools.partial(
+            create_protocol,
             ws_handler,
             ws_server,
             host=host,
@@ -917,17 +919,18 @@ class Serve:
             # serve(..., host, port) must specify host and port parameters.
             # host can be None to listen on all interfaces; port cannot be None.
             assert port is not None
-            # https://github.com/python/typeshed/pull/2763
-            creating_server = loop.create_server(  # type: ignore
-                factory, host, port, **kwargs
+            create_server = functools.partial(
+                loop.create_server, factory, host, port, **kwargs
             )
         else:
             # unix_serve(path) must not specify host and port parameters.
             assert host is None and port is None
-            creating_server = loop.create_unix_server(factory, path, **kwargs)
+            create_server = functools.partial(
+                loop.create_unix_server, factory, path, **kwargs
+            )
 
-        # This is a coroutine object.
-        self._creating_server = creating_server
+        # This is a coroutine function.
+        self._create_server = create_server
         self.ws_server = ws_server
 
     # async with serve(...)
@@ -951,7 +954,7 @@ class Serve:
         return self.__await_impl__().__await__()
 
     async def __await_impl__(self) -> WebSocketServer:
-        server = await self._creating_server
+        server = await self._create_server()
         self.ws_server.wrap(server)
         return self.ws_server
 


Test patch:
diff --git a/tests/test_client_server.py b/tests/test_client_server.py
index a880023..7281ec6 100644
--- a/tests/test_client_server.py
+++ b/tests/test_client_server.py
@@ -31,6 +31,7 @@ from websockets.handshake import build_response
 from websockets.http import USER_AGENT, Headers, read_response
 from websockets.protocol import State
 from websockets.server import *
+from websockets.uri import parse_uri
 
 from .test_protocol import MS
 from .utils import AsyncioTestCase
@@ -151,9 +152,6 @@ def get_server_uri(server, secure=False, resource_name="/", user_info=None):
         host = f"[{host}]"
     elif server_socket.family == socket.AF_INET:
         host, port = server_socket.getsockname()
-    elif server_socket.family == socket.AF_UNIX:
-        # The host and port are ignored when connecting to a Unix socket.
-        host, port = "localhost", 0
     else:  # pragma: no cover
         raise ValueError("expected an IPv6, IPv4, or Unix socket")
 
@@ -287,8 +285,11 @@ class ClientServerTestsMixin:
         # Disable pings by default in tests.
         kwargs.setdefault("ping_interval", None)
         secure = kwargs.get("ssl") is not None
-        server = self.redirecting_server if self.redirecting_server else self.server
-        server_uri = get_server_uri(server, secure, resource_name, user_info)
+        try:
+            server_uri = kwargs.pop("uri")
+        except KeyError:
+            server = self.redirecting_server if self.redirecting_server else self.server
+            server_uri = get_server_uri(server, secure, resource_name, user_info)
 
         with warnings.catch_warnings(record=True) as recorded_warnings:
             start_client = connect(server_uri, **kwargs)
@@ -440,6 +441,21 @@ class CommonClientServerTests:
                 reply = self.loop.run_until_complete(self.client.recv())
                 self.assertEqual(reply, "Hello!")
 
+    @with_server()
+    def test_explicit_host_port(self):
+        uri = get_server_uri(self.server, self.secure)
+        wsuri = parse_uri(uri)
+
+        # Change host and port to invalid values.
+        changed_uri = uri.replace(wsuri.host, "example.com").replace(
+            str(wsuri.port), str(65535 - wsuri.port)
+        )
+
+        with self.temp_client(uri=changed_uri, host=wsuri.host, port=wsuri.port):
+            self.loop.run_until_complete(self.client.send("Hello!"))
+            reply = self.loop.run_until_complete(self.client.recv())
+            self.assertEqual(reply, "Hello!")
+
     @with_server()
     def test_explicit_socket(self):
         class TrackedSocket(socket.socket):
@@ -489,18 +505,17 @@ class CommonClientServerTests:
             # Like self.start_server() but with unix_serve().
             unix_server = unix_serve(handler, path)
             self.server = self.loop.run_until_complete(unix_server)
-
-            client_socket = socket.socket(socket.AF_UNIX)
-            client_socket.connect(path)
-
             try:
-                with self.temp_client(sock=client_socket):
+                # Like self.start_client() but with unix_connect()
+                unix_client = unix_connect(path)
+                self.client = self.loop.run_until_complete(unix_client)
+                try:
                     self.loop.run_until_complete(self.client.send("Hello!"))
                     reply = self.loop.run_until_complete(self.client.recv())
                     self.assertEqual(reply, "Hello!")
-
+                finally:
+                    self.stop_client()
             finally:
-                client_socket.close()
                 self.stop_server()
 
     async def process_request_OK(path, request_headers):
diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
index fbc06e5..2cbd786 100644
--- a/tests/test_exceptions.py
+++ b/tests/test_exceptions.py
@@ -16,6 +16,11 @@ class ExceptionsTests(unittest.TestCase):
                 AbortHandshake(200, Headers(), b"OK\n"),
                 "HTTP 200, 0 headers, 3 bytes",
             ),
+            (
+                SecurityError("redirect from WSS to WS"),
+                "redirect from WSS to WS",
+
+            ),
             (
                 RedirectHandshake("wss://example.com"),
                 "redirect to wss://example.com",
diff --git a/tests/test_http.py b/tests/test_http.py
index cff97fc..41b522c 100644
--- a/tests/test_http.py
+++ b/tests/test_http.py
@@ -1,6 +1,7 @@
 import asyncio
 import unittest
 
+from websockets.exceptions import SecurityError
 from websockets.http import *
 from websockets.http import read_headers
 
@@ -120,13 +121,13 @@ class HTTPAsyncTests(AsyncioTestCase):
 
     async def test_headers_limit(self):
         self.stream.feed_data(b"foo: bar\r\n" * 257 + b"\r\n")
-        with self.assertRaises(ValueError):
+        with self.assertRaises(SecurityError):
             await read_headers(self.stream)
 
     async def test_line_limit(self):
         # Header line contains 5 + 4090 + 2 = 4097 bytes.
         self.stream.feed_data(b"foo: " + b"a" * 4090 + b"\r\n\r\n")
-        with self.assertRaises(ValueError):
+        with self.assertRaises(SecurityError):
             await read_headers(self.stream)
 
     async def test_line_ending(self):