From b14ca9d86218aef16f371e38d06e78c30e56f09a Mon Sep 17 00:00:00 2001
From: alemidev <me@alemi.dev>
Date: Thu, 25 Aug 2022 19:42:53 +0200
Subject: [PATCH] feat: automatically attempt reauth when timed out

sometimes MS times out our refresh request, making clients disconnect.
Implemented a retry loop, which will attempt to authenticate up to 5
times (configurable) waiting 60 seconds in between (configurable).
Only timeouts will be handled, any other exception will make it crash
immediately.
---
 src/treepuncher/storage.py     |  2 +-
 src/treepuncher/treepuncher.py | 29 +++++++++++++++++++----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/treepuncher/storage.py b/src/treepuncher/storage.py
index 51e348c..855fe53 100644
--- a/src/treepuncher/storage.py
+++ b/src/treepuncher/storage.py
@@ -41,7 +41,7 @@ class AddonStorage:
 		cur.execute(f"INSERT INTO documents_{self.name} VALUES (?, ?)", (key, json.dumps(val, default=str),))
 		self.db.commit()
 
-class Storage:
+class StorageDriver:
 	name : str
 	db : sqlite3.Connection
 
diff --git a/src/treepuncher/treepuncher.py b/src/treepuncher/treepuncher.py
index 163988c..d183d7b 100644
--- a/src/treepuncher/treepuncher.py
+++ b/src/treepuncher/treepuncher.py
@@ -14,7 +14,7 @@ from aiocraft.mc.packet import Packet
 from aiocraft.mc.auth import AuthInterface, AuthException, MojangAuthenticator, MicrosoftAuthenticator, OfflineAuthenticator
 from aiocraft.mc.auth.microsoft import InvalidStateError
 
-from .storage import Storage, SystemState, AuthenticatorState
+from .storage import StorageDriver, SystemState, AuthenticatorState
 from .game import GameState, GameChat, GameInventory, GameTablist, GameWorld, GameContainer
 from .addon import Addon
 from .notifier import Notifier, Provider
@@ -37,7 +37,7 @@ class Treepuncher(
 	GameWorld
 ):
 	name: str
-	storage: Storage
+	storage: StorageDriver
 
 	notifier: Notifier
 	scheduler: AsyncIOScheduler
@@ -99,7 +99,7 @@ class Treepuncher(
 			resolve_srv=opt('resolve_srv', default=True, t=bool),
 		)
 
-		self.storage = Storage(opt('session_file') or f"data/{name}.session")  # TODO wrap with pathlib
+		self.storage = StorageDriver(opt('session_file') or f"data/{name}.session")  # TODO wrap with pathlib
 
 		self.notifier = Notifier(self)
 
@@ -130,13 +130,22 @@ class Treepuncher(
 		return self.authenticator.selectedProfile.name
 
 	async def authenticate(self):
-		await super().authenticate()
-		state = AuthenticatorState(
-			date=datetime.datetime.now(),
-			token=self.authenticator.serialize(),
-			legacy=isinstance(self.authenticator, MojangAuthenticator)
-		)
-		self.storage._set_auth(state)
+		sleep_interval = self.cfg.getfloat("auth_retry_interval", fallback=60.0)
+		for _ in range(self.cfg.getint("auth_retry_count", fallback=5)):
+			try:
+				await super().authenticate()
+				state = AuthenticatorState(
+					date=datetime.datetime.now(),
+					token=self.authenticator.serialize(),
+					legacy=isinstance(self.authenticator, MojangAuthenticator)
+				)
+				self.storage._set_auth(state)
+				return
+			except AuthException as e:
+				if e.data["error"] == "request timed out":
+					await asyncio.sleep(sleep_interval)
+					continue
+				raise e  # retrying won't help anyway
 
 	async def start(self):
 		# if self.started: # TODO readd check