From 2dca5aeac20b2e3a6762c01e0b28a375ad97e6a6 Mon Sep 17 00:00:00 2001 From: Dan <14043624+delivrance@users.noreply.github.com> Date: Wed, 28 Dec 2022 00:14:04 +0100 Subject: [PATCH] Handle proxy socket connections using thread executors --- pyrogram/connection/transport/tcp/tcp.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py index 22b367ae..168f06a5 100644 --- a/pyrogram/connection/transport/tcp/tcp.py +++ b/pyrogram/connection/transport/tcp/tcp.py @@ -20,6 +20,7 @@ import asyncio import ipaddress import logging import socket +from concurrent.futures import ThreadPoolExecutor import socks @@ -38,6 +39,8 @@ class TCP: self.lock = asyncio.Lock() self.loop = asyncio.get_event_loop() + self.proxy = proxy + if proxy: hostname = proxy.get("hostname") @@ -71,10 +74,14 @@ class TCP: self.socket.setblocking(False) async def connect(self, address: tuple): - try: - await asyncio.wait_for(asyncio.get_event_loop().sock_connect(self.socket, address), TCP.TIMEOUT) - except asyncio.TimeoutError: # Re-raise as TimeoutError. asyncio.TimeoutError is deprecated in 3.11 - raise TimeoutError("Connection timed out") + if self.proxy: + with ThreadPoolExecutor(1) as executor: + await self.loop.run_in_executor(executor, self.socket.connect, address) + else: + try: + await asyncio.wait_for(asyncio.get_event_loop().sock_connect(self.socket, address), TCP.TIMEOUT) + except asyncio.TimeoutError: # Re-raise as TimeoutError. asyncio.TimeoutError is deprecated in 3.11 + raise TimeoutError("Connection timed out") self.reader, self.writer = await asyncio.open_connection(sock=self.socket)