From a8e209705a6b711872cf7b2e2fa2ba83515f1fcb Mon Sep 17 00:00:00 2001 From: Adhiraj Singh Date: Sun, 22 May 2022 20:52:21 +0530 Subject: [PATCH] feat: add retry capability to SignalKeyStore --- src/Defaults/index.ts | 1 + src/Socket/socket.ts | 5 +++-- src/Types/Auth.ts | 5 +++++ src/Types/index.ts | 4 +++- src/Utils/auth-utils.ts | 43 ++++++++++++++++++++++++++++------------- 5 files changed, 42 insertions(+), 16 deletions(-) diff --git a/src/Defaults/index.ts b/src/Defaults/index.ts index fda44dd..639e994 100644 --- a/src/Defaults/index.ts +++ b/src/Defaults/index.ts @@ -45,6 +45,7 @@ export const DEFAULT_CONNECTION_CONFIG: SocketConfig = { ...BASE_CONNECTION_CONFIG, downloadHistory: true, linkPreviewImageThumbnailWidth: 192, + transactionOpts: { maxCommitRetries: 10, delayBetweenTriesMs: 3000 }, getMessage: async() => undefined } diff --git a/src/Socket/socket.ts b/src/Socket/socket.ts index 7e7bc6a..68ce521 100644 --- a/src/Socket/socket.ts +++ b/src/Socket/socket.ts @@ -24,7 +24,8 @@ export const makeSocket = ({ browser, auth: initialAuthState, printQRInTerminal, - defaultQueryTimeoutMs + defaultQueryTimeoutMs, + transactionOpts }: SocketConfig) => { const ws = new WebSocket(waWebSocketUrl, undefined, { origin: DEFAULT_ORIGIN, @@ -51,7 +52,7 @@ export const makeSocket = ({ const { creds } = authState // add transaction capability - const keys = addTransactionCapability(authState.keys, logger) + const keys = addTransactionCapability(authState.keys, logger, transactionOpts) let lastDateRecv: Date let epoch = 1 diff --git a/src/Types/Auth.ts b/src/Types/Auth.ts index 8392461..015f5d5 100644 --- a/src/Types/Auth.ts +++ b/src/Types/Auth.ts @@ -72,6 +72,11 @@ export type SignalKeyStoreWithTransaction = SignalKeyStore & { prefetch(type: T, ids: string[]): Promise } +export type TransactionCapabilityOptions = { + maxCommitRetries: number + delayBetweenTriesMs: number +} + export type SignalAuthState = { creds: SignalCreds keys: SignalKeyStore diff --git a/src/Types/index.ts b/src/Types/index.ts index 8788d0e..dcae594 100644 --- a/src/Types/index.ts +++ b/src/Types/index.ts @@ -12,7 +12,7 @@ export * from './Call' import type NodeCache from 'node-cache' import { proto } from '../../WAProto' -import { AuthenticationState } from './Auth' +import { AuthenticationState, TransactionCapabilityOptions } from './Auth' import { CommonSocketConfig } from './Socket' export type MessageRetryMap = { [msgId: string]: number } @@ -20,6 +20,8 @@ export type MessageRetryMap = { [msgId: string]: number } export type SocketConfig = CommonSocketConfig & { /** By default true, should history messages be downloaded and processed */ downloadHistory: boolean + /** transaction capability options for SignalKeyStore */ + transactionOpts: TransactionCapabilityOptions /** provide a cache to store a user's device list */ userDevicesCache?: NodeCache /** diff --git a/src/Utils/auth-utils.ts b/src/Utils/auth-utils.ts index a5cb8c8..db72c0e 100644 --- a/src/Utils/auth-utils.ts +++ b/src/Utils/auth-utils.ts @@ -2,18 +2,9 @@ import { Boom } from '@hapi/boom' import { randomBytes } from 'crypto' import type { Logger } from 'pino' import { proto } from '../../WAProto' -import type { AuthenticationCreds, AuthenticationState, SignalDataSet, SignalDataTypeMap, SignalKeyStore, SignalKeyStoreWithTransaction } from '../Types' +import type { AuthenticationCreds, AuthenticationState, SignalDataSet, SignalDataTypeMap, SignalKeyStore, SignalKeyStoreWithTransaction, TransactionCapabilityOptions } from '../Types' import { Curve, signedKeyPair } from './crypto' -import { BufferJSON, generateRegistrationId } from './generics' - -const KEY_MAP: { [T in keyof SignalDataTypeMap]: string } = { - 'pre-key': 'preKeys', - 'session': 'sessions', - 'sender-key': 'senderKeys', - 'app-state-sync-key': 'appStateSyncKeys', - 'app-state-sync-version': 'appStateVersions', - 'sender-key-memory': 'senderKeyMemory' -} +import { BufferJSON, delay, generateRegistrationId } from './generics' /** * Adds DB like transaction capability (https://en.wikipedia.org/wiki/Database_transaction) to the SignalKeyStore, @@ -22,11 +13,15 @@ const KEY_MAP: { [T in keyof SignalDataTypeMap]: string } = { * @param logger logger to log events * @returns SignalKeyStore with transaction capability */ -export const addTransactionCapability = (state: SignalKeyStore, logger: Logger): SignalKeyStoreWithTransaction => { +export const addTransactionCapability = (state: SignalKeyStore, logger: Logger, { maxCommitRetries, delayBetweenTriesMs }: TransactionCapabilityOptions): SignalKeyStoreWithTransaction => { let inTransaction = false let transactionCache: SignalDataSet = { } let mutations: SignalDataSet = { } + /** + * prefetches some data and stores in memory, + * useful if these data points will be used together often + * */ const prefetch = async(type: keyof SignalDataTypeMap, ids: string[]) => { if(!inTransaction) { throw new Boom('Cannot prefetch without transaction') @@ -90,7 +85,19 @@ export const addTransactionCapability = (state: SignalKeyStore, logger: Logger): await work() if(Object.keys(mutations).length) { logger.debug('committing transaction') - await state.set(mutations) + // retry mechanism to ensure we've some recovery + // in case a transaction fails in the first attempt + let tries = maxCommitRetries + while(tries) { + tries -= 1 + try { + await state.set(mutations) + break + } catch(error) { + logger.warn(`failed to commit ${Object.keys(mutations).length} mutations, tries left=${tries}`) + await delay(delayBetweenTriesMs) + } + } } else { logger.debug('no mutations in transaction') } @@ -121,6 +128,16 @@ export const initAuthCreds = (): AuthenticationCreds => { } } +// useless key map only there to maintain backwards compatibility +// do not use in your own systems please +const KEY_MAP: { [T in keyof SignalDataTypeMap]: string } = { + 'pre-key': 'preKeys', + 'session': 'sessions', + 'sender-key': 'senderKeys', + 'app-state-sync-key': 'appStateSyncKeys', + 'app-state-sync-version': 'appStateVersions', + 'sender-key-memory': 'senderKeyMemory' +} /** stores the full authentication state in a single JSON file */ export const useSingleFileAuthState = (filename: string, logger?: Logger): { state: AuthenticationState, saveState: () => void } => { // require fs here so that in case "fs" is not available -- the app does not crash