LonamiWebs/Telethon
1
1
Fork 0
mirror of https://github.com/LonamiWebs/Telethon.git synced 2025-01-24 08:14:14 +03:00
Code Issues Packages Projects Releases Wiki Activity

Handle .connect() method more gracefully

Browse Source
This commit is contained in:
Lonami Exo 2017-09-21 12:37:05 +02:00
parent 446174c7de
commit 4777b8dad4
2 changed files with 32 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
telethon/extensions/tcp_client.py
View File

@ -30,19 +30,33 @@ class TcpClient:
else: # tuple, list, etc.
self._socket.set_proxy(*self._proxy)
self._socket.settimeout(self._timeout)
def connect(self, ip, port):
"""Connects to the specified IP and port number.
'timeout' must be given in seconds
"""
if not self.connected:
if ':' in ip: # IPv6
mode, address = socket.AF_INET6, (ip, port, 0, 0)
else:
mode, address = socket.AF_INET, (ip, port)
if ':' in ip: # IPv6
mode, address = socket.AF_INET6, (ip, port, 0, 0)
else:
mode, address = socket.AF_INET, (ip, port)
self._recreate_socket(mode)
self._socket.settimeout(self._timeout)
self._socket.connect(address)
while True:
try:
if not self._socket:
self._recreate_socket(mode)
self._socket.connect(address)
break # Successful connection, stop retrying to connect
except OSError as e:
# There are some errors that we know how to handle, and
# the loop will allow us to retry
if e.errno == errno.EBADF:
# Bad file descriptor, i.e. socket was closed, set it
# to none to recreate it on the next iteration
self._socket = None
else:
raise
def _get_connected(self):
return self._socket is not None

12
telethon/network/connection.py
View File

@ -3,6 +3,8 @@ from datetime import timedelta
from zlib import crc32
from enum import Enum
import errno
from ..crypto import AESModeCTR
from ..extensions import BinaryWriter, TcpClient
from ..errors import InvalidChecksumError
@ -75,9 +77,15 @@ class Connection:
setattr(self, 'read', self._read_plain)
def connect(self):
self._send_counter = 0
self.conn.connect(self.ip, self.port)
try:
self.conn.connect(self.ip, self.port)
except OSError as e:
if e.errno == errno.EISCONN:
return # Already connected, no need to re-set everything up
else:
raise
self._send_counter = 0
if self._mode == ConnectionMode.TCP_ABRIDGED:
self.conn.write(b'\xef')
elif self._mode == ConnectionMode.TCP_INTERMEDIATE: