diff --git a/src/Signal/lid-mapping.ts b/src/Signal/lid-mapping.ts index 6ad1b1c1..2e3176fb 100644 --- a/src/Signal/lid-mapping.ts +++ b/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 { isHostedPnUser, isLidUser, isPnUser, jidDecode, jidNormalizedUser, WAJIDDomains } from '../WABinary' +import { isHostedLidUser, isHostedPnUser, isLidUser, isPnUser, jidDecode, jidNormalizedUser, WAJIDDomains } from '../WABinary' // ============================================================================ // CONFIGURATION @@ -450,7 +450,8 @@ export class LIDMappingStore { const usyncFetch: { [_: string]: number[] } = {} const successfulPairs: { [_: string]: LIDMapping } = {} - const failedPns: string[] = [] + const failedPns = new Set() + const pendingByPnUser = new Map }>>() for (const pn of pns) { if (!isPnUser(pn) && !isHostedPnUser(pn)) continue @@ -459,69 +460,112 @@ export class LIDMappingStore { if (!decoded) continue const pnUser = decoded.user - let lidUser = this.mappingCache.get(`pn:${pnUser}`) + const cachedLidUser = this.mappingCache.get(`pn:${pnUser}`) - if (lidUser) { + if (cachedLidUser) { this.stats.cacheHits++ - } else { - this.stats.cacheMisses++ - - // Cache miss - check database - try { - const stored = await this.retryOperation( - () => this.keys.get('lid-mapping', [pnUser]), - 'get-lid-for-pn' - ) - lidUser = stored[pnUser] - - if (lidUser) { - this.stats.dbHits++ - this.mappingCache.set(`pn:${pnUser}`, lidUser) - this.mappingCache.set(`lid:${lidUser}`, pnUser) - } else { - this.stats.dbMisses++ - - // Need to fetch from USync - if (this.config.debugLogging) { - this.logger.trace({ pnUser }, 'No LID mapping found, queuing for USync') - } - - const device = decoded.device || 0 - let normalizedPn = jidNormalizedUser(pn) - if (isHostedPnUser(normalizedPn)) { - normalizedPn = `${pnUser}@s.whatsapp.net` - } - - if (!usyncFetch[normalizedPn]) { - usyncFetch[normalizedPn] = [device] - } else { - usyncFetch[normalizedPn]?.push(device) - } - continue - } - } catch (error) { - this.logger.error({ error, pn }, 'Failed to get LID mapping from database') - this.stats.failedOperations++ - failedPns.push(pn) + const lidUser = cachedLidUser.toString() + if (!lidUser) { + this.logger.warn({ pn, lidUser }, 'Invalid or empty LID user') continue } - } - lidUser = lidUser?.toString() - if (!lidUser) { - this.logger.warn({ pn, lidUser }, 'Invalid or empty LID user') + const pnDevice = decoded.device ?? 0 + const deviceSpecificLid = this.buildDeviceSpecificJid( + lidUser, + pnDevice, + decoded.server === 'hosted' ? 'hosted.lid' : 'lid' + ) + + if (this.config.debugLogging) { + this.logger.trace({ pn, deviceSpecificLid, pnDevice }, 'getLIDForPN: mapping found') + } + + successfulPairs[pn] = { lid: deviceSpecificLid, pn } continue } - // Push the PN device ID to the LID to maintain device separation - const pnDevice = decoded.device ?? 0 - const deviceSpecificLid = this.buildDeviceSpecificJid(lidUser, pnDevice, decoded.server === 'hosted' ? 'hosted.lid' : 'lid') + this.stats.cacheMisses++ + const pendingForUser = pendingByPnUser.get(pnUser) ?? [] + pendingForUser.push({ pn, decoded }) + pendingByPnUser.set(pnUser, pendingForUser) + } - if (this.config.debugLogging) { - this.logger.trace({ pn, deviceSpecificLid, pnDevice }, 'getLIDForPN: mapping found') + if (pendingByPnUser.size > 0) { + const pnUsers = [...pendingByPnUser.keys()] + const dbFailedPnUsers = new Set() + + for (const batch of this.chunkArray(pnUsers, this.config.batchSize)) { + try { + const stored = await this.retryOperation( + () => this.keys.get('lid-mapping', batch), + 'get-lid-for-pn' + ) + + for (const pnUser of batch) { + const lidUser = stored[pnUser] + if (lidUser) { + this.stats.dbHits++ + this.mappingCache.set(`pn:${pnUser}`, lidUser) + this.mappingCache.set(`lid:${lidUser}`, pnUser) + } else { + this.stats.dbMisses++ + } + } + } catch (error) { + this.logger.error({ error, batch }, 'Failed to get LID mapping batch from database') + this.stats.failedOperations += batch.length + batch.forEach(pnUser => dbFailedPnUsers.add(pnUser)) + } } - successfulPairs[pn] = { lid: deviceSpecificLid, pn } + for (const [pnUser, items] of pendingByPnUser.entries()) { + const lidUser = this.mappingCache.get(`pn:${pnUser}`) + for (const { pn, decoded } of items) { + if (lidUser) { + const lidUserString = lidUser.toString() + if (!lidUserString) { + this.logger.warn({ pn, lidUser }, 'Invalid or empty LID user') + continue + } + + const pnDevice = decoded.device ?? 0 + const deviceSpecificLid = this.buildDeviceSpecificJid( + lidUserString, + pnDevice, + decoded.server === 'hosted' ? 'hosted.lid' : 'lid' + ) + + if (this.config.debugLogging) { + this.logger.trace({ pn, deviceSpecificLid, pnDevice }, 'getLIDForPN: mapping found') + } + + successfulPairs[pn] = { lid: deviceSpecificLid, pn } + continue + } + + if (dbFailedPnUsers.has(pnUser)) { + failedPns.add(pn) + } + + // Need to fetch from USync + if (this.config.debugLogging) { + this.logger.trace({ pnUser }, 'No LID mapping found, queuing for USync') + } + + const device = decoded.device || 0 + let normalizedPn = jidNormalizedUser(pn) + if (isHostedPnUser(normalizedPn)) { + normalizedPn = `${pnUser}@s.whatsapp.net` + } + + if (!usyncFetch[normalizedPn]) { + usyncFetch[normalizedPn] = [device] + } else { + usyncFetch[normalizedPn]?.push(device) + } + } + } } // Fetch from USync if needed @@ -530,9 +574,9 @@ export class LIDMappingStore { } // Log warning if some PNs failed lookup - if (failedPns.length > 0) { + if (failedPns.size > 0) { this.logger.warn( - { failedCount: failedPns.length, totalRequested: pns.length }, + { failedCount: failedPns.size, totalRequested: pns.length }, 'Some PNs failed during getLIDsForPNs - results may be incomplete' ) } @@ -545,59 +589,112 @@ export class LIDMappingStore { * Get PN for LID - USER LEVEL with device construction */ async getPNForLID(lid: string): Promise { + const results = await this.getPNsForLIDs([lid]) + return results?.[0]?.pn || null + } + + /** + * Get PNs for multiple LIDs - Optimized batch operation + */ + async getPNsForLIDs(lids: string[]): Promise { this.checkDestroyed() this.stats.totalOperations++ this.stats.lastOperationAt = Date.now() - if (!isLidUser(lid)) return null + const successfulPairs: { [_: string]: LIDMapping } = {} + const failedLids = new Set() + const pendingByLidUser = new Map }>>() - const decoded = jidDecode(lid) - if (!decoded) return null + const addResolvedPair = (lid: string, decoded: ReturnType, pnUser: string): void => { + const lidDevice = decoded?.device ?? 0 + const server = decoded?.domainType === WAJIDDomains.HOSTED_LID ? 'hosted' : 's.whatsapp.net' + const pnJid = this.buildDeviceSpecificJid(pnUser, lidDevice, server) - const lidUser = decoded.user - let pnUser = this.mappingCache.get(`lid:${lidUser}`) + if (this.config.debugLogging) { + this.logger.trace({ lid, pnJid }, 'Found reverse mapping') + } + + successfulPairs[lid] = { lid, pn: pnJid } + } + + for (const lid of lids) { + if (!isLidUser(lid) && !isHostedLidUser(lid)) continue + + const decoded = jidDecode(lid) + if (!decoded) continue + + const lidUser = decoded.user + const cachedPnUser = this.mappingCache.get(`lid:${lidUser}`) + + if (cachedPnUser && typeof cachedPnUser === 'string') { + this.stats.cacheHits++ + addResolvedPair(lid, decoded, cachedPnUser) + continue + } - if (pnUser) { - this.stats.cacheHits++ - } else { this.stats.cacheMisses++ + const pendingForUser = pendingByLidUser.get(lidUser) ?? [] + pendingForUser.push({ lid, decoded }) + pendingByLidUser.set(lidUser, pendingForUser) + } - // Cache miss - check database - try { - const stored = await this.retryOperation( - () => this.keys.get('lid-mapping', [`${lidUser}_reverse`]), - 'get-pn-for-lid' - ) - pnUser = stored[`${lidUser}_reverse`] + if (pendingByLidUser.size > 0) { + const reverseKeys = [...pendingByLidUser.keys()].map(lidUser => `${lidUser}_reverse`) + const dbFailedReverseKeys = new Set() + + for (const batch of this.chunkArray(reverseKeys, this.config.batchSize)) { + try { + const stored = await this.retryOperation( + () => this.keys.get('lid-mapping', batch), + 'get-pn-for-lid' + ) + + for (const reverseKey of batch) { + const lidUser = reverseKey.replace(/_reverse$/, '') + const pnUser = stored[reverseKey] + if (pnUser && typeof pnUser === 'string') { + this.stats.dbHits++ + this.mappingCache.set(`lid:${lidUser}`, pnUser) + this.mappingCache.set(`pn:${pnUser}`, lidUser) + } else { + this.stats.dbMisses++ + } + } + } catch (error) { + this.logger.error({ error, batch }, 'Failed to get PN mapping batch from database') + this.stats.failedOperations += batch.length + batch.forEach(reverseKey => dbFailedReverseKeys.add(reverseKey)) + } + } + + for (const [lidUser, items] of pendingByLidUser.entries()) { + const pnUser = this.mappingCache.get(`lid:${lidUser}`) + for (const { lid, decoded } of items) { + if (pnUser && typeof pnUser === 'string') { + addResolvedPair(lid, decoded, pnUser) + continue + } + + if (dbFailedReverseKeys.has(`${lidUser}_reverse`)) { + failedLids.add(lid) + } - if (!pnUser || typeof pnUser !== 'string') { - this.stats.dbMisses++ if (this.config.debugLogging) { this.logger.trace({ lidUser }, 'No reverse mapping found') } - return null } - - this.stats.dbHits++ - this.mappingCache.set(`lid:${lidUser}`, pnUser) - } catch (error) { - this.logger.error({ error, lid }, 'Failed to get PN for LID') - this.stats.failedOperations++ - return null } } - // Construct device-specific PN JID - const lidDevice = decoded.device ?? 0 - const server = decoded.domainType === WAJIDDomains.HOSTED_LID ? 'hosted' : 's.whatsapp.net' - const pnJid = this.buildDeviceSpecificJid(pnUser, lidDevice, server) - - if (this.config.debugLogging) { - this.logger.trace({ lid, pnJid }, 'Found reverse mapping') + if (failedLids.size > 0) { + this.logger.warn( + { failedCount: failedLids.size, totalRequested: lids.length }, + 'Some LIDs failed during getPNsForLIDs - results may be incomplete' + ) } - this.recordMetrics('get-pn', 1) - return pnJid + this.recordMetrics('get-pn', Object.keys(successfulPairs).length) + return Object.keys(successfulPairs).length > 0 ? Object.values(successfulPairs) : null } /** diff --git a/src/Socket/chats.ts b/src/Socket/chats.ts index 8e2de602..4aede1b1 100644 --- a/src/Socket/chats.ts +++ b/src/Socket/chats.ts @@ -1187,7 +1187,17 @@ export const makeChatsSocket = (config: SocketConfig) => { ev.on('lid-mapping.update', async ({ lid, pn }) => { try { - await signalRepository.lidMapping.storeLIDPNMappings([{ lid, pn }]) + const result = await signalRepository.lidMapping.storeLIDPNMappings([{ lid, pn }]) + logger.debug( + { lid, pn, stored: result.stored, skipped: result.skipped, errors: result.errors }, + 'stored LID-PN mapping from update event' + ) + if (result.stored > 0) { + logger.info( + { lid, pn }, + 'fallback LID mapping is now available from update event' + ) + } } catch (error) { logger.warn({ lid, pn, error }, 'Failed to store LID-PN mapping') } diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts index 0ed3da54..c7844c1d 100644 --- a/src/Socket/messages-recv.ts +++ b/src/Socket/messages-recv.ts @@ -21,6 +21,7 @@ import { aesDecryptCTR, aesEncryptGCM, cleanMessage, + normalizeMessageJids, Curve, decodeMediaRetryNode, decodeMessageNode, @@ -1316,6 +1317,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { } } + await normalizeMessageJids(msg, signalRepository, logger) cleanMessage(msg, authState.creds.me!.id, authState.creds.me!.lid!) await upsertMessage(msg, node.attrs.offline ? 'append' : 'notify') diff --git a/src/Utils/process-message.ts b/src/Utils/process-message.ts index 6817fb3b..969ec3f2 100644 --- a/src/Utils/process-message.ts +++ b/src/Utils/process-message.ts @@ -107,6 +107,40 @@ export const cleanMessage = (message: WAMessage, meId: string, meLid: string) => } } +export const normalizeMessageJids = async ( + message: WAMessage, + signalRepository: SignalRepositoryWithLIDStore, + logger?: ILogger +): Promise => { + const resolveLidToPn = async (jid: string | undefined | null): Promise => { + if (!jid) { + return undefined + } + + if (isLidUser(jid) || isHostedLidUser(jid)) { + const pn = await signalRepository.lidMapping.getPNForLID(jid) + if (pn) { + logger?.debug({ lid: jid, pn }, 'Resolved LID to PN for inbound message') + } else { + logger?.debug({ lid: jid }, 'PN not found for inbound LID, keeping LID') + } + return pn || jid + } + + return jid + } + + const resolvedRemoteJid = await resolveLidToPn(message.key.remoteJid) + if (resolvedRemoteJid) { + message.key.remoteJid = resolvedRemoteJid + } + + const resolvedParticipant = await resolveLidToPn(message.key.participant) + if (resolvedParticipant) { + message.key.participant = resolvedParticipant + } +} + // TODO: target:audit AUDIT THIS FUNCTION AGAIN export const isRealMessage = (message: WAMessage) => { const normalizedContent = normalizeMessageContent(message.message) @@ -291,6 +325,22 @@ const processMessage = async ( // This is how WhatsApp Web learns mappings for chats with non-contacts if (data.lidPnMappings?.length) { logger?.debug({ count: data.lidPnMappings.length }, 'processing LID-PN mappings from history sync') + try { + const result = await signalRepository.lidMapping.storeLIDPNMappings(data.lidPnMappings) + logger?.debug( + { stored: result.stored, skipped: result.skipped, errors: result.errors }, + 'stored LID-PN mappings from history sync' + ) + if (result.stored > 0) { + logger?.info( + { stored: result.stored }, + 'fallback LID mappings are now available from history sync' + ) + } + } catch (error) { + logger?.warn({ error }, 'Failed to store LID-PN mappings from history sync') + } + // eslint-disable-next-line max-depth for (const mapping of data.lidPnMappings) { ev.emit('lid-mapping.update', mapping) diff --git a/src/__tests__/Signal/lid-mapping.test.ts b/src/__tests__/Signal/lid-mapping.test.ts index 68e49e9d..762829af 100644 --- a/src/__tests__/Signal/lid-mapping.test.ts +++ b/src/__tests__/Signal/lid-mapping.test.ts @@ -44,4 +44,42 @@ describe('LIDMappingStore', () => { expect(result).toBeNull() }) }) + + describe('getLIDsForPNs', () => { + it('should resolve multiple PNs in a single batch', async () => { + const pnOne = '11111@s.whatsapp.net' + const pnTwo = '22222:5@s.whatsapp.net' + + // @ts-ignore + mockKeys.get.mockResolvedValue({ '11111': 'aaaaa', '22222': 'bbbbb' } as SignalDataTypeMap['lid-mapping']) + + const result = await lidMappingStore.getLIDsForPNs([pnOne, pnTwo]) + + expect(result).toEqual( + expect.arrayContaining([ + { pn: pnOne, lid: 'aaaaa@lid' }, + { pn: pnTwo, lid: 'bbbbb:5@lid' } + ]) + ) + }) + }) + + describe('getPNsForLIDs', () => { + it('should resolve multiple LIDs in a single batch', async () => { + const lidOne = '33333@lid' + const lidTwo = '44444:99@hosted.lid' + + // @ts-ignore + mockKeys.get.mockResolvedValue({ '33333_reverse': '77777', '44444_reverse': '88888' } as SignalDataTypeMap['lid-mapping']) + + const result = await lidMappingStore.getPNsForLIDs([lidOne, lidTwo]) + + expect(result).toEqual( + expect.arrayContaining([ + { lid: lidOne, pn: '77777@s.whatsapp.net' }, + { lid: lidTwo, pn: '88888:99@hosted' } + ]) + ) + }) + }) }) diff --git a/src/__tests__/Utils/process-message.test.ts b/src/__tests__/Utils/process-message.test.ts index 12d5ea76..e4922dfb 100644 --- a/src/__tests__/Utils/process-message.test.ts +++ b/src/__tests__/Utils/process-message.test.ts @@ -1,5 +1,6 @@ import type { WAMessage } from '../../Types' -import { cleanMessage } from '../../Utils/process-message' +import type { SignalRepositoryWithLIDStore } from '../../Types' +import { cleanMessage, normalizeMessageJids } from '../../Utils/process-message' const createBaseMessage = (key: Partial, message?: Partial): WAMessage => { return { @@ -136,3 +137,51 @@ describe('cleanMessage', () => { }) }) }) + +describe('normalizeMessageJids', () => { + const signalRepository = { + lidMapping: { + getPNForLID: jest.fn() + } + } as unknown as SignalRepositoryWithLIDStore + + beforeEach(() => { + jest.clearAllMocks() + }) + + it('should resolve remoteJid LID to PN when mapping exists', async () => { + const message = createBaseMessage({ + remoteJid: '123456789@lid' + }) + + signalRepository.lidMapping.getPNForLID = jest.fn().mockResolvedValue('5511999999999@s.whatsapp.net') + + await normalizeMessageJids(message, signalRepository) + + expect(message.key.remoteJid).toBe('5511999999999@s.whatsapp.net') + }) + + it('should keep LID when no PN mapping exists', async () => { + const message = createBaseMessage({ + remoteJid: '123456789@lid' + }) + + signalRepository.lidMapping.getPNForLID = jest.fn().mockResolvedValue(null) + + await normalizeMessageJids(message, signalRepository) + + expect(message.key.remoteJid).toBe('123456789@lid') + }) + + it('should resolve participant LID to PN when mapping exists', async () => { + const message = createBaseMessage({ + participant: '999999999@lid' + }) + + signalRepository.lidMapping.getPNForLID = jest.fn().mockResolvedValue('5511888888888@s.whatsapp.net') + + await normalizeMessageJids(message, signalRepository) + + expect(message.key.participant).toBe('5511888888888@s.whatsapp.net') + }) +})