From 10175f7a41ea1cf17aff7983d7eace3c2e4da5c0 Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Mon, 27 Jan 2025 22:07:21 +0100 Subject: [PATCH 1/2] Refactor SOCKS proxy implementation. --- src/websockets/asyncio/client.py | 109 +++++++++++++++----------- src/websockets/sync/client.py | 127 ++++++++++++++++++------------- 2 files changed, 141 insertions(+), 95 deletions(-) diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index f76095ead..7052ca85a 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -3,6 +3,7 @@ import asyncio import logging import os +import socket import traceback import urllib.parse from collections.abc import AsyncIterator, Generator, Sequence @@ -357,15 +358,12 @@ async def create_connection(self) -> ClientConnection: ws_uri = parse_uri(self.uri) 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) @@ -381,48 +379,14 @@ def factory() -> ClientConnection: if kwargs.pop("unix", False): _, connection = await loop.create_unix_connection(factory, **kwargs) + elif proxy is not None: + kwargs["sock"] = await connect_proxy( + parse_proxy(proxy), + ws_uri, + local_addr=kwargs.pop("local_addr", None), + ) + _, connection = await loop.create_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) @@ -624,3 +588,60 @@ def unix_connect( else: uri = "wss://localhost/" return connect(uri=uri, unix=True, path=path, **kwargs) + + +try: + from python_socks import ProxyType + from python_socks.async_.asyncio import Proxy as SocksProxy + + SOCKS_PROXY_TYPES = { + "socks5h": ProxyType.SOCKS5, + "socks5": ProxyType.SOCKS5, + "socks4a": ProxyType.SOCKS4, + "socks4": ProxyType.SOCKS4, + } + + SOCKS_PROXY_RDNS = { + "socks5h": True, + "socks5": False, + "socks4a": True, + "socks4": False, + } + + async def connect_socks_proxy( + proxy: Proxy, + ws_uri: WebSocketURI, + **kwargs: Any, + ) -> socket.socket: + """Connect via a SOCKS proxy and return the socket.""" + socks_proxy = SocksProxy( + SOCKS_PROXY_TYPES[proxy.scheme], + proxy.host, + proxy.port, + proxy.username, + proxy.password, + SOCKS_PROXY_RDNS[proxy.scheme], + ) + return await socks_proxy.connect(ws_uri.host, ws_uri.port, **kwargs) + +except ImportError: + + async def connect_socks_proxy( + proxy: Proxy, + ws_uri: WebSocketURI, + **kwargs: Any, + ) -> socket.socket: + raise ImportError("python-socks is required to use a SOCKS proxy") + + +async def connect_proxy( + proxy: Proxy, + ws_uri: WebSocketURI, + **kwargs: Any, +) -> socket.socket: + """Connect via a proxy and return the socket.""" + # parse_proxy() validates proxy.scheme. + if proxy.scheme[:5] == "socks": + return await connect_socks_proxy(proxy, ws_uri, **kwargs) + else: + raise AssertionError("unsupported proxy") diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 96f62edab..5fbec67ad 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -15,7 +15,7 @@ from ..http11 import USER_AGENT, Response from ..protocol import CONNECTING, Event from ..typing import LoggerLike, Origin, Subprotocol -from ..uri import Proxy, get_proxy, parse_proxy, parse_uri +from ..uri import Proxy, WebSocketURI, get_proxy, parse_proxy, parse_uri from .connection import Connection from .utils import Deadline @@ -258,15 +258,12 @@ def connect( elif compression is not None: raise ValueError(f"unsupported compression: {compression}") - proxy_uri: Proxy | 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(proxy) # Calculate timeouts on the TCP, TLS, and WebSocket handshakes. # The TCP and TLS timeouts must be set on the socket, then removed @@ -285,54 +282,21 @@ def connect( sock.settimeout(deadline.timeout()) assert path is not None # mypy cannot figure this out sock.connect(path) + elif proxy is not None: + sock = connect_proxy( + parse_proxy(proxy), + ws_uri, + deadline, + # websockets is consistent with the socket module while + # python_socks is consistent across implementations. + local_addr=kwargs.pop("source_address", None), + ) else: - 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 == "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, - ) - sock = socks_proxy.connect( - ws_uri.host, - ws_uri.port, - timeout=deadline.timeout(), - local_addr=kwargs.pop("local_addr", None), - ) - # Proxy types are enforced in parse_proxy(). - else: - raise AssertionError("unsupported proxy") - else: - kwargs.setdefault("timeout", deadline.timeout()) - sock = socket.create_connection( - (ws_uri.host, ws_uri.port), **kwargs - ) + kwargs.setdefault("timeout", deadline.timeout()) + sock = socket.create_connection( + (ws_uri.host, ws_uri.port), + **kwargs, + ) sock.settimeout(None) # Disable Nagle algorithm @@ -420,3 +384,64 @@ def unix_connect( else: uri = "wss://localhost/" return connect(uri=uri, unix=True, path=path, **kwargs) + + +try: + from python_socks import ProxyType + from python_socks.sync import Proxy as SocksProxy + + SOCKS_PROXY_TYPES = { + "socks5h": ProxyType.SOCKS5, + "socks5": ProxyType.SOCKS5, + "socks4a": ProxyType.SOCKS4, + "socks4": ProxyType.SOCKS4, + } + + SOCKS_PROXY_RDNS = { + "socks5h": True, + "socks5": False, + "socks4a": True, + "socks4": False, + } + + def connect_socks_proxy( + proxy: Proxy, + ws_uri: WebSocketURI, + deadline: Deadline, + **kwargs: Any, + ) -> socket.socket: + """Connect via a SOCKS proxy and return the socket.""" + socks_proxy = SocksProxy( + SOCKS_PROXY_TYPES[proxy.scheme], + proxy.host, + proxy.port, + proxy.username, + proxy.password, + SOCKS_PROXY_RDNS[proxy.scheme], + ) + kwargs.setdefault("timeout", deadline.timeout()) + return socks_proxy.connect(ws_uri.host, ws_uri.port, **kwargs) + +except ImportError: + + def connect_socks_proxy( + proxy: Proxy, + ws_uri: WebSocketURI, + deadline: Deadline, + **kwargs: Any, + ) -> socket.socket: + raise ImportError("python-socks is required to use a SOCKS proxy") + + +def connect_proxy( + proxy: Proxy, + ws_uri: WebSocketURI, + deadline: Deadline, + **kwargs: Any, +) -> socket.socket: + """Connect via a proxy and return the socket.""" + # parse_proxy() validates proxy.scheme. + if proxy.scheme[:5] == "socks": + return connect_socks_proxy(proxy, ws_uri, deadline, **kwargs) + else: + raise AssertionError("unsupported proxy") From 321be894176262a74a0b020aa1327f2aaca728ca Mon Sep 17 00:00:00 2001 From: Aymeric Augustin Date: Tue, 28 Jan 2025 22:42:11 +0100 Subject: [PATCH 2/2] Improve error handling for SOCKS proxy. --- docs/reference/exceptions.rst | 8 ++- src/websockets/__init__.py | 3 + src/websockets/asyncio/client.py | 17 +++++- src/websockets/exceptions.py | 41 ++++++++------ src/websockets/sync/client.py | 10 +++- tests/asyncio/test_client.py | 94 +++++++++++++++++++++++--------- tests/asyncio/test_server.py | 6 +- tests/sync/test_client.py | 94 +++++++++++++++++++++++--------- tests/sync/test_server.py | 4 +- tests/test_exceptions.py | 4 ++ 10 files changed, 203 insertions(+), 78 deletions(-) diff --git a/docs/reference/exceptions.rst b/docs/reference/exceptions.rst index e0c2efdd1..6c09a13fa 100644 --- a/docs/reference/exceptions.rst +++ b/docs/reference/exceptions.rst @@ -34,14 +34,16 @@ also reported by :func:`~websockets.asyncio.server.serve` in logs. .. autoexception:: SecurityError -.. autoexception:: InvalidMessage - -.. autoexception:: InvalidStatus +.. autoexception:: ProxyError .. autoexception:: InvalidProxyMessage .. autoexception:: InvalidProxyStatus +.. autoexception:: InvalidMessage + +.. autoexception:: InvalidStatus + .. autoexception:: InvalidHeader .. autoexception:: InvalidHeaderFormat diff --git a/src/websockets/__init__.py b/src/websockets/__init__.py index 8bf282a73..28a10910b 100644 --- a/src/websockets/__init__.py +++ b/src/websockets/__init__.py @@ -49,6 +49,7 @@ "NegotiationError", "PayloadTooBig", "ProtocolError", + "ProxyError", "SecurityError", "WebSocketException", # .frames @@ -112,6 +113,7 @@ NegotiationError, PayloadTooBig, ProtocolError, + ProxyError, SecurityError, WebSocketException, ) @@ -173,6 +175,7 @@ "NegotiationError": ".exceptions", "PayloadTooBig": ".exceptions", "ProtocolError": ".exceptions", + "ProxyError": ".exceptions", "SecurityError": ".exceptions", "WebSocketException": ".exceptions", # .frames diff --git a/src/websockets/asyncio/client.py b/src/websockets/asyncio/client.py index 7052ca85a..9582a4bb9 100644 --- a/src/websockets/asyncio/client.py +++ b/src/websockets/asyncio/client.py @@ -12,7 +12,7 @@ from ..client import ClientProtocol, backoff from ..datastructures import HeadersLike -from ..exceptions import InvalidMessage, InvalidStatus, SecurityError +from ..exceptions import InvalidMessage, InvalidStatus, ProxyError, SecurityError from ..extensions.base import ClientExtensionFactory from ..extensions.permessage_deflate import enable_client_permessage_deflate from ..headers import validate_subprotocols @@ -148,7 +148,9 @@ def process_exception(exc: Exception) -> Exception | None: That exception will be raised, breaking out of the retry loop. """ - if isinstance(exc, (OSError, asyncio.TimeoutError)): + # This catches python-socks' ProxyConnectionError and ProxyTimeoutError. + # Remove asyncio.TimeoutError when dropping Python < 3.11. + if isinstance(exc, (OSError, TimeoutError, asyncio.TimeoutError)): return None if isinstance(exc, InvalidMessage) and isinstance(exc.__cause__, EOFError): return None @@ -266,6 +268,7 @@ class connect: 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. @@ -622,7 +625,15 @@ async def connect_socks_proxy( proxy.password, SOCKS_PROXY_RDNS[proxy.scheme], ) - return await socks_proxy.connect(ws_uri.host, ws_uri.port, **kwargs) + # connect() is documented to raise OSError. + # socks_proxy.connect() doesn't raise TimeoutError; it gets canceled. + # Wrap other exceptions in ProxyError, a subclass of InvalidHandshake. + try: + return await socks_proxy.connect(ws_uri.host, ws_uri.port, **kwargs) + except OSError: + raise + except Exception as exc: + raise ProxyError("failed to connect to SOCKS proxy") from exc except ImportError: diff --git a/src/websockets/exceptions.py b/src/websockets/exceptions.py index e70aac92e..ab1a15ca8 100644 --- a/src/websockets/exceptions.py +++ b/src/websockets/exceptions.py @@ -9,11 +9,12 @@ * :exc:`InvalidProxy` * :exc:`InvalidHandshake` * :exc:`SecurityError` + * :exc:`ProxyError` + * :exc:`InvalidProxyMessage` + * :exc:`InvalidProxyStatus` * :exc:`InvalidMessage` * :exc:`InvalidStatus` * :exc:`InvalidStatusCode` (legacy) - * :exc:`InvalidProxyMessage` - * :exc:`InvalidProxyStatus` * :exc:`InvalidHeader` * :exc:`InvalidHeaderFormat` * :exc:`InvalidHeaderValue` @@ -48,10 +49,11 @@ "InvalidProxy", "InvalidHandshake", "SecurityError", - "InvalidMessage", - "InvalidStatus", + "ProxyError", "InvalidProxyMessage", "InvalidProxyStatus", + "InvalidMessage", + "InvalidStatus", "InvalidHeader", "InvalidHeaderFormat", "InvalidHeaderValue", @@ -206,16 +208,23 @@ class SecurityError(InvalidHandshake): """ -class InvalidMessage(InvalidHandshake): +class ProxyError(InvalidHandshake): """ - Raised when a handshake request or response is malformed. + Raised when failing to connect to a proxy. """ -class InvalidStatus(InvalidHandshake): +class InvalidProxyMessage(ProxyError): """ - Raised when a handshake response rejects the WebSocket upgrade. + Raised when an HTTP proxy response is malformed. + + """ + + +class InvalidProxyStatus(ProxyError): + """ + Raised when an HTTP proxy rejects the connection. """ @@ -223,21 +232,19 @@ def __init__(self, response: http11.Response) -> None: self.response = response def __str__(self) -> str: - return ( - f"server rejected WebSocket connection: HTTP {self.response.status_code:d}" - ) + return f"proxy rejected connection: HTTP {self.response.status_code:d}" -class InvalidProxyMessage(InvalidHandshake): +class InvalidMessage(InvalidHandshake): """ - Raised when a proxy response is malformed. + Raised when a handshake request or response is malformed. """ -class InvalidProxyStatus(InvalidHandshake): +class InvalidStatus(InvalidHandshake): """ - Raised when a proxy rejects the connection. + Raised when a handshake response rejects the WebSocket upgrade. """ @@ -245,7 +252,9 @@ 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}" + return ( + f"server rejected WebSocket connection: HTTP {self.response.status_code:d}" + ) class InvalidHeader(InvalidHandshake): diff --git a/src/websockets/sync/client.py b/src/websockets/sync/client.py index 5fbec67ad..e2a287648 100644 --- a/src/websockets/sync/client.py +++ b/src/websockets/sync/client.py @@ -9,6 +9,7 @@ from ..client import ClientProtocol from ..datastructures import HeadersLike +from ..exceptions import ProxyError from ..extensions.base import ClientExtensionFactory from ..extensions.permessage_deflate import enable_client_permessage_deflate from ..headers import validate_subprotocols @@ -420,7 +421,14 @@ def connect_socks_proxy( SOCKS_PROXY_RDNS[proxy.scheme], ) kwargs.setdefault("timeout", deadline.timeout()) - return socks_proxy.connect(ws_uri.host, ws_uri.port, **kwargs) + # connect() is documented to raise OSError and TimeoutError. + # Wrap other exceptions in ProxyError, a subclass of InvalidHandshake. + try: + return socks_proxy.connect(ws_uri.host, ws_uri.port, **kwargs) + except (OSError, TimeoutError, socket.timeout): + raise + except Exception as exc: + raise ProxyError("failed to connect to SOCKS proxy") from exc except ImportError: diff --git a/tests/asyncio/test_client.py b/tests/asyncio/test_client.py index cb2b8ede6..bdd519fb8 100644 --- a/tests/asyncio/test_client.py +++ b/tests/asyncio/test_client.py @@ -17,6 +17,7 @@ InvalidProxy, InvalidStatus, InvalidURI, + ProxyError, SecurityError, ) from websockets.extensions.permessage_deflate import PerMessageDeflate @@ -379,24 +380,16 @@ def remove_accept_header(self, request, response): async def test_timeout_during_handshake(self): """Client times out before receiving handshake response from server.""" - gate = asyncio.get_running_loop().create_future() - - async def stall_connection(self, request): - await gate - - # The connection will be open for the server but failed for the client. - # Use a connection handler that exits immediately to avoid an exception. - async with serve(*args, process_request=stall_connection) as server: - try: - with self.assertRaises(TimeoutError) as raised: - async with connect(get_uri(server) + "/no-op", open_timeout=2 * MS): - self.fail("did not raise") - self.assertEqual( - str(raised.exception), - "timed out during handshake", - ) - finally: - gate.set_result(None) + # Replace the WebSocket server with a TCP server that does't respond. + with socket.create_server(("localhost", 0)) as sock: + host, port = sock.getsockname() + with self.assertRaises(TimeoutError) as raised: + async with connect(f"ws://{host}:{port}", open_timeout=MS): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "timed out during handshake", + ) async def test_connection_closed_during_handshake(self): """Client reads EOF before receiving handshake response from server.""" @@ -570,11 +563,14 @@ class ProxyClientTests(unittest.IsolatedAsyncioTestCase): async def socks_proxy(self, auth=None): if auth: proxyauth = "hello:iloveyou" - proxy_uri = "http://hello:iloveyou@localhost:1080" + proxy_uri = "http://hello:iloveyou@localhost:51080" else: proxyauth = None - proxy_uri = "http://localhost:1080" - async with async_proxy(mode=["socks5"], proxyauth=proxyauth) as record_flows: + proxy_uri = "http://localhost:51080" + async with async_proxy( + mode=["socks5@51080"], + proxyauth=proxyauth, + ) as record_flows: with patch_environ({"socks_proxy": proxy_uri}): yield record_flows @@ -602,14 +598,62 @@ async def test_authenticated_socks_proxy(self): self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(len(proxy.get_flows()), 1) + async def test_socks_proxy_connection_error(self): + """Client receives an error when connecting to the SOCKS5 proxy.""" + from python_socks import ProxyError as SocksProxyError + + async with self.socks_proxy(auth=True) as proxy: + with self.assertRaises(ProxyError) as raised: + async with connect( + "ws://example.com/", + proxy="socks5h://localhost:51080", # remove credentials + ): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "failed to connect to SOCKS proxy", + ) + self.assertIsInstance(raised.exception.__cause__, SocksProxyError) + self.assertEqual(len(proxy.get_flows()), 0) + + async def test_socks_proxy_connection_fails(self): + """Client fails to connect to the SOCKS5 proxy.""" + from python_socks import ProxyConnectionError as SocksProxyConnectionError + + with self.assertRaises(OSError) as raised: + async with connect( + "ws://example.com/", + proxy="socks5h://localhost:51080", # nothing at this address + ): + self.fail("did not raise") + # Don't test str(raised.exception) because we don't control it. + self.assertIsInstance(raised.exception, SocksProxyConnectionError) + + async def test_socks_proxy_connection_timeout(self): + """Client times out while connecting to the SOCKS5 proxy.""" + # Replace the proxy with a TCP server that does't respond. + with socket.create_server(("localhost", 0)) as sock: + host, port = sock.getsockname() + with self.assertRaises(TimeoutError) as raised: + async with connect( + "ws://example.com/", + proxy=f"socks5h://{host}:{port}/", + open_timeout=MS, + ): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "timed out during handshake", + ) + async def test_explicit_proxy(self): """Client connects to server through a proxy set explicitly.""" - async with async_proxy(mode=["socks5"]) as proxy: + async with async_proxy(mode=["socks5@51080"]) as proxy: async with serve(*args) as server: async with connect( get_uri(server), # Take this opportunity to test socks5 instead of socks5h. - proxy="socks5://localhost:1080", + proxy="socks5://localhost:51080", ) as client: self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(len(proxy.get_flows()), 1) @@ -626,13 +670,13 @@ async def test_ignore_proxy_with_existing_socket(self): async def test_unsupported_proxy(self): """Client connects to server through an unsupported proxy.""" - with patch_environ({"ws_proxy": "other://localhost:1080"}): + with patch_environ({"ws_proxy": "other://localhost:51080"}): with self.assertRaises(InvalidProxy) as raised: async with connect("ws://example.com/"): self.fail("did not raise") self.assertEqual( str(raised.exception), - "other://localhost:1080 isn't a valid proxy: scheme other isn't supported", + "other://localhost:51080 isn't a valid proxy: scheme other isn't supported", ) diff --git a/tests/asyncio/test_server.py b/tests/asyncio/test_server.py index 38c0315a1..6adfff8e9 100644 --- a/tests/asyncio/test_server.py +++ b/tests/asyncio/test_server.py @@ -65,9 +65,9 @@ async def test_connection_handler_raises_exception(self): async def test_existing_socket(self): """Server receives connection using a pre-existing socket.""" with socket.create_server(("localhost", 0)) as sock: - async with serve(handler, sock=sock, host=None, port=None): - uri = "ws://{}:{}/".format(*sock.getsockname()) - async with connect(uri) as client: + host, port = sock.getsockname() + async with serve(handler, sock=sock): + async with connect(f"ws://{host}:{port}/") as client: await self.assertEval(client, "ws.protocol.state.name", "OPEN") async def test_select_subprotocol(self): diff --git a/tests/sync/test_client.py b/tests/sync/test_client.py index 2f62dd34d..dbecadcac 100644 --- a/tests/sync/test_client.py +++ b/tests/sync/test_client.py @@ -15,6 +15,7 @@ InvalidProxy, InvalidStatus, InvalidURI, + ProxyError, ) from websockets.extensions.permessage_deflate import PerMessageDeflate from websockets.sync.client import * @@ -148,24 +149,16 @@ def remove_accept_header(self, request, response): def test_timeout_during_handshake(self): """Client times out before receiving handshake response from server.""" - gate = threading.Event() - - def stall_connection(self, request): - gate.wait() - - # The connection will be open for the server but failed for the client. - # Use a connection handler that exits immediately to avoid an exception. - with run_server(process_request=stall_connection) as server: - try: - with self.assertRaises(TimeoutError) as raised: - with connect(get_uri(server) + "/no-op", open_timeout=2 * MS): - self.fail("did not raise") - self.assertEqual( - str(raised.exception), - "timed out during handshake", - ) - finally: - gate.set() + # Replace the WebSocket server with a TCP server that does't respond. + with socket.create_server(("localhost", 0)) as sock: + host, port = sock.getsockname() + with self.assertRaises(TimeoutError) as raised: + with connect(f"ws://{host}:{port}", open_timeout=MS): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "timed out during handshake", + ) def test_connection_closed_during_handshake(self): """Client reads EOF before receiving handshake response from server.""" @@ -311,12 +304,15 @@ class ProxyClientTests(unittest.TestCase): def socks_proxy(self, auth=None): if auth: proxyauth = "hello:iloveyou" - proxy_uri = "http://hello:iloveyou@localhost:1080" + proxy_uri = "http://hello:iloveyou@localhost:51080" else: proxyauth = None - proxy_uri = "http://localhost:1080" + proxy_uri = "http://localhost:51080" - with sync_proxy(mode=["socks5"], proxyauth=proxyauth) as record_flows: + with sync_proxy( + mode=["socks5@51080"], + proxyauth=proxyauth, + ) as record_flows: with patch_environ({"socks_proxy": proxy_uri}): yield record_flows @@ -344,14 +340,62 @@ def test_authenticated_socks_proxy(self): self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(len(proxy.get_flows()), 1) + def test_socks_proxy_connection_error(self): + """Client receives an error when connecting to the SOCKS5 proxy.""" + from python_socks import ProxyError as SocksProxyError + + with self.socks_proxy(auth=True) as proxy: + with self.assertRaises(ProxyError) as raised: + with connect( + "ws://example.com/", + proxy="socks5h://localhost:51080", # remove credentials + ): + self.fail("did not raise") + self.assertEqual( + str(raised.exception), + "failed to connect to SOCKS proxy", + ) + self.assertIsInstance(raised.exception.__cause__, SocksProxyError) + self.assertEqual(len(proxy.get_flows()), 0) + + def test_socks_proxy_connection_fails(self): + """Client fails to connect to the SOCKS5 proxy.""" + from python_socks import ProxyConnectionError as SocksProxyConnectionError + + with self.assertRaises(OSError) as raised: + with connect( + "ws://example.com/", + proxy="socks5h://localhost:51080", # nothing at this address + ): + self.fail("did not raise") + # Don't test str(raised.exception) because we don't control it. + self.assertIsInstance(raised.exception, SocksProxyConnectionError) + + def test_socks_proxy_timeout(self): + """Client times out before connecting to the SOCKS5 proxy.""" + from python_socks import ProxyTimeoutError as SocksProxyTimeoutError + + # Replace the proxy with a TCP server that does't respond. + with socket.create_server(("localhost", 0)) as sock: + host, port = sock.getsockname() + with self.assertRaises(TimeoutError) as raised: + with connect( + "ws://example.com/", + proxy=f"socks5h://{host}:{port}/", + open_timeout=MS, + ): + self.fail("did not raise") + # Don't test str(raised.exception) because we don't control it. + self.assertIsInstance(raised.exception, SocksProxyTimeoutError) + def test_explicit_proxy(self): """Client connects to server through a proxy set explicitly.""" - with sync_proxy(mode=["socks5"]) as proxy: + with sync_proxy(mode=["socks5@51080"]) as proxy: with run_server() as server: with connect( get_uri(server), # Take this opportunity to test socks5 instead of socks5h. - proxy="socks5://localhost:1080", + proxy="socks5://localhost:51080", ) as client: self.assertEqual(client.protocol.state.name, "OPEN") self.assertEqual(len(proxy.get_flows()), 1) @@ -368,13 +412,13 @@ def test_ignore_proxy_with_existing_socket(self): def test_unsupported_proxy(self): """Client connects to server through an unsupported proxy.""" - with patch_environ({"ws_proxy": "other://localhost:1080"}): + with patch_environ({"ws_proxy": "other://localhost:51080"}): with self.assertRaises(InvalidProxy) as raised: with connect("ws://example.com/"): self.fail("did not raise") self.assertEqual( str(raised.exception), - "other://localhost:1080 isn't a valid proxy: scheme other isn't supported", + "other://localhost:51080 isn't a valid proxy: scheme other isn't supported", ) diff --git a/tests/sync/test_server.py b/tests/sync/test_server.py index f59671efd..d04d1859a 100644 --- a/tests/sync/test_server.py +++ b/tests/sync/test_server.py @@ -64,9 +64,9 @@ def test_connection_handler_raises_exception(self): def test_existing_socket(self): """Server receives connection using a pre-existing socket.""" with socket.create_server(("localhost", 0)) as sock: + host, port = sock.getsockname() with run_server(sock=sock): - uri = "ws://{}:{}/".format(*sock.getsockname()) - with connect(uri) as client: + with connect(f"ws://{host}:{port}/") as client: self.assertEval(client, "ws.protocol.state.name", "OPEN") def test_select_subprotocol(self): diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py index 8b437ab5e..b4e7acee7 100644 --- a/tests/test_exceptions.py +++ b/tests/test_exceptions.py @@ -95,6 +95,10 @@ def test_str(self): SecurityError("redirect from WSS to WS"), "redirect from WSS to WS", ), + ( + ProxyError("failed to connect to SOCKS proxy"), + "failed to connect to SOCKS proxy", + ), ( InvalidMessage("malformed HTTP message"), "malformed HTTP message",