fix: auth storage fixes

- .reset() no longer resets auth keys by default
- auth keys are stored immediately in sqlite
- update loop fixes for logout
- tests for sqlite storage

likely closes #13 (?)
This commit is contained in:
alina 🌸 2023-11-12 07:42:51 +03:00
parent 38de001e8d
commit f525c12f83
Signed by: teidesu
SSH key fingerprint: SHA256:uNeCpw6aTSU4aIObXLvHfLkDa82HWH9EiOj9AXOIRpI
18 changed files with 254 additions and 67 deletions

View file

@ -421,7 +421,7 @@ async function main() {
}
for await (const file of getFiles(path.join(__dirname, '../src/methods'))) {
if (!file.startsWith('.') && file.endsWith('.ts')) {
if (!file.startsWith('.') && file.endsWith('.ts') && !file.endsWith('.web.ts')) {
await addSingleMethod(state, file)
}
}

View file

@ -6,6 +6,8 @@ import { TelegramClient } from '../client.js'
// @copy
import { Conversation } from '../types/conversation.js'
// @copy
import { logOut } from './auth/log-out.js'
// @copy
import { start } from './auth/start.js'
// @copy
import {

View file

@ -5,6 +5,7 @@ import { assertTypeIs } from '@mtcute/core/utils.js'
import { User } from '../../types/peers/user.js'
const STATE_SYMBOL = Symbol('authState')
/** @exported */
export interface AuthState {
// local copy of "self" in storage,
@ -70,7 +71,11 @@ export function getAuthState(client: BaseTelegramClient): AuthState {
}
/** @internal */
export function _onAuthorization(client: BaseTelegramClient, auth: tl.auth.TypeAuthorization, bot = false): User {
export async function _onAuthorization(
client: BaseTelegramClient,
auth: tl.auth.TypeAuthorization,
bot = false,
): Promise<User> {
if (auth._ === 'auth.authorizationSignUpRequired') {
throw new MtUnsupportedError(
'Signup is no longer supported by Telegram for non-official clients. Please use your mobile device to sign up.',
@ -86,6 +91,7 @@ export function _onAuthorization(client: BaseTelegramClient, auth: tl.auth.TypeA
state.selfChanged = true
client.notifyLoggedIn(auth)
await client.saveStorage()
// telegram ignores invokeWithoutUpdates for auth methods
if (client.network.params.disableUpdates) client.network.resetSessions()

View file

@ -861,6 +861,14 @@ function fetchDifferenceLater(
0,
fetchDifference(client, state, requestedDiff)
.catch((err) => {
if (tl.RpcError.is(err, 'AUTH_KEY_UNREGISTERED')) {
// for some reason, when logging out telegram may send updatesTooLong
// in any case, we need to stop updates loop
stopUpdatesLoop(client)
return
}
state.log.warn('error fetching common difference: %s', err)
})
.then(() => {

View file

@ -249,12 +249,10 @@ export class DcConnectionManager {
this.manager._log.debug('key change for dc %d from connection %d', this.dcId, idx)
// send key to other connections
Promise.all([
this.manager._storage.setAuthKeyFor(this.dcId, key),
this.upload.setAuthKey(key),
this.download.setAuthKey(key),
this.downloadSmall.setAuthKey(key),
])
this.upload.setAuthKey(key)
this.download.setAuthKey(key)
this.downloadSmall.setAuthKey(key)
Promise.resolve(this.manager._storage.setAuthKeyFor(this.dcId, key))
.then(() => {
this.upload.notifyKeyChange()
this.download.notifyKeyChange()
@ -272,12 +270,11 @@ export class DcConnectionManager {
this.manager._log.debug('temp key change for dc %d from connection %d', this.dcId, idx)
// send key to other connections
Promise.all([
this.manager._storage.setTempAuthKeyFor(this.dcId, idx, key, expires * 1000),
this.upload.setAuthKey(key, true),
this.download.setAuthKey(key, true),
this.downloadSmall.setAuthKey(key, true),
])
this.upload.setAuthKey(key, true)
this.download.setAuthKey(key, true)
this.downloadSmall.setAuthKey(key, true)
Promise.resolve(this.manager._storage.setTempAuthKeyFor(this.dcId, idx, key, expires * 1000))
.then(() => {
this.upload.notifyKeyChange()
this.download.notifyKeyChange()

View file

@ -78,9 +78,11 @@ export interface ITelegramStorage {
destroy?(): MaybeAsync<void>
/**
* Reset session to its default state
* Reset session to its default state, optionally resetting auth keys
*
* @param [withAuthKeys=false] Whether to also reset auth keys
*/
reset(): void
reset(withAuthKeys?: boolean): void
/**
* Set default datacenter to use with this session.
@ -128,6 +130,9 @@ export interface ITelegramStorage {
/**
* Update local database of input peers from the peer info list
*
* Client will call `.save()` after all updates-related methods
* are called, so you can safely batch these updates
*/
updatePeers(peers: ITelegramStorage.PeerInfo[]): MaybeAsync<void>
/**
@ -151,18 +156,30 @@ export interface ITelegramStorage {
/**
* Set common `pts` value
*
* Client will call `.save()` after all updates-related methods
* are called, so you can safely batch these updates
*/
setUpdatesPts(val: number): MaybeAsync<void>
/**
* Set common `qts` value
*
* Client will call `.save()` after all updates-related methods
* are called, so you can safely batch these updates
*/
setUpdatesQts(val: number): MaybeAsync<void>
/**
* Set updates `date` value
*
* Client will call `.save()` after all updates-related methods
* are called, so you can safely batch these updates
*/
setUpdatesDate(val: number): MaybeAsync<void>
/**
* Set updates `seq` value
*
* Client will call `.save()` after all updates-related methods
* are called, so you can safely batch these updates
*/
setUpdatesSeq(val: number): MaybeAsync<void>
@ -172,8 +189,12 @@ export interface ITelegramStorage {
getChannelPts(entityId: number): MaybeAsync<number | null>
/**
* Set channels `pts` values in batch.
*
* Storage is supposed to replace stored channel `pts` values
* with given in the object (key is unmarked peer id, value is the `pts`)
*
* Client will call `.save()` after all updates-related methods
* are called, so you can safely batch these updates
*/
setManyChannelPts(values: Map<number, number>): MaybeAsync<void>

View file

@ -87,7 +87,7 @@ export class MemoryStorage implements ITelegramStorage {
*/
vacuumInterval?: number
}) {
this.reset()
this.reset(true)
this._cachedFull = new LruMap(params?.cacheSize ?? 100)
this._vacuumInterval = params?.vacuumInterval ?? 300_000
}
@ -100,13 +100,13 @@ export class MemoryStorage implements ITelegramStorage {
clearInterval(this._vacuumTimeout)
}
reset(): void {
reset(withAuthKeys = false): void {
this._state = {
$version: CURRENT_VERSION,
defaultDcs: null,
authKeys: new Map(),
authKeysTemp: new Map(),
authKeysTempExpiry: new Map(),
authKeys: withAuthKeys ? new Map<number, Uint8Array>() : this._state.authKeys,
authKeysTemp: withAuthKeys ? new Map<string, Uint8Array>() : this._state.authKeysTemp,
authKeysTempExpiry: withAuthKeys ? new Map<string, number>() : this._state.authKeysTempExpiry,
entities: new Map(),
phoneIndex: new Map(),
usernameIndex: new Map(),

View file

@ -8,7 +8,7 @@ import { __tlWriterMap } from '@mtcute/tl/binary/writer.js'
import { MaybeAsync } from '../types/index.js'
import { defaultProductionDc } from '../utils/default-dcs.js'
import { LogManager } from '../utils/index.js'
import { hexEncode, Logger, LogManager, TlReaderMap, TlWriterMap } from '../utils/index.js'
import { ITelegramStorage } from './abstract.js'
export const stubPeerUser: ITelegramStorage.PeerInfo = {
@ -39,22 +39,34 @@ const peerChannelInput: tl.TypeInputPeer = {
accessHash: Long.fromBits(666, 555),
}
export function testStorage(s: ITelegramStorage): void {
beforeAll(async () => {
await s.load?.()
function maybeHexEncode(x: Uint8Array | null): string | null {
if (x == null) return null
return hexEncode(x)
}
export function testStorage<T extends ITelegramStorage>(
s: T,
params?: {
skipEntityOverwrite?: boolean
customTests?: (s: T) => void
},
): void {
beforeAll(async () => {
const logger = new LogManager()
logger.level = 0
s.setup?.(logger, __tlReaderMap, __tlWriterMap)
await s.load?.()
})
afterAll(() => s.destroy?.())
beforeEach(() => s.reset?.())
beforeEach(() => s.reset(true))
describe('default dc', () => {
it('should store', async () => {
await s.setDefaultDcs(defaultProductionDc)
expect(await s.getDefaultDcs()).toBe(defaultProductionDc)
expect(await s.getDefaultDcs()).toEqual(defaultProductionDc)
})
it('should remove', async () => {
@ -79,8 +91,8 @@ export function testStorage(s: ITelegramStorage): void {
await s.setAuthKeyFor(2, key2)
await s.setAuthKeyFor(3, key3)
expect(await s.getAuthKeyFor(2)).toEqual(key2)
expect(await s.getAuthKeyFor(3)).toEqual(key3)
expect(maybeHexEncode(await s.getAuthKeyFor(2))).toEqual(hexEncode(key2))
expect(maybeHexEncode(await s.getAuthKeyFor(3))).toEqual(hexEncode(key3))
})
it('should store temp auth keys', async () => {
@ -91,10 +103,10 @@ export function testStorage(s: ITelegramStorage): void {
await s.setTempAuthKeyFor(3, 0, key3i0, expire)
await s.setTempAuthKeyFor(3, 1, key3i1, expire)
expect(await s.getAuthKeyFor(2, 0)).toEqual(key2i0)
expect(await s.getAuthKeyFor(2, 1)).toEqual(key2i1)
expect(await s.getAuthKeyFor(3, 0)).toEqual(key3i0)
expect(await s.getAuthKeyFor(3, 1)).toEqual(key3i1)
expect(maybeHexEncode(await s.getAuthKeyFor(2, 0))).toEqual(hexEncode(key2i0))
expect(maybeHexEncode(await s.getAuthKeyFor(2, 1))).toEqual(hexEncode(key2i1))
expect(maybeHexEncode(await s.getAuthKeyFor(3, 0))).toEqual(hexEncode(key3i0))
expect(maybeHexEncode(await s.getAuthKeyFor(3, 1))).toEqual(hexEncode(key3i1))
})
it('should expire temp auth keys', async () => {
@ -128,7 +140,7 @@ export function testStorage(s: ITelegramStorage): void {
expect(await s.getAuthKeyFor(2)).toBeNull()
expect(await s.getAuthKeyFor(2, 0)).toBeNull()
expect(await s.getAuthKeyFor(2, 1)).toBeNull()
expect(await s.getAuthKeyFor(3)).toEqual(key3) // should not be removed
expect(maybeHexEncode(await s.getAuthKeyFor(3))).toEqual(hexEncode(key3)) // should not be removed
})
it('should remove all auth keys with dropAuthKeysFor', async () => {
@ -144,13 +156,32 @@ export function testStorage(s: ITelegramStorage): void {
expect(await s.getAuthKeyFor(2)).toBeNull()
expect(await s.getAuthKeyFor(2, 0)).toBeNull()
expect(await s.getAuthKeyFor(2, 1)).toBeNull()
expect(await s.getAuthKeyFor(3)).toEqual(key3) // should not be removed
expect(maybeHexEncode(await s.getAuthKeyFor(3))).toEqual(hexEncode(key3)) // should not be removed
})
it('should not reset auth keys on reset()', async () => {
await s.setAuthKeyFor(2, key2)
await s.setAuthKeyFor(3, key3)
s.reset()
expect(maybeHexEncode(await s.getAuthKeyFor(2))).toEqual(hexEncode(key2))
expect(maybeHexEncode(await s.getAuthKeyFor(3))).toEqual(hexEncode(key3))
})
it('should reset auth keys on reset(true)', async () => {
await s.setAuthKeyFor(2, key2)
await s.setAuthKeyFor(3, key3)
s.reset(true)
expect(await s.getAuthKeyFor(2)).toBeNull()
expect(await s.getAuthKeyFor(3)).toBeNull()
})
})
describe('peers', () => {
it('should cache and return peers', async () => {
await s.updatePeers([stubPeerUser, peerChannel])
await s.save?.() // update-related methods are batched, so we need to save
expect(await s.getPeerById(stubPeerUser.id)).toEqual(peerUserInput)
expect(await s.getPeerById(peerChannel.id)).toEqual(peerChannelInput)
@ -158,6 +189,7 @@ export function testStorage(s: ITelegramStorage): void {
it('should cache and return peers by username', async () => {
await s.updatePeers([stubPeerUser, peerChannel])
await s.save?.() // update-related methods are batched, so we need to save
expect(await s.getPeerByUsername(stubPeerUser.username!)).toEqual(peerUserInput)
expect(await s.getPeerByUsername(peerChannel.username!)).toEqual(peerChannelInput)
@ -165,21 +197,26 @@ export function testStorage(s: ITelegramStorage): void {
it('should cache and return peers by phone', async () => {
await s.updatePeers([stubPeerUser])
await s.save?.() // update-related methods are batched, so we need to save
expect(await s.getPeerByPhone(stubPeerUser.phone!)).toEqual(peerUserInput)
})
it('should overwrite existing cached peers', async () => {
await s.updatePeers([stubPeerUser])
await s.updatePeers([{ ...stubPeerUser, username: 'whatever' }])
if (!params?.skipEntityOverwrite) {
it('should overwrite existing cached peers', async () => {
await s.updatePeers([stubPeerUser])
await s.updatePeers([{ ...stubPeerUser, username: 'whatever' }])
await s.save?.() // update-related methods are batched, so we need to save
expect(await s.getPeerById(stubPeerUser.id)).toEqual(peerUserInput)
expect(await s.getPeerByUsername(stubPeerUser.username!)).toBeNull()
expect(await s.getPeerByUsername('whatever')).toEqual(peerUserInput)
})
expect(await s.getPeerById(stubPeerUser.id)).toEqual(peerUserInput)
expect(await s.getPeerByUsername(stubPeerUser.username!)).toBeNull()
expect(await s.getPeerByUsername('whatever')).toEqual(peerUserInput)
})
}
it('should cache full peer info', async () => {
await s.updatePeers([stubPeerUser, peerChannel])
await s.save?.() // update-related methods are batched, so we need to save
expect(await s.getFullPeerById(stubPeerUser.id)).toEqual(stubPeerUser.full)
expect(await s.getFullPeerById(peerChannel.id)).toEqual(peerChannel.full)
@ -210,6 +247,8 @@ export function testStorage(s: ITelegramStorage): void {
await s.setUpdatesQts(2)
await s.setUpdatesDate(3)
await s.setUpdatesSeq(4)
await s.save?.() // update-related methods are batched, so we need to save
expect(await s.getUpdatesState()).toEqual([1, 2, 3, 4])
})
@ -220,6 +259,7 @@ export function testStorage(s: ITelegramStorage): void {
[3, 4],
]),
)
await s.save?.() // update-related methods are batched, so we need to save
expect(await s.getChannelPts(1)).toEqual(2)
expect(await s.getChannelPts(3)).toEqual(4)
@ -230,9 +270,16 @@ export function testStorage(s: ITelegramStorage): void {
expect(await s.getUpdatesState()).toBeNull()
})
})
params?.customTests?.(s)
}
interface IStateStorage {
setup?(log: Logger, readerMap: TlReaderMap, writerMap: TlWriterMap): void
load?(): MaybeAsync<void>
save?(): MaybeAsync<void>
destroy?(): MaybeAsync<void>
reset(): MaybeAsync<void>
getState(key: string): MaybeAsync<unknown>
setState(key: string, state: unknown, ttl?: number): MaybeAsync<void>
deleteState(key: string): MaybeAsync<void>
@ -244,6 +291,17 @@ interface IStateStorage {
}
export function testStateStorage(s: IStateStorage) {
beforeAll(async () => {
const logger = new LogManager()
logger.level = 0
s.setup?.(logger, __tlReaderMap, __tlWriterMap)
await s.load?.()
})
afterAll(() => s.destroy?.())
beforeEach(() => s.reset())
describe('key-value state', () => {
beforeAll(() => void vi.useFakeTimers())
afterAll(() => void vi.useRealTimers())

View file

@ -1,3 +1,7 @@
export type ThrottledFunction = (() => void) & {
reset: () => void
}
/**
* Throttle a function with a given delay.
* Similar to lodash.
@ -10,10 +14,10 @@
* @param func Function to throttle
* @param delay Throttle delay
*/
export function throttle(func: () => void, delay: number): () => void {
export function throttle(func: () => void, delay: number): ThrottledFunction {
let timeout: NodeJS.Timeout | null
return function () {
const res: ThrottledFunction = function () {
if (timeout) {
return
}
@ -24,4 +28,13 @@ export function throttle(func: () => void, delay: number): () => void {
}
timeout = setTimeout(later, delay)
}
res.reset = () => {
if (timeout) {
clearTimeout(timeout)
timeout = null
}
}
return res
}

View file

@ -135,4 +135,11 @@ export class LruMap<K extends string | number, V> {
const item = this._map.get(key)
if (item) this._remove(item)
}
clear(): void {
this._map.clear()
this._first = undefined
this._last = undefined
this._size = 0
}
}

View file

@ -1,7 +1,7 @@
import { describe } from 'vitest'
// eslint-disable-next-line import/no-relative-packages
import { testCryptoProvider } from '../../core/src/utils/crypto/crypto.test-utils.js'
import { testCryptoProvider } from '@mtcute/core/src/utils/crypto/crypto.test-utils.js'
import { NodeNativeCryptoProvider } from '../src/index.js'
describe('NodeNativeCryptoProvider', () => {

View file

@ -5,7 +5,7 @@
"description": "SQLite-based storage for mtcute",
"author": "Alina Sireneva <alina@tei.su>",
"license": "MIT",
"main": "index.ts",
"main": "src/index.ts",
"type": "module",
"distOnlyFields": {
"exports": {

View file

@ -9,6 +9,7 @@ import {
longToFastString,
LruMap,
throttle,
ThrottledFunction,
TlBinaryReader,
TlBinaryWriter,
TlReaderMap,
@ -98,10 +99,13 @@ const SCHEMA = `
const RESET = `
delete from kv where key <> 'ver';
delete from state;
delete from auth_keys;
delete from pts;
delete from entities
`
const RESET_AUTH_KEYS = `
delete from auth_keys;
delete from temp_auth_keys;
`
const USERNAME_TTL = 86400000 // 24 hours
@ -179,7 +183,7 @@ export class SqliteStorage implements ITelegramStorage /*, IStateStorage*/ {
private _reader!: TlBinaryReader
private _saveUnimportantLater: () => void
private _saveUnimportantLater: ThrottledFunction
private _vacuumTimeout?: NodeJS.Timeout
private _vacuumInterval: number
@ -435,7 +439,7 @@ export class SqliteStorage implements ITelegramStorage /*, IStateStorage*/ {
load(): void {
this._db = sqlite3(this._filename, {
verbose: this.log.mgr.level === 5 ? (this.log.verbose as Options['verbose']) : undefined,
verbose: this.log.mgr.level >= 5 ? (this.log.verbose as Options['verbose']) : undefined,
})
this._initialize()
@ -475,12 +479,20 @@ export class SqliteStorage implements ITelegramStorage /*, IStateStorage*/ {
clearInterval(this._vacuumTimeout)
}
reset(): void {
reset(withAuthKeys = false): void {
this._db.exec(RESET)
if (withAuthKeys) this._db.exec(RESET_AUTH_KEYS)
this._pending = []
this._pendingUnimportant = {}
this._cache?.clear()
this._fsmCache?.clear()
this._rlCache?.clear()
this._saveUnimportantLater.reset()
}
setDefaultDcs(dc: ITelegramStorage.DcOptions | null): void {
return this._setToKv('def_dc', dc)
return this._setToKv('def_dc', dc, true)
}
getDefaultDcs(): ITelegramStorage.DcOptions | null {
@ -500,21 +512,24 @@ export class SqliteStorage implements ITelegramStorage /*, IStateStorage*/ {
}
setAuthKeyFor(dcId: number, key: Uint8Array | null): void {
this._pending.push([
key === null ? this._statements.delAuth : this._statements.setAuth,
key === null ? [dcId] : [dcId, key],
])
if (key !== null) {
this._statements.setAuth.run(dcId, key)
} else {
this._statements.delAuth.run(dcId)
}
}
setTempAuthKeyFor(dcId: number, index: number, key: Uint8Array | null, expires: number): void {
this._pending.push([
key === null ? this._statements.delAuthTemp : this._statements.setAuthTemp,
key === null ? [dcId, index] : [dcId, index, key, expires],
])
if (key !== null) {
this._statements.setAuthTemp.run(dcId, index, key, expires)
} else {
this._statements.delAuthTemp.run(dcId, index)
}
}
dropAuthKeysFor(dcId: number): void {
this._pending.push([this._statements.delAuth, [dcId]], [this._statements.delAllAuthTemp, [dcId]])
this._statements.delAuth.run(dcId)
this._statements.delAllAuthTemp.run(dcId)
}
getSelf(): ITelegramStorage.SelfInfo | null {
@ -522,7 +537,7 @@ export class SqliteStorage implements ITelegramStorage /*, IStateStorage*/ {
}
setSelf(self: ITelegramStorage.SelfInfo | null): void {
return this._setToKv('self', self)
return this._setToKv('self', self, true)
}
getUpdatesState(): [number, number, number, number] | null {

View file

@ -0,0 +1,50 @@
import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
import { stubPeerUser, testStateStorage, testStorage } from '@mtcute/core/src/storage/storage.test-utils.js'
import { SqliteStorage } from '../src/index.js'
describe('SqliteStorage', () => {
testStorage(new SqliteStorage(), {
// sqlite implements "unimportant" updates, which are batched once every 30sec (tested below)
skipEntityOverwrite: true,
customTests: (s) => {
describe('batching', () => {
beforeAll(() => void vi.useFakeTimers())
afterAll(() => void vi.useRealTimers())
it('should batch entity writes', async () => {
s.updatePeers([stubPeerUser])
s.updatePeers([{ ...stubPeerUser, username: 'test123' }])
s.save()
// eslint-disable-next-line
expect(Object.keys(s['_pendingUnimportant'])).toEqual([String(stubPeerUser.id)])
// not yet updated
expect(s.getPeerByUsername(stubPeerUser.username!)).not.toBeNull()
expect(s.getPeerByUsername('test123')).toBeNull()
await vi.advanceTimersByTimeAsync(30001)
expect(s.getPeerByUsername(stubPeerUser.username!)).toBeNull()
expect(s.getPeerByUsername('test123')).not.toBeNull()
})
it('should batch update state writes', () => {
s.setUpdatesPts(123)
s.setUpdatesQts(456)
s.setUpdatesDate(789)
s.setUpdatesSeq(999)
// not yet updated
expect(s.getUpdatesState()).toBeNull()
s.save()
expect(s.getUpdatesState()).toEqual([123, 456, 789, 999])
})
})
},
})
testStateStorage(new SqliteStorage())
})

View file

@ -0,0 +1,9 @@
{
"extends": "../../../tsconfig.json",
"include": [
"."
],
"references": [
{ "path": "../" }
]
}

View file

@ -1,10 +1,11 @@
{
"extends": "../../tsconfig.json",
"compilerOptions": {
"outDir": "./dist/esm"
"outDir": "./dist/esm",
"rootDir": "./src"
},
"include": [
"./index.ts"
"./src"
],
"references": [
{ "path": "../core" }

View file

@ -1,4 +1,4 @@
module.exports = {
extends: ['../../typedoc.base.cjs'],
entryPoints: ['./index.ts'],
entryPoints: ['./src/index.ts'],
}

View file

@ -62,9 +62,9 @@ export class StubMemoryTelegramStorage extends MemoryStorage implements ITelegra
super.destroy()
}
reset(): void {
reset(withKeys = false): void {
this.params?.onReset?.()
super.reset()
super.reset(withKeys)
}
decryptOutgoingMessage(crypto: ICryptoProvider, data: Uint8Array, dcId: number, tempIndex?: number | undefined) {