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

fix: dont export everything on first connection

Browse source
This commit is contained in:
alina 🌸 2023-09-18 03:16:29 +03:00
parent 55edbde3e7
commit 7bf63b2507
Signed by: teidesu
SSH key fingerprint: SHA256:uNeCpw6aTSU4aIObXLvHfLkDa82HWH9EiOj9AXOIRpI
11 changed files with 256 additions and 140 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
packages/client/src/methods/auth/start-test.ts
View file

@ -76,12 +76,16 @@ export async function startTest(
}
const id = parseInt(phone[5])
if (!availableDcs.find((dc) => dc.id === id)) { throw new MtArgumentError(`${phone} has invalid DC ID (${id})`) }
if (!availableDcs.find((dc) => dc.id === id)) {
throw new MtArgumentError(`${phone} has invalid DC ID (${id})`)
}
} else {
let dcId = this._defaultDc.id
let dcId = this._defaultDcs.main.id
if (params.dcId) {
if (!availableDcs.find((dc) => dc.id === params!.dcId)) { throw new MtArgumentError(`DC ID is invalid (${dcId})`) }
if (!availableDcs.find((dc) => dc.id === params!.dcId)) {
throw new MtArgumentError(`DC ID is invalid (${dcId})`)
}
dcId = params.dcId
}

2
packages/client/src/methods/files/download-iterable.ts
View file

@ -74,7 +74,7 @@ export async function* downloadAsIterable(
const isWeb = tl.isAnyInputWebFileLocation(location)
// we will receive a FileMigrateError in case this is invalid
if (!dcId) dcId = this._defaultDc.id
if (!dcId) dcId = this._defaultDcs.main.id
const partSizeKb =
params.partSize ?? (fileSize ? determinePartSize(fileSize) : 64)

37
packages/core/src/base-client.ts
View file

@ -77,7 +77,7 @@ export interface BaseTelegramClientOptions {
* When session already contains primary DC, this parameter is ignored.
* Defaults to Production DC 2.
*/
defaultDc?: tl.RawDcOption
defaultDcs?: ITelegramStorage.DcOptions
/**
* Whether to connect to test servers.
@ -211,10 +211,10 @@ export class BaseTelegramClient extends EventEmitter {
protected readonly _testMode: boolean
/**
* Primary DC taken from {@link BaseTelegramClientOptions.defaultDc},
* Primary DCs taken from {@link BaseTelegramClientOptions.defaultDcs},
* loaded from session or changed by other means (like redirecting).
*/
protected _defaultDc: tl.RawDcOption
protected _defaultDcs: ITelegramStorage.DcOptions
private _niceStacks: boolean
readonly _layer: number
@ -264,7 +264,7 @@ export class BaseTelegramClient extends EventEmitter {
this._useIpv6 = Boolean(opts.useIpv6)
this._testMode = Boolean(opts.testMode)
let dc = opts.defaultDc
let dc = opts.defaultDcs
if (!dc) {
if (this._testMode) {
@ -276,7 +276,7 @@ export class BaseTelegramClient extends EventEmitter {
}
}
this._defaultDc = dc
this._defaultDcs = dc
this._niceStacks = opts.niceStacks ?? true
this._layer = opts.overrideLayer ?? tl.LAYER
@ -348,11 +348,11 @@ export class BaseTelegramClient extends EventEmitter {
const promise = (this._connected = createControllablePromise())
await this._loadStorage()
const primaryDc = await this.storage.getDefaultDc()
if (primaryDc !== null) this._defaultDc = primaryDc
const primaryDc = await this.storage.getDefaultDcs()
if (primaryDc !== null) this._defaultDcs = primaryDc
const defaultDcAuthKey = await this.storage.getAuthKeyFor(
this._defaultDc.id,
this._defaultDcs.main.id,
)
if ((this._importForce || !defaultDcAuthKey) && this._importFrom) {
@ -369,21 +369,24 @@ export class BaseTelegramClient extends EventEmitter {
)
}
this._defaultDc = data.primaryDc
await this.storage.setDefaultDc(data.primaryDc)
this._defaultDcs = data.primaryDcs
await this.storage.setDefaultDcs(data.primaryDcs)
if (data.self) {
await this.storage.setSelf(data.self)
}
// await this.primaryConnection.setupKeys(data.authKey)
await this.storage.setAuthKeyFor(data.primaryDc.id, data.authKey)
await this.storage.setAuthKeyFor(
data.primaryDcs.main.id,
data.authKey,
)
await this._saveStorage(true)
}
this.network
.connect(this._defaultDc)
.connect(this._defaultDcs)
.then(() => {
promise.resolve()
this._connected = true
@ -574,17 +577,17 @@ export class BaseTelegramClient extends EventEmitter {
* > with [@BotFather](//t.me/botfather)
*/
async exportSession(): Promise<string> {
const primaryDc = await this.storage.getDefaultDc()
if (!primaryDc) throw new Error('No default DC set')
const primaryDcs = await this.storage.getDefaultDcs()
if (!primaryDcs) throw new Error('No default DC set')
const authKey = await this.storage.getAuthKeyFor(primaryDc.id)
const authKey = await this.storage.getAuthKeyFor(primaryDcs.main.id)
if (!authKey) throw new Error('Auth key is not ready yet')
return writeStringSession(this._writerMap, {
version: 1,
version: 2,
self: await this.storage.getSelf(),
testMode: this._testMode,
primaryDc,
primaryDcs,
authKey,
})
}

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

@ -155,6 +155,7 @@ export class MultiSessionConnection extends EventEmitter {
isMainConnection: this.params.isMainConnection && i === 0,
withUpdates:
this.params.isMainConnection &&
this.params.isMainDcConnection &&
!this.params.disableUpdates,
},
session,

165
packages/core/src/network/network-manager.ts
View file

@ -164,7 +164,7 @@ export class DcConnectionManager {
crypto: this.manager.params.crypto,
initConnection: this.manager._initConnectionParams,
transportFactory: this.manager._transportFactory,
dc: this._dc,
dc: this._dcs.media,
testMode: this.manager.params.testMode,
reconnectionStrategy: this.manager._reconnectionStrategy,
layer: this.manager.params.layer,
@ -173,6 +173,7 @@ export class DcConnectionManager {
writerMap: this.manager.params.writerMap,
usePfs: this.manager.params.usePfs,
isMainConnection: false,
isMainDcConnection: this.isPrimary,
inactivityTimeout: this.manager.params.inactivityTimeout ?? 60_000,
})
@ -184,7 +185,7 @@ export class DcConnectionManager {
this.__baseConnectionParams(),
this.manager._connectionCount(
'upload',
this._dc.id,
this.dcId,
this.manager.params.isPremium,
),
this._log,
@ -195,7 +196,7 @@ export class DcConnectionManager {
this.__baseConnectionParams(),
this.manager._connectionCount(
'download',
this._dc.id,
this.dcId,
this.manager.params.isPremium,
),
this._log,
@ -206,7 +207,7 @@ export class DcConnectionManager {
this.__baseConnectionParams(),
this.manager._connectionCount(
'downloadSmall',
this._dc.id,
this.dcId,
this.manager.params.isPremium,
),
this._log,
@ -222,13 +223,14 @@ export class DcConnectionManager {
constructor(
readonly manager: NetworkManager,
readonly dcId: number,
readonly _dc: tl.RawDcOption,
readonly _dcs: ITelegramStorage.DcOptions,
public isPrimary = false,
) {
this._log.prefix = `[DC ${dcId}] `
const mainParams = this.__baseConnectionParams()
mainParams.isMainConnection = true
mainParams.dc = _dcs.main
if (isPrimary) {
mainParams.inactivityTimeout = undefined
@ -361,17 +363,13 @@ export class DcConnectionManager {
setIsPremium(isPremium: boolean): void {
this.upload.setCount(
this.manager._connectionCount('upload', this._dc.id, isPremium),
this.manager._connectionCount('upload', this.dcId, isPremium),
)
this.download.setCount(
this.manager._connectionCount('download', this._dc.id, isPremium),
this.manager._connectionCount('download', this.dcId, isPremium),
)
this.downloadSmall.setCount(
this.manager._connectionCount(
'downloadSmall',
this._dc.id,
isPremium,
),
this.manager._connectionCount('downloadSmall', this.dcId, isPremium),
)
}
@ -481,6 +479,33 @@ export class NetworkManager {
config.onConfigUpdate(this._onConfigChanged)
}
private async _findDcOptions(
dcId: number,
): Promise<ITelegramStorage.DcOptions> {
const main = await this.config.findOption({
dcId,
allowIpv6: this.params.useIpv6,
preferIpv6: this.params.useIpv6,
allowMedia: false,
cdn: false,
})
const media = await this.config.findOption({
dcId,
allowIpv6: this.params.useIpv6,
preferIpv6: this.params.useIpv6,
allowMedia: true,
preferMedia: true,
cdn: false,
})
if (!main || !media) {
throw new Error(`Could not find DC ${dcId}`)
}
return { main, media }
}
private _switchPrimaryDc(dc: DcConnectionManager) {
if (this._primaryDc && this._primaryDc !== dc) {
this._primaryDc.setIsPrimary(false)
@ -536,19 +561,9 @@ export class NetworkManager {
this._log.debug('creating new DC %d', dcId)
try {
const dcOption = await this.config.findOption({
dcId,
allowIpv6: this.params.useIpv6,
preferIpv6: this.params.useIpv6,
allowMedia: true,
preferMedia: true,
cdn: false,
})
const dcOptions = await this._findDcOptions(dcId)
if (!dcOption) {
throw new Error(`Could not find DC ${dcId}`)
}
const dc = new DcConnectionManager(this, dcId, dcOption)
const dc = new DcConnectionManager(this, dcId, dcOptions)
if (!(await dc.loadKeys())) {
dc.main.requestAuth()
@ -567,56 +582,66 @@ export class NetworkManager {
/**
* Perform initial connection to the default DC
*
* @param defaultDc Default DC to connect to
* @param defaultDcs Default DCs to connect to
*/
async connect(defaultDc: tl.RawDcOption): Promise<void> {
if (this._dcConnections[defaultDc.id]) {
async connect(defaultDcs: ITelegramStorage.DcOptions): Promise<void> {
if (defaultDcs.main.id !== defaultDcs.media.id) {
throw new Error('Default DCs must be the same')
}
if (this._dcConnections[defaultDcs.main.id]) {
// shouldn't happen
throw new Error('DC manager already exists')
}
const dc = new DcConnectionManager(this, defaultDc.id, defaultDc)
this._dcConnections[defaultDc.id] = dc
const dc = new DcConnectionManager(this, defaultDcs.main.id, defaultDcs)
this._dcConnections[defaultDcs.main.id] = dc
await this._switchPrimaryDc(dc)
}
private _pendingExports: Record<number, Promise<void>> = {}
private async _exportAuthTo(manager: DcConnectionManager): Promise<void> {
const auth = await this.call({
_: 'auth.exportAuthorization',
dcId: manager.dcId,
})
if (manager.dcId in this._pendingExports) {
this._log.debug('waiting for auth export to dc %d', manager.dcId)
const res = await this.call(
{
_: 'auth.importAuthorization',
id: auth.id,
bytes: auth.bytes,
},
{ manager },
)
return this._pendingExports[manager.dcId]
}
if (res._ !== 'auth.authorization') {
throw new Error(
`Unexpected response from auth.importAuthorization: ${res._}`,
this._log.debug('exporting auth to dc %d', manager.dcId)
const promise = createControllablePromise<void>()
this._pendingExports[manager.dcId] = promise
try {
const auth = await this.call({
_: 'auth.exportAuthorization',
dcId: manager.dcId,
})
const res = await this.call(
{
_: 'auth.importAuthorization',
id: auth.id,
bytes: auth.bytes,
},
{ manager },
)
}
}
async exportAuth(): Promise<void> {
const dcs: Record<number, number> = {}
const config = await this.config.get()
if (res._ !== 'auth.authorization') {
throw new Error(
`Unexpected response from auth.importAuthorization: ${res._}`,
)
}
for (const dc of config.dcOptions) {
if (dc.cdn) continue
dcs[dc.id] = dc.id
}
for (const dc of Object.values(dcs)) {
if (dc === this._primaryDc!.dcId) continue
this._log.debug('exporting auth for dc %d', dc)
const manager = await this._getOtherDc(dc)
await this._exportAuthTo(manager)
promise.resolve()
delete this._pendingExports[manager.dcId]
} catch (e) {
this._log.warn(
'failed to export auth to dc %d: %s',
manager.dcId,
e,
)
promise.reject(e)
throw e
}
}
@ -628,6 +653,8 @@ export class NetworkManager {
})
}
// future-proofing. should probably remove once the implementation is stable
// eslint-disable-next-line @typescript-eslint/require-await
async notifyLoggedIn(auth: tl.auth.TypeAuthorization): Promise<void> {
if (
auth._ === 'auth.authorizationSignUpRequired' ||
@ -642,7 +669,7 @@ export class NetworkManager {
this.setIsPremium(auth.user.premium!)
await this.exportAuth()
// await this.exportAuth()
}
resetSessions(): void {
@ -664,27 +691,17 @@ export class NetworkManager {
async changePrimaryDc(newDc: number): Promise<void> {
if (newDc === this._primaryDc?.dcId) return
const option = await this.config.findOption({
dcId: newDc,
allowIpv6: this.params.useIpv6,
preferIpv6: this.params.useIpv6,
cdn: false,
allowMedia: false,
})
if (!option) {
throw new Error(`DC ${newDc} not found`)
}
const options = await this._findDcOptions(newDc)
if (!this._dcConnections[newDc]) {
this._dcConnections[newDc] = new DcConnectionManager(
this,
newDc,
option,
options,
)
}
await this._storage.setDefaultDc(option)
await this._storage.setDefaultDcs(options)
await this._switchPrimaryDc(this._dcConnections[newDc])
}

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

@ -41,6 +41,7 @@ export interface SessionConnectionParams extends PersistentConnectionParams {
disableUpdates?: boolean
withUpdates?: boolean
isMainConnection: boolean
isMainDcConnection: boolean
usePfs?: boolean
readerMap: TlReaderMap

16
packages/core/src/storage/abstract.ts
View file

@ -21,6 +21,11 @@ export namespace ITelegramStorage {
isBot: boolean
userId: number
}
export interface DcOptions {
main: tl.RawDcOption
media: tl.RawDcOption
}
}
/**
@ -68,12 +73,12 @@ export interface ITelegramStorage {
/**
* Set default datacenter to use with this session.
*/
setDefaultDc(dc: tl.RawDcOption | null): MaybeAsync<void>
setDefaultDcs(dcs: ITelegramStorage.DcOptions | null): MaybeAsync<void>
/**
* Get default datacenter for this session
* (by default should return null)
*/
getDefaultDc(): MaybeAsync<tl.RawDcOption | null>
getDefaultDcs(): MaybeAsync<ITelegramStorage.DcOptions | null>
/**
* Get auth_key for a given DC
@ -92,7 +97,12 @@ export interface ITelegramStorage {
* Set temp_auth_key for a given DC
* expiresAt is unix time in ms
*/
setTempAuthKeyFor(dcId: number, index: number, key: Buffer | null, expiresAt: number): MaybeAsync<void>
setTempAuthKeyFor(
dcId: number,
index: number,
key: Buffer | null,
expiresAt: number
): MaybeAsync<void>
/**
* Remove all saved auth keys (both temp and perm)
* for the given DC. Used when perm_key becomes invalid,

12
packages/core/src/storage/memory.ts
View file

@ -13,7 +13,7 @@ export interface MemorySessionState {
// forwards compatibility for persistent storages
$version: typeof CURRENT_VERSION
defaultDc: tl.RawDcOption | null
defaultDcs: ITelegramStorage.DcOptions | null
authKeys: Record<number, Buffer | null>
authKeysTemp: Record<string, Buffer | null>
authKeysTempExpiry: Record<string, number>
@ -110,7 +110,7 @@ export class MemoryStorage implements ITelegramStorage, IStateStorage {
reset(): void {
this._state = {
$version: CURRENT_VERSION,
defaultDc: null,
defaultDcs: null,
authKeys: {},
authKeysTemp: {},
authKeysTempExpiry: {},
@ -183,12 +183,12 @@ export class MemoryStorage implements ITelegramStorage, IStateStorage {
})
}
getDefaultDc(): tl.RawDcOption | null {
return this._state.defaultDc
getDefaultDcs(): ITelegramStorage.DcOptions | null {
return this._state.defaultDcs
}
setDefaultDc(dc: tl.RawDcOption | null): void {
this._state.defaultDc = dc
setDefaultDcs(dcs: ITelegramStorage.DcOptions | null): void {
this._state.defaultDcs = dcs
}
setTempAuthKeyFor(

82
packages/core/src/utils/default-dcs.ts
View file

@ -1,37 +1,73 @@
import { tl } from '@mtcute/tl'
import { ITelegramStorage } from '../storage'
/** @internal */
export const defaultProductionDc: tl.RawDcOption = {
_: 'dcOption',
ipAddress: '149.154.167.50',
port: 443,
id: 2,
export const defaultProductionDc: ITelegramStorage.DcOptions = {
main: {
_: 'dcOption',
ipAddress: '149.154.167.50',
port: 443,
id: 2,
},
media: {
_: 'dcOption',
ipAddress: '149.154.167.222',
port: 443,
id: 2,
mediaOnly: true,
},
}
/** @internal */
export const defaultProductionIpv6Dc: tl.RawDcOption = {
_: 'dcOption',
ipAddress: '2001:67c:4e8:f002::a',
ipv6: true,
port: 443,
id: 2,
export const defaultProductionIpv6Dc: ITelegramStorage.DcOptions = {
main: {
_: 'dcOption',
ipAddress: '2001:067c:04e8:f002:0000:0000:0000:000a',
ipv6: true,
port: 443,
id: 2,
},
media: {
_: 'dcOption',
ipAddress: '2001:067c:04e8:f002:0000:0000:0000:000b',
ipv6: true,
mediaOnly: true,
port: 443,
id: 2,
},
}
/** @internal */
export const defaultTestDc: tl.RawDcOption = {
_: 'dcOption',
ipAddress: '149.154.167.40',
port: 443,
id: 2,
export const defaultTestDc: ITelegramStorage.DcOptions = {
main: {
_: 'dcOption',
ipAddress: '149.154.167.40',
port: 443,
id: 2,
},
media: {
_: 'dcOption',
ipAddress: '149.154.167.40',
port: 443,
id: 2,
},
}
/** @internal */
export const defaultTestIpv6Dc: tl.RawDcOption = {
_: 'dcOption',
ipAddress: '2001:67c:4e8:f002::e',
port: 443,
ipv6: true,
id: 2,
export const defaultTestIpv6Dc: ITelegramStorage.DcOptions = {
main: {
_: 'dcOption',
ipAddress: '2001:67c:4e8:f002::e',
port: 443,
ipv6: true,
id: 2,
},
media: {
_: 'dcOption',
ipAddress: '2001:67c:4e8:f002::e',
port: 443,
ipv6: true,
id: 2,
},
}
export const defaultDcs = {

42
packages/core/src/utils/string-session.ts
View file

@ -1,5 +1,10 @@
import { tl } from '@mtcute/tl'
import { TlBinaryReader, TlBinaryWriter, TlReaderMap, TlWriterMap } from '@mtcute/tl-runtime'
import {
TlBinaryReader,
TlBinaryWriter,
TlReaderMap,
TlWriterMap,
} from '@mtcute/tl-runtime'
import { ITelegramStorage } from '../storage'
import { encodeUrlSafeBase64, parseUrlSafeBase64 } from './buffer-utils'
@ -7,7 +12,7 @@ import { encodeUrlSafeBase64, parseUrlSafeBase64 } from './buffer-utils'
export interface StringSessionData {
version: number
testMode: boolean
primaryDc: tl.TypeDcOption
primaryDcs: ITelegramStorage.DcOptions
self?: ITelegramStorage.SelfInfo | null
authKey: Buffer
}
@ -20,7 +25,7 @@ export function writeStringSession(
const version = data.version
if (version !== 1) {
if (version !== 1 && version !== 2) {
throw new Error(`Unsupported string session version: ${version}`)
}
@ -38,7 +43,12 @@ export function writeStringSession(
writer.pos += 1
writer.int(flags)
writer.object(data.primaryDc)
writer.object(data.primaryDcs.main)
if (version >= 2 && data.primaryDcs.media !== data.primaryDcs.main) {
flags |= 4
writer.object(data.primaryDcs.media)
}
if (data.self) {
writer.int53(data.self.userId)
@ -50,23 +60,32 @@ export function writeStringSession(
return encodeUrlSafeBase64(writer.result())
}
export function readStringSession(readerMap: TlReaderMap, data: string): StringSessionData {
export function readStringSession(
readerMap: TlReaderMap,
data: string,
): StringSessionData {
const buf = parseUrlSafeBase64(data)
if (buf[0] !== 1) { throw new Error(`Invalid session string (version = ${buf[0]})`) }
const version = buf[0]
if (version !== 1 && version !== 2) {
throw new Error(`Invalid session string (version = ${version})`)
}
const reader = new TlBinaryReader(readerMap, buf, 1)
const flags = reader.int()
const hasSelf = flags & 1
const testMode = Boolean(flags & 2)
const hasMedia = version >= 2 && Boolean(flags & 4)
const primaryDc = reader.object() as tl.TypeDcOption
const primaryMediaDc = hasMedia ?
(reader.object() as tl.TypeDcOption) :
primaryDc
if (primaryDc._ !== 'dcOption') {
throw new Error(
`Invalid session string (dc._ = ${primaryDc._})`,
)
throw new Error(`Invalid session string (dc._ = ${primaryDc._})`)
}
let self: ITelegramStorage.SelfInfo | null = null
@ -86,7 +105,10 @@ export function readStringSession(readerMap: TlReaderMap, data: string): StringS
return {
version: 1,
testMode,
primaryDc,
primaryDcs: {
main: primaryDc,
media: primaryMediaDc,
},
self,
authKey: key,
}

28
packages/sqlite/index.ts
View file

@ -54,7 +54,7 @@ function getInputPeer(
throw new Error(`Invalid peer type: ${row.type}`)
}
const CURRENT_VERSION = 3
const CURRENT_VERSION = 4
// language=SQLite format=false
const TEMP_AUTH_TABLE = `
@ -392,6 +392,28 @@ export class SqliteStorage implements ITelegramStorage, IStateStorage {
from = 3
}
if (from === 3) {
// media dc support added
const oldDc = this._db
.prepare("select value from kv where key = 'def_dc'")
.get()
if (oldDc) {
const oldDcValue = JSON.parse(
(oldDc as { value: string }).value,
) as tl.RawDcOption
this._db
.prepare("update kv set value = ? where key = 'def_dc'")
.run([
JSON.stringify({
main: oldDcValue,
media: oldDcValue,
}),
])
}
from = 4
}
if (from !== CURRENT_VERSION) {
// an assertion just in case i messed up
throw new Error('Migration incomplete')
@ -499,11 +521,11 @@ export class SqliteStorage implements ITelegramStorage, IStateStorage {
this._db.exec(RESET)
}
setDefaultDc(dc: tl.RawDcOption | null): void {
setDefaultDcs(dc: ITelegramStorage.DcOptions | null): void {
return this._setToKv('def_dc', dc)
}
getDefaultDc(): tl.RawDcOption | null {
getDefaultDcs(): ITelegramStorage.DcOptions | null {
return this._getFromKv('def_dc')
}