GramJS: Store test server flag with session (#5084)

src/api/gramjs/apiBuilders/misc.ts

@@ -220,6 +220,7 @@ export function buildApiUrlAuthResult(result: GramJs.TypeUrlAuthResult): ApiUrlA
 export function buildApiConfig(config: GramJs.Config): ApiConfig {
   const defaultReaction = config.reactionsDefault && buildApiReaction(config.reactionsDefault);
   return {
+    isTestServer: config.testMode,
     expiresAt: config.expires,
     gifSearchUsername: config.gifSearchUsername,
     defaultReaction,

src/api/gramjs/methods/client.ts

@@ -69,10 +69,11 @@ export async function init(initialArgs: ApiInitialArgs) {
   }
 
   const {
-    userAgent, platform, sessionData, isTest, isWebmSupported, maxBufferSize, webAuthToken, dcId,
+    userAgent, platform, sessionData, isWebmSupported, maxBufferSize, webAuthToken, dcId,
     mockScenario, shouldForceHttpTransport, shouldAllowHttpTransport,
-    shouldDebugExportedSenders, langCode,
+    shouldDebugExportedSenders, langCode, isTestServerRequested,
   } = initialArgs;
+
   const session = new sessions.CallbackSession(sessionData, onSessionUpdate);
 
   // eslint-disable-next-line no-restricted-globals
@@ -93,9 +94,9 @@ export async function init(initialArgs: ApiInitialArgs) {
       shouldDebugExportedSenders,
       shouldForceHttpTransport,
       shouldAllowHttpTransport,
-      testServers: isTest,
       dcId,
       langCode,
+      isTestServerRequested,
     } as any,
   );
 

src/api/types/misc.ts

@@ -17,6 +17,7 @@ export interface ApiInitialArgs {
   shouldForceHttpTransport?: boolean;
   shouldDebugExportedSenders?: boolean;
   langCode: string;
+  isTestServerRequested?: boolean;
 }
 
 export interface ApiOnProgress {
@@ -96,6 +97,7 @@ export interface ApiWebSession {
 
 export interface ApiSessionData {
   mainDcId: number;
+  isTest?: true;
   keys: Record<number, string | number[]>;
   hashes: Record<number, string | number[]>;
 }
@@ -242,6 +244,7 @@ export interface ApiConfig {
   gifSearchUsername?: string;
   maxGroupSize: number;
   autologinToken?: string;
+  isTestServer?: boolean;
 }
 
 export type ApiPeerColorSet = string[];

src/components/App.tsx

@@ -38,6 +38,7 @@ type StateProps = {
   hasPasscode?: boolean;
   isInactiveAuth?: boolean;
   hasWebAuthTokenFailed?: boolean;
+  isTestServer?: boolean;
   theme: ThemeKey;
 };
 
@@ -57,6 +58,7 @@ const App: FC<StateProps> = ({
   hasPasscode,
   isInactiveAuth,
   hasWebAuthTokenFailed,
+  isTestServer,
   theme,
 }) => {
   const { disconnect } = getActions();
@@ -221,6 +223,7 @@ const App: FC<StateProps> = ({
       >
         {renderContent}
       </Transition>
+      {activeKey === AppScreens.auth && isTestServer && <div className="test-server-badge">Test server</div>}
     </UiLoader>
   );
 };
@@ -234,6 +237,7 @@ export default withGlobal(
       isInactiveAuth: selectTabState(global).isInactive,
       hasWebAuthTokenFailed: global.hasWebAuthTokenFailed || global.hasWebAuthTokenPasswordRequired,
       theme: selectTheme(global),
+      isTestServer: global.config?.isTestServer,
     };
   },
 )(App);

src/components/auth/Auth.scss

@@ -219,6 +219,12 @@
   word-break: normal !important;
 }
 
+.test-server-badge {
+  position: fixed;
+  bottom: 0.5rem;
+  right: 0.5rem;
+}
+
 @keyframes qr-show {
   0% {
     opacity: 0;

src/components/main/Main.tsx

@@ -397,6 +397,8 @@ const Main = ({
   // Parse deep link
   useEffect(() => {
     if (!isSynced) return;
+    updatePageTitle();
+
     const parsedInitialLocationHash = parseInitialLocationHash();
     if (parsedInitialLocationHash?.tgaddr) {
       processDeepLink(decodeURIComponent(parsedInitialLocationHash.tgaddr));

src/global/actions/api/initial.ts

@@ -40,11 +40,12 @@ import {
 addActionHandler('initApi', (global, actions): ActionReturnType => {
   const initialLocationHash = parseInitialLocationHash();
 
+  const hasTestParam = window.location.search.includes('test') || initialLocationHash?.tgWebAuthTest === '1';
+
   void initApi(actions.apiUpdate, {
     userAgent: navigator.userAgent,
     platform: PLATFORM_ENV,
     sessionData: loadStoredSession(),
-    isTest: window.location.search.includes('test') || initialLocationHash?.tgWebAuthTest === '1',
     isWebmSupported: IS_WEBM_SUPPORTED,
     maxBufferSize: MAX_BUFFER_SIZE,
     webAuthToken: initialLocationHash?.tgWebAuthToken,
@@ -54,6 +55,7 @@ addActionHandler('initApi', (global, actions): ActionReturnType => {
     shouldForceHttpTransport: global.settings.byKey.shouldForceHttpTransport,
     shouldDebugExportedSenders: global.settings.byKey.shouldDebugExportedSenders,
     langCode: global.settings.byKey.language,
+    isTestServerRequested: hasTestParam,
   });
 
   void setShouldEnableDebugLog(Boolean(global.settings.byKey.shouldCollectDebugLogs));

src/global/actions/apiUpdaters/initial.ts

@@ -243,10 +243,21 @@ function onUpdateConnectionState<T extends GlobalState>(
 
 function onUpdateSession<T extends GlobalState>(global: T, actions: RequiredGlobalActions, update: ApiUpdateSession) {
   const { sessionData } = update;
-  global = getGlobal();
   const { authRememberMe, authState } = global;
   const isEmpty = !sessionData || !sessionData.mainDcId;
 
+  const isTest = sessionData?.isTest;
+  if (isTest) {
+    global = {
+      ...global,
+      config: {
+        ...global.config,
+        isTestServer: isTest,
+      },
+    };
+    setGlobal(global);
+  }
+
   if (!authRememberMe || authState !== 'authorizationStateReady' || isEmpty) {
     return;
   }

src/global/actions/ui/misc.ts

@@ -749,10 +749,12 @@ addActionHandler('updatePageTitle', (global, actions, payload): ActionReturnType
   const { tabId = getCurrentTabId() } = payload || {};
   const { canDisplayChatInTitle } = global.settings.byKey;
   const currentUserId = global.currentUserId;
+  const isTestServer = global.config?.isTestServer;
+  const prefix = isTestServer ? '[T] ' : '';
 
   if (document.title.includes(INACTIVE_MARKER)) {
     updateIcon(false);
-    setPageTitleInstant(`${PAGE_TITLE} ${INACTIVE_MARKER}`);
+    setPageTitleInstant(`${prefix}${PAGE_TITLE} ${INACTIVE_MARKER}`);
     return;
   }
 
@@ -762,7 +764,7 @@ addActionHandler('updatePageTitle', (global, actions, payload): ActionReturnType
     const newUnread = notificationCount - global.initialUnreadNotifications;
 
     if (newUnread > 0) {
-      setPageTitleInstant(`${newUnread} notification${newUnread > 1 ? 's' : ''}`);
+      setPageTitleInstant(`${prefix}${newUnread} notification${newUnread > 1 ? 's' : ''}`);
       updateIcon(true);
       return;
     }
@@ -779,16 +781,16 @@ addActionHandler('updatePageTitle', (global, actions, payload): ActionReturnType
       const title = getChatTitle(langProvider.oldTranslate, currentChat, chatId === currentUserId);
       const topic = selectTopic(global, chatId, threadId);
       if (currentChat.isForum && topic) {
-        setPageTitle(`${title} › ${topic.title}`);
+        setPageTitle(`${prefix}${title} › ${topic.title}`);
         return;
       }
 
-      setPageTitle(title);
+      setPageTitle(`${prefix}${title}`);
       return;
     }
   }
 
-  setPageTitleInstant(IS_ELECTRON ? '' : PAGE_TITLE);
+  setPageTitleInstant(IS_ELECTRON ? '' : `${prefix}${PAGE_TITLE}`);
 });
 
 addActionHandler('closeInviteViaLinkModal', (global, actions, payload): ActionReturnType => {

src/global/cache.ts

@@ -291,6 +291,7 @@ function reduceGlobal<T extends GlobalState>(global: T) {
     ...INITIAL_GLOBAL_STATE,
     ...pick(global, [
       'appConfig',
+      'config',
       'authState',
       'authPhoneNumber',
       'authRememberMe',

src/lib/gramjs/client/TelegramClient.js

@@ -77,8 +77,8 @@ class TelegramClient {
         baseLogger: 'gramjs',
         useWSS: false,
         additionalDcsDisabled: false,
-        testServers: false,
         dcId: DEFAULT_DC_ID,
+        isTestServerRequested: false,
         shouldAllowHttpTransport: false,
         shouldForceHttpTransport: false,
         shouldDebugExportedSenders: false,
@@ -218,10 +218,10 @@ class TelegramClient {
         this._sender._disconnected = true;
 
         const connection = new this._connection(
-            this.session.serverAddress, this.session.port, this.session.dcId, this._log, this._args.testServers,
+            this.session.serverAddress, this.session.port, this.session.dcId, this._log, this.session.isTestServer,
         );
         const fallbackConnection = new this._fallbackConnection(
-            this.session.serverAddress, this.session.port, this.session.dcId, this._log, this._args.testServers,
+            this.session.serverAddress, this.session.port, this.session.dcId, this._log, this.session.isTestServer,
         );
 
         const newConnection = await this._sender.connect(connection, undefined, fallbackConnection);
@@ -257,7 +257,9 @@ class TelegramClient {
         if (!this.session.serverAddress || (this.session.serverAddress.includes(':') !== this._useIPV6)) {
             const DC = utils.getDC(this.defaultDcId);
             // TODO Fill IP addresses for when `this._useIPV6` is used
-            this.session.setDC(this.defaultDcId, DC.ipAddress, this._args.useWSS ? 443 : 80);
+            this.session.setDC(
+                this.defaultDcId, DC.ipAddress, this._args.useWSS ? 443 : 80, this._args.isTestServerRequested,
+            );
         }
     }
 
@@ -417,7 +419,8 @@ class TelegramClient {
     async _switchDC(newDc) {
         this._log.info(`Reconnecting to new data center ${newDc}`);
         const DC = utils.getDC(newDc);
-        this.session.setDC(newDc, DC.ipAddress, DC.port);
+        const isTestServer = this.session.isTestServer || this._args.isTestServerRequested;
+        this.session.setDC(newDc, DC.ipAddress, DC.port, isTestServer);
         // authKey's are associated with a server, which has now changed
         // so it's not valid anymore. Set to None to force recreating it.
         await this._sender.authKey.setKey(undefined);
@@ -495,7 +498,7 @@ class TelegramClient {
                     dc.port,
                     dcId,
                     this._log,
-                    this._args.testServers,
+                    this.session.isTestServer,
                     // Premium DCs are not stable for obtaining auth keys, so need to we first connect to regular ones
                     hasAuthKey ? isPremium : false,
                 ), undefined, new this._fallbackConnection(
@@ -503,7 +506,7 @@ class TelegramClient {
                     dc.port,
                     dcId,
                     this._log,
-                    this._args.testServers,
+                    this.session.isTestServer,
                     hasAuthKey ? isPremium : false,
                 ));
 

src/lib/gramjs/network/connection/Connection.js

@@ -22,6 +22,7 @@ class Connection {
         this._dcId = dcId;
         this._log = loggers;
         this._testServers = testServers;
+
         this._isPremium = isPremium;
         this._connected = false;
         this._sendTask = undefined;

src/lib/gramjs/sessions/CallbackSession.js

@@ -29,13 +29,14 @@ class CallbackSession extends MemorySession {
             mainDcId,
             keys,
             hashes,
+            isTest,
         } = this._sessionData;
         const {
             ipAddress,
             port,
         } = utils.getDC(mainDcId);
 
-        this.setDC(mainDcId, ipAddress, port, true);
+        this.setDC(mainDcId, ipAddress, port, isTest, true);
 
         await Promise.all(Object.keys(keys)
             .map(async (dcId) => {
@@ -56,10 +57,11 @@ class CallbackSession extends MemorySession {
             }));
     }
 
-    setDC(dcId, serverAddress, port, skipOnUpdate = false) {
+    setDC(dcId, serverAddress, port, isTestServer, skipOnUpdate = false) {
         this._dcId = dcId;
         this._serverAddress = serverAddress;
         this._port = port;
+        this._isTestServer = isTestServer;
 
         delete this._authKeys[dcId];
 
@@ -83,6 +85,7 @@ class CallbackSession extends MemorySession {
             mainDcId: this._dcId,
             keys: {},
             hashes: {},
+            isTest: this._isTestServer || undefined,
         };
 
         Object

src/lib/gramjs/sessions/Memory.js

@@ -8,6 +8,7 @@ class MemorySession extends Session {
         this._dcId = 0;
         this._port = undefined;
         this._takeoutId = undefined;
+        this._isTestServer = false;
 
         this._entities = new Set();
         this._updateStates = {};
@@ -33,222 +34,16 @@ class MemorySession extends Session {
         this._authKey = value;
     }
 
-    setDC(dcId, serverAddress, port) {
+    get isTestServer() {
+        return this._isTestServer;
+    }
+
+    setDC(dcId, serverAddress, port, isTestServer) {
         this._dcId = dcId | 0;
         this._serverAddress = serverAddress;
         this._port = port;
+        this._isTestServer = isTestServer;
     }
-
-    /* CONTEST
-    get takeoutId() {
-        return this._takeoutId
-    }
-
-    set takeoutId(value) {
-        this._takeoutId = value
-    }
-
-    getUpdateState(entityId) {
-        return this._updateStates[entityId]
-    }
-
-    setUpdateState(entityId, state) {
-        return this._updateStates[entityId] = state
-    }
-
-    close() {
-    }
-
-    save() {
-    }
-
-    async load() {
-
-    }
-
-    delete() {
-    }
-
-    _entityValuesToRow(id, hash, username, phone, name) {
-        // While this is a simple implementation it might be overrode by,
-        // other classes so they don't need to implement the plural form
-        // of the method. Don't remove.
-        return [id, hash, username, phone, name]
-    }
-
-    _entityToRow(e) {
-        if (!(e.classType === "constructor")) {
-            return
-        }
-        let p
-        let markedId
-        try {
-            p = utils.getInputPeer(e, false)
-            markedId = utils.getPeerId(p)
-        } catch (e) {
-            // Note: `get_input_peer` already checks for non-zero `accessHash`.
-            // See issues #354 and #392. It also checks that the entity
-            // is not `min`, because its `accessHash` cannot be used
-            // anywhere (since layer 102, there are two access hashes).
-            return
-        }
-        let pHash
-        if (p instanceof types.InputPeerUser || p instanceof types.InputPeerChannel) {
-            pHash = p.accessHash
-        } else if (p instanceof types.InputPeerChat) {
-            pHash = 0
-        } else {
-            return
-        }
-
-        let username = e.username
-        if (username) {
-            username = username.toLowerCase()
-        }
-        const phone = e.phone
-        const name = utils.getDisplayName(e)
-        return this._entityValuesToRow(markedId, pHash, username, phone, name)
-    }
-
-    _entitiesToRows(tlo) {
-        let entities = []
-        if (tlo.classType === "constructor" && utils.isListLike(tlo)) {
-            // This may be a list of users already for instance
-            entities = tlo
-        } else {
-            if (tlo instanceof Object) {
-                if ('user' in tlo) {
-                    entities.push(tlo.user)
-                }
-                if ('chats' in tlo && utils.isListLike(tlo.chats)) {
-                    entities.concat(tlo.chats)
-                }
-                if ('users' in tlo && utils.isListLike(tlo.users)) {
-                    entities.concat(tlo.users)
-                }
-            }
-        }
-        const rows = [] // Rows to add (id, hash, username, phone, name)
-        for (const e of entities) {
-            const row = this._entityToRow(e)
-            if (row) {
-                rows.push(row)
-            }
-        }
-        return rows
-    }
-
-    processEntities(tlo) {
-        const entitiesSet = this._entitiesToRows(tlo)
-        for (const e of entitiesSet) {
-            this._entities.add(e)
-        }
-    }
-
-    getEntityRowsByPhone(phone) {
-        for (const e of this._entities) { // id, hash, username, phone, name
-            if (e[3] === phone) {
-                return [e[0], e[1]]
-            }
-        }
-    }
-
-    getEntityRowsByUsername(username) {
-        for (const e of this._entities) { // id, hash, username, phone, name
-            if (e[2] === username) {
-                return [e[0], e[1]]
-            }
-        }
-    }
-
-    getEntityRowsByName(name) {
-        for (const e of this._entities) { // id, hash, username, phone, name
-            if (e[4] === name) {
-                return [e[0], e[1]]
-            }
-        }
-    }
-
-    getEntityRowsById(id, exact = true) {
-        if (exact) {
-            for (const e of this._entities) { // id, hash, username, phone, name
-                if (e[0] === id) {
-                    return [e[0], e[1]]
-                }
-            }
-        } else {
-            const ids = [utils.getPeerId(new types.PeerUser({ userId: id })),
-                utils.getPeerId(new types.PeerChat({ chatId: id })),
-                utils.getPeerId(new types.PeerChannel({ channelId: id })),
-            ]
-            for (const e of this._entities) { // id, hash, username, phone, name
-                if (ids.includes(e[0])) {
-                    return [e[0], e[1]]
-                }
-            }
-        }
-    }
-
-    getInputEntity(key) {
-        let exact
-        if (key.SUBCLASS_OF_ID !== undefined) {
-            if ([0xc91c90b6, 0xe669bf46, 0x40f202fd].includes(key.SUBCLASS_OF_ID)) {
-                // hex(crc32(b'InputPeer', b'InputUser' and b'InputChannel'))
-                // We already have an Input version, so nothing else required
-                return key
-            }
-            // Try to early return if this key can be casted as input peer
-            return utils.getInputPeer(key)
-        } else {
-            // Not a TLObject or can't be cast into InputPeer
-            if (key.classType === 'constructor') {
-                key = utils.getPeerId(key)
-                exact = true
-            } else {
-                exact = !(typeof key == 'number') || key < 0
-            }
-        }
-        let result = null
-        if (typeof key === 'string') {
-            const phone = utils.parsePhone(key)
-            if (phone) {
-                result = this.getEntityRowsByPhone(phone)
-            } else {
-                const { username, isInvite } = utils.parseUsername(key)
-                if (username && !isInvite) {
-                    result = this.getEntityRowsByUsername(username)
-                } else {
-                    const tup = utils.resolveInviteLink(key)[1]
-                    if (tup) {
-                        result = this.getEntityRowsById(tup, false)
-                    }
-                }
-            }
-        } else if (typeof key === 'number') {
-            result = this.getEntityRowsById(key, exact)
-        }
-        if (!result && typeof key === 'string') {
-            result = this.getEntityRowsByName(key)
-        }
-
-        if (result) {
-            let entityId = result[0] // unpack resulting tuple
-            const entityHash = result[1]
-            const resolved = utils.resolveId(entityId)
-            entityId = resolved[0]
-            const kind = resolved[1]
-            // removes the mark and returns type of entity
-            if (kind === types.PeerUser) {
-                return new types.InputPeerUser({ userId: entityId, accessHash: entityHash })
-            } else if (kind === types.PeerChat) {
-                return new types.InputPeerChat({ chatId: entityId })
-            } else if (kind === types.PeerChannel) {
-                return new types.InputPeerChannel({ channelId: entityId, accessHash: entityHash })
-            }
-        } else {
-            throw new Error('Could not find input entity with key ' + key)
-        }
-    } */
 }
 
 module.exports = MemorySession;

src/util/sessions.ts

@@ -27,9 +27,15 @@ export function hasStoredSession() {
 }
 
 export function storeSession(sessionData: ApiSessionData, currentUserId?: string) {
-  const { mainDcId, keys, hashes } = sessionData;
+  const {
+    mainDcId, keys, hashes, isTest,
+  } = sessionData;
 
-  localStorage.setItem(SESSION_USER_KEY, JSON.stringify({ dcID: mainDcId, id: currentUserId }));
+  localStorage.setItem(SESSION_USER_KEY, JSON.stringify({
+    dcID: mainDcId,
+    id: currentUserId,
+    test: isTest,
+  }));
   localStorage.setItem('dc', String(mainDcId));
   Object.keys(keys).map(Number).forEach((dcId) => {
     localStorage.setItem(`dc${dcId}_auth_key`, JSON.stringify(keys[dcId]));
@@ -64,6 +70,7 @@ export function loadStoredSession(): ApiSessionData | undefined {
     return undefined;
   }
   const mainDcId = Number(userAuth.dcID);
+  const isTest = userAuth.test;
   const keys: Record<number, string> = {};
   const hashes: Record<number, string> = {};
 
@@ -93,6 +100,7 @@ export function loadStoredSession(): ApiSessionData | undefined {
     mainDcId,
     keys,
     hashes,
+    isTest,
   };
 }