diff --git a/receptor/connection/ws.py b/receptor/connection/ws.py index 5d04f656..f22a9a5f 100644 --- a/receptor/connection/ws.py +++ b/receptor/connection/ws.py @@ -3,6 +3,8 @@ import aiohttp import aiohttp.web import asyncio +from aiohttp.helpers import proxies_from_env +from urllib.parse import urlparse from .base import Transport, log_ssl_detail @@ -37,8 +39,19 @@ async def connect( worker = factory() try: - async with aiohttp.ClientSession().ws_connect(uri, ssl=ssl_context, - headers=ws_extra_headers, heartbeat=ws_heartbeat) as ws: + proxy_scheme = {'ws': 'http', 'wss': 'https'}[urlparse(uri).scheme] + proxies = proxies_from_env() + if proxy_scheme in proxies: + proxy = proxies[proxy_scheme].proxy + proxy_auth = proxies[proxy_scheme].proxy_auth + else: + proxy = None + proxy_auth = None + async with aiohttp.ClientSession().ws_connect( + uri, ssl=ssl_context, headers=ws_extra_headers, + heartbeat=ws_heartbeat, + proxy=proxy, proxy_auth=proxy_auth + ) as ws: log_ssl_detail(ws) t = WebSocket(ws) await worker.client(t)