diff --git a/mautrix_telegram/web/provisioning/__init__.py b/mautrix_telegram/web/provisioning/__init__.py index c52101d6..293a23c8 100644 --- a/mautrix_telegram/web/provisioning/__init__.py +++ b/mautrix_telegram/web/provisioning/__init__.py @@ -14,6 +14,7 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . from __future__ import annotations +import datetime from typing import TYPE_CHECKING, Awaitable, Callable import asyncio @@ -542,19 +543,33 @@ class ProvisioningAPI(AuthAPI): ws = web.WebSocketResponse(protocols=["net.maunium.telegram.login"]) await ws.prepare(request) + retries = 0 user_info = None - try: - await qr_login.recreate() - await ws.send_json({"code": qr_login.url, "timeout": 30}) - user_info = await qr_login.wait() - except asyncio.TimeoutError: + while retries < 4: + try: + await qr_login.recreate() + await ws.send_json( + { + "code": qr_login.url, + "timeout": int( + ( + qr_login.expires - datetime.datetime.now(tz=datetime.timezone.utc) + ).total_seconds() + ), + } + ) + user_info = await qr_login.wait() + break + except asyncio.TimeoutError: + retries += 1 + except SessionPasswordNeededError: + await ws.send_json({"success": False, "error": "password-needed"}) + await ws.close() + return ws + else: await ws.send_json({"success": False, "error": "timeout"}) await ws.close() return ws - except SessionPasswordNeededError: - await ws.send_json({"success": False, "error": "password-needed"}) - await ws.close() - return ws await self.postprocess_login(user, user_info) await ws.send_json({"success": True})