Add helpers to locate proxy for client connections.

Author: aaugustin

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

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/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.

src/websockets/uri.py

@@ -2,13 +2,18 @@
 
 import dataclasses
 import urllib.parse
+import urllib.request
 
-from .exceptions import InvalidURI
+from .exceptions import InvalidProxy, InvalidURI
 
 
 __all__ = ["parse_uri", "WebSocketURI"]
 
 
+# All characters from the gen-delims and sub-delims sets in RFC 3987.
+DELIMS = ":/?#[]@!$&'()*+,;="
+
+
 @dataclasses.dataclass
 class WebSocketURI:
     """
@@ -53,10 +58,6 @@ def user_info(self) -> tuple[str, str] | None:
         return (self.username, self.password)
 
 
-# All characters from the gen-delims and sub-delims sets in RFC 3987.
-DELIMS = ":/?#[]@!$&'()*+,;="
-
-
 def parse_uri(uri: str) -> WebSocketURI:
     """
     Parse and validate a WebSocket URI.
@@ -105,3 +106,120 @@ def parse_uri(uri: str) -> WebSocketURI:
             password = urllib.parse.quote(password, safe=DELIMS)
 
     return WebSocketURI(secure, host, port, path, query, username, password)
+
+
+@dataclasses.dataclass
+class Proxy:
+    """
+    Proxy.
+
+    Attributes:
+        scheme: ``"socks5h"``, ``"socks5"``, ``"socks4a"``, ``"socks4"``,
+            ``"https"``, or ``"http"``.
+        host: Normalized to lower case.
+        port: Always set even if it's the default.
+        username: Available when the proxy address contains `User Information`_.
+        password: Available when the proxy address contains `User Information`_.
+
+    .. _User Information: https://datatracker.ietf.org/doc/html/rfc3986#section-3.2.1
+
+    """
+
+    scheme: str
+    host: str
+    port: int
+    username: str | None = None
+    password: str | None = None
+
+    @property
+    def user_info(self) -> tuple[str, str] | None:
+        if self.username is None:
+            return None
+        assert self.password is not None
+        return (self.username, self.password)
+
+
+def parse_proxy(proxy: str) -> Proxy:
+    """
+    Parse and validate a proxy.
+
+    Args:
+        proxy: proxy.
+
+    Returns:
+        Parsed proxy.
+
+    Raises:
+        InvalidProxy: If ``proxy`` isn't a valid proxy.
+
+    """
+    parsed = urllib.parse.urlparse(proxy)
+    if parsed.scheme not in ["socks5h", "socks5", "socks4a", "socks4", "https", "http"]:
+        raise InvalidProxy(proxy, f"scheme {parsed.scheme} isn't supported")
+    if parsed.hostname is None:
+        raise InvalidProxy(proxy, "hostname isn't provided")
+    if parsed.path not in ["", "/"]:
+        raise InvalidProxy(proxy, "path is meaningless")
+    if parsed.query != "":
+        raise InvalidProxy(proxy, "query is meaningless")
+    if parsed.fragment != "":
+        raise InvalidProxy(proxy, "fragment is meaningless")
+
+    scheme = parsed.scheme
+    host = parsed.hostname
+    port = parsed.port or (443 if parsed.scheme == "https" else 80)
+    username = parsed.username
+    password = parsed.password
+    # urllib.parse.urlparse accepts URLs with a username but without a
+    # password. This doesn't make sense for HTTP Basic Auth credentials.
+    if username is not None and password is None:
+        raise InvalidProxy(proxy, "username provided without password")
+
+    try:
+        proxy.encode("ascii")
+    except UnicodeEncodeError:
+        # Input contains non-ASCII characters.
+        # It must be an IRI. Convert it to a URI.
+        host = host.encode("idna").decode()
+        if username is not None:
+            assert password is not None
+            username = urllib.parse.quote(username, safe=DELIMS)
+            password = urllib.parse.quote(password, safe=DELIMS)
+
+    return Proxy(scheme, host, port, username, password)
+
+
+def get_proxy(uri: WebSocketURI) -> str | None:
+    """
+    Return the proxy to use for connecting to the given WebSocket URI, if any.
+
+    """
+    if urllib.request.proxy_bypass(f"{uri.host}:{uri.port}"):
+        return None
+
+    # According to the _Proxy Usage_ section of RFC 6455, use a SOCKS5 proxy if
+    # available, else favor the proxy for HTTPS connections over the proxy for
+    # HTTP connections.
+
+    # The priority of a proxy for WebSocket connections is unspecified. We give
+    # it the highest priority. This makes it easy to configure a specific proxy
+    # for websockets.
+
+    # getproxies() may return SOCKS proxies as {"socks": "http://host:port"} or
+    # as {"https": "socks5h://host:port"} depending on whether they're declared
+    # in the operating system or in environment variables.
+
+    proxies = urllib.request.getproxies()
+    if uri.secure:
+        schemes = ["wss", "socks", "https"]
+    else:
+        schemes = ["ws", "socks", "https", "http"]
+
+    for scheme in schemes:
+        proxy = proxies.get(scheme)
+        if proxy is not None:
+            if scheme == "socks" and proxy.startswith("http://"):
+                proxy = "socks5h://" + proxy[7:]
+            return proxy
+    else:
+        return None

tests/test_exceptions.py

@@ -83,6 +83,10 @@ def test_str(self):
                 InvalidURI("|", "not at all!"),
                 "| isn't a valid URI: not at all!",
             ),
+            (
+                InvalidProxy("|", "not at all!"),
+                "| isn't a valid proxy: not at all!",
+            ),
             (
                 InvalidHandshake("invalid request"),
                 "invalid request",
@@ -99,6 +103,14 @@ def test_str(self):
                 InvalidStatus(Response(401, "Unauthorized", Headers())),
                 "server rejected WebSocket connection: HTTP 401",
             ),
+            (
+                InvalidProxyMessage("malformed HTTP message"),
+                "malformed HTTP message",
+            ),
+            (
+                InvalidProxyStatus(Response(401, "Unauthorized", Headers())),
+                "proxy rejected connection: HTTP 401",
+            ),
             (
                 InvalidHeader("Name"),
                 "missing Name header",