From 6e77f583f1660944903e3ac71824e8fcb5fc5fdb Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Fri, 12 Oct 2018 19:47:40 +0200
Subject: [PATCH] Detect arbitrary negative HTTP error codes

---
 telethon/errors/__init__.py            |  2 +-
 telethon/errors/common.py              | 21 +++++++++++++++------
 telethon/network/mtprotoplainsender.py |  7 +++----
 telethon/network/mtprotosender.py      | 10 +++++++---
 telethon/network/mtprotostate.py       |  8 ++------
 5 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/telethon/errors/__init__.py b/telethon/errors/__init__.py
index e3106631..03125dc5 100644
--- a/telethon/errors/__init__.py
+++ b/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
diff --git a/telethon/errors/common.py b/telethon/errors/common.py
index 9af19c9f..65cd363c 100644
--- a/telethon/errors/common.py
+++ b/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):
diff --git a/telethon/network/mtprotoplainsender.py b/telethon/network/mtprotoplainsender.py
index b53ac24d..f476dda5 100644
--- a/telethon/network/mtprotoplainsender.py
+++ b/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()
diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py
index 86f5b739..d0248680 100644
--- a/telethon/network/mtprotosender.py
+++ b/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
diff --git a/telethon/network/mtprotostate.py b/telethon/network/mtprotostate.py
index 3495ae8a..61642620 100644
--- a/telethon/network/mtprotostate.py
+++ b/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]