lid-mapping, messages-send: Batch getLIDForPN calls

Author: purpshell

src/Signal/lid-mapping.ts

@@ -1,7 +1,7 @@
 import { LRUCache } from 'lru-cache'
 import type { LIDMapping, SignalKeyStoreWithTransaction } from '../Types'
 import type { ILogger } from '../Utils/logger'
-import { isLidUser, isPnUser, jidDecode } from '../WABinary'
+import { isLidUser, isPnUser, jidDecode, jidNormalizedUser } from '../WABinary'
 
 export class LIDMappingStore {
 	private readonly mappingCache = new LRUCache<string, string>({
@@ -89,7 +89,7 @@ export class LIDMappingStore {
 	}
 
 	async getLIDsForPNs(pns: string[]): Promise<LIDMapping[] | null> {
-		const usyncFetch: string[] = []
+		const usyncFetch: { [_: string]: number[] } = {}
 		// mapped from pn to lid mapping to prevent duplication in results later
 		const successfulPairs: { [_: string]: LIDMapping } = {}
 		for (const pn of pns) {
@@ -112,7 +112,14 @@ export class LIDMappingStore {
 					this.mappingCache.set(`lid:${lidUser}`, pnUser)
 				} else {
 					this.logger.trace(`No LID mapping found for PN user ${pnUser}; batch getting from USync`)
-					usyncFetch.push(pn)
+					const device = decoded.device || 0
+					const normalizedPn = jidNormalizedUser(pn)
+					if (!usyncFetch[normalizedPn]) {
+						usyncFetch[normalizedPn] = [device]
+					} else {
+						usyncFetch[normalizedPn].push(device)
+					}
+
 					continue
 				}
 			}
@@ -125,20 +132,32 @@ export class LIDMappingStore {
 
 			// Push the PN device ID to the LID to maintain device separation
 			const pnDevice = decoded.device !== undefined ? decoded.device : 0
-			const deviceSpecificLid = `${lidUser}:${pnDevice}@lid`
+			const deviceSpecificLid = `${lidUser}${!!pnDevice ? `:${pnDevice}` : ``}@lid`
 
 			this.logger.trace(`getLIDForPN: ${pn} → ${deviceSpecificLid} (user mapping with device ${pnDevice})`)
 			successfulPairs[pn] = { lid: deviceSpecificLid, pn }
 		}
 
-		if (usyncFetch.length > 0) {
-			const result = await this.pnToLIDFunc?.(usyncFetch) // this function already adds LIDs to mapping
+		if (Object.keys(usyncFetch).length > 0) {
+			const result = await this.pnToLIDFunc?.(Object.keys(usyncFetch)) // this function already adds LIDs to mapping
 			if (result && result.length > 0) {
 				this.storeLIDPNMappings(result)
 				for (const pair of result) {
+					const pnUser = jidDecode(pair.pn)?.user
+					if (!pnUser) continue
 					const lidUser = jidDecode(pair.lid)?.user
-						if (lidUser) {
-						successfulPairs[pair.pn] = pair
+					if (!lidUser) continue
+
+					for (const device of usyncFetch[pair.pn]!) {
+						const deviceSpecificLid = `${lidUser}${!!device ? `:${device}` : ``}@lid`
+
+						this.logger.trace(
+							`getLIDForPN: USYNC success for ${pair.pn} → ${deviceSpecificLid} (user mapping with device ${device})`
+						)
+
+						const deviceSpecificPn = `${pnUser}${!!device ? `:${device}` : ``}@s.whatsapp.net`
+
+						successfulPairs[deviceSpecificPn] = { lid: deviceSpecificLid, pn: deviceSpecificPn }
 					}
 				}
 			} else {

src/Socket/messages-send.ts

@@ -392,16 +392,14 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 		if (jidsRequiringFetch.length) {
 			// LID if mapped, otherwise original
-			const wireJids = await Promise.all(
-				jidsRequiringFetch.map(async jid => {
-					if (jid.includes('@s.whatsapp.net')) {
-						const lid = await signalRepository.lidMapping.getLIDForPN(jid)
-						return lid ? lid : jid
-					}
-
-					return jid
-				})
-			)
+			const wireJids = [
+				...jidsRequiringFetch.filter(jid => !!jid.includes('@lid')),
+				...(
+					(await signalRepository.lidMapping.getLIDsForPNs(
+						jidsRequiringFetch.filter(jid => !!jid.includes('@s.whatsapp.net'))
+					)) || []
+				).map(a => a.lid)
+			]
 
 			logger.debug({ jidsRequiringFetch, wireJids }, 'fetching sessions')
 			const result = await query({