Detect arbitrary negative HTTP error codes

telethon/errors/__init__.py

@@ -8,7 +8,7 @@ from threading import Thread
 
 from .common import (
     ReadCancelledError, TypeNotFoundError, InvalidChecksumError,
-    BrokenAuthKeyError, SecurityError, CdnFileTamperedError, MultiError
+    InvalidBufferError, SecurityError, CdnFileTamperedError, MultiError
 )
 
 # This imports the base errors too, as they're imported there

telethon/errors/common.py

@@ -1,4 +1,6 @@
 """Errors not related to the Telegram API itself"""
+import struct
+
 from ..tl import TLRequest
 
 
@@ -39,14 +41,21 @@ class InvalidChecksumError(Exception):
         self.valid_checksum = valid_checksum
 
 
-class BrokenAuthKeyError(Exception):
+class InvalidBufferError(BufferError):
     """
-    Occurs when the authorization key for a data center is not valid.
+    Occurs when the buffer is invalid, and may contain an HTTP error code.
+    For instance, 404 means "forgotten/broken authorization key", while
     """
-    def __init__(self):
-        super().__init__(
-            'The authorization key is broken, and it must be reset.'
-        )
+    def __init__(self, payload):
+        self.payload = payload
+        if len(payload) == 4:
+            self.code = -struct.unpack('<i', payload)[0]
+            super().__init__(
+                'Invalid response buffer (HTTP code {})'.format(self.code))
+        else:
+            self.code = None
+            super().__init__(
+                'Invalid response buffer (too short {})'.format(self.payload))
 
 
 class SecurityError(Exception):

telethon/network/mtprotoplainsender.py

@@ -5,7 +5,7 @@ in plain text, when no authorization key has been created yet.
 import struct
 
 from .mtprotostate import MTProtoState
-from ..errors import BrokenAuthKeyError
+from ..errors import InvalidBufferError
 from ..extensions import BinaryReader
 
 
@@ -34,9 +34,8 @@ class MTProtoPlainSender:
         )
 
         body = await self._connection.recv()
-        if body == b'l\xfe\xff\xff':  # -404 little endian signed
-            # Broken authorization, must reset the auth key
-            raise BrokenAuthKeyError()
+        if len(body) < 8:
+            raise InvalidBufferError(body)
 
         with BinaryReader(body) as reader:
             auth_key_id = reader.read_long()

telethon/network/mtprotosender.py

@@ -9,7 +9,7 @@ from .requeststate import RequestState
 from ..tl.tlobject import TLRequest
 from .. import utils
 from ..errors import (
-    BadMessageError, BrokenAuthKeyError, SecurityError, TypeNotFoundError,
+    BadMessageError, InvalidBufferError, SecurityError, TypeNotFoundError,
     InvalidChecksumError, rpc_message_to_error
 )
 from ..extensions import BinaryReader
@@ -382,8 +382,12 @@ class MTProtoSender:
                 )
             except asyncio.CancelledError:
                 return
-            except (BrokenAuthKeyError, BufferError):
-                __log__.info('Broken authorization key; resetting')
+            except BufferError as e:
+                if isinstance(e, InvalidBufferError) and e.code == 404:
+                    __log__.info('Broken authorization key; resetting')
+                else:
+                    __log__.warning('Invalid buffer %s', e)
+
                 self._connection._state.auth_key = None
                 self._start_reconnect()
                 return

telethon/network/mtprotostate.py

@@ -5,7 +5,7 @@ import time
 from hashlib import sha256
 
 from ..crypto import AES
-from ..errors import SecurityError, BrokenAuthKeyError
+from ..errors import SecurityError, InvalidBufferError
 from ..extensions import BinaryReader
 from ..tl.core import TLMessage
 from ..tl.functions import InvokeAfterMsgRequest
@@ -114,11 +114,7 @@ class MTProtoState:
         Inverse of `encrypt_message_data` for incoming server messages.
         """
         if len(body) < 8:
-            # TODO If len == 4, raise HTTPErrorCode(-little endian int)
-            if body == b'l\xfe\xff\xff':
-                raise BrokenAuthKeyError()
-            else:
-                raise BufferError("Can't decode packet ({})".format(body))
+            raise InvalidBufferError(body)
 
         # TODO Check salt, session_id and sequence_number
         key_id = struct.unpack('<Q', body[:8])[0]