From e89abb6acb3745795ea969685df164fc48d5209c Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 12:58:55 +0100 Subject: [PATCH 01/11] Experimental support for https proxies from Codex --- docs/source/guides.rst | 9 +++-- picows/api.py | 59 +++++++++++++++++++++++++---- picows/picows.pyi | 1 + tests/test_redirects_and_proxies.py | 33 ++++++++++++++-- 4 files changed, 89 insertions(+), 13 deletions(-) diff --git a/docs/source/guides.rst b/docs/source/guides.rst index 5b2cdbb..2e8c6c0 100644 --- a/docs/source/guides.rst +++ b/docs/source/guides.rst @@ -236,17 +236,20 @@ Using proxies :any:`ws_connect` supports HTTP, SOCKS4 and SOCKS5 proxies via `python-socks `_. -Use the ``proxy`` argument with a proxy URL. HTTPS proxy URLs (``https://...``) -are not currently supported: +Use the ``proxy`` argument with a proxy URL: .. code-block:: python transport, listener = await ws_connect( ClientListener, "ws://127.0.0.1:9000/", - proxy="socks5://user:password@127.0.0.1:1080", + proxy_ssl_context=ssl.create_default_context(), + proxy="https://user:password@127.0.0.1:1080", ) +When ``https://`` proxy URL scheme is used, TLS is established with the proxy +first. ``proxy_ssl_context`` can be used to customize certificate verification. + When connecting to ``wss://`` URLs through a proxy, **picows** establishes a tunnel through the proxy and then performs the TLS handshake with the websocket server. diff --git a/picows/api.py b/picows/api.py index 3a2b9e3..dd8fd54 100644 --- a/picows/api.py +++ b/picows/api.py @@ -1,10 +1,12 @@ import asyncio +import ssl import urllib.parse from logging import getLogger from ssl import SSLContext from typing import Callable, Optional, Union from python_socks.async_.asyncio import Proxy +from python_socks.async_.asyncio.v2 import Proxy as ProxyV2 from .types import (WSHeadersLike, WSUpgradeRequest, WSUpgradeResponseWithListener, WSError) @@ -12,6 +14,8 @@ WSProtocol) from .url import parse_url, ParsedURL +_proxy_stream_lifecycle_guards: set[object] = set() + def _maybe_handle_redirect(exc: WSError, old_parsed_url: ParsedURL, max_redirects: int) -> ParsedURL: if max_redirects <= 0: @@ -38,10 +42,23 @@ def _maybe_handle_redirect(exc: WSError, old_parsed_url: ParsedURL, max_redirect return parsed_url +def _hold_proxy_stream_until_disconnect(proxy_stream: object, ws_transport: WSTransport) -> None: + _proxy_stream_lifecycle_guards.add(proxy_stream) + + async def _cleanup() -> None: + try: + await ws_transport.wait_disconnected() + finally: + _proxy_stream_lifecycle_guards.discard(proxy_stream) + + asyncio.get_running_loop().create_task(_cleanup()) + + async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: ignore [no-untyped-def] url: str, *, ssl_context: Optional[SSLContext] = None, + proxy_ssl_context: Optional[SSLContext] = None, disconnect_on_exception: bool = True, websocket_handshake_timeout: float = 5, logger_name: str = "client", @@ -68,6 +85,8 @@ async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: igno :param url: Destination URL :param ssl_context: optional SSLContext to override default one when the wss scheme is used + :param proxy_ssl_context: optional SSLContext to override default one when + https proxy scheme is used :param disconnect_on_exception: Indicates whether the client should initiate disconnect on any exception thrown from WSListener.on_ws_frame callbacks @@ -104,8 +123,7 @@ async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: igno * How many times we can follow HTTP redirects. Set to 0 in order to disable redirects. :param proxy: Optional proxy URL. Supported schemes are ``http://``, ``socks4://`` - and ``socks5://`` (including authenticated variants). - HTTPS proxy scheme (``https://``) is currently not supported. + ``https://`` and ``socks5://`` (including authenticated variants). :return: :any:`WSTransport` object and a user handler returned by `ws_listener_factory()` """ @@ -147,16 +165,43 @@ def ws_protocol_factory() -> WSProtocol: try: loop = asyncio.get_running_loop() conn_kwargs = dict(kwargs) - if proxy is not None and urllib.parse.urlsplit(proxy).scheme.lower() == "https": - raise ValueError("HTTPS proxy URL scheme is not supported, use http://, socks4:// or socks5://") proxy_socket = None host = None port = None if proxy is not None: - proxy_socket = await Proxy.from_url(proxy).connect( - dest_host=parsed_url.host, - dest_port=parsed_url.port) + proxy_url = urllib.parse.urlsplit(proxy) + proxy_scheme = proxy_url.scheme.lower() + if proxy_scheme == "https": + if proxy_ssl_context is None: + current_proxy_ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) + else: + current_proxy_ssl_context = proxy_ssl_context + + if ssl is None: + destination_ssl_context = None + elif isinstance(ssl, SSLContext): + destination_ssl_context = ssl + else: + destination_ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) + + http_proxy_url = urllib.parse.urlunsplit( + ("http", proxy_url.netloc, "", "", "") + ) + stream = await ProxyV2.from_url(http_proxy_url, proxy_ssl=current_proxy_ssl_context).connect( + dest_host=parsed_url.host, + dest_port=parsed_url.port, + dest_ssl=destination_ssl_context) + ws_protocol = ws_protocol_factory() + stream.writer.transport.set_protocol(ws_protocol) + ws_protocol.connection_made(stream.writer.transport) + await ws_protocol.wait_until_handshake_complete() + _hold_proxy_stream_until_disconnect(stream, ws_protocol.transport) + return ws_protocol.transport, ws_protocol.listener + else: + proxy_socket = await Proxy.from_url(proxy).connect( + dest_host=parsed_url.host, + dest_port=parsed_url.port) if ssl is not None and "server_hostname" not in conn_kwargs: conn_kwargs["server_hostname"] = parsed_url.host diff --git a/picows/picows.pyi b/picows/picows.pyi index 74c4711..5a3716f 100644 --- a/picows/picows.pyi +++ b/picows/picows.pyi @@ -177,6 +177,7 @@ async def ws_connect( url: str, *args: Any, ssl_context: Union[SSLContext, None] = None, + proxy_ssl_context: Union[SSLContext, None] = None, disconnect_on_exception: bool = True, websocket_handshake_timeout: float = 5, logger_name: str = "client", diff --git a/tests/test_redirects_and_proxies.py b/tests/test_redirects_and_proxies.py index acc9b03..5f8ac5a 100644 --- a/tests/test_redirects_and_proxies.py +++ b/tests/test_redirects_and_proxies.py @@ -5,12 +5,13 @@ import anyio import pytest +from anyio.streams.tls import TLSListener from tiny_proxy import HttpProxyHandler, Socks4ProxyHandler, Socks5ProxyHandler import picows from tests.utils import ClientAsyncContext, AsyncClient, \ create_client_ssl_context, echo_server, multiloop_event_loop_policy, \ - ServerAsyncContext + ServerAsyncContext, create_server_ssl_context event_loop_policy = multiloop_event_loop_policy() @@ -21,6 +22,10 @@ def _create_proxy_handler(proxy_type: str): return HttpProxyHandler() if proxy_type == "http_auth": return HttpProxyHandler(username="user", password="password") + if proxy_type == "https": + return HttpProxyHandler() + if proxy_type == "https_auth": + return HttpProxyHandler(username="user", password="password") if proxy_type == "socks4": return Socks4ProxyHandler() if proxy_type == "socks5": @@ -32,6 +37,8 @@ def _create_proxy_handler(proxy_type: str): _proxy_url_templates = { "http": "http://127.0.0.1:{port}", "http_auth": "http://user:password@127.0.0.1:{port}", + "https": "https://127.0.0.1:{port}", + "https_auth": "https://user:password@127.0.0.1:{port}", "socks4": "socks4://127.0.0.1:{port}", "socks5": "socks5://user:password@127.0.0.1:{port}" } @@ -45,10 +52,11 @@ async def ProxyServer(proxy_type: str): url_template = _proxy_url_templates[proxy_type] handler = _create_proxy_handler(proxy_type) listener = await anyio.create_tcp_listener(local_host="127.0.0.1") + proxy_listener = TLSListener(listener, create_server_ssl_context()) if proxy_type.startswith("https") else listener task_group = anyio.create_task_group() await task_group.__aenter__() - task_group.start_soon(listener.serve, handler.handle) + task_group.start_soon(proxy_listener.serve, handler.handle) try: proxy_port = listener.listeners[0].extra(anyio.abc.SocketAttribute.local_port) @@ -56,7 +64,7 @@ async def ProxyServer(proxy_type: str): finally: task_group.cancel_scope.cancel() await task_group.__aexit__(None, None, None) - await listener.aclose() + await proxy_listener.aclose() @pytest.fixture() @@ -125,3 +133,22 @@ async def test_proxy_dns_resolution(proxy_type): frame = await listener.get_message() assert frame.msg_type == picows.WSMsgType.BINARY assert frame.payload_as_bytes == b"hello over proxy" + + +@pytest.mark.parametrize("proxy_type", ["https", "https_auth"]) +async def test_https_proxy(echo_server, proxy_type): + client_ssl_ctx = create_client_ssl_context() + proxy_ssl_ctx = create_client_ssl_context() + + async with ProxyServer(proxy_type) as proxy_url: + async with ClientAsyncContext( + AsyncClient, + echo_server, + ssl_context=client_ssl_ctx, + proxy=proxy_url, + proxy_ssl_context=proxy_ssl_ctx, + ) as (transport, listener): + transport.send(picows.WSMsgType.BINARY, b"hello over https proxy") + frame = await listener.get_message() + assert frame.msg_type == picows.WSMsgType.BINARY + assert frame.payload_as_bytes == b"hello over https proxy" From 2f52dda4f8b647e0c7afc451dfab9d7cd7d3e2b2 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 13:05:43 +0100 Subject: [PATCH 02/11] Fix mypy errors --- AGENTS.md | 5 +++++ picows/api.py | 14 +++++++------- 2 files changed, 12 insertions(+), 7 deletions(-) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..b269e04 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,5 @@ +# Run lint after updating code, and fix all errors +flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics + +# Run mypy after updating code, and fix all errors. Disable errors that seems to be mypy quirks with #ignore comments +mypy picows \ No newline at end of file diff --git a/picows/api.py b/picows/api.py index dd8fd54..1710a58 100644 --- a/picows/api.py +++ b/picows/api.py @@ -141,9 +141,9 @@ async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: igno logger.warning("Basic authentication was requested in URL, but it is not currently supported, ignore username and password") if parsed_url.secure: - ssl = ssl_context if ssl_context is not None else True + ssl_arg = ssl_context if ssl_context is not None else True else: - ssl = None + ssl_arg = None def ws_protocol_factory() -> WSProtocol: return WSProtocol( @@ -178,10 +178,10 @@ def ws_protocol_factory() -> WSProtocol: else: current_proxy_ssl_context = proxy_ssl_context - if ssl is None: + if ssl_arg is None: destination_ssl_context = None - elif isinstance(ssl, SSLContext): - destination_ssl_context = ssl + elif isinstance(ssl_arg, SSLContext): + destination_ssl_context = ssl_arg else: destination_ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) @@ -203,14 +203,14 @@ def ws_protocol_factory() -> WSProtocol: dest_host=parsed_url.host, dest_port=parsed_url.port) - if ssl is not None and "server_hostname" not in conn_kwargs: + if ssl_arg is not None and "server_hostname" not in conn_kwargs: conn_kwargs["server_hostname"] = parsed_url.host else: host = parsed_url.host port = parsed_url.port (_, ws_protocol) = await loop.create_connection( - ws_protocol_factory, host, port, ssl=ssl, sock=proxy_socket, **conn_kwargs) # type: ignore[arg-type] + ws_protocol_factory, host, port, ssl=ssl_arg, sock=proxy_socket, **conn_kwargs) # type: ignore[arg-type] await ws_protocol.wait_until_handshake_complete() return ws_protocol.transport, ws_protocol.listener From 3b04791123ff9671821454766c9d9b000ed5b904 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 13:28:37 +0100 Subject: [PATCH 03/11] Update tests --- picows/api.py | 8 +++++++- tests/test_redirects_and_proxies.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/picows/api.py b/picows/api.py index 1710a58..fd57e6d 100644 --- a/picows/api.py +++ b/picows/api.py @@ -1,5 +1,6 @@ import asyncio import ssl +import sys import urllib.parse from logging import getLogger from ssl import SSLContext @@ -12,7 +13,7 @@ WSUpgradeResponseWithListener, WSError) from .picows import (WSListener, WSTransport, WSAutoPingStrategy, # type: ignore [attr-defined] WSProtocol) -from .url import parse_url, ParsedURL +from .url import parse_url, ParsedURL, WSInvalidURL _proxy_stream_lifecycle_guards: set[object] = set() @@ -173,6 +174,11 @@ def ws_protocol_factory() -> WSProtocol: proxy_url = urllib.parse.urlsplit(proxy) proxy_scheme = proxy_url.scheme.lower() if proxy_scheme == "https": + if sys.version_info < (3, 11): + raise WSInvalidURL( + proxy, + "https proxy requires Python 3.11+ (asyncio StreamWriter.start_tls support)" + ) if proxy_ssl_context is None: current_proxy_ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) else: diff --git a/tests/test_redirects_and_proxies.py b/tests/test_redirects_and_proxies.py index 5f8ac5a..6113ee8 100644 --- a/tests/test_redirects_and_proxies.py +++ b/tests/test_redirects_and_proxies.py @@ -1,4 +1,6 @@ +import asyncio import ssl +import sys from contextlib import asynccontextmanager from http import HTTPStatus from logging import getLogger @@ -9,6 +11,7 @@ from tiny_proxy import HttpProxyHandler, Socks4ProxyHandler, Socks5ProxyHandler import picows +import picows.api as picows_api from tests.utils import ClientAsyncContext, AsyncClient, \ create_client_ssl_context, echo_server, multiloop_event_loop_policy, \ ServerAsyncContext, create_server_ssl_context @@ -137,10 +140,25 @@ async def test_proxy_dns_resolution(proxy_type): @pytest.mark.parametrize("proxy_type", ["https", "https_auth"]) async def test_https_proxy(echo_server, proxy_type): + loop = asyncio.get_running_loop() client_ssl_ctx = create_client_ssl_context() proxy_ssl_ctx = create_client_ssl_context() async with ProxyServer(proxy_type) as proxy_url: + if sys.version_info < (3, 11) and isinstance(loop, asyncio.AbstractEventLoop): + pytest.skip("HTTPS proxy using asyncio requires Python 3.11+") + return + # with pytest.raises(picows.WSInvalidURL, + # match="https proxy requires Python 3.11\\+"): + # await picows.ws_connect( + # AsyncClient, + # echo_server, + # ssl_context=client_ssl_ctx, + # proxy=proxy_url, + # proxy_ssl_context=proxy_ssl_ctx, + # ) + # return + async with ClientAsyncContext( AsyncClient, echo_server, @@ -152,3 +170,4 @@ async def test_https_proxy(echo_server, proxy_type): frame = await listener.get_message() assert frame.msg_type == picows.WSMsgType.BINARY assert frame.payload_as_bytes == b"hello over https proxy" + From 2e9246aa3900d661d558fc4ac0f1a26df8dda5d1 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 13:37:48 +0100 Subject: [PATCH 04/11] Cleanup --- tests/test_redirects_and_proxies.py | 20 +++++--------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/tests/test_redirects_and_proxies.py b/tests/test_redirects_and_proxies.py index 6113ee8..fe4738c 100644 --- a/tests/test_redirects_and_proxies.py +++ b/tests/test_redirects_and_proxies.py @@ -140,25 +140,15 @@ async def test_proxy_dns_resolution(proxy_type): @pytest.mark.parametrize("proxy_type", ["https", "https_auth"]) async def test_https_proxy(echo_server, proxy_type): - loop = asyncio.get_running_loop() + is_asyncio_loop = isinstance(asyncio.get_event_loop_policy(), asyncio.DefaultEventLoopPolicy) + if sys.version_info < (3, 11) and is_asyncio_loop: + pytest.skip("HTTPS proxy using asyncio requires Python 3.11+") + return + client_ssl_ctx = create_client_ssl_context() proxy_ssl_ctx = create_client_ssl_context() async with ProxyServer(proxy_type) as proxy_url: - if sys.version_info < (3, 11) and isinstance(loop, asyncio.AbstractEventLoop): - pytest.skip("HTTPS proxy using asyncio requires Python 3.11+") - return - # with pytest.raises(picows.WSInvalidURL, - # match="https proxy requires Python 3.11\\+"): - # await picows.ws_connect( - # AsyncClient, - # echo_server, - # ssl_context=client_ssl_ctx, - # proxy=proxy_url, - # proxy_ssl_context=proxy_ssl_ctx, - # ) - # return - async with ClientAsyncContext( AsyncClient, echo_server, From 5c1776a56380c417480e5adac58a25454239fc1b Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 13:42:32 +0100 Subject: [PATCH 05/11] Fix tests --- picows/api.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/picows/api.py b/picows/api.py index fd57e6d..57dfad7 100644 --- a/picows/api.py +++ b/picows/api.py @@ -174,10 +174,13 @@ def ws_protocol_factory() -> WSProtocol: proxy_url = urllib.parse.urlsplit(proxy) proxy_scheme = proxy_url.scheme.lower() if proxy_scheme == "https": - if sys.version_info < (3, 11): + is_asyncio_loop = isinstance( + asyncio.get_event_loop_policy(), + asyncio.DefaultEventLoopPolicy) + if sys.version_info < (3, 11) and is_asyncio_loop: raise WSInvalidURL( proxy, - "https proxy requires Python 3.11+ (asyncio StreamWriter.start_tls support)" + "HTTPS proxy with asyncio requires Python 3.11+ (asyncio StreamWriter.start_tls support)" ) if proxy_ssl_context is None: current_proxy_ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) From 42d822b8227c1f1eaa83c9ac5bd63a535b332045 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 13:44:28 +0100 Subject: [PATCH 06/11] Fix mypy --- picows/api.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/picows/api.py b/picows/api.py index 57dfad7..b663bbc 100644 --- a/picows/api.py +++ b/picows/api.py @@ -174,9 +174,7 @@ def ws_protocol_factory() -> WSProtocol: proxy_url = urllib.parse.urlsplit(proxy) proxy_scheme = proxy_url.scheme.lower() if proxy_scheme == "https": - is_asyncio_loop = isinstance( - asyncio.get_event_loop_policy(), - asyncio.DefaultEventLoopPolicy) + is_asyncio_loop = isinstance(asyncio.get_event_loop_policy(), asyncio.DefaultEventLoopPolicy) # type: ignore [attr-defined] if sys.version_info < (3, 11) and is_asyncio_loop: raise WSInvalidURL( proxy, From a5f5a15089db4f158bbafa61e40251f3bcfeb701 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 14:21:09 +0100 Subject: [PATCH 07/11] Fix tests --- picows/api.py | 91 ++++++++++++++--------------- tests/test_redirects_and_proxies.py | 41 +++++-------- 2 files changed, 57 insertions(+), 75 deletions(-) diff --git a/picows/api.py b/picows/api.py index b663bbc..63e8817 100644 --- a/picows/api.py +++ b/picows/api.py @@ -55,11 +55,37 @@ async def _cleanup() -> None: asyncio.get_running_loop().create_task(_cleanup()) +async def _connect_through_https_proxy(ws_protocol_factory, ssl_context, proxy: str, proxy_ssl_context, proxy_parsed_url, parsed_url): + is_asyncio_loop = isinstance(asyncio.get_event_loop_policy(), + asyncio.DefaultEventLoopPolicy) # type: ignore [attr-defined] + if sys.version_info < (3, 11) and is_asyncio_loop: + raise WSInvalidURL( + proxy, + "HTTPS proxy with asyncio requires Python 3.11+ (asyncio StreamWriter.start_tls support)" + ) + proxy_ssl_context = proxy_ssl_context or ssl.create_default_context( + ssl.Purpose.SERVER_AUTH) + + http_proxy_url = urllib.parse.urlunsplit( + ("http", proxy_parsed_url.netloc, "", "", "") + ) + stream = await ProxyV2.from_url(http_proxy_url, + proxy_ssl=proxy_ssl_context).connect( + dest_host=parsed_url.host, + dest_port=parsed_url.port, + dest_ssl=ssl_context) + ws_protocol = ws_protocol_factory() + stream.writer.transport.set_protocol(ws_protocol) + ws_protocol.connection_made(stream.writer.transport) + await ws_protocol.wait_until_handshake_complete() + _hold_proxy_stream_until_disconnect(stream, ws_protocol.transport) + return ws_protocol.transport, ws_protocol.listener + + async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: ignore [no-untyped-def] url: str, *, ssl_context: Optional[SSLContext] = None, - proxy_ssl_context: Optional[SSLContext] = None, disconnect_on_exception: bool = True, websocket_handshake_timeout: float = 5, logger_name: str = "client", @@ -72,6 +98,7 @@ async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: igno extra_headers: Optional[WSHeadersLike] = None, max_redirects: int = 5, proxy: Optional[str] = None, + proxy_ssl_context: Optional[SSLContext] = None, **kwargs ) -> tuple[WSTransport, WSListener]: """ @@ -86,8 +113,6 @@ async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: igno :param url: Destination URL :param ssl_context: optional SSLContext to override default one when the wss scheme is used - :param proxy_ssl_context: optional SSLContext to override default one when - https proxy scheme is used :param disconnect_on_exception: Indicates whether the client should initiate disconnect on any exception thrown from WSListener.on_ws_frame callbacks @@ -125,6 +150,8 @@ async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: igno :param proxy: Optional proxy URL. Supported schemes are ``http://``, ``socks4://`` ``https://`` and ``socks5://`` (including authenticated variants). + :param proxy_ssl_context: optional SSLContext to override default one when + https proxy scheme is used :return: :any:`WSTransport` object and a user handler returned by `ws_listener_factory()` """ @@ -141,11 +168,6 @@ async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: igno if parsed_url.username is not None or parsed_url.password is not None: logger.warning("Basic authentication was requested in URL, but it is not currently supported, ignore username and password") - if parsed_url.secure: - ssl_arg = ssl_context if ssl_context is not None else True - else: - ssl_arg = None - def ws_protocol_factory() -> WSProtocol: return WSProtocol( parsed_url.netloc, @@ -163,61 +185,36 @@ def ws_protocol_factory() -> WSProtocol: max_frame_size, extra_headers) - try: - loop = asyncio.get_running_loop() - conn_kwargs = dict(kwargs) + current_ssl_context = ssl_context if parsed_url.secure else None + + loop = asyncio.get_running_loop() + conn_kwargs = dict(kwargs) - proxy_socket = None - host = None - port = None + proxy_socket = None + host = None + port = None + + try: if proxy is not None: proxy_url = urllib.parse.urlsplit(proxy) proxy_scheme = proxy_url.scheme.lower() if proxy_scheme == "https": - is_asyncio_loop = isinstance(asyncio.get_event_loop_policy(), asyncio.DefaultEventLoopPolicy) # type: ignore [attr-defined] - if sys.version_info < (3, 11) and is_asyncio_loop: - raise WSInvalidURL( - proxy, - "HTTPS proxy with asyncio requires Python 3.11+ (asyncio StreamWriter.start_tls support)" - ) - if proxy_ssl_context is None: - current_proxy_ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) - else: - current_proxy_ssl_context = proxy_ssl_context - - if ssl_arg is None: - destination_ssl_context = None - elif isinstance(ssl_arg, SSLContext): - destination_ssl_context = ssl_arg - else: - destination_ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) - - http_proxy_url = urllib.parse.urlunsplit( - ("http", proxy_url.netloc, "", "", "") - ) - stream = await ProxyV2.from_url(http_proxy_url, proxy_ssl=current_proxy_ssl_context).connect( - dest_host=parsed_url.host, - dest_port=parsed_url.port, - dest_ssl=destination_ssl_context) - ws_protocol = ws_protocol_factory() - stream.writer.transport.set_protocol(ws_protocol) - ws_protocol.connection_made(stream.writer.transport) - await ws_protocol.wait_until_handshake_complete() - _hold_proxy_stream_until_disconnect(stream, ws_protocol.transport) - return ws_protocol.transport, ws_protocol.listener + return await _connect_through_https_proxy( + ws_protocol_factory, current_ssl_context, proxy, + proxy_ssl_context, proxy_url, parsed_url) else: proxy_socket = await Proxy.from_url(proxy).connect( dest_host=parsed_url.host, dest_port=parsed_url.port) - if ssl_arg is not None and "server_hostname" not in conn_kwargs: + if parsed_url.secure and "server_hostname" not in conn_kwargs: conn_kwargs["server_hostname"] = parsed_url.host else: host = parsed_url.host port = parsed_url.port (_, ws_protocol) = await loop.create_connection( - ws_protocol_factory, host, port, ssl=ssl_arg, sock=proxy_socket, **conn_kwargs) # type: ignore[arg-type] + ws_protocol_factory, host, port, ssl=current_ssl_context, sock=proxy_socket, **conn_kwargs) # type: ignore[arg-type] await ws_protocol.wait_until_handshake_complete() return ws_protocol.transport, ws_protocol.listener diff --git a/tests/test_redirects_and_proxies.py b/tests/test_redirects_and_proxies.py index fe4738c..258a8ca 100644 --- a/tests/test_redirects_and_proxies.py +++ b/tests/test_redirects_and_proxies.py @@ -98,7 +98,7 @@ def listener_factory(r): yield server_urls.tcp_url -@pytest.mark.parametrize("proxy_type", ["direct", "http", "http_auth", "socks4", "socks5"]) +@pytest.mark.parametrize("proxy_type", ["direct", "socks4", "socks5", "http", "http_auth", "https", "https_auth"]) async def test_redirect_through_proxy(redirect_server_2, proxy_type: str): # This is an absolute masterpiece! Best test I wrote ever! # @@ -108,20 +108,29 @@ async def test_redirect_through_proxy(redirect_server_2, proxy_type: str): # echo server, send request and validate response. # # God bless pytest! + + is_https = proxy_type in ("https", "https_auth") + is_asyncio_loop = isinstance(asyncio.get_event_loop_policy(), asyncio.DefaultEventLoopPolicy) + + if sys.version_info < (3, 11) and is_asyncio_loop and is_https: + pytest.skip("HTTPS proxy using asyncio requires Python 3.11+") + return + client_ssl_ctx = create_client_ssl_context() + proxy_ssl_ctx = create_client_ssl_context() if is_https else None async with ProxyServer(proxy_type) as proxy_url: - async with ClientAsyncContext(AsyncClient, redirect_server_2, ssl_context=client_ssl_ctx, proxy=proxy_url) as (transport, listener): + async with ClientAsyncContext(AsyncClient, redirect_server_2, ssl_context=client_ssl_ctx, proxy=proxy_url, proxy_ssl_context=proxy_ssl_ctx) as (transport, listener): transport.send(picows.WSMsgType.BINARY, b"hello over proxy") frame = await listener.get_message() assert frame.msg_type == picows.WSMsgType.BINARY assert frame.payload_as_bytes == b"hello over proxy" with pytest.raises(picows.WSError, match="status 101"): - await picows.ws_connect(AsyncClient, redirect_server_2, max_redirects=0, proxy=proxy_url) + await picows.ws_connect(AsyncClient, redirect_server_2, max_redirects=0, proxy=proxy_url, proxy_ssl_context=proxy_ssl_ctx) with pytest.raises(picows.WSError, match="status 101"): - await picows.ws_connect(AsyncClient, redirect_server_2, max_redirects=1, proxy=proxy_url) + await picows.ws_connect(AsyncClient, redirect_server_2, max_redirects=1, proxy=proxy_url, proxy_ssl_context=proxy_ssl_ctx) @pytest.mark.parametrize("proxy_type", ["direct", "http", "http_auth", "socks4", "socks5"]) @@ -137,27 +146,3 @@ async def test_proxy_dns_resolution(proxy_type): assert frame.msg_type == picows.WSMsgType.BINARY assert frame.payload_as_bytes == b"hello over proxy" - -@pytest.mark.parametrize("proxy_type", ["https", "https_auth"]) -async def test_https_proxy(echo_server, proxy_type): - is_asyncio_loop = isinstance(asyncio.get_event_loop_policy(), asyncio.DefaultEventLoopPolicy) - if sys.version_info < (3, 11) and is_asyncio_loop: - pytest.skip("HTTPS proxy using asyncio requires Python 3.11+") - return - - client_ssl_ctx = create_client_ssl_context() - proxy_ssl_ctx = create_client_ssl_context() - - async with ProxyServer(proxy_type) as proxy_url: - async with ClientAsyncContext( - AsyncClient, - echo_server, - ssl_context=client_ssl_ctx, - proxy=proxy_url, - proxy_ssl_context=proxy_ssl_ctx, - ) as (transport, listener): - transport.send(picows.WSMsgType.BINARY, b"hello over https proxy") - frame = await listener.get_message() - assert frame.msg_type == picows.WSMsgType.BINARY - assert frame.payload_as_bytes == b"hello over https proxy" - From ac12096a1f11235ebe8aea44d75351cb5de242f0 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 14:38:19 +0100 Subject: [PATCH 08/11] Use a little cleaner approach --- picows/api.py | 48 ++++++++++++++++++++++++++++-------------------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/picows/api.py b/picows/api.py index 63e8817..1eecf77 100644 --- a/picows/api.py +++ b/picows/api.py @@ -4,7 +4,7 @@ import urllib.parse from logging import getLogger from ssl import SSLContext -from typing import Callable, Optional, Union +from typing import Any, Callable, Optional, Union, cast from python_socks.async_.asyncio import Proxy from python_socks.async_.asyncio.v2 import Proxy as ProxyV2 @@ -15,9 +15,6 @@ WSProtocol) from .url import parse_url, ParsedURL, WSInvalidURL -_proxy_stream_lifecycle_guards: set[object] = set() - - def _maybe_handle_redirect(exc: WSError, old_parsed_url: ParsedURL, max_redirects: int) -> ParsedURL: if max_redirects <= 0: raise exc @@ -43,21 +40,31 @@ def _maybe_handle_redirect(exc: WSError, old_parsed_url: ParsedURL, max_redirect return parsed_url -def _hold_proxy_stream_until_disconnect(proxy_stream: object, ws_transport: WSTransport) -> None: - _proxy_stream_lifecycle_guards.add(proxy_stream) +class _DetachedWriterTransport: + def is_closing(self) -> bool: + return True + + def close(self) -> None: + return - async def _cleanup() -> None: - try: - await ws_transport.wait_disconnected() - finally: - _proxy_stream_lifecycle_guards.discard(proxy_stream) - asyncio.get_running_loop().create_task(_cleanup()) +def _detach_stream_writer_transport(stream: Any) -> asyncio.Transport: + transport = cast(asyncio.Transport, stream.writer.transport) + # Prevent StreamWriter.__del__ from closing a transport we hand over to WSProtocol. + stream.writer._transport = _DetachedWriterTransport() + return transport -async def _connect_through_https_proxy(ws_protocol_factory, ssl_context, proxy: str, proxy_ssl_context, proxy_parsed_url, parsed_url): - is_asyncio_loop = isinstance(asyncio.get_event_loop_policy(), - asyncio.DefaultEventLoopPolicy) # type: ignore [attr-defined] +async def _connect_through_https_proxy( + ws_protocol_factory: Callable[[], WSProtocol], + ssl_context: Optional[SSLContext], + proxy: str, + proxy_ssl_context: Optional[SSLContext], + proxy_parsed_url: urllib.parse.SplitResult, + parsed_url: ParsedURL +) -> tuple[WSTransport, WSListener]: + loop = asyncio.get_running_loop() + is_asyncio_loop = loop.__class__.__module__.startswith("asyncio") if sys.version_info < (3, 11) and is_asyncio_loop: raise WSInvalidURL( proxy, @@ -69,16 +76,15 @@ async def _connect_through_https_proxy(ws_protocol_factory, ssl_context, proxy: http_proxy_url = urllib.parse.urlunsplit( ("http", proxy_parsed_url.netloc, "", "", "") ) - stream = await ProxyV2.from_url(http_proxy_url, - proxy_ssl=proxy_ssl_context).connect( + stream = await ProxyV2.from_url(http_proxy_url, proxy_ssl=proxy_ssl_context).connect( dest_host=parsed_url.host, dest_port=parsed_url.port, dest_ssl=ssl_context) ws_protocol = ws_protocol_factory() - stream.writer.transport.set_protocol(ws_protocol) - ws_protocol.connection_made(stream.writer.transport) + transport = _detach_stream_writer_transport(stream) + transport.set_protocol(ws_protocol) + ws_protocol.connection_made(transport) await ws_protocol.wait_until_handshake_complete() - _hold_proxy_stream_until_disconnect(stream, ws_protocol.transport) return ws_protocol.transport, ws_protocol.listener @@ -164,6 +170,8 @@ async def ws_connect(ws_listener_factory: Callable[[], WSListener], # type: igno logger = getLogger(f"picows.{logger_name}") parsed_url = parse_url(url) + # Loop in order to follow redirects + # Break loop if we were able to upgrade while True: if parsed_url.username is not None or parsed_url.password is not None: logger.warning("Basic authentication was requested in URL, but it is not currently supported, ignore username and password") From 5641435270154cbcfc1a9aa231fdce7a084e76c8 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 14:45:12 +0100 Subject: [PATCH 09/11] Update tests --- tests/test_redirects_and_proxies.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/test_redirects_and_proxies.py b/tests/test_redirects_and_proxies.py index 258a8ca..26e9152 100644 --- a/tests/test_redirects_and_proxies.py +++ b/tests/test_redirects_and_proxies.py @@ -133,12 +133,21 @@ async def test_redirect_through_proxy(redirect_server_2, proxy_type: str): await picows.ws_connect(AsyncClient, redirect_server_2, max_redirects=1, proxy=proxy_url, proxy_ssl_context=proxy_ssl_ctx) -@pytest.mark.parametrize("proxy_type", ["direct", "http", "http_auth", "socks4", "socks5"]) +@pytest.mark.parametrize("proxy_type", ["direct", "socks4", "socks5", "http", "http_auth", "https", "https_auth"]) +@pytest.mark.skip(reason="echo server may respond with 429 (too many requests if we spam it a lot)") async def test_proxy_dns_resolution(proxy_type): + is_https = proxy_type in ("https", "https_auth") + is_asyncio_loop = isinstance(asyncio.get_event_loop_policy(), asyncio.DefaultEventLoopPolicy) + + if sys.version_info < (3, 11) and is_asyncio_loop and is_https: + pytest.skip("HTTPS proxy using asyncio requires Python 3.11+") + return + client_ssl_ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) + proxy_ssl_ctx = create_client_ssl_context() if is_https else None async with ProxyServer(proxy_type) as proxy_url: - async with ClientAsyncContext(AsyncClient, "wss://echo.websocket.org", ssl_context=client_ssl_ctx, proxy=proxy_url) as (transport, listener): + async with ClientAsyncContext(AsyncClient, "wss://echo.websocket.org", ssl_context=client_ssl_ctx, proxy=proxy_url, proxy_ssl_context=proxy_ssl_ctx) as (transport, listener): frame = await listener.get_message() _logger.debug("Welcome frame from echo.websocket.org: %s", frame.payload_as_ascii_text) transport.send(picows.WSMsgType.BINARY, b"hello over proxy") From c785160da1099bfbd8586ed2e30bafdcc75d6e66 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 14:46:15 +0100 Subject: [PATCH 10/11] Increate timeouts --- tests/test_redirects_and_proxies.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_redirects_and_proxies.py b/tests/test_redirects_and_proxies.py index 26e9152..604d5e9 100644 --- a/tests/test_redirects_and_proxies.py +++ b/tests/test_redirects_and_proxies.py @@ -122,7 +122,7 @@ async def test_redirect_through_proxy(redirect_server_2, proxy_type: str): async with ProxyServer(proxy_type) as proxy_url: async with ClientAsyncContext(AsyncClient, redirect_server_2, ssl_context=client_ssl_ctx, proxy=proxy_url, proxy_ssl_context=proxy_ssl_ctx) as (transport, listener): transport.send(picows.WSMsgType.BINARY, b"hello over proxy") - frame = await listener.get_message() + frame = await listener.get_message(1.0) assert frame.msg_type == picows.WSMsgType.BINARY assert frame.payload_as_bytes == b"hello over proxy" From fc91a85f750caf39e8a8812ee97283a8f58e53c5 Mon Sep 17 00:00:00 2001 From: taras Date: Tue, 17 Feb 2026 14:50:37 +0100 Subject: [PATCH 11/11] Update tests --- tests/test_redirects_and_proxies.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_redirects_and_proxies.py b/tests/test_redirects_and_proxies.py index 604d5e9..eae72b3 100644 --- a/tests/test_redirects_and_proxies.py +++ b/tests/test_redirects_and_proxies.py @@ -133,7 +133,7 @@ async def test_redirect_through_proxy(redirect_server_2, proxy_type: str): await picows.ws_connect(AsyncClient, redirect_server_2, max_redirects=1, proxy=proxy_url, proxy_ssl_context=proxy_ssl_ctx) -@pytest.mark.parametrize("proxy_type", ["direct", "socks4", "socks5", "http", "http_auth", "https", "https_auth"]) +@pytest.mark.parametrize("proxy_type", ["socks4", "socks5", "http"]) @pytest.mark.skip(reason="echo server may respond with 429 (too many requests if we spam it a lot)") async def test_proxy_dns_resolution(proxy_type): is_https = proxy_type in ("https", "https_auth")