Async + Multi Cache (#1741)

* 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 <devastro0010@gmail.com>

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)

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<number>(key) || 0
+			let retryCount = (await msgRetryCache.get<number>(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<number>(key) || 1
+		const retryCount = (await msgRetryCache.get<number>(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<number>(key) || 0
-		return retryCount <= maxMsgRetryCount
+		const retryCount = (await msgRetryCache.get<number>(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<number>(key) || 0) + 1
-		msgRetryCache.set(key, newValue)
+		const newValue = ((await msgRetryCache.get<number>(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<WACallEvent>(call.id)
+		const existingCall = await callOfferCache.get<WACallEvent>(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])

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<string, JidWithDevice[] | undefined>
 
+		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])
+				}
 			}
 		}
 

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<T>(key: string): T | undefined
+	get<T>(key: string): Promise<T> | T | undefined
 	/** set a key in the cache */
-	set<T>(key: string, value: T): void
+	set<T>(key: string, value: T): Promise<void> | void | number | boolean
 	/** delete a key from the cache */
-	del(key: string): void
+	del(key: string): void | Promise<void> | number | boolean
 	/** flush all data */
-	flushAll(): void
+	flushAll(): void | Promise<void>
+}
+
+export type PossiblyExtendedCacheStore = CacheStore & {
+	mget?: <T>(keys: string[]) => Promise<Record<string, T | undefined>>
+	mset?: <T>(entries: { key: string; value: T }[]) => Promise<void> | void | number | boolean
+	mdel?: (keys: string[]) => void | Promise<void> | 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 */

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?.()
 		}
 	}

src/Utils/messages.ts

@@ -152,7 +152,7 @@ export const prepareWAMessageMedia = async (
 	}
 
 	if (cacheableKey) {
-		const mediaBuff = options.mediaCache!.get<Buffer>(cacheableKey)
+		const mediaBuff = await options.mediaCache!.get<Buffer>(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

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