From 2ffac2dcdb1919a6486b10deaf28465d564aea7a Mon Sep 17 00:00:00 2001 From: Lonami Exo Date: Fri, 21 Oct 2022 15:52:00 +0200 Subject: [PATCH] Handle DestroyAuthKey result more gracefully --- telethon/network/mtprotosender.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py index 8a117de1..b7eee784 100644 --- a/telethon/network/mtprotosender.py +++ b/telethon/network/mtprotosender.py @@ -16,11 +16,12 @@ from ..errors import ( from ..extensions import BinaryReader from ..tl.core import RpcResult, MessageContainer, GzipPacked from ..tl.functions.auth import LogOutRequest -from ..tl.functions import PingRequest, DestroySessionRequest +from ..tl.functions import PingRequest, DestroySessionRequest, DestroyAuthKeyRequest from ..tl.types import ( MsgsAck, Pong, BadServerSalt, BadMsgNotification, FutureSalts, MsgNewDetailedInfo, NewSessionCreated, MsgDetailedInfo, MsgsStateReq, MsgsStateInfo, MsgsAllInfo, MsgResendReq, upload, DestroySessionOk, DestroySessionNone, + DestroyAuthKeyOk, DestroyAuthKeyNone, DestroyAuthKeyFail ) from ..tl import types as _tl from ..crypto import AuthKey @@ -111,6 +112,9 @@ class MTProtoSender: MsgsAllInfo.CONSTRUCTOR_ID: self._handle_msg_all, DestroySessionOk.CONSTRUCTOR_ID: self._handle_destroy_session, DestroySessionNone.CONSTRUCTOR_ID: self._handle_destroy_session, + DestroyAuthKeyOk.CONSTRUCTOR_ID: self._handle_destroy_auth_key, + DestroyAuthKeyNone.CONSTRUCTOR_ID: self._handle_destroy_auth_key, + DestroyAuthKeyFail.CONSTRUCTOR_ID: self._handle_destroy_auth_key, } # Public API @@ -849,3 +853,26 @@ class MTProtoSender: del self._pending_state[msg_id] if not state.future.cancelled(): state.future.set_result(message.obj) + + async def _handle_destroy_auth_key(self, message): + """ + Handles :tl:`DestroyAuthKeyFail`, :tl:`DestroyAuthKeyNone`, and :tl:`DestroyAuthKeyOk`. + + :tl:`DestroyAuthKey` is not intended for users to use, but they still + might, and the response won't come in `rpc_result`, so thhat's worked + around here. + """ + self._log.debug('Handling destroy auth key %s', message.obj) + for msg_id, state in list(self._pending_state.items()): + if isinstance(state.request, DestroyAuthKeyRequest): + del self._pending_state[msg_id] + if not state.future.cancelled(): + state.future.set_result(message.obj) + + # If the auth key has been destroyed, that pretty much means the + # library can't continue as our auth key will no longer be found + # on the server. + # Even if the library didn't disconnect, the server would (and then + # the library would reconnect and learn about auth key being invalid). + if isinstance(message.obj, DestroyAuthKeyOk): + await self._disconnect(error=AuthKeyNotFound())