fix(core): fixes for highload bots

packages/core/src/network/multi-session-connection.ts

@@ -2,7 +2,7 @@ import EventEmitter from 'events'
 
 import { tl } from '@mtcute/tl'
 
-import { Logger } from '../utils/index.js'
+import { createControllablePromise, Logger } from '../utils/index.js'
 import { MtprotoSession } from './mtproto-session.js'
 import { SessionConnection, SessionConnectionParams } from './session-connection.js'
 import { TransportFactory } from './transports/index.js'
@@ -124,9 +124,21 @@ export class MultiSessionConnection extends EventEmitter {
         }
 
         if (enforcePfsChanged) {
-            this._connections.forEach((conn) => {
-                conn.setUsePfs(this.params.usePfs || this._enforcePfs)
-            })
+            // we need to fetch new auth keys first
+            const promise = createControllablePromise<void>()
+            this.emit('request-keys', promise)
+
+            promise
+                .then(() => {
+                    this._connections.forEach((conn) => {
+                        conn.setUsePfs(this.params.usePfs || this._enforcePfs)
+
+                        if (connect) conn.connect()
+                    })
+                })
+                .catch((err) => {
+                    this.emit('error', err)
+                })
         }
 
         // create new connections
@@ -144,7 +156,7 @@ export class MultiSessionConnection extends EventEmitter {
                 session,
             )
 
-            if (this.params.isMainConnection) {
+            if (this.params.isMainConnection && this.params.isMainDcConnection) {
                 conn.on('update', (update) => this.emit('update', update))
             }
             conn.on('error', (err) => this.emit('error', err, conn))
@@ -178,7 +190,8 @@ export class MultiSessionConnection extends EventEmitter {
             })
 
             this._connections.push(conn)
-            if (connect) conn.connect()
+            // if enforcePfsChanged, we need to connect after setting the new auth key
+            if (connect && !enforcePfsChanged) conn.connect()
         }
     }
 

packages/core/src/network/network-manager.ts

@@ -3,7 +3,7 @@ import { TlReaderMap, TlWriterMap } from '@mtcute/tl-runtime'
 
 import { ITelegramStorage } from '../storage/index.js'
 import { MtArgumentError, MtcuteError } from '../types/index.js'
-import { createControllablePromise, ICryptoProvider, Logger, sleep } from '../utils/index.js'
+import { ControllablePromise, createControllablePromise, ICryptoProvider, Logger, sleep } from '../utils/index.js'
 import { assertTypeIs } from '../utils/type-assertions.js'
 import { ConfigManager } from './config-manager.js'
 import { MultiSessionConnection } from './multi-session-connection.js'
@@ -303,6 +303,13 @@ export class DcConnectionManager {
             this.main.requestAuth()
         })
 
+        // fucking awesome architecture, but whatever
+        connection.on('request-keys', (promise: ControllablePromise<void>) => {
+            this.loadKeys(true)
+                .then(() => promise.resolve())
+                .catch((e: Error) => promise.reject(e))
+        })
+
         connection.on('error', (err: Error, conn: SessionConnection) => {
             this.manager.params._emitError(err, conn)
         })
@@ -311,6 +318,7 @@ export class DcConnectionManager {
     setIsPrimary(isPrimary: boolean): void {
         if (this.isPrimary === isPrimary) return
         this.isPrimary = isPrimary
+        this.main.params.isMainDcConnection = isPrimary
 
         if (isPrimary) {
             this.main.setInactivityTimeout(undefined)
@@ -325,7 +333,7 @@ export class DcConnectionManager {
         this.downloadSmall.setCount(this.manager._connectionCount('downloadSmall', this.dcId, isPremium))
     }
 
-    async loadKeys(): Promise<boolean> {
+    async loadKeys(forcePfs = false): Promise<boolean> {
         const permanent = await this.manager._storage.getAuthKeyFor(this.dcId)
 
         this.main.setAuthKey(permanent)
@@ -337,7 +345,7 @@ export class DcConnectionManager {
             return false
         }
 
-        if (this.manager.params.usePfs) {
+        if (this.manager.params.usePfs || forcePfs) {
             await Promise.all(
                 this.main._sessions.map(async (_, i) => {
                     const temp = await this.manager._storage.getAuthKeyFor(this.dcId, i)

packages/core/src/network/session-connection.ts

@@ -312,6 +312,7 @@ export class SessionConnection extends PersistentConnection {
             this._isPfsBindingPending = true
         }
 
+        process.exit(0)
         doAuthorization(this, this._crypto, TEMP_AUTH_KEY_EXPIRY)
             .then(async ([tempAuthKey, tempServerSalt]) => {
                 if (!this._usePfs) {

packages/core/src/utils/controllable-promise.ts

@@ -6,11 +6,6 @@ export type ControllablePromise<T = unknown> = Promise<T> & {
     reject(err?: unknown): void
 }
 
-/**
- * The promise was cancelled
- */
-export class PromiseCancelledError extends Error {}
-
 /**
  * Creates a promise that can be resolved or rejected from outside.
  */