TeamPGM/pyrogram
2
0
Fork 0
mirror of https://github.com/TeamPGM/pyrogram.git synced 2024-11-28 00:56:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Handle proxy socket connections using thread executors

Browse Source
This commit is contained in:
Dan 2022-12-28 00:14:04 +01:00
parent 1fa6f3b924
commit 2dca5aeac2
1 changed files with 11 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
pyrogram/connection/transport/tcp/tcp.py
View File

@ -20,6 +20,7 @@ import asyncio
import ipaddress import ipaddress
import logging import logging
import socket import socket
from concurrent.futures import ThreadPoolExecutor
import socks import socks
@ -38,6 +39,8 @@ class TCP:
self.lock = asyncio.Lock() self.lock = asyncio.Lock()
self.loop = asyncio.get_event_loop() self.loop = asyncio.get_event_loop()
self.proxy = proxy
if proxy: if proxy:
hostname = proxy.get("hostname") hostname = proxy.get("hostname")
@ -71,10 +74,14 @@ class TCP:
self.socket.setblocking(False) self.socket.setblocking(False)
async def connect(self, address: tuple): async def connect(self, address: tuple):
try: if self.proxy:
await asyncio.wait_for(asyncio.get_event_loop().sock_connect(self.socket, address), TCP.TIMEOUT) with ThreadPoolExecutor(1) as executor:
except asyncio.TimeoutError: # Re-raise as TimeoutError. asyncio.TimeoutError is deprecated in 3.11 await self.loop.run_in_executor(executor, self.socket.connect, address)
raise TimeoutError("Connection timed out") 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) self.reader, self.writer = await asyncio.open_connection(sock=self.socket)