teidesu/mtcute
1
0
Fork 0
Code Issues 15 Pull requests 1 Projects Releases Packages Wiki Activity Actions

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

Browse source 
holy shit i actually made it 🫣
This commit is contained in:
teidesu 2022-11-14 03:24:18 +03:00 committed by Alina Sireneva
parent f68d83da06
commit 0b4edbf70e
Signed by: teidesu
SSH key fingerprint: SHA256:uNeCpw6aTSU4aIObXLvHfLkDa82HWH9EiOj9AXOIRpI
2 changed files with 65 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
packages/core/src/network/multi-session-connection.ts
View file

@ -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()
}
}

32
packages/core/src/network/session-connection.ts
View file

@ -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',