Add mutex-based transaction safety to Signal key store (re-reworked) (#1697)

* mutex reimplementation

* lint

* lint

* lint fix

* Add cleanup for expired sender key mutexes

Introduces a mechanism to periodically clean up unused sender key mutexes in addTransactionCapability, reducing memory usage by removing mutexes that have not been used for over an hour and are not locked. A timer is started when the first mutex is created, and cleanup runs every 30 minutes.

* Update auth-utils.ts

* Refactor Signal key transaction usage

Introduces a local variable for parsed Signal keys in libsignal.ts to avoid repeated type assertions and improve code clarity.

* Refactor transaction handling for key operations

Introduces per-entity mutexes and passes key identifiers to transaction calls for finer-grained concurrency control in Signal key storage and socket operations. Updates transaction signatures and usage across Signal, Socket, and Utils modules to improve atomicity and performance, especially for system and sync messages.

* Improve SignalKeyStore transaction handling

Refactored transaction logic in SignalKeyStore to add commit retries, cleanup of transaction state, and improved mutex management for sender keys. Enhanced validation and batching for pre-key deletions and updates, improving concurrency and reliability.

* Replace custom mutex cleanup with LRU cache

Switched from manual mutex cleanup logic to using the lru-cache package for managing mutexes in addTransactionCapability. This simplifies resource management and ensures expired mutexes are automatically purged. Added lru-cache as a dependency.

* Lint fix

* Update src/Signal/libsignal.ts

* Update libsignal.ts

---------

Co-authored-by: João Lucas <jlucaso@hotmail.com>
Co-authored-by: Rajeh Taher <rajeh@reforward.dev>

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<SenderKeyRecord>
+
 	storeSenderKey(senderKeyName: SenderKeyName, record: SenderKeyRecord): Promise<void>
 }
 
@@ -20,64 +20,56 @@ export class GroupCipher {
 		this.senderKeyName = senderKeyName
 	}
 
-	private queueJob<T>(awaitable: () => Promise<T>): Promise<T> {
-		return queueJob(this.senderKeyName.toString(), awaitable)
-	}
-
 	public async encrypt(paddedPlaintext: Uint8Array | string): Promise<Uint8Array> {
-		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<Uint8Array> {
-		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) {

src/Signal/Group/queue-job.ts

@@ -1,70 +0,0 @@
-interface QueueJob<T> {
-	awaitable: () => Promise<T>
-	resolve: (value: T | PromiseLike<T>) => void
-	reject: (reason?: unknown) => void
-}
-
-const _queueAsyncBuckets = new Map<string | number, Array<QueueJob<any>>>()
-
-export function cleanupQueues() {
-	_queueAsyncBuckets.clear()
-}
-
-const _gcLimit = 10000
-
-async function _asyncQueueExecutor(queue: Array<QueueJob<any>>, cleanup: () => void): Promise<void> {
-	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<T>(bucket: string | number, awaitable: () => Promise<T>): Promise<T> {
-	// 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<T>((resolve, reject) => {
-		queue.push({
-			awaitable,
-			resolve: resolve as (value: any) => void,
-			reject
-		})
-	})
-
-	if (inactive) {
-		_asyncQueueExecutor(queue, () => _queueAsyncBuckets.delete(bucket))
-	}
-
-	return job
-}

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<string, boolean>({
 		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 }) {

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) {

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) => {

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
 	}

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 {

src/Types/Auth.ts

@@ -88,7 +88,7 @@ export type SignalKeyStore = {
 
 export type SignalKeyStoreWithTransaction = SignalKeyStore & {
 	isInTransaction: () => boolean
-	transaction<T>(exec: () => Promise<T>): Promise<T>
+	transaction<T>(exec: () => Promise<T>, key: string): Promise<T>
 }
 
 export type TransactionCapabilityOptions = {

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<SignalDataTypeMap[typeof type]>(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<SignalDataTypeMap[typeof type]>(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<void> {
+	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<void> {
+	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<T>(
+	keyTypes: string[],
+	getKeyTypeMutex: (type: string) => Mutex,
+	fn: () => Promise<T>
+): Promise<T> {
+	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<string, Mutex>({
+		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<T>(senderKeyName: string, operation: () => Promise<T>): Promise<T> {
+		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<void> {
+		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<any>): Promise<any> {
+		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<ReturnType<typeof work>>
-			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 => {

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))
+}
+

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 {