diff --git a/src/Signal/Group/group_cipher.ts b/src/Signal/Group/group_cipher.ts index ef881eab..b0839e0b 100644 --- a/src/Signal/Group/group_cipher.ts +++ b/src/Signal/Group/group_cipher.ts @@ -1,6 +1,5 @@ /* @ts-ignore */ import { decrypt, encrypt } from 'libsignal/src/crypto' -import queueJob from './queue-job' import { SenderKeyMessage } from './sender-key-message' import { SenderKeyName } from './sender-key-name' import { SenderKeyRecord } from './sender-key-record' @@ -8,6 +7,7 @@ import { SenderKeyState } from './sender-key-state' export interface SenderKeyStore { loadSenderKey(senderKeyName: SenderKeyName): Promise + storeSenderKey(senderKeyName: SenderKeyName, record: SenderKeyRecord): Promise } @@ -20,64 +20,56 @@ export class GroupCipher { this.senderKeyName = senderKeyName } - private queueJob(awaitable: () => Promise): Promise { - return queueJob(this.senderKeyName.toString(), awaitable) - } - public async encrypt(paddedPlaintext: Uint8Array | string): Promise { - return await this.queueJob(async () => { - const record = await this.senderKeyStore.loadSenderKey(this.senderKeyName) - if (!record) { - throw new Error('No SenderKeyRecord found for encryption') - } + const record = await this.senderKeyStore.loadSenderKey(this.senderKeyName) + if (!record) { + throw new Error('No SenderKeyRecord found for encryption') + } - const senderKeyState = record.getSenderKeyState() - if (!senderKeyState) { - throw new Error('No session to encrypt message') - } + const senderKeyState = record.getSenderKeyState() + if (!senderKeyState) { + throw new Error('No session to encrypt message') + } - const iteration = senderKeyState.getSenderChainKey().getIteration() - const senderKey = this.getSenderKey(senderKeyState, iteration === 0 ? 0 : iteration + 1) + const iteration = senderKeyState.getSenderChainKey().getIteration() + const senderKey = this.getSenderKey(senderKeyState, iteration === 0 ? 0 : iteration + 1) - const ciphertext = await this.getCipherText(senderKey.getIv(), senderKey.getCipherKey(), paddedPlaintext) + const ciphertext = await this.getCipherText(senderKey.getIv(), senderKey.getCipherKey(), paddedPlaintext) - const senderKeyMessage = new SenderKeyMessage( - senderKeyState.getKeyId(), - senderKey.getIteration(), - ciphertext, - senderKeyState.getSigningKeyPrivate() - ) + const senderKeyMessage = new SenderKeyMessage( + senderKeyState.getKeyId(), + senderKey.getIteration(), + ciphertext, + senderKeyState.getSigningKeyPrivate() + ) - await this.senderKeyStore.storeSenderKey(this.senderKeyName, record) - return senderKeyMessage.serialize() - }) + await this.senderKeyStore.storeSenderKey(this.senderKeyName, record) + return senderKeyMessage.serialize() } public async decrypt(senderKeyMessageBytes: Uint8Array): Promise { - return await this.queueJob(async () => { - const record = await this.senderKeyStore.loadSenderKey(this.senderKeyName) - if (!record) { - throw new Error('No SenderKeyRecord found for decryption') - } + const record = await this.senderKeyStore.loadSenderKey(this.senderKeyName) + if (!record) { + throw new Error('No SenderKeyRecord found for decryption') + } - const senderKeyMessage = new SenderKeyMessage(null, null, null, null, senderKeyMessageBytes) - const senderKeyState = record.getSenderKeyState(senderKeyMessage.getKeyId()) - if (!senderKeyState) { - throw new Error('No session found to decrypt message') - } + const senderKeyMessage = new SenderKeyMessage(null, null, null, null, senderKeyMessageBytes) + const senderKeyState = record.getSenderKeyState(senderKeyMessage.getKeyId()) + if (!senderKeyState) { + throw new Error('No session found to decrypt message') + } - senderKeyMessage.verifySignature(senderKeyState.getSigningKeyPublic()) - const senderKey = this.getSenderKey(senderKeyState, senderKeyMessage.getIteration()) + senderKeyMessage.verifySignature(senderKeyState.getSigningKeyPublic()) + const senderKey = this.getSenderKey(senderKeyState, senderKeyMessage.getIteration()) - const plaintext = await this.getPlainText( - senderKey.getIv(), - senderKey.getCipherKey(), - senderKeyMessage.getCipherText() - ) + const plaintext = await this.getPlainText( + senderKey.getIv(), + senderKey.getCipherKey(), + senderKeyMessage.getCipherText() + ) - await this.senderKeyStore.storeSenderKey(this.senderKeyName, record) - return plaintext - }) + await this.senderKeyStore.storeSenderKey(this.senderKeyName, record) + return plaintext } private getSenderKey(senderKeyState: SenderKeyState, iteration: number) { diff --git a/src/Signal/Group/queue-job.ts b/src/Signal/Group/queue-job.ts deleted file mode 100644 index 4719f3dd..00000000 --- a/src/Signal/Group/queue-job.ts +++ /dev/null @@ -1,70 +0,0 @@ -interface QueueJob { - awaitable: () => Promise - resolve: (value: T | PromiseLike) => void - reject: (reason?: unknown) => void -} - -const _queueAsyncBuckets = new Map>>() - -export function cleanupQueues() { - _queueAsyncBuckets.clear() -} - -const _gcLimit = 10000 - -async function _asyncQueueExecutor(queue: Array>, cleanup: () => void): Promise { - let offt = 0 - - while (true) { - const limit = Math.min(queue.length, _gcLimit) - for (let i = offt; i < limit; i++) { - const job = queue[i]! - try { - job.resolve(await job.awaitable()) - } catch (e) { - job.reject(e) - } - } - - if (limit < queue.length) { - if (limit >= _gcLimit) { - queue.splice(0, limit) - offt = 0 - } else { - offt = limit - } - } else { - break - } - } - - cleanup() -} - -export default function queueJob(bucket: string | number, awaitable: () => Promise): Promise { - // Skip name assignment since it's readonly in strict mode - if (typeof bucket !== 'string') { - console.warn('Unhandled bucket type (for naming):', typeof bucket, bucket) - } - - let inactive = false - if (!_queueAsyncBuckets.has(bucket)) { - _queueAsyncBuckets.set(bucket, []) - inactive = true - } - - const queue = _queueAsyncBuckets.get(bucket)! - const job = new Promise((resolve, reject) => { - queue.push({ - awaitable, - resolve: resolve as (value: any) => void, - reject - }) - }) - - if (inactive) { - _asyncQueueExecutor(queue, () => _queueAsyncBuckets.delete(bucket)) - } - - return job -} diff --git a/src/Signal/libsignal.ts b/src/Signal/libsignal.ts index c601c1be..0c338c0a 100644 --- a/src/Signal/libsignal.ts +++ b/src/Signal/libsignal.ts @@ -16,6 +16,22 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository const lidMapping = new LIDMappingStore(auth.keys as SignalKeyStoreWithTransaction) const storage = signalStorage(auth, lidMapping) + + + const parsedKeys = auth.keys as SignalKeyStoreWithTransaction + + function isLikelySyncMessage(addr: any): boolean { + const key = addr.toString() + + // Only bypass for WhatsApp system addresses, not regular user contacts + // Be very specific about sync service patterns + return ( + key.includes('@lid.whatsapp.net') || // WhatsApp system messages + key.includes('@broadcast') || // Broadcast messages + key.includes('@newsletter') + ) + } + // Simple operation-level deduplication (5 minutes) const recentMigrations = new LRUCache({ max: 500, @@ -23,11 +39,15 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository }) const repository: SignalRepository = { + decryptGroupMessage({ group, authorJid, msg }) { const senderName = jidToSignalSenderKeyName(group, authorJid) const cipher = new GroupCipher(storage, senderName) - return cipher.decrypt(msg) + // Use transaction to ensure atomicity + return parsedKeys.transaction(async () => { + return cipher.decrypt(msg) + }, group) }, async processSenderKeyDistributionMessage({ item, authorJid }) { const builder = new GroupSessionBuilder(storage) @@ -50,24 +70,44 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository await storage.storeSenderKey(senderName, new SenderKeyRecord()) } - await builder.process(senderName, senderMsg) + return parsedKeys.transaction(async () => { + const { [senderNameStr]: senderKey } = await auth.keys.get('sender-key', [senderNameStr]) + if (!senderKey) { + await storage.storeSenderKey(senderName, new SenderKeyRecord()) + } + + await builder.process(senderName, senderMsg) + }, item.groupId) }, async decryptMessage({ jid, type, ciphertext }) { const addr = jidToSignalProtocolAddress(jid) const session = new libsignal.SessionCipher(storage, addr) - let result: Buffer - switch (type) { - case 'pkmsg': - result = await session.decryptPreKeyWhisperMessage(ciphertext) - break - case 'msg': - result = await session.decryptWhisperMessage(ciphertext) - break - default: - throw new Error(`Unknown message type: ${type}`) + + async function doDecrypt() { + let result: Buffer + switch (type) { + case 'pkmsg': + result = await session.decryptPreKeyWhisperMessage(ciphertext) + break + case 'msg': + result = await session.decryptWhisperMessage(ciphertext) + break + } + + return result } - return result + if (isLikelySyncMessage(addr)) { + // If it's a sync message, we can skip the transaction + // as it is likely to be a system message that doesn't require strict atomicity + return await doDecrypt() + } + + // If it's not a sync message, we need to ensure atomicity + // For regular messages, we use a transaction to ensure atomicity + return parsedKeys.transaction(async () => { + return await doDecrypt() + }, jid) }, async encryptMessage({ jid, data }) { @@ -101,32 +141,40 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository const addr = jidToSignalProtocolAddress(encryptionJid) const cipher = new libsignal.SessionCipher(storage, addr) - const { type: sigType, body } = await cipher.encrypt(data) - const type = sigType === 3 ? 'pkmsg' : 'msg' - return { type, ciphertext: Buffer.from(body, 'binary') } + // Use transaction to ensure atomicity + return parsedKeys.transaction(async () => { + const { type: sigType, body } = await cipher.encrypt(data) + const type = sigType === 3 ? 'pkmsg' : 'msg' + return { type, ciphertext: Buffer.from(body, 'binary') } + }, jid) }, async encryptGroupMessage({ group, meId, data }) { const senderName = jidToSignalSenderKeyName(group, meId) const builder = new GroupSessionBuilder(storage) const senderNameStr = senderName.toString() - const { [senderNameStr]: senderKey } = await auth.keys.get('sender-key', [senderNameStr]) - if (!senderKey) { - await storage.storeSenderKey(senderName, new SenderKeyRecord()) - } - const senderKeyDistributionMessage = await builder.create(senderName) - const session = new GroupCipher(storage, senderName) - const ciphertext = await session.encrypt(data) + return parsedKeys.transaction(async () => { + const { [senderNameStr]: senderKey } = await auth.keys.get('sender-key', [senderNameStr]) + if (!senderKey) { + await storage.storeSenderKey(senderName, new SenderKeyRecord()) + } - return { - ciphertext, - senderKeyDistributionMessage: senderKeyDistributionMessage.serialize() - } + const senderKeyDistributionMessage = await builder.create(senderName) + const session = new GroupCipher(storage, senderName) + const ciphertext = await session.encrypt(data) + + return { + ciphertext, + senderKeyDistributionMessage: senderKeyDistributionMessage.serialize() + } + }, group) }, async injectE2ESession({ jid, session }) { const cipher = new libsignal.SessionBuilder(storage, jidToSignalProtocolAddress(jid)) - await cipher.initOutgoing(session) + return parsedKeys.transaction(async () => { + await cipher.initOutgoing(session) + }, jid) }, async validateSession(jid: string) { try { @@ -180,9 +228,9 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository async deleteSession(jid: string) { const addr = jidToSignalProtocolAddress(jid) - return (auth.keys as SignalKeyStoreWithTransaction).transaction(async () => { + return parsedKeys.transaction(async () => { await auth.keys.set({ session: { [addr.toString()]: null } }) - }) + }, jid) }, async migrateSession(fromJid: string, toJid: string) { @@ -211,7 +259,7 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository return } - return (auth.keys as SignalKeyStoreWithTransaction).transaction(async () => { + return parsedKeys.transaction(async () => { // Store mapping await lidMapping.storeLIDPNMapping(toJid, fromJid) @@ -232,7 +280,7 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository } recentMigrations.set(migrationKey, true) - }) + }, fromJid) }, async encryptMessageWithWire({ encryptionJid, wireJid, data }) { diff --git a/src/Socket/chats.ts b/src/Socket/chats.ts index ea743452..f84f91ec 100644 --- a/src/Socket/chats.ts +++ b/src/Socket/chats.ts @@ -601,7 +601,7 @@ export const makeChatsSocket = (config: SocketConfig) => { } } } - }) + }, authState?.creds?.me?.id || 'resync-app-state') const { onMutation } = newAppStateChunkHandler(isInitialSync) for (const key in globalMutationMap) { @@ -807,7 +807,7 @@ export const makeChatsSocket = (config: SocketConfig) => { await query(node) await authState.keys.set({ 'app-state-sync-version': { [name]: state } }) - }) + }, authState?.creds?.me?.id || 'app-patch') }) if (config.emitOwnEvents) { diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts index 4f4e6db6..1d86367a 100644 --- a/src/Socket/messages-recv.ts +++ b/src/Socket/messages-recv.ts @@ -307,8 +307,8 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { await sendNode(receipt) - logger.info({ msgAttrs: node.attrs, retryCount, shouldRecreateSession, recreateReason }, 'sent retry receipt') - }) + logger.info({ msgAttrs: node.attrs, retryCount }, 'sent retry receipt') + }, authState?.creds?.me?.id || 'sendRetryRequest') } const handleEncryptNotification = async (node: BinaryNode) => { diff --git a/src/Socket/messages-send.ts b/src/Socket/messages-send.ts index be81042a..5cb6507f 100644 --- a/src/Socket/messages-send.ts +++ b/src/Socket/messages-send.ts @@ -1128,11 +1128,11 @@ export const makeMessagesSocket = (config: SocketConfig) => { await sendNode(stanza) - // Add message to retry cache if enabled + // Add message to retry cache if enabled if (messageRetryManager && !participant) { messageRetryManager.addRecentMessage(destinationJid, msgId, message) } - }) + }, meId) return msgId } diff --git a/src/Socket/socket.ts b/src/Socket/socket.ts index 76263b50..495d8028 100644 --- a/src/Socket/socket.ts +++ b/src/Socket/socket.ts @@ -331,7 +331,7 @@ export const makeSocket = (config: SocketConfig) => { // Update credentials immediately to prevent duplicate IDs on retry ev.emit('creds.update', update) return node // Only return node since update is already used - }) + }, creds?.me?.id || 'upload-pre-keys') // Upload to server (outside transaction, can fail without affecting local keys) try { diff --git a/src/Types/Auth.ts b/src/Types/Auth.ts index 187ba87b..a8918bdd 100644 --- a/src/Types/Auth.ts +++ b/src/Types/Auth.ts @@ -88,7 +88,7 @@ export type SignalKeyStore = { export type SignalKeyStoreWithTransaction = SignalKeyStore & { isInTransaction: () => boolean - transaction(exec: () => Promise): Promise + transaction(exec: () => Promise, key: string): Promise } export type TransactionCapabilityOptions = { diff --git a/src/Utils/auth-utils.ts b/src/Utils/auth-utils.ts index 7544538f..1dc18ae1 100644 --- a/src/Utils/auth-utils.ts +++ b/src/Utils/auth-utils.ts @@ -1,5 +1,7 @@ import NodeCache from '@cacheable/node-cache' +import { Mutex } from 'async-mutex' import { randomBytes } from 'crypto' +import { LRUCache } from 'lru-cache' import { DEFAULT_CACHE_TTLS } from '../Defaults' import type { AuthenticationCreds, @@ -33,49 +35,56 @@ export function makeCacheableSignalKeyStore( deleteOnExpire: true }) + // Mutex for protecting cache operations + const cacheMutex = new Mutex() + function getUniqueId(type: string, id: string) { return `${type}.${id}` } return { async get(type, ids) { - const data: { [_: string]: SignalDataTypeMap[typeof type] } = {} - const idsToFetch: string[] = [] - for (const id of ids) { - const item = cache.get(getUniqueId(type, id)) as any - if (typeof item !== 'undefined') { - data[id] = item - } else { - idsToFetch.push(id) - } - } - - if (idsToFetch.length) { - logger?.trace({ items: idsToFetch.length }, 'loading from store') - const fetched = await store.get(type, idsToFetch) - for (const id of idsToFetch) { - const item = fetched[id] - if (item) { + return cacheMutex.runExclusive(async () => { + const data: { [_: string]: SignalDataTypeMap[typeof type] } = {} + const idsToFetch: string[] = [] + for (const id of ids) { + const item = cache.get(getUniqueId(type, id)) as any + if (typeof item !== 'undefined') { data[id] = item - cache.set(getUniqueId(type, id), item) + } else { + idsToFetch.push(id) } } - } - return data + if (idsToFetch.length) { + logger?.trace({ items: idsToFetch.length }, 'loading from store') + const fetched = await store.get(type, idsToFetch) + for (const id of idsToFetch) { + const item = fetched[id] + if (item) { + data[id] = item + cache.set(getUniqueId(type, id), item) + } + } + } + + return data + }) }, async set(data) { - let keys = 0 - for (const type in data) { - for (const id in data[type as keyof SignalDataTypeMap]) { - cache.set(getUniqueId(type, id), data[type as keyof SignalDataTypeMap]![id]!) - keys += 1 + return cacheMutex.runExclusive(async () => { + let keys = 0 + for (const type in data) { + for (const id in data[type as keyof SignalDataTypeMap]) { + cache.set(getUniqueId(type, id), data[type as keyof SignalDataTypeMap]![id]!) + keys += 1 + } } - } - logger?.trace({ keys }, 'updated cache') + logger?.trace({ keys }, 'updated cache') - await store.set(data) + await store.set(data) + }) }, async clear() { cache.flushAll() @@ -84,6 +93,181 @@ export function makeCacheableSignalKeyStore( } } +// Module-level specialized mutexes for pre-key operations +const preKeyMutex = new Mutex() +const signedPreKeyMutex = new Mutex() + +/** + * Get the appropriate mutex for the key type + */ +const getPreKeyMutex = (keyType: string): Mutex => { + return keyType === 'signed-pre-key' ? signedPreKeyMutex : preKeyMutex +} + +/** + * Handles pre-key operations with mutex protection + */ +async function handlePreKeyOperations( + data: SignalDataSet, + keyType: keyof SignalDataTypeMap, + transactionCache: SignalDataSet, + mutations: SignalDataSet, + logger: ILogger, + isInTransaction: boolean, + state?: SignalKeyStore +): Promise { + const mutex = getPreKeyMutex(keyType) + + await mutex.runExclusive(async () => { + const keyData = data[keyType] + if (!keyData) return + + // Ensure structures exist + transactionCache[keyType] = transactionCache[keyType] || ({} as any) + mutations[keyType] = mutations[keyType] || ({} as any) + + // Separate deletions from updates for batch processing + const deletionKeys: string[] = [] + const updateKeys: string[] = [] + + for (const keyId in keyData) { + if (keyData[keyId] === null) { + deletionKeys.push(keyId) + } else { + updateKeys.push(keyId) + } + } + + // Process updates first (no validation needed) + for (const keyId of updateKeys) { + if (transactionCache[keyType]) { + transactionCache[keyType][keyId] = keyData[keyId]! + } + + if (mutations[keyType]) { + mutations[keyType][keyId] = keyData[keyId]! + } + } + + // Process deletions with validation + if (deletionKeys.length === 0) return + + if (isInTransaction) { + // In transaction, only allow deletion if key exists in cache + for (const keyId of deletionKeys) { + if (transactionCache[keyType]) { + transactionCache[keyType][keyId] = null + if (mutations[keyType]) { + // Mark for deletion in mutations + mutations[keyType][keyId] = null + } + } else { + logger.warn(`Skipping deletion of non-existent ${keyType} in transaction: ${keyId}`) + } + } + + return + } + + // Outside transaction, batch validate all deletions + if (!state) return + + const existingKeys = await state.get(keyType, deletionKeys) + for (const keyId of deletionKeys) { + if (existingKeys[keyId]) { + if (transactionCache[keyType]) transactionCache[keyType][keyId] = null + + if (mutations[keyType]) mutations[keyType][keyId] = null + } else { + logger.warn(`Skipping deletion of non-existent ${keyType}: ${keyId}`) + } + } + }) +} + +/** + * Handles normal key operations for transactions + */ +function handleNormalKeyOperations( + data: SignalDataSet, + key: keyof SignalDataTypeMap, + transactionCache: SignalDataSet, + mutations: SignalDataSet +) { + Object.assign(transactionCache[key]!, data[key]) + mutations[key] = mutations[key] || ({} as any) + Object.assign(mutations[key]!, data[key]) +} + +/** + * Process pre-key deletions with validation + */ +async function processPreKeyDeletions( + data: SignalDataSet, + keyType: keyof SignalDataTypeMap, + state: SignalKeyStore, + logger: ILogger +): Promise { + const mutex = getPreKeyMutex(keyType) + + await mutex.runExclusive(async () => { + const keyData = data[keyType] + if (!keyData) return + + // Validate deletions + for (const keyId in keyData) { + if (keyData[keyId] === null) { + const existingKeys = await state.get(keyType, [keyId]) + if (!existingKeys[keyId]) { + logger.warn(`Skipping deletion of non-existent ${keyType}: ${keyId}`) + + if (data[keyType]) delete data[keyType][keyId] + } + } + } + }) +} + +/** + * Executes a function with mutexes acquired for given key types + * Uses async-mutex's runExclusive with efficient batching + */ +async function withMutexes( + keyTypes: string[], + getKeyTypeMutex: (type: string) => Mutex, + fn: () => Promise +): Promise { + if (keyTypes.length === 0) { + return fn() + } + + if (keyTypes.length === 1) { + return getKeyTypeMutex(keyTypes[0]!).runExclusive(fn) + } + + // For multiple mutexes, sort by key type to prevent deadlocks + // Then acquire all mutexes in order using Promise.all for better efficiency + const sortedKeyTypes = [...keyTypes].sort() + const mutexes = sortedKeyTypes.map(getKeyTypeMutex) + + // Acquire all mutexes in order to prevent deadlocks + const releases: (() => void)[] = [] + + try { + for (const mutex of mutexes) { + releases.push(await mutex.acquire()) + } + + return await fn() + } finally { + // Release in reverse order + while (releases.length > 0) { + const release = releases.pop() + if (release) release() + } + } +} + /** * Adds DB like transaction capability (https://en.wikipedia.org/wiki/Database_transaction) to the SignalKeyStore, * this allows batch read & write operations & improves the performance of the lib @@ -102,8 +286,96 @@ export const addTransactionCapability = ( let transactionCache: SignalDataSet = {} let mutations: SignalDataSet = {} + // LRU Cache to hold mutexes for different key types + const mutexCache = new LRUCache({ + ttl: 60 * 60 * 1000, // 1 hour + ttlAutopurge: true, + updateAgeOnGet: true + }) + let transactionsInProgress = 0 + function getKeyTypeMutex(type: string): Mutex { + return getMutex(`keytype:${type}`) + } + + function getSenderKeyMutex(senderKeyName: string): Mutex { + return getMutex(`senderkey:${senderKeyName}`) + } + + function getTransactionMutex(key: string): Mutex { + return getMutex(`transaction:${key}`) + } + + // Get or create a mutex for a specific key name + function getMutex(key: string): Mutex { + let mutex = mutexCache.get(key) + if (!mutex) { + mutex = new Mutex() + mutexCache.set(key, mutex) + logger.info({ key }, 'created new mutex') + } + + return mutex + } + + // Sender key operations with proper mutex sequencing + function queueSenderKeyOperation(senderKeyName: string, operation: () => Promise): Promise { + return getSenderKeyMutex(senderKeyName).runExclusive(operation) + } + + // Check if we are currently in a transaction + function isInTransaction() { + return transactionsInProgress > 0 + } + + // Helper function to handle transaction commit with retries + async function commitTransaction(): Promise { + if (!Object.keys(mutations).length) { + logger.trace('no mutations in transaction') + return + } + + logger.trace('committing transaction') + let tries = maxCommitRetries + + while (tries > 0) { + tries -= 1 + try { + await state.set(mutations) + logger.trace({ dbQueriesInTransaction }, 'committed transaction') + return + } catch (error) { + logger.warn(`failed to commit ${Object.keys(mutations).length} mutations, tries left=${tries}`) + if (tries > 0) { + await delay(delayBetweenTriesMs) + } + } + } + } + + // Helper function to clean up transaction state + function cleanupTransactionState(): void { + transactionsInProgress -= 1 + if (transactionsInProgress === 0) { + transactionCache = {} + mutations = {} + dbQueriesInTransaction = 0 + } + } + + // Helper function to execute work within transaction + async function executeTransactionWork(work: () => Promise): Promise { + const result = await work() + + // commit if this is the outermost transaction + if (transactionsInProgress === 1) { + await commitTransaction() + } + + return result + } + return { get: async (type, ids) => { if (isInTransaction()) { @@ -112,10 +384,31 @@ export const addTransactionCapability = ( // only fetch if there are any items to fetch if (idsRequiringFetch.length) { dbQueriesInTransaction += 1 - const result = await state.get(type, idsRequiringFetch) - transactionCache[type] ||= {} - Object.assign(transactionCache[type]!, result) + // Use per-sender-key queue for sender-key operations when possible + if (type === 'sender-key') { + logger.info({ idsRequiringFetch }, 'processing sender keys in transaction') + // For sender keys, process each one with queued operations to maintain serialization + for (const senderKeyName of idsRequiringFetch) { + await queueSenderKeyOperation(senderKeyName, async () => { + logger.info({ senderKeyName }, 'fetching sender key in transaction') + const result = await state.get(type, [senderKeyName]) + // Update transaction cache + transactionCache[type] ||= {} + Object.assign(transactionCache[type]!, result) + logger.info({ senderKeyName, hasResult: !!result[senderKeyName] }, 'sender key fetch complete') + }) + } + } else { + // Use runExclusive for cleaner mutex handling + await getKeyTypeMutex(type as string).runExclusive(async () => { + const result = await state.get(type, idsRequiringFetch) + + // Update transaction cache + transactionCache[type] ||= {} + Object.assign(transactionCache[type]!, result) + }) + } } return ids.reduce((dict: { [T in string]: any }, id) => { @@ -127,73 +420,123 @@ export const addTransactionCapability = ( return dict }, {}) } else { - return state.get(type, ids) + // Not in transaction, fetch directly with queue protection + if (type === 'sender-key') { + // For sender keys, use individual queues to maintain per-key serialization + const results: { [key: string]: SignalDataTypeMap[typeof type] } = {} + for (const senderKeyName of ids) { + const result = await queueSenderKeyOperation( + senderKeyName, + async () => await state.get(type, [senderKeyName]) + ) + Object.assign(results, result) + } + + return results + } else { + return await getKeyTypeMutex(type as string).runExclusive(() => state.get(type, ids)) + } } }, - set: data => { + set: async data => { if (isInTransaction()) { logger.trace({ types: Object.keys(data) }, 'caching in transaction') for (const key_ in data) { const key = key_ as keyof SignalDataTypeMap transactionCache[key] = transactionCache[key] || ({} as any) - Object.assign(transactionCache[key]!, data[key]) - mutations[key] = mutations[key] || ({} as any) - Object.assign(mutations[key]!, data[key]) + // Special handling for pre-keys and signed-pre-keys + if (key === 'pre-key') { + await handlePreKeyOperations(data, key, transactionCache, mutations, logger, true) + } else { + // Normal handling for other key types + handleNormalKeyOperations(data, key, transactionCache, mutations) + } } } else { - return state.set(data) + // Not in transaction, apply directly with mutex protection + const hasSenderKeys = 'sender-key' in data + const senderKeyNames = hasSenderKeys ? Object.keys(data['sender-key'] || {}) : [] + + if (hasSenderKeys) { + logger.info({ senderKeyNames }, 'processing sender key set operations') + // Handle sender key operations with per-key queues + for (const senderKeyName of senderKeyNames) { + await queueSenderKeyOperation(senderKeyName, async () => { + // Create data subset for this specific sender key + const senderKeyData = { + 'sender-key': { + [senderKeyName]: data['sender-key']![senderKeyName] + } + } + + logger.trace({ senderKeyName }, 'storing sender key') + // Apply changes to the store + await state.set(senderKeyData as SignalDataSet) + logger.trace({ senderKeyName }, 'sender key stored') + }) + } + + // Handle any non-sender-key data with regular mutexes + const nonSenderKeyData = { ...data } + delete nonSenderKeyData['sender-key'] + + if (Object.keys(nonSenderKeyData).length > 0) { + await withMutexes(Object.keys(nonSenderKeyData), getKeyTypeMutex, async () => { + // Process pre-keys and signed-pre-keys separately with specialized mutexes + for (const key_ in nonSenderKeyData) { + const keyType = key_ as keyof SignalDataTypeMap + if (keyType === 'pre-key') { + await processPreKeyDeletions(nonSenderKeyData, keyType, state, logger) + } + } + + // Apply changes to the store + await state.set(nonSenderKeyData) + }) + } + } else { + // No sender keys - use original logic + await withMutexes(Object.keys(data), getKeyTypeMutex, async () => { + // Process pre-keys and signed-pre-keys separately with specialized mutexes + for (const key_ in data) { + const keyType = key_ as keyof SignalDataTypeMap + if (keyType === 'pre-key') { + await processPreKeyDeletions(data, keyType, state, logger) + } + } + + // Apply changes to the store + await state.set(data) + }) + } } }, isInTransaction, - async transaction(work) { - let result: Awaited> - transactionsInProgress += 1 - if (transactionsInProgress === 1) { - logger.trace('entering transaction') - } + async transaction(work, key) { + const releaseTxMutex = await getTransactionMutex(key).acquire() try { - result = await work() - // commit if this is the outermost transaction + transactionsInProgress += 1 if (transactionsInProgress === 1) { - if (Object.keys(mutations).length) { - logger.trace('committing transaction') - // retry mechanism to ensure we've some recovery - // in case a transaction fails in the first attempt - let tries = maxCommitRetries - while (tries) { - tries -= 1 - //eslint-disable-next-line max-depth - try { - await state.set(mutations) - logger.trace({ dbQueriesInTransaction }, 'committed transaction') - break - } catch (error) { - logger.warn(`failed to commit ${Object.keys(mutations).length} mutations, tries left=${tries}`) - await delay(delayBetweenTriesMs) - } - } - } else { - logger.trace('no mutations in transaction') - } + logger.trace('entering transaction') } - } finally { - transactionsInProgress -= 1 - if (transactionsInProgress === 0) { - transactionCache = {} - mutations = {} - dbQueriesInTransaction = 0 + + // Release the transaction mutex now that we've updated the counter + // This allows other transactions to start preparing + releaseTxMutex() + + try { + return await executeTransactionWork(work) + } finally { + cleanupTransactionState() } + } catch (error) { + releaseTxMutex() + throw error } - - return result } } - - function isInTransaction() { - return transactionsInProgress > 0 - } } export const initAuthCreds = (): AuthenticationCreds => { diff --git a/src/Utils/decode-wa-message.ts b/src/Utils/decode-wa-message.ts index 90f61839..685518d1 100644 --- a/src/Utils/decode-wa-message.ts +++ b/src/Utils/decode-wa-message.ts @@ -58,6 +58,13 @@ const storeMappingFromEnvelope = async ( export const NO_MESSAGE_FOUND_ERROR_TEXT = 'Message absent from node' export const MISSING_KEYS_ERROR_TEXT = 'Key used already or never filled' +// Retry configuration for failed decryption +export const DECRYPTION_RETRY_CONFIG = { + maxRetries: 3, + baseDelayMs: 100, + sessionRecordErrors: ['No session record', 'SessionError: No session record'] +} + export const NACK_REASONS = { ParsingError: 487, UnrecognizedStanza: 488, @@ -238,6 +245,7 @@ export const decryptMessageNode = ( try { const e2eType = tag === 'plaintext' ? 'plaintext' : attrs.type + switch (e2eType) { case 'skmsg': msgBuffer = await repository.decryptGroupMessage({ @@ -278,7 +286,7 @@ export const decryptMessageNode = ( item: msg.senderKeyDistributionMessage }) } catch (err) { - logger.error({ key: fullMessage.key, err }, 'failed to decrypt message') + logger.error({ key: fullMessage.key, err }, 'failed to process sender key distribution message') } } @@ -288,7 +296,17 @@ export const decryptMessageNode = ( fullMessage.message = msg } } catch (err: any) { - logger.error({ key: fullMessage.key, err }, 'failed to decrypt message') + const errorContext = { + key: fullMessage.key, + err, + messageType: tag === 'plaintext' ? 'plaintext' : attrs.type, + sender, + author, + isSessionRecordError: isSessionRecordError(err) + } + + logger.error(errorContext, 'failed to decrypt message') + fullMessage.messageStubType = proto.WebMessageInfo.StubType.CIPHERTEXT fullMessage.messageStubParameters = [err.message] } @@ -303,3 +321,12 @@ export const decryptMessageNode = ( } } } + +/** + * Utility function to check if an error is related to missing session record + */ +function isSessionRecordError(error: any): boolean { + const errorMessage = error?.message || error?.toString() || '' + return DECRYPTION_RETRY_CONFIG.sessionRecordErrors.some(errorPattern => errorMessage.includes(errorPattern)) +} + diff --git a/src/Utils/process-message.ts b/src/Utils/process-message.ts index 4951f352..dfd62211 100644 --- a/src/Utils/process-message.ts +++ b/src/Utils/process-message.ts @@ -225,7 +225,7 @@ const processMessage = async ( } logger?.info({ newAppStateSyncKeyId, newKeys }, 'injecting new app state sync keys') - }) + }, meId) ev.emit('creds.update', { myAppStateKeyId: newAppStateSyncKeyId }) } else {