Set auth_key on connection

This commit is contained in:
Lonami Exo 2018-10-01 14:20:50 +02:00
parent 21ffa2f26b
commit cf7e5d5592
2 changed files with 9 additions and 13 deletions

View File

@ -11,7 +11,6 @@ from .. import version
from ..crypto import rsa from ..crypto import rsa
from ..extensions import markdown from ..extensions import markdown
from ..network import MTProtoSender, ConnectionTcpFull from ..network import MTProtoSender, ConnectionTcpFull
from ..network.mtprotostate import MTProtoState
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
@ -226,7 +225,7 @@ class TelegramBaseClient(abc.ABC):
self._connection = connection self._connection = connection
self._sender = MTProtoSender( self._sender = MTProtoSender(
self.session.auth_key, self._loop, self._loop,
retries=self._connection_retries, retries=self._connection_retries,
auto_reconnect=self._auto_reconnect, auto_reconnect=self._auto_reconnect,
update_callback=self._handle_update, update_callback=self._handle_update,
@ -305,7 +304,7 @@ class TelegramBaseClient(abc.ABC):
""" """
Connects to Telegram. Connects to Telegram.
""" """
await self._sender.connect(self._connection( await self._sender.connect(self.session.auth_key, self._connection(
self.session.server_address, self.session.port, loop=self._loop)) self.session.server_address, self.session.port, loop=self._loop))
await self._sender.send(self._init_with( await self._sender.send(self._init_with(
@ -369,7 +368,7 @@ 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.session.auth_key = self._sender._connection._state.auth_key = None self.session.auth_key = None
self.session.save() self.session.save()
await self._disconnect() await self._disconnect()
return await self.connect() return await self.connect()
@ -416,8 +415,8 @@ class TelegramBaseClient(abc.ABC):
# #
# If one were to do that, Telegram would reset the connection # If one were to do that, Telegram would reset the connection
# with no further clues. # with no further clues.
sender = MTProtoSender(None, self._loop) sender = MTProtoSender(self._loop)
await sender.connect(self._connection( await sender.connect(None, self._connection(
dc.ip_address, dc.port, loop=self._loop)) dc.ip_address, dc.port, loop=self._loop))
__log__.info('Exporting authorization for data center %s', dc) __log__.info('Exporting authorization for data center %s', dc)
auth = await self(functions.auth.ExportAuthorizationRequest(dc_id)) auth = await self(functions.auth.ExportAuthorizationRequest(dc_id))

View File

@ -38,10 +38,9 @@ class MTProtoSender:
A new authorization key will be generated on connection if no other A new authorization key will be generated on connection if no other
key exists yet. key exists yet.
""" """
def __init__(self, auth_key, loop, *, def __init__(self, loop, *,
retries=5, auto_reconnect=True, update_callback=None, retries=5, auto_reconnect=True, update_callback=None,
auth_key_callback=None, auto_reconnect_callback=None): auth_key_callback=None, auto_reconnect_callback=None):
self._auth_key = auth_key
self._connection = None # MTProtoLayer, a.k.a. encrypted connection self._connection = None # MTProtoLayer, a.k.a. encrypted connection
self._loop = loop self._loop = loop
self._retries = retries self._retries = retries
@ -100,17 +99,15 @@ class MTProtoSender:
# Public API # Public API
async def connect(self, connection): async def connect(self, auth_key, connection):
""" """
Connects to the specified ``ip:port``, and generates a new Connects to the specified given connection using the given auth key.
authorization key for the `MTProtoSender.session` if it does
not exist yet.
""" """
if self._user_connected: if self._user_connected:
__log__.info('User is already connected!') __log__.info('User is already connected!')
return return
self._connection = MTProtoLayer(connection, self._auth_key) self._connection = MTProtoLayer(connection, auth_key)
self._user_connected = True self._user_connected = True
await self._connect() await self._connect()