From 0b4edbf70e058517d9dc6019a45993a9f640bd14 Mon Sep 17 00:00:00 2001 From: teidesu <86301490+teidesu@users.noreply.github.com> Date: Mon, 14 Nov 2022 03:24:18 +0300 Subject: [PATCH] feat(core): update connection count and pfs usage on the go MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit holy shit i actually made it 🫣 --- .../src/network/multi-session-connection.ts | 38 ++++++++++++++++++- .../core/src/network/session-connection.ts | 32 ++++++++++++++-- 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/packages/core/src/network/multi-session-connection.ts b/packages/core/src/network/multi-session-connection.ts index e5388d76..be4b3892 100644 --- a/packages/core/src/network/multi-session-connection.ts +++ b/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() } } diff --git a/packages/core/src/network/session-connection.ts b/packages/core/src/network/session-connection.ts index ccc279d8..320736b2 100644 --- a/packages/core/src/network/session-connection.ts +++ b/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',