mtcute/packages/core/src/network/auth-key.ts

import Long from 'long'

import { tl } from '@mtcute/tl'
import { TlBinaryReader, TlReaderMap } from '@mtcute/tl-runtime'

import { MtcuteError } from '../types'
import { buffersEqual, ICryptoProvider, Logger, randomBytes } from '../utils'
import { createAesIgeForMessage } from '../utils/crypto/mtproto'

export class AuthKey {
    ready = false

    key!: Buffer
    id!: Buffer
    clientSalt!: Buffer
    serverSalt!: Buffer

    constructor(readonly _crypto: ICryptoProvider, readonly log: Logger, readonly _readerMap: TlReaderMap) {}

    match(keyId: Buffer): boolean {
        return this.ready && buffersEqual(keyId, this.id)
    }

    async setup(authKey?: Buffer | null): Promise<void> {
        if (!authKey) return this.reset()

        this.ready = true
        this.key = authKey
        this.clientSalt = authKey.slice(88, 120)
        this.serverSalt = authKey.slice(96, 128)
        this.id = (await this._crypto.sha1(authKey)).slice(-8)

        this.log.verbose('auth key set up, id = %h', this.id)
    }

    async encryptMessage(message: Buffer, serverSalt: Long, sessionId: Long): Promise<Buffer> {
        if (!this.ready) throw new MtcuteError('Keys are not set up!')

        let padding = (16 /* header size */ + message.length + 12) /* min padding */ % 16
        padding = 12 + (padding ? 16 - padding : 0)

        const buf = Buffer.alloc(16 + message.length + padding)

        buf.writeInt32LE(serverSalt.low)
        buf.writeInt32LE(serverSalt.high, 4)
        buf.writeInt32LE(sessionId.low, 8)
        buf.writeInt32LE(sessionId.high, 12)
        message.copy(buf, 16)
        randomBytes(padding).copy(buf, 16 + message.length)

        const messageKey = (await this._crypto.sha256(Buffer.concat([this.clientSalt, buf]))).slice(8, 24)
        const ige = await createAesIgeForMessage(this._crypto, this.key, messageKey, true)
        const encryptedData = await ige.encrypt(buf)

        return Buffer.concat([this.id, messageKey, encryptedData])
    }

    async decryptMessage(
        data: Buffer,
        sessionId: Long,
        callback: (msgId: tl.Long, seqNo: number, data: TlBinaryReader) => void,
    ): Promise<void> {
        const messageKey = data.slice(8, 24)
        const encryptedData = data.slice(24)

        const ige = await createAesIgeForMessage(this._crypto, this.key, messageKey, false)
        const innerData = await ige.decrypt(encryptedData)

        const expectedMessageKey = (await this._crypto.sha256(Buffer.concat([this.serverSalt, innerData]))).slice(8, 24)

        if (!buffersEqual(messageKey, expectedMessageKey)) {
            this.log.warn(
                '[%h] received message with invalid messageKey = %h (expected %h)',
                messageKey,
                expectedMessageKey,
            )

            return
        }

        const innerReader = new TlBinaryReader(this._readerMap, innerData)
        innerReader.seek(8) // skip salt
        const sessionId_ = innerReader.long()
        const messageId = innerReader.long(true)

        if (sessionId_.neq(sessionId)) {
            this.log.warn('ignoring message with invalid sessionId = %h', sessionId_)

            return
        }

        const seqNo = innerReader.uint()
        const length = innerReader.uint()

        if (length > innerData.length - 32 /* header size */) {
            this.log.warn('ignoring message with invalid length: %d > %d', length, innerData.length - 32)

            return
        }

        if (length % 4 !== 0) {
            this.log.warn('ignoring message with invalid length: %d is not a multiple of 4', length)

            return
        }

        const paddingSize = innerData.length - length - 32 // header size

        if (paddingSize < 12 || paddingSize > 1024) {
            this.log.warn('ignoring message with invalid padding size: %d', paddingSize)

            return
        }

        callback(messageId, seqNo, innerReader)
    }

    copyFrom(authKey: AuthKey): void {
        this.ready = authKey.ready
        this.key = authKey.key
        this.id = authKey.id
        this.serverSalt = authKey.serverSalt
        this.clientSalt = authKey.clientSalt
    }

    reset(): void {
        this.ready = false
    }
}