Add smarter auth import to deal with race conditions by multi sessions

delivrance · delivrance · commit c1321a4c0115 · 2019-08-10T22:37:07.000+02:00
- Add a retry mechanism (up to three times) - Narrow the window in which export+import executes - Remove a line of duplicated code Fixes #299
diff --git a/pyrogram/client/client.py b/pyrogram/client/client.py
@@ -45,7 +45,7 @@
     PhoneCodeExpired, PhoneCodeEmpty, SessionPasswordNeeded,
     PasswordHashInvalid, FloodWait, PeerIdInvalid, FirstnameInvalid, PhoneNumberBanned,
     VolumeLocNotFound, UserMigrate, ChannelPrivate, PhoneNumberOccupied,
-    PasswordRecoveryNa, PasswordEmpty
+    PasswordRecoveryNa, PasswordEmpty, AuthBytesInvalid
 )
 from pyrogram.session import Auth, Session
 from .ext import utils, Syncer, BaseClient, Dispatcher
@@ -1229,7 +1229,7 @@ def send(self, data: TLObject, retries: int = Session.MAX_RETRIES, timeout: floa
     def load_config(self):
         parser = ConfigParser()
         parser.read(str(self.config_file))
-        
+
         if self.bot_token:
             pass
         else:
@@ -1720,30 +1720,35 @@ def get_file(
 
             if session is None:
                 if dc_id != self.storage.dc_id:
-                    exported_auth = self.send(
-                        functions.auth.ExportAuthorization(
-                            dc_id=dc_id
-                        )
-                    )
-
                     session = Session(self, dc_id, Auth(self, dc_id).create(), is_media=True)
-
                     session.start()
 
-                    self.media_sessions[dc_id] = session
-
-                    session.send(
-                        functions.auth.ImportAuthorization(
-                            id=exported_auth.id,
-                            bytes=exported_auth.bytes
+                    for _ in range(3):
+                        exported_auth = self.send(
+                            functions.auth.ExportAuthorization(
+                                dc_id=dc_id
+                            )
                         )
-                    )
+
+                        try:
+                            session.send(
+                                functions.auth.ImportAuthorization(
+                                    id=exported_auth.id,
+                                    bytes=exported_auth.bytes
+                                )
+                            )
+                        except AuthBytesInvalid:
+                            continue
+                        else:
+                            break
+                    else:
+                        session.stop()
+                        raise AuthBytesInvalid
                 else:
                     session = Session(self, dc_id, self.storage.auth_key, is_media=True)
-
                     session.start()
 
-                    self.media_sessions[dc_id] = session
+                self.media_sessions[dc_id] = session
 
         if media_type == 1:
             location = types.InputPeerPhotoFileLocation(
