feat(core): update connection count and pfs usage on the go

holy shit i actually made it 🫣

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

@@ -11,6 +11,7 @@ import { MtprotoSession } from './mtproto-session'
 export class MultiSessionConnection extends EventEmitter {
     private _log: Logger
     readonly _sessions: MtprotoSession[]
+    private _enforcePfs = false
 
     constructor(
         readonly params: SessionConnectionParams,
@@ -19,6 +20,7 @@ export class MultiSessionConnection extends EventEmitter {
     ) {
         super()
         this._log = log.create('multi')
+        this._enforcePfs = _count > 1 && params.isMainConnection
 
         this._sessions = []
         this._updateConnections()
@@ -28,10 +30,17 @@ export class MultiSessionConnection extends EventEmitter {
 
     setCount(count: number, doUpdate = true): void {
         this._count = count
-        if (doUpdate) this._updateConnections()
+
+        if (doUpdate) this._updateConnections(true)
     }
 
     private _updateSessions(): void {
+        this._log.debug(
+            'updating sessions count: %d -> %d',
+            this._sessions.length,
+            this._count
+        )
+
         // there are two cases
         // 1. this msc is main, in which case every connection should have its own session
         // 2. this msc is not main, in which case all connections should share the same session
@@ -85,10 +94,27 @@ export class MultiSessionConnection extends EventEmitter {
         }
     }
 
-    private _updateConnections(): void {
+    private _updateConnections(active = false): void {
         this._updateSessions()
         if (this._connections.length === this._count) return
 
+        this._log.debug(
+            'updating connections count: %d -> %d',
+            this._connections.length,
+            this._count
+        )
+
+        const newEnforcePfs = this._count > 1 && this.params.isMainConnection
+        const enforcePfsChanged = newEnforcePfs !== this._enforcePfs
+        if (enforcePfsChanged) {
+            this._log.debug(
+                'enforcePfs changed: %s -> %s',
+                this._enforcePfs,
+                newEnforcePfs
+            )
+            this._enforcePfs = newEnforcePfs
+        }
+
         if (this._connections.length > this._count) {
             // destroy extra connections
             for (let i = this._connections.length - 1; i >= this._count; i--) {
@@ -100,6 +126,12 @@ export class MultiSessionConnection extends EventEmitter {
             return
         }
 
+        if (enforcePfsChanged) {
+            this._connections.forEach((conn) => {
+                conn.setUsePfs(this.params.usePfs || this._enforcePfs)
+            })
+        }
+
         // create new connections
         for (let i = this._connections.length; i < this._count; i++) {
             const session = this.params.isMainConnection
@@ -108,6 +140,7 @@ export class MultiSessionConnection extends EventEmitter {
             const conn = new SessionConnection(
                 {
                     ...this.params,
+                    usePfs: this.params.usePfs || this._enforcePfs,
                     isMainConnection: this.params.isMainConnection && i === 0,
                 },
                 session
@@ -142,6 +175,7 @@ export class MultiSessionConnection extends EventEmitter {
             conn.on('request-auth', () => this.emit('request-auth', i))
 
             this._connections.push(conn)
+            if (active) conn.connect()
         }
     }
 

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

@@ -92,6 +92,7 @@ export class SessionConnection extends PersistentConnection {
     private _lastPingMsgId = Long.ZERO
     private _lastSessionCreatedUid = Long.ZERO
 
+    private _usePfs = this.params.usePfs ?? false
     private _isPfsBindingPending = false
     private _isPfsBindingPendingInBackground = false
     private _pfsUpdateTimeout?: NodeJS.Timeout
@@ -124,6 +125,21 @@ export class SessionConnection extends PersistentConnection {
         return key.key
     }
 
+    setUsePfs(usePfs: boolean): void {
+        if (this._usePfs === usePfs) return
+
+        this.log.debug('use pfs changed to %s', usePfs)
+        this._usePfs = usePfs
+        if (!usePfs) {
+            this._isPfsBindingPending = false
+            this._isPfsBindingPendingInBackground = false
+            this._session._authKeyTemp.reset()
+            clearTimeout(this._pfsUpdateTimeout!)
+        }
+
+        this._resetSession()
+    }
+
     onTransportClose(): void {
         super.onTransportClose()
 
@@ -168,7 +184,7 @@ export class SessionConnection extends PersistentConnection {
             return
         }
 
-        if (this.params.usePfs && !this._session._authKeyTemp.ready) {
+        if (this._usePfs && !this._session._authKeyTemp.ready) {
             this.log.info('no temp auth key but using pfs, authorizing')
             this._authorizePfs()
             return
@@ -185,7 +201,7 @@ export class SessionConnection extends PersistentConnection {
                 // if we are using pfs, this could be due to the server
                 // forgetting our temp key (which is kinda weird but expected)
 
-                if (this.params.usePfs) {
+                if (this._usePfs) {
                     if (
                         !this._isPfsBindingPending &&
                         this._session._authKeyTemp.ready
@@ -283,7 +299,7 @@ export class SessionConnection extends PersistentConnection {
 
                 this.emit('key-change', authKey)
 
-                if (this.params.usePfs) {
+                if (this._usePfs) {
                     return this._authorizePfs()
                 } else {
                     this.onConnectionUsable()
@@ -321,6 +337,11 @@ export class SessionConnection extends PersistentConnection {
 
         doAuthorization(this, this.params.crypto, TEMP_AUTH_KEY_EXPIRY)
             .then(async ([tempAuthKey, tempServerSalt]) => {
+                if (!this._usePfs) {
+                    this.log.info('pfs has been disabled while generating temp key')
+                    return
+                }
+
                 const tempKey = await this._session._authKeyTempSecondary
                 await tempKey.setup(tempAuthKey)
 
@@ -426,6 +447,11 @@ export class SessionConnection extends PersistentConnection {
 
                 this._session.pendingMessages.delete(msgId)
 
+                if (!this._usePfs) {
+                    this.log.info('pfs has been disabled while binding temp key')
+                    return
+                }
+
                 if (typeof res === 'object') {
                     this.log.error(
                         'failed to bind temp key: %s:%s',