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