Handle DestroyAuthKey result more gracefully

This commit is contained in:
Lonami Exo 2022-10-21 15:52:00 +02:00
parent f902c9293a
commit 2ffac2dcdb

View File

@ -16,11 +16,12 @@ from ..errors import (
from ..extensions import BinaryReader from ..extensions import BinaryReader
from ..tl.core import RpcResult, MessageContainer, GzipPacked from ..tl.core import RpcResult, MessageContainer, GzipPacked
from ..tl.functions.auth import LogOutRequest from ..tl.functions.auth import LogOutRequest
from ..tl.functions import PingRequest, DestroySessionRequest from ..tl.functions import PingRequest, DestroySessionRequest, DestroyAuthKeyRequest
from ..tl.types import ( from ..tl.types import (
MsgsAck, Pong, BadServerSalt, BadMsgNotification, FutureSalts, MsgsAck, Pong, BadServerSalt, BadMsgNotification, FutureSalts,
MsgNewDetailedInfo, NewSessionCreated, MsgDetailedInfo, MsgsStateReq, MsgNewDetailedInfo, NewSessionCreated, MsgDetailedInfo, MsgsStateReq,
MsgsStateInfo, MsgsAllInfo, MsgResendReq, upload, DestroySessionOk, DestroySessionNone, MsgsStateInfo, MsgsAllInfo, MsgResendReq, upload, DestroySessionOk, DestroySessionNone,
DestroyAuthKeyOk, DestroyAuthKeyNone, DestroyAuthKeyFail
) )
from ..tl import types as _tl from ..tl import types as _tl
from ..crypto import AuthKey from ..crypto import AuthKey
@ -111,6 +112,9 @@ class MTProtoSender:
MsgsAllInfo.CONSTRUCTOR_ID: self._handle_msg_all, MsgsAllInfo.CONSTRUCTOR_ID: self._handle_msg_all,
DestroySessionOk.CONSTRUCTOR_ID: self._handle_destroy_session, DestroySessionOk.CONSTRUCTOR_ID: self._handle_destroy_session,
DestroySessionNone.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 # Public API
@ -849,3 +853,26 @@ class MTProtoSender:
del self._pending_state[msg_id] del self._pending_state[msg_id]
if not state.future.cancelled(): if not state.future.cancelled():
state.future.set_result(message.obj) 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())