Spike SOCKS proxy implementation.

Author: aaugustin

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 ProxyURI, WebSocketURI, get_proxy, parse_proxy_uri, 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/proxy>` 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: ProxyURI | None = None
+        if kwargs.pop("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_uri(proxy)
 
         def factory() -> ClientConnection:
             return self.protocol_factory(ws_uri)
@@ -365,6 +382,38 @@ 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[:6] == "socks5":
+                        proxy_type = ProxyType.SOCKS5
+                    elif proxy_uri.scheme[:6] == "socks4":
+                        proxy_type = ProxyType.SOCKS4
+                    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.pop("rdns", None),
+                    )
+                    kwargs["sock"] = await socks_proxy.connect(
+                        ws_uri.host,
+                        ws_uri.port,
+                        local_addr=kwargs.pop("local_addr", None),
+                    )
+                else:
+                    raise NotImplementedError(
+                        f"proxy scheme not implemented yet: {proxy_uri.scheme}"
+                    )
             if kwargs.get("sock") is None:
                 kwargs.setdefault("host", ws_uri.host)
                 kwargs.setdefault("port", ws_uri.port)

src/websockets/sync/client.py

@@ -5,7 +5,7 @@
 import threading
 import warnings
 from collections.abc import Sequence
-from typing import Any
+from typing import Any, Literal
 
 from ..client import ClientProtocol
 from ..datastructures import HeadersLike
@@ -15,7 +15,7 @@
 from ..http11 import USER_AGENT, Response
 from ..protocol import CONNECTING, Event
 from ..typing import LoggerLike, Origin, Subprotocol
-from ..uri import parse_uri
+from ..uri import ProxyURI, get_proxy, parse_proxy_uri, parse_uri
 from .connection import Connection
 from .utils import Deadline
 
@@ -139,6 +139,7 @@ def connect(
     # HTTP
     additional_headers: HeadersLike | None = None,
     user_agent_header: str | None = USER_AGENT,
+    proxy: str | Literal[True] | None = True,
     # Timeouts
     open_timeout: float | None = 10,
     ping_interval: float | None = 20,
@@ -189,6 +190,10 @@ def 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/proxy>` for details.
         open_timeout: Timeout for opening the connection in seconds.
             :obj:`None` disables the timeout.
         ping_interval: Interval between keepalive pings in seconds.
@@ -253,6 +258,16 @@ def connect(
     elif compression is not None:
         raise ValueError(f"unsupported compression: {compression}")
 
+    proxy_uri: ProxyURI | None = None
+    if unix:
+        proxy = None
+    if sock is not None:
+        proxy = None
+    if proxy is True:
+        proxy = get_proxy(ws_uri)
+    if proxy is not None:
+        proxy_uri = parse_proxy_uri(proxy)
+
     # 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().
@@ -271,14 +286,49 @@ def connect(
                 assert path is not None  # mypy cannot figure this out
                 sock.connect(path)
             else:
-                kwargs.setdefault("timeout", deadline.timeout())
-                sock = socket.create_connection((ws_uri.host, ws_uri.port), **kwargs)
+                if proxy_uri is not None:
+                    if proxy_uri.scheme[:5] == "socks":
+                        try:
+                            from python_socks import ProxyType
+                            from python_socks.sync import Proxy
+                        except ImportError:
+                            raise ImportError(
+                                "python-socks is required to use a SOCKS proxy"
+                            )
+                        if proxy_uri.scheme[:6] == "socks5":
+                            proxy_type = ProxyType.SOCKS5
+                        elif proxy_uri.scheme[:6] == "socks4":
+                            proxy_type = ProxyType.SOCKS4
+                        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.pop("rdns", None),
+                        )
+                        sock = socks_proxy.connect(
+                            ws_uri.host,
+                            ws_uri.port,
+                            timeout=deadline.timeout(),
+                            local_addr=kwargs.pop("local_addr", None),
+                        )
+                    else:
+                        raise NotImplementedError(
+                            f"proxy scheme not implemented yet: {proxy_uri.scheme}"
+                        )
+                else:
+                    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)
+        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, True)
 
         # Initialize TLS wrapper and perform TLS handshake
 

src/websockets/uri.py

@@ -188,7 +188,7 @@ def parse_proxy_uri(uri: str) -> ProxyURI:
 
     """
     parsed = urllib.parse.urlparse(uri)
-    if parsed.scheme not in ["socks5", "socks4", "https", "http"]:
+    if parsed.scheme not in ["socks5h", "socks5", "socks4a", "socks4", "https", "http"]:
         raise InvalidURI(uri, f"proxy scheme isn't supported: {parsed.scheme}")
     if parsed.hostname is None:
         raise InvalidURI(uri, "proxy hostname isn't provided")