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>
This commit is contained in:
Paulo Victor Lund
2025-09-07 14:22:47 -03:00
committed by GitHub
parent ae0cb89714
commit f83a1c2aff
11 changed files with 576 additions and 236 deletions
+38 -46
View File
@@ -1,6 +1,5 @@
/* @ts-ignore */ /* @ts-ignore */
import { decrypt, encrypt } from 'libsignal/src/crypto' import { decrypt, encrypt } from 'libsignal/src/crypto'
import queueJob from './queue-job'
import { SenderKeyMessage } from './sender-key-message' import { SenderKeyMessage } from './sender-key-message'
import { SenderKeyName } from './sender-key-name' import { SenderKeyName } from './sender-key-name'
import { SenderKeyRecord } from './sender-key-record' import { SenderKeyRecord } from './sender-key-record'
@@ -8,6 +7,7 @@ import { SenderKeyState } from './sender-key-state'
export interface SenderKeyStore { export interface SenderKeyStore {
loadSenderKey(senderKeyName: SenderKeyName): Promise<SenderKeyRecord> loadSenderKey(senderKeyName: SenderKeyName): Promise<SenderKeyRecord>
storeSenderKey(senderKeyName: SenderKeyName, record: SenderKeyRecord): Promise<void> storeSenderKey(senderKeyName: SenderKeyName, record: SenderKeyRecord): Promise<void>
} }
@@ -20,64 +20,56 @@ export class GroupCipher {
this.senderKeyName = senderKeyName 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> { public async encrypt(paddedPlaintext: Uint8Array | string): Promise<Uint8Array> {
return await this.queueJob(async () => { const record = await this.senderKeyStore.loadSenderKey(this.senderKeyName)
const record = await this.senderKeyStore.loadSenderKey(this.senderKeyName) if (!record) {
if (!record) { throw new Error('No SenderKeyRecord found for encryption')
throw new Error('No SenderKeyRecord found for encryption') }
}
const senderKeyState = record.getSenderKeyState() const senderKeyState = record.getSenderKeyState()
if (!senderKeyState) { if (!senderKeyState) {
throw new Error('No session to encrypt message') throw new Error('No session to encrypt message')
} }
const iteration = senderKeyState.getSenderChainKey().getIteration() const iteration = senderKeyState.getSenderChainKey().getIteration()
const senderKey = this.getSenderKey(senderKeyState, iteration === 0 ? 0 : iteration + 1) 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( const senderKeyMessage = new SenderKeyMessage(
senderKeyState.getKeyId(), senderKeyState.getKeyId(),
senderKey.getIteration(), senderKey.getIteration(),
ciphertext, ciphertext,
senderKeyState.getSigningKeyPrivate() senderKeyState.getSigningKeyPrivate()
) )
await this.senderKeyStore.storeSenderKey(this.senderKeyName, record) await this.senderKeyStore.storeSenderKey(this.senderKeyName, record)
return senderKeyMessage.serialize() return senderKeyMessage.serialize()
})
} }
public async decrypt(senderKeyMessageBytes: Uint8Array): Promise<Uint8Array> { public async decrypt(senderKeyMessageBytes: Uint8Array): Promise<Uint8Array> {
return await this.queueJob(async () => { const record = await this.senderKeyStore.loadSenderKey(this.senderKeyName)
const record = await this.senderKeyStore.loadSenderKey(this.senderKeyName) if (!record) {
if (!record) { throw new Error('No SenderKeyRecord found for decryption')
throw new Error('No SenderKeyRecord found for decryption') }
}
const senderKeyMessage = new SenderKeyMessage(null, null, null, null, senderKeyMessageBytes) const senderKeyMessage = new SenderKeyMessage(null, null, null, null, senderKeyMessageBytes)
const senderKeyState = record.getSenderKeyState(senderKeyMessage.getKeyId()) const senderKeyState = record.getSenderKeyState(senderKeyMessage.getKeyId())
if (!senderKeyState) { if (!senderKeyState) {
throw new Error('No session found to decrypt message') throw new Error('No session found to decrypt message')
} }
senderKeyMessage.verifySignature(senderKeyState.getSigningKeyPublic()) senderKeyMessage.verifySignature(senderKeyState.getSigningKeyPublic())
const senderKey = this.getSenderKey(senderKeyState, senderKeyMessage.getIteration()) const senderKey = this.getSenderKey(senderKeyState, senderKeyMessage.getIteration())
const plaintext = await this.getPlainText( const plaintext = await this.getPlainText(
senderKey.getIv(), senderKey.getIv(),
senderKey.getCipherKey(), senderKey.getCipherKey(),
senderKeyMessage.getCipherText() senderKeyMessage.getCipherText()
) )
await this.senderKeyStore.storeSenderKey(this.senderKeyName, record) await this.senderKeyStore.storeSenderKey(this.senderKeyName, record)
return plaintext return plaintext
})
} }
private getSenderKey(senderKeyState: SenderKeyState, iteration: number) { private getSenderKey(senderKeyState: SenderKeyState, iteration: number) {
-70
View File
@@ -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
}
+80 -32
View File
@@ -16,6 +16,22 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository
const lidMapping = new LIDMappingStore(auth.keys as SignalKeyStoreWithTransaction) const lidMapping = new LIDMappingStore(auth.keys as SignalKeyStoreWithTransaction)
const storage = signalStorage(auth, lidMapping) 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) // Simple operation-level deduplication (5 minutes)
const recentMigrations = new LRUCache<string, boolean>({ const recentMigrations = new LRUCache<string, boolean>({
max: 500, max: 500,
@@ -23,11 +39,15 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository
}) })
const repository: SignalRepository = { const repository: SignalRepository = {
decryptGroupMessage({ group, authorJid, msg }) { decryptGroupMessage({ group, authorJid, msg }) {
const senderName = jidToSignalSenderKeyName(group, authorJid) const senderName = jidToSignalSenderKeyName(group, authorJid)
const cipher = new GroupCipher(storage, senderName) 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 }) { async processSenderKeyDistributionMessage({ item, authorJid }) {
const builder = new GroupSessionBuilder(storage) const builder = new GroupSessionBuilder(storage)
@@ -50,24 +70,44 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository
await storage.storeSenderKey(senderName, new SenderKeyRecord()) 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 }) { async decryptMessage({ jid, type, ciphertext }) {
const addr = jidToSignalProtocolAddress(jid) const addr = jidToSignalProtocolAddress(jid)
const session = new libsignal.SessionCipher(storage, addr) const session = new libsignal.SessionCipher(storage, addr)
let result: Buffer
switch (type) { async function doDecrypt() {
case 'pkmsg': let result: Buffer
result = await session.decryptPreKeyWhisperMessage(ciphertext) switch (type) {
break case 'pkmsg':
case 'msg': result = await session.decryptPreKeyWhisperMessage(ciphertext)
result = await session.decryptWhisperMessage(ciphertext) break
break case 'msg':
default: result = await session.decryptWhisperMessage(ciphertext)
throw new Error(`Unknown message type: ${type}`) 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 }) { async encryptMessage({ jid, data }) {
@@ -101,32 +141,40 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository
const addr = jidToSignalProtocolAddress(encryptionJid) const addr = jidToSignalProtocolAddress(encryptionJid)
const cipher = new libsignal.SessionCipher(storage, addr) const cipher = new libsignal.SessionCipher(storage, addr)
const { type: sigType, body } = await cipher.encrypt(data) // Use transaction to ensure atomicity
const type = sigType === 3 ? 'pkmsg' : 'msg' return parsedKeys.transaction(async () => {
return { type, ciphertext: Buffer.from(body, 'binary') } 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 }) { async encryptGroupMessage({ group, meId, data }) {
const senderName = jidToSignalSenderKeyName(group, meId) const senderName = jidToSignalSenderKeyName(group, meId)
const builder = new GroupSessionBuilder(storage) const builder = new GroupSessionBuilder(storage)
const senderNameStr = senderName.toString() 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) return parsedKeys.transaction(async () => {
const session = new GroupCipher(storage, senderName) const { [senderNameStr]: senderKey } = await auth.keys.get('sender-key', [senderNameStr])
const ciphertext = await session.encrypt(data) if (!senderKey) {
await storage.storeSenderKey(senderName, new SenderKeyRecord())
}
return { const senderKeyDistributionMessage = await builder.create(senderName)
ciphertext, const session = new GroupCipher(storage, senderName)
senderKeyDistributionMessage: senderKeyDistributionMessage.serialize() const ciphertext = await session.encrypt(data)
}
return {
ciphertext,
senderKeyDistributionMessage: senderKeyDistributionMessage.serialize()
}
}, group)
}, },
async injectE2ESession({ jid, session }) { async injectE2ESession({ jid, session }) {
const cipher = new libsignal.SessionBuilder(storage, jidToSignalProtocolAddress(jid)) 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) { async validateSession(jid: string) {
try { try {
@@ -180,9 +228,9 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository
async deleteSession(jid: string) { async deleteSession(jid: string) {
const addr = jidToSignalProtocolAddress(jid) const addr = jidToSignalProtocolAddress(jid)
return (auth.keys as SignalKeyStoreWithTransaction).transaction(async () => { return parsedKeys.transaction(async () => {
await auth.keys.set({ session: { [addr.toString()]: null } }) await auth.keys.set({ session: { [addr.toString()]: null } })
}) }, jid)
}, },
async migrateSession(fromJid: string, toJid: string) { async migrateSession(fromJid: string, toJid: string) {
@@ -211,7 +259,7 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository
return return
} }
return (auth.keys as SignalKeyStoreWithTransaction).transaction(async () => { return parsedKeys.transaction(async () => {
// Store mapping // Store mapping
await lidMapping.storeLIDPNMapping(toJid, fromJid) await lidMapping.storeLIDPNMapping(toJid, fromJid)
@@ -232,7 +280,7 @@ export function makeLibSignalRepository(auth: SignalAuthState): SignalRepository
} }
recentMigrations.set(migrationKey, true) recentMigrations.set(migrationKey, true)
}) }, fromJid)
}, },
async encryptMessageWithWire({ encryptionJid, wireJid, data }) { async encryptMessageWithWire({ encryptionJid, wireJid, data }) {
+2 -2
View File
@@ -601,7 +601,7 @@ export const makeChatsSocket = (config: SocketConfig) => {
} }
} }
} }
}) }, authState?.creds?.me?.id || 'resync-app-state')
const { onMutation } = newAppStateChunkHandler(isInitialSync) const { onMutation } = newAppStateChunkHandler(isInitialSync)
for (const key in globalMutationMap) { for (const key in globalMutationMap) {
@@ -807,7 +807,7 @@ export const makeChatsSocket = (config: SocketConfig) => {
await query(node) await query(node)
await authState.keys.set({ 'app-state-sync-version': { [name]: state } }) await authState.keys.set({ 'app-state-sync-version': { [name]: state } })
}) }, authState?.creds?.me?.id || 'app-patch')
}) })
if (config.emitOwnEvents) { if (config.emitOwnEvents) {
+2 -2
View File
@@ -307,8 +307,8 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
await sendNode(receipt) 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) => { const handleEncryptNotification = async (node: BinaryNode) => {
+2 -2
View File
@@ -1128,11 +1128,11 @@ export const makeMessagesSocket = (config: SocketConfig) => {
await sendNode(stanza) await sendNode(stanza)
// Add message to retry cache if enabled // Add message to retry cache if enabled
if (messageRetryManager && !participant) { if (messageRetryManager && !participant) {
messageRetryManager.addRecentMessage(destinationJid, msgId, message) messageRetryManager.addRecentMessage(destinationJid, msgId, message)
} }
}) }, meId)
return msgId return msgId
} }
+1 -1
View File
@@ -331,7 +331,7 @@ export const makeSocket = (config: SocketConfig) => {
// Update credentials immediately to prevent duplicate IDs on retry // Update credentials immediately to prevent duplicate IDs on retry
ev.emit('creds.update', update) ev.emit('creds.update', update)
return node // Only return node since update is already used 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) // Upload to server (outside transaction, can fail without affecting local keys)
try { try {
+1 -1
View File
@@ -88,7 +88,7 @@ export type SignalKeyStore = {
export type SignalKeyStoreWithTransaction = SignalKeyStore & { export type SignalKeyStoreWithTransaction = SignalKeyStore & {
isInTransaction: () => boolean isInTransaction: () => boolean
transaction<T>(exec: () => Promise<T>): Promise<T> transaction<T>(exec: () => Promise<T>, key: string): Promise<T>
} }
export type TransactionCapabilityOptions = { export type TransactionCapabilityOptions = {
+420 -77
View File
@@ -1,5 +1,7 @@
import NodeCache from '@cacheable/node-cache' import NodeCache from '@cacheable/node-cache'
import { Mutex } from 'async-mutex'
import { randomBytes } from 'crypto' import { randomBytes } from 'crypto'
import { LRUCache } from 'lru-cache'
import { DEFAULT_CACHE_TTLS } from '../Defaults' import { DEFAULT_CACHE_TTLS } from '../Defaults'
import type { import type {
AuthenticationCreds, AuthenticationCreds,
@@ -33,49 +35,56 @@ export function makeCacheableSignalKeyStore(
deleteOnExpire: true deleteOnExpire: true
}) })
// Mutex for protecting cache operations
const cacheMutex = new Mutex()
function getUniqueId(type: string, id: string) { function getUniqueId(type: string, id: string) {
return `${type}.${id}` return `${type}.${id}`
} }
return { return {
async get(type, ids) { async get(type, ids) {
const data: { [_: string]: SignalDataTypeMap[typeof type] } = {} return cacheMutex.runExclusive(async () => {
const idsToFetch: string[] = [] const data: { [_: string]: SignalDataTypeMap[typeof type] } = {}
for (const id of ids) { const idsToFetch: string[] = []
const item = cache.get<SignalDataTypeMap[typeof type]>(getUniqueId(type, id)) as any for (const id of ids) {
if (typeof item !== 'undefined') { const item = cache.get<SignalDataTypeMap[typeof type]>(getUniqueId(type, id)) as any
data[id] = item if (typeof item !== 'undefined') {
} 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) {
data[id] = item 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) { async set(data) {
let keys = 0 return cacheMutex.runExclusive(async () => {
for (const type in data) { let keys = 0
for (const id in data[type as keyof SignalDataTypeMap]) { for (const type in data) {
cache.set(getUniqueId(type, id), data[type as keyof SignalDataTypeMap]![id]!) for (const id in data[type as keyof SignalDataTypeMap]) {
keys += 1 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() { async clear() {
cache.flushAll() 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, * 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 * this allows batch read & write operations & improves the performance of the lib
@@ -102,8 +286,96 @@ export const addTransactionCapability = (
let transactionCache: SignalDataSet = {} let transactionCache: SignalDataSet = {}
let mutations: 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 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 { return {
get: async (type, ids) => { get: async (type, ids) => {
if (isInTransaction()) { if (isInTransaction()) {
@@ -112,10 +384,31 @@ export const addTransactionCapability = (
// only fetch if there are any items to fetch // only fetch if there are any items to fetch
if (idsRequiringFetch.length) { if (idsRequiringFetch.length) {
dbQueriesInTransaction += 1 dbQueriesInTransaction += 1
const result = await state.get(type, idsRequiringFetch)
transactionCache[type] ||= {} // Use per-sender-key queue for sender-key operations when possible
Object.assign(transactionCache[type]!, result) 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) => { return ids.reduce((dict: { [T in string]: any }, id) => {
@@ -127,73 +420,123 @@ export const addTransactionCapability = (
return dict return dict
}, {}) }, {})
} else { } 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()) { if (isInTransaction()) {
logger.trace({ types: Object.keys(data) }, 'caching in transaction') logger.trace({ types: Object.keys(data) }, 'caching in transaction')
for (const key_ in data) { for (const key_ in data) {
const key = key_ as keyof SignalDataTypeMap const key = key_ as keyof SignalDataTypeMap
transactionCache[key] = transactionCache[key] || ({} as any) transactionCache[key] = transactionCache[key] || ({} as any)
Object.assign(transactionCache[key]!, data[key])
mutations[key] = mutations[key] || ({} as any) // Special handling for pre-keys and signed-pre-keys
Object.assign(mutations[key]!, data[key]) 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 { } 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, isInTransaction,
async transaction(work) { async transaction(work, key) {
let result: Awaited<ReturnType<typeof work>> const releaseTxMutex = await getTransactionMutex(key).acquire()
transactionsInProgress += 1
if (transactionsInProgress === 1) {
logger.trace('entering transaction')
}
try { try {
result = await work() transactionsInProgress += 1
// commit if this is the outermost transaction
if (transactionsInProgress === 1) { if (transactionsInProgress === 1) {
if (Object.keys(mutations).length) { logger.trace('entering transaction')
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')
}
} }
} finally {
transactionsInProgress -= 1 // Release the transaction mutex now that we've updated the counter
if (transactionsInProgress === 0) { // This allows other transactions to start preparing
transactionCache = {} releaseTxMutex()
mutations = {}
dbQueriesInTransaction = 0 try {
return await executeTransactionWork(work)
} finally {
cleanupTransactionState()
} }
} catch (error) {
releaseTxMutex()
throw error
} }
return result
} }
} }
function isInTransaction() {
return transactionsInProgress > 0
}
} }
export const initAuthCreds = (): AuthenticationCreds => { export const initAuthCreds = (): AuthenticationCreds => {
+29 -2
View File
@@ -58,6 +58,13 @@ const storeMappingFromEnvelope = async (
export const NO_MESSAGE_FOUND_ERROR_TEXT = 'Message absent from node' export const NO_MESSAGE_FOUND_ERROR_TEXT = 'Message absent from node'
export const MISSING_KEYS_ERROR_TEXT = 'Key used already or never filled' 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 = { export const NACK_REASONS = {
ParsingError: 487, ParsingError: 487,
UnrecognizedStanza: 488, UnrecognizedStanza: 488,
@@ -238,6 +245,7 @@ export const decryptMessageNode = (
try { try {
const e2eType = tag === 'plaintext' ? 'plaintext' : attrs.type const e2eType = tag === 'plaintext' ? 'plaintext' : attrs.type
switch (e2eType) { switch (e2eType) {
case 'skmsg': case 'skmsg':
msgBuffer = await repository.decryptGroupMessage({ msgBuffer = await repository.decryptGroupMessage({
@@ -278,7 +286,7 @@ export const decryptMessageNode = (
item: msg.senderKeyDistributionMessage item: msg.senderKeyDistributionMessage
}) })
} catch (err) { } 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 fullMessage.message = msg
} }
} catch (err: any) { } 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.messageStubType = proto.WebMessageInfo.StubType.CIPHERTEXT
fullMessage.messageStubParameters = [err.message] 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))
}
+1 -1
View File
@@ -225,7 +225,7 @@ const processMessage = async (
} }
logger?.info({ newAppStateSyncKeyId, newKeys }, 'injecting new app state sync keys') logger?.info({ newAppStateSyncKeyId, newKeys }, 'injecting new app state sync keys')
}) }, meId)
ev.emit('creds.update', { myAppStateKeyId: newAppStateSyncKeyId }) ev.emit('creds.update', { myAppStateKeyId: newAppStateSyncKeyId })
} else { } else {