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