From e4a2bd2f69200aab8d92b235ab45912e432d1e21 Mon Sep 17 00:00:00 2001
From: Tulir Asokan <tulir@maunium.net>
Date: Wed, 20 Oct 2021 20:02:09 +0300
Subject: [PATCH] Catch authorization errors in get_me()

---
 mautrix_telegram/commands/telegram/auth.py    | 11 +++++---
 mautrix_telegram/user.py                      | 23 +++++++++++++---
 mautrix_telegram/web/provisioning/__init__.py | 26 +++++++++++--------
 3 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/mautrix_telegram/commands/telegram/auth.py b/mautrix_telegram/commands/telegram/auth.py
index 54e142ce..10320124 100644
--- a/mautrix_telegram/commands/telegram/auth.py
+++ b/mautrix_telegram/commands/telegram/auth.py
@@ -46,10 +46,13 @@
                  help_section=SECTION_AUTH,
                  help_text="Check if you're logged into Telegram.")
 async def ping(evt: CommandEvent) -> EventID:
-    me = await evt.sender.client.get_me() if await evt.sender.is_logged_in() else None
-    if me:
-        human_tg_id = f"@{me.username}" if me.username else f"+{me.phone}"
-        return await evt.reply(f"You're logged in as {human_tg_id}")
+    if await evt.sender.is_logged_in():
+        me = await evt.sender.get_me()
+        if me:
+            human_tg_id = f"@{me.username}" if me.username else f"+{me.phone}"
+            return await evt.reply(f"You're logged in as {human_tg_id}")
+        else:
+            return await evt.reply("You were logged in, but there appears to have been an error.")
     else:
         return await evt.reply("You're not logged in.")
 
diff --git a/mautrix_telegram/user.py b/mautrix_telegram/user.py
index 220b35b6..57a166ab 100644
--- a/mautrix_telegram/user.py
+++ b/mautrix_telegram/user.py
@@ -22,12 +22,14 @@
 from telethon.tl.types import (TypeUpdate, UpdateNewMessage, UpdateNewChannelMessage,
                                UpdateShortChatMessage, UpdateShortMessage, User as TLUser, Chat,
                                ChatForbidden, UpdateFolderPeers, UpdatePinnedDialogs,
-                               UpdateNotifySettings, NotifyPeer)
+                               UpdateNotifySettings, NotifyPeer, InputUserSelf)
 from telethon.tl.custom import Dialog
 from telethon.tl.types.contacts import ContactsNotModified
 from telethon.tl.functions.contacts import GetContactsRequest, SearchRequest
 from telethon.tl.functions.account import UpdateStatusRequest
-from telethon.errors import AuthKeyDuplicatedError
+from telethon.tl.functions.users import GetUsersRequest
+from telethon.errors import (AuthKeyDuplicatedError, UserDeactivatedError, UserDeactivatedBanError,
+                             SessionRevokedError, UnauthorizedError)
 
 from mautrix.client import Client
 from mautrix.errors import MatrixRequestError, MNotFound
@@ -333,8 +335,22 @@ async def set_presence(self, online: bool = True) -> None:
         if not self.is_bot:
             await self.client(UpdateStatusRequest(offline=not online))
 
+    async def get_me(self) -> Optional[TLUser]:
+        try:
+            return (await self.client(GetUsersRequest([InputUserSelf()])))[0]
+        except UnauthorizedError as e:
+            self.log.error(f"Authorization error in get_me(): {e}")
+            await self.push_bridge_state(BridgeStateEvent.BAD_CREDENTIALS, error="tg-auth-error",
+                                         message=str(e), ttl=3600)
+            await self.stop()
+            return None
+
     async def update_info(self, info: TLUser = None) -> None:
-        info = info or await self.client.get_me()
+        if not info:
+            info = await self.get_me()
+            if not info:
+                self.log.warning("get_me() returned None, aborting update_info()")
+                return
         changed = False
         if self.is_bot != info.bot:
             self.is_bot = info.bot
@@ -350,6 +366,7 @@ async def update_info(self, info: TLUser = None) -> None:
             self.by_tgid[self.tgid] = self
         if changed:
             await self.save()
+        return info
 
     async def log_out(self) -> bool:
         puppet = pu.Puppet.get(self.tgid)
diff --git a/mautrix_telegram/web/provisioning/__init__.py b/mautrix_telegram/web/provisioning/__init__.py
index 2ec4cae0..a7308761 100644
--- a/mautrix_telegram/web/provisioning/__init__.py
+++ b/mautrix_telegram/web/provisioning/__init__.py
@@ -21,7 +21,10 @@
 from aiohttp import web
 
 from telethon.utils import get_peer_id, resolve_id
-from telethon.tl.types import ChatForbidden, ChannelForbidden, TypeChat
+from telethon.tl.types import ChatForbidden, ChannelForbidden, TypeChat, InputUserSelf
+from telethon.tl.functions.users import GetUsersRequest
+from telethon.errors import (UserDeactivatedError, UserDeactivatedBanError, SessionRevokedError,
+                             UnauthorizedError)
 
 from mautrix.appservice import AppService
 from mautrix.errors import MatrixRequestError, IntentError
@@ -294,16 +297,17 @@ async def get_user_info(self, request: web.Request) -> web.Response:
 
         user_data = None
         if await user.is_logged_in():
-            me = await user.client.get_me()
-            await user.update_info(me)
-            user_data = {
-                "id": user.tgid,
-                "username": user.username,
-                "first_name": me.first_name,
-                "last_name": me.last_name,
-                "phone": me.phone,
-                "is_bot": user.is_bot,
-            }
+            me = await user.get_me()
+            if me:
+                await user.update_info(me)
+                user_data = {
+                    "id": user.tgid,
+                    "username": user.username,
+                    "first_name": me.first_name,
+                    "last_name": me.last_name,
+                    "phone": me.phone,
+                    "is_bot": user.is_bot,
+                }
         return web.json_response({
             "telegram": user_data,
             "mxid": user.mxid,