Fix valid auth_key never being saved after switching DC

This commit is contained in:
Lonami Exo 2018-11-03 18:53:26 +01:00
parent 0fc97a3e8c
commit cdbd1f6193
3 changed files with 16 additions and 3 deletions

View File

@ -13,7 +13,6 @@ from ..network import MTProtoSender, ConnectionTcpFull
from ..sessions import Session, SQLiteSession, MemorySession from ..sessions import Session, SQLiteSession, MemorySession
from ..tl import TLObject, functions, types from ..tl import TLObject, functions, types
from ..tl.alltlobjects import LAYER from ..tl.alltlobjects import LAYER
from ..crypto import AuthKey
DEFAULT_DC_ID = 4 DEFAULT_DC_ID = 4
DEFAULT_IPV4_IP = '149.154.167.51' DEFAULT_IPV4_IP = '149.154.167.51'
@ -245,6 +244,7 @@ class TelegramBaseClient(abc.ABC):
delay=self._retry_delay, delay=self._retry_delay,
auto_reconnect=self._auto_reconnect, auto_reconnect=self._auto_reconnect,
connect_timeout=self._timeout, connect_timeout=self._timeout,
auth_key_callback=self._auth_key_callback,
update_callback=self._handle_update, update_callback=self._handle_update,
auto_reconnect_callback=self._handle_auto_reconnect auto_reconnect_callback=self._handle_auto_reconnect
) )
@ -386,12 +386,13 @@ class TelegramBaseClient(abc.ABC):
self.session.set_dc(dc.id, dc.ip_address, dc.port) self.session.set_dc(dc.id, dc.ip_address, dc.port)
# auth_key's are associated with a server, which has now changed # auth_key's are associated with a server, which has now changed
# so it's not valid anymore. Set to None to force recreating it. # so it's not valid anymore. Set to None to force recreating it.
self._sender.auth_key.key = None
self.session.auth_key = None self.session.auth_key = None
self.session.save() self.session.save()
self._disconnect() self._disconnect()
return await self.connect() return await self.connect()
async def _auth_key_callback(self, auth_key): def _auth_key_callback(self, auth_key):
""" """
Callback from the sender whenever it needed to generate a Callback from the sender whenever it needed to generate a
new authorization key. This means we are not authorized. new authorization key. This means we are not authorized.

View File

@ -61,6 +61,7 @@ class MTProtoSender:
""" """
def __init__(self, auth_key, loop, *, def __init__(self, auth_key, loop, *,
retries=5, delay=1, auto_reconnect=True, connect_timeout=None, retries=5, delay=1, auto_reconnect=True, connect_timeout=None,
auth_key_callback=None,
update_callback=None, auto_reconnect_callback=None): update_callback=None, auto_reconnect_callback=None):
self._connection = None self._connection = None
self._loop = loop self._loop = loop
@ -68,6 +69,7 @@ class MTProtoSender:
self._delay = delay self._delay = delay
self._auto_reconnect = auto_reconnect self._auto_reconnect = auto_reconnect
self._connect_timeout = connect_timeout self._connect_timeout = connect_timeout
self._auth_key_callback = auth_key_callback
self._update_callback = update_callback self._update_callback = update_callback
self._auto_reconnect_callback = auto_reconnect_callback self._auto_reconnect_callback = auto_reconnect_callback
@ -232,6 +234,13 @@ class MTProtoSender:
self.auth_key.key, self._state.time_offset =\ self.auth_key.key, self._state.time_offset =\
await authenticator.do_authentication(plain) await authenticator.do_authentication(plain)
# This is *EXTREMELY* important since we don't control
# external references to the authorization key, we must
# notify whenever we change it. This is crucial when we
# switch to different data centers.
if self._auth_key_callback:
self._auth_key_callback(self.auth_key)
break break
except (SecurityError, AssertionError) as e: except (SecurityError, AssertionError) as e:
__log__.warning('Attempt {} at new auth_key failed: {}' __log__.warning('Attempt {} at new auth_key failed: {}'
@ -420,6 +429,9 @@ class MTProtoSender:
__log__.warning('Invalid buffer %s', e) __log__.warning('Invalid buffer %s', e)
self.auth_key.key = None self.auth_key.key = None
if self._auth_key_callback:
self._auth_key_callback(None)
self._start_reconnect() self._start_reconnect()
return return
except Exception: except Exception:

View File

@ -1,3 +1,3 @@
# Versions should comply with PEP440. # Versions should comply with PEP440.
# This line is parsed in setup.py: # This line is parsed in setup.py:
__version__ = '1.4' __version__ = '1.4.1'