From 8029446511e4fd4eea7d19710e1b88ae905ee9d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mart=C3=ADn=20Schere?= <61954243+MartinSchere@users.noreply.github.com> Date: Sun, 14 Sep 2025 16:13:11 +0200 Subject: [PATCH] Async + Multi Cache (#1741) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: handle invalid signatureKeyPublic types in sender-key-state - added extra check to ensure signatureKeyPublic is either a base64 string or a Buffer - fallback to empty buffer when unexpected object type is received - prevents ERR_INVALID_ARG_TYPE error in Buffer.from * adjusted based on @jlucaso1 recommmendation * fix: handle chain key objects for skmsg group message decryption previously, some sender chain keys were stored or retrieved as plain objects (e.g. { '0': 85, '1': 100, ... }) instead of Buffers. this caused "ERR_INVALID_ARG_TYPE" during initial decryption of skmsg group messages. this fixes any chain key object is converted to a proper Buffer before being passed to SenderChainKey. * fix: add more robust checks and fallbacks * feat: async cache feat: async caching * fix: linting issues * fix: mget logic --------- Co-authored-by: αѕтяσχ11 --- src/Signal/Group/sender-chain-key.ts | 1 - src/Socket/messages-recv.ts | 36 +++++++++--------- src/Socket/messages-send.ts | 55 ++++++++++++++++++---------- src/Types/Socket.ts | 16 +++++--- src/Utils/auth-utils.ts | 4 +- src/Utils/messages.ts | 6 +-- src/Utils/process-message.ts | 2 +- 7 files changed, 70 insertions(+), 50 deletions(-) diff --git a/src/Signal/Group/sender-chain-key.ts b/src/Signal/Group/sender-chain-key.ts index 604c3aa8..ffe11b1a 100644 --- a/src/Signal/Group/sender-chain-key.ts +++ b/src/Signal/Group/sender-chain-key.ts @@ -10,7 +10,6 @@ export class SenderChainKey { constructor(iteration: number, chainKey: Uint8Array | Buffer) { this.iteration = iteration if (Buffer.isBuffer(chainKey)) { - // backported from @MartinSchere's PR this.chainKey = chainKey } else if (chainKey instanceof Uint8Array) { this.chainKey = Buffer.from(chainKey) diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts index c835b31b..603f154c 100644 --- a/src/Socket/messages-recv.ts +++ b/src/Socket/messages-recv.ts @@ -400,7 +400,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { } else { // Fallback to old system const key = `${msgId}:${msgKey?.participant}` - let retryCount = msgRetryCache.get(key) || 0 + let retryCount = (await msgRetryCache.get(key)) || 0 if (retryCount >= maxMsgRetryCount) { logger.debug({ retryCount, msgId }, 'reached retry limit, clearing') msgRetryCache.del(key) @@ -408,11 +408,11 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { } retryCount += 1 - msgRetryCache.set(key, retryCount) + await msgRetryCache.set(key, retryCount) } const key = `${msgId}:${msgKey?.participant}` - const retryCount = msgRetryCache.get(key) || 1 + const retryCount = (await msgRetryCache.get(key)) || 1 const { account, signedPreKey, signedIdentityKey: identityKey } = authState.creds const fromJid = node.attrs.from! @@ -862,16 +862,16 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { return data instanceof Buffer ? data : Buffer.from(data) } - const willSendMessageAgain = (id: string, participant: string) => { + const willSendMessageAgain = async (id: string, participant: string) => { const key = `${id}:${participant}` - const retryCount = msgRetryCache.get(key) || 0 - return retryCount <= maxMsgRetryCount + const retryCount = (await msgRetryCache.get(key)) || 0 + return retryCount < maxMsgRetryCount } - const updateSendMessageAgainCount = (id: string, participant: string) => { + const updateSendMessageAgainCount = async (id: string, participant: string) => { const key = `${id}:${participant}` - const newValue = (msgRetryCache.get(key) || 0) + 1 - msgRetryCache.set(key, newValue) + const newValue = ((await msgRetryCache.get(key)) || 0) + 1 + await msgRetryCache.set(key, newValue) } const sendMessagesAgain = async (key: proto.IMessageKey, ids: string[], retryNode: BinaryNode) => { @@ -950,7 +950,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { for (const [i, msg] of msgs.entries()) { if (!ids[i]) continue - if (msg && willSendMessageAgain(ids[i], participant)) { + if (msg && (await willSendMessageAgain(ids[i], participant))) { updateSendMessageAgainCount(ids[i], participant) const msgRelayOpts: MessageRelayOptions = { messageId: ids[i] } @@ -1039,7 +1039,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { // correctly set who is asking for the retry key.participant = key.participant || attrs.from const retryNode = getBinaryNodeChild(node, 'retry') - if (ids[0] && key.participant && willSendMessageAgain(ids[0], key.participant)) { + if (ids[0] && key.participant && (await willSendMessageAgain(ids[0], key.participant))) { if (key.fromMe) { try { updateSendMessageAgainCount(ids[0], key.participant) @@ -1127,8 +1127,8 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { logger.debug('received unavailable message, acked and requested resend from phone') } else { - if (placeholderResendCache.get(node.attrs.id!)) { - placeholderResendCache.del(node.attrs.id!) + if (await placeholderResendCache.get(node.attrs.id!)) { + await placeholderResendCache.del(node.attrs.id!) } } @@ -1292,7 +1292,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { status = getCallStatusFromNode(infoChild) if (isLidUser(from) && infoChild.tag === 'relaylatency') { - const verify = callOfferCache.get(callId) + const verify = await callOfferCache.get(callId) if (!verify) { status = 'offer' const callLid: WACallEvent = { @@ -1303,7 +1303,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { offline: !!attrs.offline, status } - callOfferCache.set(callId, callLid) + await callOfferCache.set(callId, callLid) } } @@ -1320,10 +1320,10 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { call.isVideo = !!getBinaryNodeChild(infoChild, 'video') call.isGroup = infoChild.attrs.type === 'group' || !!infoChild.attrs['group-jid'] call.groupJid = infoChild.attrs['group-jid'] - callOfferCache.set(call.id, call) + await callOfferCache.set(call.id, call) } - const existingCall = callOfferCache.get(call.id) + const existingCall = await callOfferCache.get(call.id) // use existing call info to populate this event if (existingCall) { @@ -1333,7 +1333,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { // delete data once call has ended if (status === 'reject' || status === 'accept' || status === 'timeout' || status === 'terminate') { - callOfferCache.del(call.id) + await callOfferCache.del(call.id) } ev.emit('call', [call]) diff --git a/src/Socket/messages-send.ts b/src/Socket/messages-send.ts index 742efb14..91677227 100644 --- a/src/Socket/messages-send.ts +++ b/src/Socket/messages-send.ts @@ -245,30 +245,40 @@ export const makeMessagesSocket = (config: SocketConfig) => { } const toFetch: string[] = [] - // Deduplicate and normalize JIDs jids = deduplicateLidPnJids(Array.from(new Set(jids))) + const jidsWithUser = jids + .map(jid => { + const decoded = jidDecode(jid) + const user = decoded?.user + const device = decoded?.device + const isExplicitDevice = typeof device === 'number' && device >= 0 - for (let jid of jids) { - const decoded = jidDecode(jid) - const user = decoded?.user - const device = decoded?.device - const isExplicitDevice = typeof device === 'number' && device >= 0 + if (isExplicitDevice && user) { + deviceResults.push({ + user, + device, + wireJid: jid + }) + return null + } - // Handle explicit device JIDs directly - if (isExplicitDevice && user) { - deviceResults.push({ - user, - device, - wireJid: jid // Preserve exact JID format for wire protocol - }) - continue - } + jid = jidNormalizedUser(jid) + return { jid, user } + }) + .filter(jid => jid !== null) - // For user JIDs, normalize and prepare for device enumeration - jid = jidNormalizedUser(jid) + let mgetDevices: undefined | Record + if (useCache && userDevicesCache.mget) { + const usersToFetch = jidsWithUser.map(j => j?.user).filter(Boolean) as string[] + mgetDevices = await userDevicesCache.mget(usersToFetch) + } + + for (const { jid, user } of jidsWithUser) { if (useCache) { - const devices = userDevicesCache.get(user!) as JidWithDevice[] + const devices = + mgetDevices?.[user!] || + (userDevicesCache.mget ? undefined : ((await userDevicesCache.get(user!)) as JidWithDevice[])) if (devices) { const isLidJid = jid.includes('@lid') const devicesWithWire = devices.map(d => ({ @@ -342,8 +352,13 @@ export const makeMessagesSocket = (config: SocketConfig) => { } } - for (const key in deviceMap) { - userDevicesCache.set(key, deviceMap[key]!) + if (userDevicesCache.mset) { + // if the cache supports mset, we can set all devices in one go + await userDevicesCache.mset(Object.entries(deviceMap).map(([key, value]) => ({ key, value }))) + } else { + for (const key in deviceMap) { + if (deviceMap[key]) await userDevicesCache.set(key, deviceMap[key]) + } } } diff --git a/src/Types/Socket.ts b/src/Types/Socket.ts index 7dcec711..09ca891c 100644 --- a/src/Types/Socket.ts +++ b/src/Types/Socket.ts @@ -13,13 +13,19 @@ export type WABrowserDescription = [string, string, string] export type CacheStore = { /** get a cached key and change the stats */ - get(key: string): T | undefined + get(key: string): Promise | T | undefined /** set a key in the cache */ - set(key: string, value: T): void + set(key: string, value: T): Promise | void | number | boolean /** delete a key from the cache */ - del(key: string): void + del(key: string): void | Promise | number | boolean /** flush all data */ - flushAll(): void + flushAll(): void | Promise +} + +export type PossiblyExtendedCacheStore = CacheStore & { + mget?: (keys: string[]) => Promise> + mset?: (entries: { key: string; value: T }[]) => Promise | void | number | boolean + mdel?: (keys: string[]) => void | Promise | number | boolean } export type PatchedMessageWithRecipientJID = proto.IMessage & { recipientJid?: string } @@ -78,7 +84,7 @@ export type SocketConfig = { * used to determine whether to retry a message or not */ msgRetryCounterCache?: CacheStore /** provide a cache to store a user's device list */ - userDevicesCache?: CacheStore + userDevicesCache?: PossiblyExtendedCacheStore /** cache to store call offers */ callOfferCache?: CacheStore /** cache to track placeholder resends */ diff --git a/src/Utils/auth-utils.ts b/src/Utils/auth-utils.ts index 1dc18ae1..d1bf44b1 100644 --- a/src/Utils/auth-utils.ts +++ b/src/Utils/auth-utils.ts @@ -76,7 +76,7 @@ export function makeCacheableSignalKeyStore( 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]!) + await cache.set(getUniqueId(type, id), data[type as keyof SignalDataTypeMap]![id]!) keys += 1 } } @@ -87,7 +87,7 @@ export function makeCacheableSignalKeyStore( }) }, async clear() { - cache.flushAll() + await cache.flushAll() await store.clear?.() } } diff --git a/src/Utils/messages.ts b/src/Utils/messages.ts index 7618b842..f7cd7236 100644 --- a/src/Utils/messages.ts +++ b/src/Utils/messages.ts @@ -152,7 +152,7 @@ export const prepareWAMessageMedia = async ( } if (cacheableKey) { - const mediaBuff = options.mediaCache!.get(cacheableKey) + const mediaBuff = await options.mediaCache!.get(cacheableKey) if (mediaBuff) { logger?.debug({ cacheableKey }, 'got media cache hit') @@ -202,7 +202,7 @@ export const prepareWAMessageMedia = async ( if (cacheableKey) { logger?.debug({ cacheableKey }, 'set cache') - options.mediaCache!.set(cacheableKey, WAProto.Message.encode(obj).finish()) + await options.mediaCache!.set(cacheableKey, WAProto.Message.encode(obj).finish()) } return obj @@ -305,7 +305,7 @@ export const prepareWAMessageMedia = async ( if (cacheableKey) { logger?.debug({ cacheableKey }, 'set cache') - options.mediaCache!.set(cacheableKey, WAProto.Message.encode(obj).finish()) + await options.mediaCache!.set(cacheableKey, WAProto.Message.encode(obj).finish()) } return obj diff --git a/src/Utils/process-message.ts b/src/Utils/process-message.ts index 73bedc3a..258d342c 100644 --- a/src/Utils/process-message.ts +++ b/src/Utils/process-message.ts @@ -264,7 +264,7 @@ const processMessage = async ( case proto.Message.ProtocolMessage.Type.PEER_DATA_OPERATION_REQUEST_RESPONSE_MESSAGE: const response = protocolMsg.peerDataOperationRequestResponseMessage! if (response) { - placeholderResendCache?.del(response.stanzaId!) + await placeholderResendCache?.del(response.stanzaId!) // TODO: IMPLEMENT HISTORY SYNC ETC (sticker uploads etc.). const { peerDataOperationResult } = response for (const result of peerDataOperationResult!) {