Add support for SOCKS proxies to clients.

docs/project/changelog.rst

@@ -25,13 +25,28 @@ fixing regressions shortly after a release.
 Only documented APIs are public. Undocumented, private APIs may change without
 notice.
 
-.. _14.3:
+.. _15.0:
 
-14.3
+15.0
 ----
 
 *In development*
 
+Backwards-incompatible changes
+..............................
+
+.. admonition:: Client connections use SOCKS proxies automatically.
+    :class: important
+
+    If a proxy is configured in the operating system or with an environment
+    variable, websockets uses it automatically when connecting to a server.
+    This feature requires installing the third-party library `python-socks`_.
+
+    If you want to disable the proxy, add ``proxy=None`` when calling
+    :func:`~asyncio.client.connect`. See :doc:`../topics/proxies` for details.
+
+    .. _python-socks: https://github.com/romis2012/python-socks
+
 New features
 ............
 

docs/reference/exceptions.rst

@@ -28,14 +28,20 @@ also reported by :func:`~websockets.asyncio.server.serve` in logs.
 
 .. autoexception:: InvalidURI
 
-.. autoexception:: InvalidHandshake
+.. autoexception:: InvalidProxy
 
-.. autoexception:: InvalidMessage
+.. autoexception:: InvalidHandshake
 
 .. autoexception:: SecurityError
 
+.. autoexception:: InvalidMessage
+
 .. autoexception:: InvalidStatus
 
+.. autoexception:: InvalidProxyMessage
+
+.. autoexception:: InvalidProxyStatus
+
 .. autoexception:: InvalidHeader
 
 .. autoexception:: InvalidHeaderFormat

docs/reference/features.rst

@@ -168,11 +168,10 @@ Client
     +------------------------------------+--------+--------+--------+--------+
     | Connect via HTTP proxy (`#364`_)   | ❌     | ❌     | —      | ❌     |
     +------------------------------------+--------+--------+--------+--------+
-    | Connect via SOCKS5 proxy (`#475`_) | ❌     | ❌     | —      | ❌     |
+    | Connect via SOCKS5 proxy           | ✅     | ✅     | —      | ❌     |
     +------------------------------------+--------+--------+--------+--------+
 
 .. _#364: https://github.com/python-websockets/websockets/issues/364
-.. _#475: https://github.com/python-websockets/websockets/issues/475
 .. _#784: https://github.com/python-websockets/websockets/issues/784
 
 Known limitations

docs/topics/index.rst

@@ -15,3 +15,4 @@ Get a deeper understanding of how websockets is built and why.
    memory
    security
    performance
+   proxies

docs/topics/proxies.rst

@@ -0,0 +1,66 @@
+Proxies
+=======
+
+.. currentmodule:: websockets
+
+If a proxy is configured in the operating system or with an environment
+variable, websockets uses it automatically when connecting to a server.
+
+Configuration
+-------------
+
+First, if the server is in the proxy bypass list of the operating system or in
+the ``no_proxy`` environment variable, websockets connects directly.
+
+Then, it looks for a proxy in the following locations:
+
+1. The ``wss_proxy`` or ``ws_proxy`` environment variables for ``wss://`` and
+   ``ws://`` connections respectively. They allow configuring a specific proxy
+   for WebSocket connections.
+2. A SOCKS proxy configured in the operating system.
+3. An HTTP proxy configured in the operating system or in the ``https_proxy``
+   environment variable, for both ``wss://`` and ``ws://`` connections.
+4. An HTTP proxy configured in the operating system or in the ``http_proxy``
+   environment variable, only for ``ws://`` connections.
+
+Finally, if no proxy is found, websockets connects directly.
+
+While environment variables are case-insensitive, the lower-case spelling is the
+most common, for `historical reasons`_, and recommended.
+
+.. _historical reasons: https://unix.stackexchange.com/questions/212894/
+
+.. admonition:: Any environment variable can configure a SOCKS proxy or an HTTP proxy.
+    :class: tip
+
+    For example, ``https_proxy=socks5h://proxy:1080/`` configures a SOCKS proxy
+    for all WebSocket connections. Likewise, ``wss_proxy=http://proxy:8080/``
+    configures an HTTP proxy only for ``wss://`` connections.
+
+.. admonition:: What if websockets doesn't select the right proxy?
+    :class: hint
+
+    websockets relies on :func:`~urllib.request.getproxies()` to read the proxy
+    configuration. Check that it returns what you expect. If it doesn't, review
+    your proxy configuration.
+
+You can override the default configuration and configure a proxy explicitly with
+the ``proxy`` argument of :func:`~asyncio.client.connect`. Set ``proxy=None`` to
+disable the proxy.
+
+SOCKS proxies
+-------------
+
+Connecting through a SOCKS proxy requires installing the third-party library
+`python-socks`_::
+
+    $ pip install python-socks\[asyncio\]
+
+.. _python-socks: https://github.com/romis2012/python-socks
+
+python-socks supports SOCKS4, SOCKS4a, SOCKS5, and SOCKS5h. The protocol version
+is configured in the address of the proxy e.g. ``socks5h://proxy:1080/``. When a
+SOCKS proxy is configured in the operating system, python-socks uses SOCKS5h.
+
+python-socks supports username/password authentication for SOCKS5 (:rfc:`1929`)
+but does not support other authentication methods such as GSSAPI (:rfc:`1961`).

src/websockets/__init__.py

@@ -39,6 +39,9 @@
     "InvalidOrigin",
     "InvalidParameterName",
     "InvalidParameterValue",
+    "InvalidProxy",
+    "InvalidProxyMessage",
+    "InvalidProxyStatus",
     "InvalidState",
     "InvalidStatus",
     "InvalidUpgrade",
@@ -99,6 +102,9 @@
         InvalidOrigin,
         InvalidParameterName,
         InvalidParameterValue,
+        InvalidProxy,
+        InvalidProxyMessage,
+        InvalidProxyStatus,
         InvalidState,
         InvalidStatus,
         InvalidUpgrade,
@@ -157,6 +163,9 @@
             "InvalidOrigin": ".exceptions",
             "InvalidParameterName": ".exceptions",
             "InvalidParameterValue": ".exceptions",
+            "InvalidProxy": ".exceptions",
+            "InvalidProxyMessage": ".exceptions",
+            "InvalidProxyStatus": ".exceptions",
             "InvalidState": ".exceptions",
             "InvalidStatus": ".exceptions",
             "InvalidUpgrade": ".exceptions",

src/websockets/asyncio/client.py

@@ -7,7 +7,7 @@
 import urllib.parse
 from collections.abc import AsyncIterator, Generator, Sequence
 from types import TracebackType
-from typing import Any, Callable
+from typing import Any, Callable, Literal
 
 from ..client import ClientProtocol, backoff
 from ..datastructures import HeadersLike
@@ -18,7 +18,7 @@
 from ..http11 import USER_AGENT, Response
 from ..protocol import CONNECTING, Event
 from ..typing import LoggerLike, Origin, Subprotocol
-from ..uri import WebSocketURI, parse_uri
+from ..uri import Proxy, WebSocketURI, get_proxy, parse_proxy, parse_uri
 from .compatibility import TimeoutError, asyncio_timeout
 from .connection import Connection
 
@@ -208,6 +208,10 @@ class connect:
         user_agent_header: Value of  the ``User-Agent`` request header.
             It defaults to ``"Python/x.y.z websockets/X.Y"``.
             Setting it to :obj:`None` removes the header.
+        proxy: If a proxy is configured, it is used by default. Set ``proxy``
+            to :obj:`None` to disable the proxy or to the address of a proxy
+            to override the system configuration. See the :doc:`proxy docs
+            <../../topics/proxies>` for details.
         process_exception: When reconnecting automatically, tell whether an
             error is transient or fatal. The default behavior is defined by
             :func:`process_exception`. Refer to its documentation for details.
@@ -279,6 +283,7 @@ def __init__(
         # HTTP
         additional_headers: HeadersLike | None = None,
         user_agent_header: str | None = USER_AGENT,
+        proxy: str | Literal[True] | None = True,
         process_exception: Callable[[Exception], Exception | None] = process_exception,
         # Timeouts
         open_timeout: float | None = 10,
@@ -333,6 +338,7 @@ def protocol_factory(uri: WebSocketURI) -> ClientConnection:
             )
             return connection
 
+        self.proxy = proxy
         self.protocol_factory = protocol_factory
         self.handshake_args = (
             additional_headers,
@@ -346,9 +352,20 @@ def protocol_factory(uri: WebSocketURI) -> ClientConnection:
     async def create_connection(self) -> ClientConnection:
         """Create TCP or Unix connection."""
         loop = asyncio.get_running_loop()
+        kwargs = self.connection_kwargs.copy()
 
         ws_uri = parse_uri(self.uri)
-        kwargs = self.connection_kwargs.copy()
+
+        proxy = self.proxy
+        proxy_uri: Proxy | None = None
+        if kwargs.get("unix", False):
+            proxy = None
+        if kwargs.get("sock") is not None:
+            proxy = None
+        if proxy is True:
+            proxy = get_proxy(ws_uri)
+        if proxy is not None:
+            proxy_uri = parse_proxy(proxy)
 
         def factory() -> ClientConnection:
             return self.protocol_factory(ws_uri)
@@ -365,6 +382,47 @@ def factory() -> ClientConnection:
         if kwargs.pop("unix", False):
             _, connection = await loop.create_unix_connection(factory, **kwargs)
         else:
+            if proxy_uri is not None:
+                if proxy_uri.scheme[:5] == "socks":
+                    try:
+                        from python_socks import ProxyType
+                        from python_socks.async_.asyncio import Proxy
+                    except ImportError:
+                        raise ImportError(
+                            "python-socks is required to use a SOCKS proxy"
+                        )
+                    if proxy_uri.scheme == "socks5h":
+                        proxy_type = ProxyType.SOCKS5
+                        rdns = True
+                    elif proxy_uri.scheme == "socks5":
+                        proxy_type = ProxyType.SOCKS5
+                        rdns = False
+                    # We use mitmproxy for testing and it doesn't support SOCKS4.
+                    elif proxy_uri.scheme == "socks4a":  # pragma: no cover
+                        proxy_type = ProxyType.SOCKS4
+                        rdns = True
+                    elif proxy_uri.scheme == "socks4":  # pragma: no cover
+                        proxy_type = ProxyType.SOCKS4
+                        rdns = False
+                    # Proxy types are enforced in parse_proxy().
+                    else:
+                        raise AssertionError("unsupported SOCKS proxy")
+                    socks_proxy = Proxy(
+                        proxy_type,
+                        proxy_uri.host,
+                        proxy_uri.port,
+                        proxy_uri.username,
+                        proxy_uri.password,
+                        rdns,
+                    )
+                    kwargs["sock"] = await socks_proxy.connect(
+                        ws_uri.host,
+                        ws_uri.port,
+                        local_addr=kwargs.pop("local_addr", None),
+                    )
+                # Proxy types are enforced in parse_proxy().
+                else:
+                    raise AssertionError("unsupported proxy")
             if kwargs.get("sock") is None:
                 kwargs.setdefault("host", ws_uri.host)
                 kwargs.setdefault("port", ws_uri.port)

src/websockets/exceptions.py

@@ -6,11 +6,14 @@
         * :exc:`ConnectionClosedOK`
         * :exc:`ConnectionClosedError`
     * :exc:`InvalidURI`
+    * :exc:`InvalidProxy`
     * :exc:`InvalidHandshake`
         * :exc:`SecurityError`
         * :exc:`InvalidMessage`
         * :exc:`InvalidStatus`
         * :exc:`InvalidStatusCode` (legacy)
+        * :exc:`InvalidProxyMessage`
+        * :exc:`InvalidProxyStatus`
         * :exc:`InvalidHeader`
             * :exc:`InvalidHeaderFormat`
             * :exc:`InvalidHeaderValue`
@@ -42,13 +45,16 @@
     "ConnectionClosedOK",
     "ConnectionClosedError",
     "InvalidURI",
+    "InvalidProxy",
     "InvalidHandshake",
     "SecurityError",
+    "InvalidMessage",
     "InvalidStatus",
+    "InvalidProxyMessage",
+    "InvalidProxyStatus",
     "InvalidHeader",
     "InvalidHeaderFormat",
     "InvalidHeaderValue",
-    "InvalidMessage",
     "InvalidOrigin",
     "InvalidUpgrade",
     "NegotiationError",
@@ -169,6 +175,20 @@ def __str__(self) -> str:
         return f"{self.uri} isn't a valid URI: {self.msg}"
 
 
+class InvalidProxy(WebSocketException):
+    """
+    Raised when connecting via a proxy that isn't valid.
+
+    """
+
+    def __init__(self, proxy: str, msg: str) -> None:
+        self.proxy = proxy
+        self.msg = msg
+
+    def __str__(self) -> str:
+        return f"{self.proxy} isn't a valid proxy: {self.msg}"
+
+
 class InvalidHandshake(WebSocketException):
     """
     Base class for exceptions raised when the opening handshake fails.
@@ -208,6 +228,26 @@ def __str__(self) -> str:
         )
 
 
+class InvalidProxyMessage(InvalidHandshake):
+    """
+    Raised when a proxy response is malformed.
+
+    """
+
+
+class InvalidProxyStatus(InvalidHandshake):
+    """
+    Raised when a proxy rejects the connection.
+
+    """
+
+    def __init__(self, response: http11.Response) -> None:
+        self.response = response
+
+    def __str__(self) -> str:
+        return f"proxy rejected connection: HTTP {self.response.status_code:d}"
+
+
 class InvalidHeader(InvalidHandshake):
     """
     Raised when an HTTP header doesn't have a valid format or value.