From 194b557b0f3f4aacfa4d5feecf5541f19c646e67 Mon Sep 17 00:00:00 2001 From: Gustavo Quadri <87215048+gusquadri@users.noreply.github.com> Date: Sun, 14 Sep 2025 15:58:59 -0300 Subject: [PATCH] fix: optimize LID migration performance and introduce cache (#1782) * chore: lid-map cache, bulk migration, fixes * fix: type error * fix: signalrepository with lid type, remove check from assertsessions, lid-mapping store level deduplication * fix: migrations logs * fix: remove migrationsops lenght check * fix: skip deletesession if cached or existing migrations * chore: remove check migrated session * fix: remove return null loadsession * fix: getLIDforPN call and jidsrequiringfetch * fix: longer lid map cache TTL * fix: lint * fix: lid type * fix: lid socket type --- src/Signal/libsignal.ts | 151 ++++++++++++++++++-------------- src/Signal/lid-mapping.ts | 89 ++++++++++++++----- src/Socket/messages-recv.ts | 15 ++-- src/Socket/messages-send.ts | 155 ++++++++++++--------------------- src/Socket/socket.ts | 11 ++- src/Types/Signal.ts | 12 +-- src/Types/Socket.ts | 4 +- src/Utils/decode-wa-message.ts | 13 +-- src/Utils/process-message.ts | 7 +- src/Utils/signal.ts | 4 +- 10 files changed, 241 insertions(+), 220 deletions(-) diff --git a/src/Signal/libsignal.ts b/src/Signal/libsignal.ts index 0678bdbe..a0965f04 100644 --- a/src/Signal/libsignal.ts +++ b/src/Signal/libsignal.ts @@ -1,9 +1,9 @@ +/* @ts-ignore */ import * as libsignal from 'libsignal' -import { LRUCache } from 'lru-cache' import type { SignalAuthState, SignalKeyStoreWithTransaction } from '../Types' -import type { SignalRepository } from '../Types/Signal' +import type { SignalRepositoryWithLIDStore } from '../Types/Signal' import { generateSignalPubKey } from '../Utils' -import { jidDecode } from '../WABinary' +import { jidDecode, transferDevice } from '../WABinary' import type { SenderKeyStore } from './Group/group_cipher' import { SenderKeyName } from './Group/sender-key-name' import { SenderKeyRecord } from './Group/sender-key-record' @@ -20,7 +20,7 @@ export function makeLibSignalRepository( }[] | undefined > -): SignalRepository { +): SignalRepositoryWithLIDStore { const lidMapping = new LIDMappingStore(auth.keys as SignalKeyStoreWithTransaction, onWhatsAppFunc) const storage = signalStorage(auth, lidMapping) @@ -38,13 +38,7 @@ export function makeLibSignalRepository( ) } - // Simple operation-level deduplication (5 minutes) - const recentMigrations = new LRUCache({ - max: 500, - ttl: 5 * 60 * 1000 - }) - - const repository: SignalRepository = { + const repository: SignalRepositoryWithLIDStore = { decryptGroupMessage({ group, authorJid, msg }) { const senderName = jidToSignalSenderKeyName(group, authorJid) const cipher = new GroupCipher(storage, senderName) @@ -136,7 +130,7 @@ export function makeLibSignalRepository( if (pnSession) { // Migrate PN to LID - await repository.migrateSession(jid, lidForPN) + await repository.migrateSession([jid], lidForPN) encryptionJid = lidForPN } } @@ -185,13 +179,8 @@ export function makeLibSignalRepository( return jidToSignalProtocolAddress(jid).toString() }, - async storeLIDPNMapping(lid: string, pn: string) { - await lidMapping.storeLIDPNMapping(lid, pn) - }, - - getLIDMappingStore() { - return lidMapping - }, + // Optimized direct access to LID mapping store + lidMapping, async validateSession(jid: string) { try { @@ -212,71 +201,97 @@ export function makeLibSignalRepository( } }, - async deleteSession(jid: string) { - const addr = jidToSignalProtocolAddress(jid) + async deleteSession(jids: string[]) { + if (!jids.length) return + // Convert JIDs to signal addresses and prepare for bulk deletion + const sessionUpdates: { [key: string]: null } = {} + jids.forEach(jid => { + const addr = jidToSignalProtocolAddress(jid) + sessionUpdates[addr.toString()] = null + }) + + // Single transaction for all deletions return parsedKeys.transaction(async () => { - await auth.keys.set({ session: { [addr.toString()]: null } }) - }, jid) + await auth.keys.set({ session: sessionUpdates }) + }, `delete-${jids.length}-sessions`) }, - async migrateSession(fromJid: string, toJid: string) { - // Only migrate PN → LID - if (!fromJid.includes('@s.whatsapp.net') || !toJid.includes('@lid')) { - return - } + async migrateSession( + fromJids: string[], + toJid: string + ): Promise<{ migrated: number; skipped: number; total: number }> { + if (!fromJids.length || !toJid.includes('@lid')) return { migrated: 0, skipped: 0, total: 0 } - const fromDecoded = jidDecode(fromJid) - const toDecoded = jidDecode(toJid) - if (!fromDecoded || !toDecoded) return + // Filter valid PN JIDs + const validJids = fromJids.filter(jid => jid.includes('@s.whatsapp.net')) + if (!validJids.length) return { migrated: 0, skipped: 0, total: fromJids.length } - const deviceId = fromDecoded.device || 0 - const migrationKey = `${fromDecoded.user}.${deviceId}→${toDecoded.user}.${deviceId}` + // Single optimized transaction for all migrations + return parsedKeys.transaction( + async (): Promise<{ migrated: number; skipped: number; total: number }> => { + // 1. Batch store all LID mappings + const mappings = validJids.map(jid => ({ + lid: transferDevice(jid, toJid), + pn: jid + })) + await lidMapping.storeLIDPNMappings(mappings) - // Check if recently migrated (5 min window) - if (recentMigrations.has(migrationKey)) { - return - } + // 2. Prepare migration operations + const migrationOps = validJids.map(jid => { + const lidWithDevice = transferDevice(jid, toJid) + const fromDecoded = jidDecode(jid)! + const toDecoded = jidDecode(lidWithDevice)! - // Check if LID session already exists - const lidAddr = jidToSignalProtocolAddress(toJid) - const { [lidAddr.toString()]: lidExists } = await auth.keys.get('session', [lidAddr.toString()]) - if (lidExists) { - recentMigrations.set(migrationKey, true) - return - } + return { + fromJid: jid, + toJid: lidWithDevice, + pnUser: fromDecoded.user, + lidUser: toDecoded.user, + deviceId: fromDecoded.device || 0, + fromAddr: jidToSignalProtocolAddress(jid), + toAddr: jidToSignalProtocolAddress(lidWithDevice) + } + }) - return parsedKeys.transaction(async () => { - // Store mapping - await lidMapping.storeLIDPNMapping(toJid, fromJid) + // 3. Batch check which LID sessions already exist + const lidAddrs = migrationOps.map(op => op.toAddr.toString()) + const existingSessions = await auth.keys.get('session', lidAddrs) - // Load and copy session - const fromAddr = jidToSignalProtocolAddress(fromJid) - const fromSession = await storage.loadSession(fromAddr.toString()) + // 4. Filter out sessions that already have LID sessions + const opsToMigrate = migrationOps.filter(op => !existingSessions[op.toAddr.toString()]) + const skippedCount = migrationOps.length - opsToMigrate.length - if (fromSession?.haveOpenSession()) { - // Deep copy session to prevent reference issues - const sessionBytes = fromSession.serialize() - const copiedSession = libsignal.SessionRecord.deserialize(sessionBytes) + if (!opsToMigrate.length) { + return { migrated: 0, skipped: skippedCount, total: validJids.length } + } - // Store at LID address - await storage.storeSession(lidAddr.toString(), copiedSession) + // 5. Execute all migrations in parallel + await Promise.all( + opsToMigrate.map(async op => { + const fromSession = await storage.loadSession(op.fromAddr.toString()) - // Delete PN session - maintain single encryption layer - await auth.keys.set({ session: { [fromAddr.toString()]: null } }) - } + if (fromSession?.haveOpenSession()) { + // Copy session to LID address + const sessionBytes = fromSession.serialize() + const copiedSession = libsignal.SessionRecord.deserialize(sessionBytes) + await storage.storeSession(op.toAddr.toString(), copiedSession) - recentMigrations.set(migrationKey, true) - }, fromJid) + // Delete PN session + await auth.keys.set({ session: { [op.fromAddr.toString()]: null } }) + } + }) + ) + + return { migrated: opsToMigrate.length, skipped: skippedCount, total: validJids.length } + }, + `migrate-${validJids.length}-sessions-${jidDecode(toJid)?.user}` + ) }, async encryptMessageWithWire({ encryptionJid, wireJid, data }) { const result = await repository.encryptMessage({ jid: encryptionJid, data }) return { ...result, wireJid } - }, - - destroy() { - recentMigrations.clear() } } @@ -287,6 +302,12 @@ const jidToSignalProtocolAddress = (jid: string): libsignal.ProtocolAddress => { const decoded = jidDecode(jid)! const { user, device, server } = decoded + if (!user) { + throw new Error( + `JID decoded but user is empty: "${jid}" -> user: "${user}", server: "${server}", device: ${device}` + ) + } + // LID addresses get _1 suffix for Signal protocol const signalUser = server === 'lid' ? `${user}_1` : user const finalDevice = device || 0 diff --git a/src/Signal/lid-mapping.ts b/src/Signal/lid-mapping.ts index 03231736..aae42d2f 100644 --- a/src/Signal/lid-mapping.ts +++ b/src/Signal/lid-mapping.ts @@ -1,9 +1,14 @@ +import { LRUCache } from 'lru-cache' import type { SignalKeyStoreWithTransaction } from '../Types' import logger from '../Utils/logger' import { isLidUser, isPnUser, jidDecode } from '../WABinary' -//TODO: Caching export class LIDMappingStore { + private readonly mappingCache = new LRUCache({ + ttl: 7 * 24 * 60 * 60 * 1000, // 7 days + ttlAutopurge: true, + updateAgeOnGet: true + }) private readonly keys: SignalKeyStoreWithTransaction private onWhatsAppFunc?: (...jids: string[]) => Promise< | { @@ -29,13 +34,6 @@ export class LIDMappingStore { this.onWhatsAppFunc = onWhatsAppFunc // needed to get LID from PN if not found } - /** - * Store LID-PN mapping - USER LEVEL - */ - async storeLIDPNMapping(lid: string, pn: string): Promise { - return this.storeLIDPNMappings([{ lid, pn }]) - } - /** * Store LID-PN mapping - USER LEVEL */ @@ -57,6 +55,25 @@ export class LIDMappingStore { const pnUser = pnDecoded.user const lidUser = lidDecoded.user + + // Check if mapping already exists (cache first, then database) + let existingLidUser = this.mappingCache.get(`pn:${pnUser}`) + if (!existingLidUser) { + // Cache miss - check database + const stored = await this.keys.get('lid-mapping', [pnUser]) + existingLidUser = stored[pnUser] + if (existingLidUser) { + // Update cache with database value + this.mappingCache.set(`pn:${pnUser}`, existingLidUser) + this.mappingCache.set(`lid:${existingLidUser}`, pnUser) + } + } + + if (existingLidUser === lidUser) { + logger.debug({ pnUser, lidUser }, 'LID mapping already exists, skipping') + continue + } + pairMap[pnUser] = lidUser } @@ -70,6 +87,10 @@ export class LIDMappingStore { [`${lidUser}_reverse`]: pnUser // "102765716062358_reverse" -> "554396160286" } }) + + // Update cache with both directions + this.mappingCache.set(`pn:${pnUser}`, lidUser) + this.mappingCache.set(`lid:${lidUser}`, pnUser) } }, 'lid-mapping') } @@ -83,22 +104,38 @@ export class LIDMappingStore { const decoded = jidDecode(pn) if (!decoded) return null - // Look up user-level mapping (whatsmeow approach) + // Check cache first for PN → LID mapping const pnUser = decoded.user - const stored = await this.keys.get('lid-mapping', [pnUser]) - let lidUser = stored[pnUser] + let lidUser = this.mappingCache.get(`pn:${pnUser}`) if (!lidUser) { - logger.trace(`No LID mapping found for PN user ${pnUser}; getting from USync`) - const { exists, lid } = (await this.onWhatsAppFunc?.(pn))?.[0]! // this function already adds LIDs to mapping - if (exists) { - lidUser = jidDecode(lid)?.user + // Cache miss - check database + const stored = await this.keys.get('lid-mapping', [pnUser]) + lidUser = stored[pnUser] + + if (lidUser) { + // Cache the database result + this.mappingCache.set(`pn:${pnUser}`, lidUser) } else { - return null + // Not in database - try USync + logger.trace(`No LID mapping found for PN user ${pnUser}; getting from USync`) + const { exists, lid } = (await this.onWhatsAppFunc?.(pn))?.[0]! // this function already adds LIDs to mapping + if (exists && lid) { + lidUser = jidDecode(lid)?.user + if (lidUser) { + // Cache the USync result + this.mappingCache.set(`pn:${pnUser}`, lidUser) + } + } else { + return null + } } } - if (typeof lidUser !== 'string') return null + if (typeof lidUser !== 'string' || !lidUser) { + logger.warn(`Invalid or empty LID user for PN ${pn}: lidUser = "${lidUser}"`) + return null + } // Push the PN device ID to the LID to maintain device separation const pnDevice = decoded.device !== undefined ? decoded.device : 0 @@ -117,15 +154,21 @@ export class LIDMappingStore { const decoded = jidDecode(lid) if (!decoded) return null - // Look up reverse user mapping + // Check cache first for LID → PN mapping const lidUser = decoded.user - // TODO: remove this style and instead load all mappings somehow, and then assign them in the map - const stored = await this.keys.get('lid-mapping', [`${lidUser}_reverse`]) - const pnUser = stored[`${lidUser}_reverse`] + let pnUser = this.mappingCache.get(`lid:${lidUser}`) if (!pnUser || typeof pnUser !== 'string') { - logger.trace(`No reverse mapping found for LID user: ${lidUser}`) - return null + // Cache miss - check database + const stored = await this.keys.get('lid-mapping', [`${lidUser}_reverse`]) + pnUser = stored[`${lidUser}_reverse`] + + if (!pnUser || typeof pnUser !== 'string') { + logger.trace(`No reverse mapping found for LID user: ${lidUser}`) + return null + } + + this.mappingCache.set(`lid:${lidUser}`, pnUser) } // Construct device-specific PN JID diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts index 603f154c..e254b246 100644 --- a/src/Socket/messages-recv.ts +++ b/src/Socket/messages-recv.ts @@ -1150,21 +1150,24 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => { const lid = jidNormalizedUser(node.attrs.from), pn = jidNormalizedUser(node.attrs.sender_pn) ev.emit('lid-mapping.update', { lid, pn }) - await signalRepository.storeLIDPNMapping(lid, pn) + await signalRepository.lidMapping.storeLIDPNMappings([{ lid, pn }]) } const alt = msg.key.participantAlt || msg.key.remoteJidAlt // store new mappings we didn't have before if (!!alt) { const altServer = jidDecode(alt)?.server - const lidMapping = signalRepository.getLIDMappingStore() if (altServer === 'lid') { - if (typeof (await lidMapping.getPNForLID(alt)) === 'string') { - await lidMapping.storeLIDPNMapping(alt, msg.key.participant || msg.key.remoteJid!) + if (typeof (await signalRepository.lidMapping.getPNForLID(alt)) === 'string') { + await signalRepository.lidMapping.storeLIDPNMappings([ + { lid: alt, pn: msg.key.participant || msg.key.remoteJid! } + ]) } } else { - if (typeof (await lidMapping.getLIDForPN(alt)) === 'string') { - await lidMapping.storeLIDPNMapping(msg.key.participant || msg.key.remoteJid!, alt) + if (typeof (await signalRepository.lidMapping.getLIDForPN(alt)) === 'string') { + await signalRepository.lidMapping.storeLIDPNMappings([ + { lid: msg.key.participant || msg.key.remoteJid!, pn: alt } + ]) } } } diff --git a/src/Socket/messages-send.ts b/src/Socket/messages-send.ts index ec80359d..703d7fbe 100644 --- a/src/Socket/messages-send.ts +++ b/src/Socket/messages-send.ts @@ -47,8 +47,7 @@ import { jidEncode, jidNormalizedUser, type JidWithDevice, - S_WHATSAPP_NET, - transferDevice + S_WHATSAPP_NET } from '../WABinary' import { USyncQuery, USyncUser } from '../WAUSync' import { makeNewsletterSocket } from './newsletter' @@ -366,19 +365,6 @@ export const makeMessagesSocket = (config: SocketConfig) => { return deviceResults } - // Helper to check if JID has migrated LID session - const checkForMigratedLidSession = async (jid: string): Promise => { - if (!jid.includes('@s.whatsapp.net')) return false - - const lidMapping = signalRepository.getLIDMappingStore() - const lidForPN = await lidMapping.getLIDForPN(jid) - if (!lidForPN?.includes('@lid')) return false - - const lidSignalId = signalRepository.jidToSignalProtocolAddress(lidForPN) - const lidSessions = await authState.keys.get('session', [lidSignalId]) - return !!lidSessions[lidSignalId] - } - const assertSessions = async (jids: string[], force: boolean) => { let didFetchNewSession = false const jidsRequiringFetch: string[] = [] @@ -391,20 +377,13 @@ export const makeMessagesSocket = (config: SocketConfig) => { const addrs = jids.map(jid => signalRepository.jidToSignalProtocolAddress(jid)) const sessions = await authState.keys.get('session', addrs) - // Helper to check session for a JID - const checkJidSession = async (jid: string) => { + // Simplified: Check session existence directly + const checkJidSession = (jid: string) => { const signalId = signalRepository.jidToSignalProtocolAddress(jid) - let hasSession = !!sessions[signalId] - - // Check for migrated LID session if PN session missing - if (!hasSession) { - hasSession = await checkForMigratedLidSession(jid) - if (hasSession) { - logger.debug({ jid }, 'Found migrated LID session during force assert, skipping PN fetch') - } - } + const hasSession = !!sessions[signalId] // Add to fetch list if no session exists + // Session type selection (LID vs PN) is handled in encryptMessage if (!hasSession) { if (jid.includes('@lid')) { logger.debug({ jid }, 'No LID session found, will create new LID session') @@ -416,10 +395,9 @@ export const makeMessagesSocket = (config: SocketConfig) => { // Process all JIDs for (const jid of jids) { - await checkJidSession(jid) + checkJidSession(jid) } } else { - const lidMapping = signalRepository.getLIDMappingStore() const addrs = jids.map(jid => signalRepository.jidToSignalProtocolAddress(jid)) const sessions = await authState.keys.get('session', addrs) @@ -441,7 +419,9 @@ export const makeMessagesSocket = (config: SocketConfig) => { } try { - const mapping = await lidMapping.getLIDForPN(user) + // Convert user to proper PN JID format for getLIDForPN + const pnJid = `${user}@s.whatsapp.net` + const mapping = await signalRepository.lidMapping.getLIDForPN(pnJid) if (mapping?.includes('@lid')) { logger.debug( { user, lidForPN: mapping, deviceCount: userJids.length }, @@ -456,27 +436,6 @@ export const makeMessagesSocket = (config: SocketConfig) => { return { shouldMigrate: false, lidForPN: undefined } } - // Helper to migrate a single device - const migrateDeviceToLid = async (jid: string, lidForPN: string) => { - if (!jid.includes('@s.whatsapp.net')) return - - try { - const lidWithDevice = transferDevice(jid, lidForPN) - await signalRepository.migrateSession(jid, lidWithDevice) - logger.debug({ fromJid: jid, toJid: lidWithDevice }, 'Migrated device session to LID') - - // Delete PN session after successful migration - try { - await signalRepository.deleteSession(jid) - logger.debug({ deletedPNSession: jid }, 'Deleted PN session after migration') - } catch (deleteError) { - logger.warn({ jid, error: deleteError }, 'Failed to delete PN session') - } - } catch (migrationError) { - logger.warn({ jid, error: migrationError }, 'Failed to migrate device session') - } - } - // Process each user group for potential bulk LID migration for (const [user, userJids] of userGroups) { const mappingResult = await checkUserLidMapping(user, userJids) @@ -485,59 +444,52 @@ export const makeMessagesSocket = (config: SocketConfig) => { // Migrate all devices for this user if LID mapping exists if (shouldMigrateUser && lidForPN) { - // Migrate each device individually - for (const jid of userJids) { - await migrateDeviceToLid(jid, lidForPN) - } + // Bulk migrate all user devices in single transaction + const migrationResult = await signalRepository.migrateSession(userJids, lidForPN) - logger.info( - { - user, - lidMapping: lidForPN, - deviceCount: userJids.length - }, - 'Completed migration attempt for user devices' - ) - } - - // Helper to check session for migrated user - const checkMigratedSession = async (jid: string) => { - const signalId = signalRepository.jidToSignalProtocolAddress(jid) - let hasSession = !!sessions[signalId] - let jidToFetch = jid - - // Check if we should use migrated LID session instead - if (shouldMigrateUser && lidForPN && jid.includes('@s.whatsapp.net')) { - const originalDecoded = jidDecode(jid) - const deviceId = originalDecoded?.device || 0 - const lidDecoded = jidDecode(lidForPN) - const lidWithDevice = jidEncode(lidDecoded?.user!, 'lid', deviceId) - - // Check if LID session exists - const lidSignalId = signalRepository.jidToSignalProtocolAddress(lidWithDevice) - const lidSessions = await authState.keys.get('session', [lidSignalId]) - hasSession = !!lidSessions[lidSignalId] - jidToFetch = lidWithDevice - - if (hasSession) { - logger.debug({ originalJid: jid, lidJid: lidWithDevice }, '✅ Found bulk-migrated LID session') - } - } - - // Add to fetch list if no session exists - if (!hasSession) { - jidsRequiringFetch.push(jidToFetch) + if (migrationResult.migrated > 0) { + logger.info( + { + user, + lidMapping: lidForPN, + migrated: migrationResult.migrated, + skipped: migrationResult.skipped, + total: migrationResult.total + }, + 'Completed bulk migration for user devices' + ) + } else { logger.debug( - { jid: jidToFetch, originalJid: jid !== jidToFetch ? jid : undefined }, - 'Adding to session fetch list' + { + user, + lidMapping: lidForPN, + skipped: migrationResult.skipped, + total: migrationResult.total + }, + 'All user device sessions already migrated' ) } } - // Now check which sessions need to be fetched for this user - for (const jid of userJids) { - await checkMigratedSession(jid) + // Direct bulk session check with LID single source of truth + const addMissingSessionsToFetchList = (jid: string) => { + const signalId = signalRepository.jidToSignalProtocolAddress(jid) + if (sessions[signalId]) return + + // Determine correct JID to fetch (LID if mapping exists, otherwise original) + if (jid.includes('@s.whatsapp.net') && shouldMigrateUser && lidForPN) { + const decoded = jidDecode(jid)! + const lidDeviceJid = + decoded.device !== undefined ? `${jidDecode(lidForPN)!.user}:${decoded.device}@lid` : lidForPN + jidsRequiringFetch.push(lidDeviceJid) + logger.debug({ pnJid: jid, lidJid: lidDeviceJid }, 'Adding LID JID to fetch list (conversion)') + } else { + jidsRequiringFetch.push(jid) + logger.debug({ jid }, 'Adding JID to fetch list') + } } + + userJids.forEach(addMissingSessionsToFetchList) } } @@ -674,8 +626,7 @@ export const makeMessagesSocket = (config: SocketConfig) => { if (!wireJid.includes('@s.whatsapp.net')) return wireJid try { - const lidMapping = signalRepository.getLIDMappingStore() - const lidForPN = await lidMapping.getLIDForPN(wireJid) + const lidForPN = await signalRepository.lidMapping.getLIDForPN(wireJid) if (!lidForPN?.includes('@lid')) return wireJid @@ -687,7 +638,7 @@ export const makeMessagesSocket = (config: SocketConfig) => { // Migrate session to LID for unified encryption layer try { - await signalRepository.migrateSession(wireJid, lidWithDevice) + const migrationResult = await signalRepository.migrateSession([wireJid], lidWithDevice) const recipientUser = jidNormalizedUser(wireJid) const ownPnUser = jidNormalizedUser(meId) const isOwnDevice = recipientUser === ownPnUser @@ -695,8 +646,10 @@ export const makeMessagesSocket = (config: SocketConfig) => { // Delete PN session after successful migration try { - await signalRepository.deleteSession(wireJid) - logger.debug({ deletedPNSession: wireJid }, 'Deleted PN session') + if (migrationResult.migrated) { + await signalRepository.deleteSession([wireJid]) + logger.debug({ deletedPNSession: wireJid }, 'Deleted PN session') + } } catch (deleteError) { logger.warn({ wireJid, error: deleteError }, 'Failed to delete PN session') } diff --git a/src/Socket/socket.ts b/src/Socket/socket.ts index 7f3bf590..ea22d8b2 100644 --- a/src/Socket/socket.ts +++ b/src/Socket/socket.ts @@ -271,10 +271,9 @@ export const makeSocket = (config: SocketConfig) => { const results = await executeUSyncQuery(usyncQuery) if (results) { - if (results.list.filter(a => a.lid).length > 0) { - const lidMapping = signalRepository.getLIDMappingStore() - const lidOnly = results.list.filter(a => a.lid) - await lidMapping.storeLIDPNMappings(lidOnly.map(a => ({ pn: a.id, lid: a.lid as string }))) + if (results.list.filter(a => !!a.lid).length > 0) { + const lidOnly = results.list.filter(a => !!a.lid) + await signalRepository.lidMapping.storeLIDPNMappings(lidOnly.map(a => ({ pn: a.id, lid: a.lid as string }))) } return results.list @@ -851,10 +850,10 @@ export const makeSocket = (config: SocketConfig) => { const myPN = authState.creds.me!.id // Store our own LID-PN mapping - await signalRepository.storeLIDPNMapping(myLID, myPN) + await signalRepository.lidMapping.storeLIDPNMappings([{ lid: myLID, pn: myPN }]) // Create LID session for ourselves (whatsmeow pattern) - await signalRepository.migrateSession(myPN, myLID) + await signalRepository.migrateSession([myPN], myLID) logger.info({ myPN, myLID }, 'Own LID session created successfully') } catch (error) { diff --git a/src/Types/Signal.ts b/src/Types/Signal.ts index f20a866b..52cadcf0 100644 --- a/src/Types/Signal.ts +++ b/src/Types/Signal.ts @@ -76,10 +76,12 @@ export type SignalRepository = { injectE2ESession(opts: E2ESessionOpts): Promise validateSession(jid: string): Promise<{ exists: boolean; reason?: string }> jidToSignalProtocolAddress(jid: string): string - storeLIDPNMapping(lid: string, pn: string): Promise - getLIDMappingStore(): LIDMappingStore - migrateSession(fromJid: string, toJid: string): Promise + migrateSession(fromJids: string[], toJid: string): Promise<{ migrated: number; skipped: number; total: number }> validateSession(jid: string): Promise<{ exists: boolean; reason?: string }> - deleteSession(jid: string): Promise - destroy(): void + deleteSession(jids: string[]): Promise +} + +// Optimized repository with pre-loaded LID mapping store +export interface SignalRepositoryWithLIDStore extends SignalRepository { + lidMapping: LIDMappingStore } diff --git a/src/Types/Socket.ts b/src/Types/Socket.ts index 09ca891c..3a14100f 100644 --- a/src/Types/Socket.ts +++ b/src/Types/Socket.ts @@ -6,7 +6,7 @@ import type { ILogger } from '../Utils/logger' import type { AuthenticationState, SignalAuthState, TransactionCapabilityOptions } from './Auth' import type { GroupMetadata } from './GroupMetadata' import { type MediaConnInfo } from './Message' -import type { SignalRepository } from './Signal' +import type { SignalRepositoryWithLIDStore } from './Signal' export type WAVersion = [number, number, number] export type WABrowserDescription = [string, string, string] @@ -153,5 +153,5 @@ export type SocketConfig = { }[] | undefined > - ) => SignalRepository + ) => SignalRepositoryWithLIDStore } diff --git a/src/Utils/decode-wa-message.ts b/src/Utils/decode-wa-message.ts index 2c7c1417..9755c5d7 100644 --- a/src/Utils/decode-wa-message.ts +++ b/src/Utils/decode-wa-message.ts @@ -1,6 +1,7 @@ import { Boom } from '@hapi/boom' import { proto } from '../../WAProto/index.js' -import type { SignalRepository, WAMessage, WAMessageKey } from '../Types' +import type { WAMessage, WAMessageKey } from '../Types' +import type { SignalRepositoryWithLIDStore } from '../Types/Signal' import { areJidsSameUser, type BinaryNode, @@ -16,26 +17,26 @@ import { import { unpadRandomMax16 } from './generics' import type { ILogger } from './logger' -const getDecryptionJid = async (sender: string, repository: SignalRepository): Promise => { +const getDecryptionJid = async (sender: string, repository: SignalRepositoryWithLIDStore): Promise => { if (!sender.includes('@s.whatsapp.net')) { return sender } - return (await repository.getLIDMappingStore().getLIDForPN(sender))! + return (await repository.lidMapping.getLIDForPN(sender))! } const storeMappingFromEnvelope = async ( stanza: BinaryNode, sender: string, decryptionJid: string, - repository: SignalRepository, + repository: SignalRepositoryWithLIDStore, logger: ILogger ): Promise => { const { senderAlt } = extractAddressingContext(stanza) if (senderAlt && isLidUser(senderAlt) && isPnUser(sender) && decryptionJid === sender) { try { - await repository.storeLIDPNMapping(senderAlt, sender) + await repository.lidMapping.storeLIDPNMappings([{ lid: senderAlt, pn: sender }]) logger.debug({ sender, senderAlt }, 'Stored LID mapping from envelope') } catch (error) { logger.warn({ sender, senderAlt, error }, 'Failed to store LID mapping') @@ -206,7 +207,7 @@ export const decryptMessageNode = ( stanza: BinaryNode, meId: string, meLid: string, - repository: SignalRepository, + repository: SignalRepositoryWithLIDStore, logger: ILogger ) => { const { fullMessage, author, sender } = decodeMessageNode(stanza, meId, meLid) diff --git a/src/Utils/process-message.ts b/src/Utils/process-message.ts index 258d342c..a6d5b6e4 100644 --- a/src/Utils/process-message.ts +++ b/src/Utils/process-message.ts @@ -10,7 +10,7 @@ import type { RequestJoinAction, RequestJoinMethod, SignalKeyStoreWithTransaction, - SignalRepository + SignalRepositoryWithLIDStore } from '../Types' import { WAMessageStubType } from '../Types' import { getContentType, normalizeMessageContent } from '../Utils/messages' @@ -28,7 +28,7 @@ type ProcessMessageContext = { ev: BaileysEventEmitter logger?: ILogger options: AxiosRequestConfig<{}> - signalRepository: SignalRepository + signalRepository: SignalRepositoryWithLIDStore } const REAL_MSG_STUB_TYPES = new Set([ @@ -304,7 +304,6 @@ const processMessage = async ( ]) break case proto.Message.ProtocolMessage.Type.LID_MIGRATION_MAPPING_SYNC: - const lidMappingStore = signalRepository.getLIDMappingStore() const encodedPayload = protocolMsg.lidMigrationMappingSyncMessage?.encodedMappingPayload! const { pnToLidMappings, chatDbMigrationTimestamp } = proto.LIDMigrationMappingSyncPayload.decode(encodedPayload) @@ -315,7 +314,7 @@ const processMessage = async ( pairs.push({ lid: `${lid}@lid`, pn: `${pn}@s.whatsapp.net` }) } - await lidMappingStore.storeLIDPNMappings(pairs) + await signalRepository.lidMapping.storeLIDPNMappings(pairs) } } else if (content?.reactionMessage) { const reaction: proto.IReaction = { diff --git a/src/Utils/signal.ts b/src/Utils/signal.ts index 8db313d5..a755ae25 100644 --- a/src/Utils/signal.ts +++ b/src/Utils/signal.ts @@ -1,5 +1,5 @@ import { KEY_BUNDLE_TYPE } from '../Defaults' -import type { SignalRepository } from '../Types' +import type { SignalRepositoryWithLIDStore } from '../Types' import type { AuthenticationCreds, AuthenticationState, @@ -85,7 +85,7 @@ export const xmppPreKey = (pair: KeyPair, id: number): BinaryNode => ({ ] }) -export const parseAndInjectE2ESessions = async (node: BinaryNode, repository: SignalRepository) => { +export const parseAndInjectE2ESessions = async (node: BinaryNode, repository: SignalRepositoryWithLIDStore) => { const extractKey = (key: BinaryNode) => key ? {