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. else: # tuple, list, etc.
self._socket.set_proxy(*self._proxy) self._socket.set_proxy(*self._proxy)
self._socket.settimeout(self._timeout)
def connect(self, ip, port): def connect(self, ip, port):
"""Connects to the specified IP and port number. """Connects to the specified IP and port number.
'timeout' must be given in seconds 'timeout' must be given in seconds
""" """
if not self.connected: if ':' in ip: # IPv6
if ':' in ip: # IPv6 mode, address = socket.AF_INET6, (ip, port, 0, 0)
mode, address = socket.AF_INET6, (ip, port, 0, 0) else:
else: mode, address = socket.AF_INET, (ip, port)
mode, address = socket.AF_INET, (ip, port)
self._recreate_socket(mode) while True:
self._socket.settimeout(self._timeout) try:
self._socket.connect(address) 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): def _get_connected(self):
return self._socket is not None 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 zlib import crc32
from enum import Enum from enum import Enum
import errno
from ..crypto import AESModeCTR from ..crypto import AESModeCTR
from ..extensions import BinaryWriter, TcpClient from ..extensions import BinaryWriter, TcpClient
from ..errors import InvalidChecksumError from ..errors import InvalidChecksumError
@ -75,9 +77,15 @@ class Connection:
setattr(self, 'read', self._read_plain) setattr(self, 'read', self._read_plain)
def connect(self): def connect(self):
self._send_counter = 0 try:
self.conn.connect(self.ip, self.port) 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: if self._mode == ConnectionMode.TCP_ABRIDGED:
self.conn.write(b'\xef') self.conn.write(b'\xef')
elif self._mode == ConnectionMode.TCP_INTERMEDIATE: elif self._mode == ConnectionMode.TCP_INTERMEDIATE: