From 592f70b81ddecb4e2914f18646a614db14ad60b9 Mon Sep 17 00:00:00 2001 From: Rajeh Taher Date: Fri, 3 Oct 2025 00:04:39 +0300 Subject: [PATCH] general: Bring back USYNC calls for LID getting, and improve --- src/Signal/libsignal.ts | 31 ++++-------- src/Signal/lid-mapping.ts | 94 +++++++++++++++++++++++++------------ src/Socket/messages-send.ts | 5 ++ src/Socket/socket.ts | 50 +++++++++++++------- src/Types/Auth.ts | 5 ++ src/Types/Socket.ts | 11 +---- 6 files changed, 119 insertions(+), 77 deletions(-) diff --git a/src/Signal/libsignal.ts b/src/Signal/libsignal.ts index ea9cd7bb..86859f3c 100644 --- a/src/Signal/libsignal.ts +++ b/src/Signal/libsignal.ts @@ -1,7 +1,7 @@ /* @ts-ignore */ import * as libsignal from 'libsignal' import { LRUCache } from 'lru-cache' -import type { SignalAuthState, SignalKeyStoreWithTransaction } from '../Types' +import type { LIDMapping, SignalAuthState, SignalKeyStoreWithTransaction } from '../Types' import type { SignalRepositoryWithLIDStore } from '../Types/Signal' import { generateSignalPubKey } from '../Utils' import type { ILogger } from '../Utils/logger' @@ -12,8 +12,12 @@ import { SenderKeyRecord } from './Group/sender-key-record' import { GroupCipher, GroupSessionBuilder, SenderKeyDistributionMessage } from './Group' import { LIDMappingStore } from './lid-mapping' -export function makeLibSignalRepository(auth: SignalAuthState, logger: ILogger): SignalRepositoryWithLIDStore { - const lidMapping = new LIDMappingStore(auth.keys as SignalKeyStoreWithTransaction, logger) +export function makeLibSignalRepository( + auth: SignalAuthState, + logger: ILogger, + pnToLIDFunc?: (jids: string[]) => Promise +): SignalRepositoryWithLIDStore { + const lidMapping = new LIDMappingStore(auth.keys as SignalKeyStoreWithTransaction, logger, pnToLIDFunc) const storage = signalStorage(auth, lidMapping) const parsedKeys = auth.keys as SignalKeyStoreWithTransaction @@ -23,18 +27,6 @@ export function makeLibSignalRepository(auth: SignalAuthState, logger: ILogger): updateAgeOnGet: true }) - function isLikelySyncMessage(addr: libsignal.ProtocolAddress): boolean { - const key = addr.toString() - - // Only bypass for WhatsApp system addresses, not regular user contacts - // Be very specific about sync service patterns - return ( - key.includes('@lid.whatsapp.net') || // WhatsApp system messages - key.includes('@broadcast') || // Broadcast messages - key.includes('@newsletter') - ) - } - const repository: SignalRepositoryWithLIDStore = { decryptGroupMessage({ group, authorJid, msg }) { const senderName = jidToSignalSenderKeyName(group, authorJid) @@ -93,12 +85,6 @@ export function makeLibSignalRepository(auth: SignalAuthState, logger: ILogger): return result } - if (isLikelySyncMessage(addr)) { - // If it's a sync message, we can skip the transaction - // as it is likely to be a system message that doesn't require strict atomicity - return await doDecrypt() - } - // If it's not a sync message, we need to ensure atomicity // For regular messages, we use a transaction to ensure atomicity return parsedKeys.transaction(async () => { @@ -117,6 +103,7 @@ export function makeLibSignalRepository(auth: SignalAuthState, logger: ILogger): return { type, ciphertext: Buffer.from(body, 'binary') } }, jid) }, + async encryptGroupMessage({ group, meId, data }) { const senderName = jidToSignalSenderKeyName(group, meId) const builder = new GroupSessionBuilder(storage) @@ -139,6 +126,7 @@ export function makeLibSignalRepository(auth: SignalAuthState, logger: ILogger): } }, group) }, + async injectE2ESession({ jid, session }) { const cipher = new libsignal.SessionBuilder(storage, jidToSignalProtocolAddress(jid)) return parsedKeys.transaction(async () => { @@ -191,6 +179,7 @@ export function makeLibSignalRepository(auth: SignalAuthState, logger: ILogger): fromJid: string, toJid: string ): Promise<{ migrated: number; skipped: number; total: number }> { + // TODO: use usync to handle this entire mess if (!fromJid || !toJid.includes('@lid')) return { migrated: 0, skipped: 0, total: 0 } // Only support PN to LID migration diff --git a/src/Signal/lid-mapping.ts b/src/Signal/lid-mapping.ts index 4f1adc53..39c8cd7a 100644 --- a/src/Signal/lid-mapping.ts +++ b/src/Signal/lid-mapping.ts @@ -1,5 +1,5 @@ import { LRUCache } from 'lru-cache' -import type { SignalKeyStoreWithTransaction } from '../Types' +import type { LIDMapping, SignalKeyStoreWithTransaction } from '../Types' import type { ILogger } from '../Utils/logger' import { isLidUser, isPnUser, jidDecode } from '../WABinary' @@ -12,15 +12,22 @@ export class LIDMappingStore { private readonly keys: SignalKeyStoreWithTransaction private readonly logger: ILogger - constructor(keys: SignalKeyStoreWithTransaction, logger: ILogger) { + private pnToLIDFunc?: (jids: string[]) => Promise + + constructor( + keys: SignalKeyStoreWithTransaction, + logger: ILogger, + pnToLIDFunc?: (jids: string[]) => Promise + ) { this.keys = keys + this.pnToLIDFunc = pnToLIDFunc this.logger = logger } /** * Store LID-PN mapping - USER LEVEL */ - async storeLIDPNMappings(pairs: { lid: string; pn: string }[]): Promise { + async storeLIDPNMappings(pairs: LIDMapping[]): Promise { // Validate inputs const pairMap: { [_: string]: string } = {} for (const { lid, pn } of pairs) { @@ -78,41 +85,70 @@ export class LIDMappingStore { * Get LID for PN - Returns device-specific LID based on user mapping */ async getLIDForPN(pn: string): Promise { - if (!isPnUser(pn)) return null + return (await this.getLIDsForPNs([pn]))?.[0]?.lid || null + } - const decoded = jidDecode(pn) - if (!decoded) return null + async getLIDsForPNs(pns: string[]): Promise { + const usyncFetch: string[] = [] + // mapped from pn to lid mapping to prevent duplication in results later + const successfulPairs: { [_: string]: LIDMapping } = {} + for (const pn of pns) { + if (!isPnUser(pn)) continue - // Check cache first for PN → LID mapping - const pnUser = decoded.user - let lidUser = this.mappingCache.get(`pn:${pnUser}`) + const decoded = jidDecode(pn) + if (!decoded) continue - if (!lidUser) { - // Cache miss - check database - const stored = await this.keys.get('lid-mapping', [pnUser]) - lidUser = stored[pnUser] + // Check cache first for PN → LID mapping + const pnUser = decoded.user + let lidUser = this.mappingCache.get(`pn:${pnUser}`) - if (lidUser) { - this.mappingCache.set(`pn:${pnUser}`, lidUser) - this.mappingCache.set(`lid:${lidUser}`, pnUser) + if (!lidUser) { + // Cache miss - check database + const stored = await this.keys.get('lid-mapping', [pnUser]) + lidUser = stored[pnUser] + + if (lidUser) { + this.mappingCache.set(`pn:${pnUser}`, lidUser) + 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) + continue; + } + } + + lidUser = lidUser.toString() + if (!lidUser) { + this.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 + const deviceSpecificLid = `${lidUser}:${pnDevice}@lid` + + this.logger.trace(`getLIDForPN: ${pn} → ${deviceSpecificLid} (user mapping with device ${pnDevice})`) + successfulPairs[pn] = { lid: deviceSpecificLid, pn } + } + + console.log(this.pnToLIDFunc) + if (usyncFetch.length > 0) { + console.log(usyncFetch) + const result = await this.pnToLIDFunc?.(usyncFetch) // this function already adds LIDs to mapping + if (result && result.length > 0) { + this.storeLIDPNMappings(result) + for (const pair of result) { + const lidUser = jidDecode(pair.lid)?.user + if (lidUser) { + successfulPairs[pair.pn] = pair + } + } } else { - this.logger.trace(`No LID mapping found for PN user ${pnUser}`) return null } } - lidUser = lidUser.toString() - if (!lidUser) { - this.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 - const deviceSpecificLid = `${lidUser}:${pnDevice}@lid` - - this.logger.trace(`getLIDForPN: ${pn} → ${deviceSpecificLid} (user mapping with device ${pnDevice})`) - return deviceSpecificLid + return Object.values(successfulPairs) } /** diff --git a/src/Socket/messages-send.ts b/src/Socket/messages-send.ts index 364716bd..a38ed728 100644 --- a/src/Socket/messages-send.ts +++ b/src/Socket/messages-send.ts @@ -292,6 +292,11 @@ export const makeMessagesSocket = (config: SocketConfig) => { if (result) { // TODO: LID MAP this stuff (lid protocol will now return lid with devices) + const lidResults = result.list.filter(a => !!a.lid) + if (lidResults.length > 0) { + logger.trace('Storing LID maps from device call') + await signalRepository.lidMapping.storeLIDPNMappings(lidResults.map(a => ({lid: a.lid as string, pn: a.id}))) + } const extracted = extractDeviceJids(result?.list, authState.creds.me!.id, ignoreZeroDevices) const deviceMap: { [_: string]: FullJid[] } = {} diff --git a/src/Socket/socket.ts b/src/Socket/socket.ts index dcbe0983..316baf3f 100644 --- a/src/Socket/socket.ts +++ b/src/Socket/socket.ts @@ -12,7 +12,7 @@ import { NOISE_WA_HEADER, UPLOAD_TIMEOUT } from '../Defaults' -import type { SocketConfig } from '../Types' +import type { LIDMapping, SocketConfig } from '../Types' import { DisconnectReason } from '../Types' import { addTransactionCapability, @@ -258,13 +258,13 @@ export const makeSocket = (config: SocketConfig) => { return usyncQuery.parseUSyncQueryResult(result) } - const onWhatsApp = async (...jids: string[]) => { - let usyncQuery = new USyncQuery().withLIDProtocol() + const onWhatsApp = async (...phoneNumber: string[]) => { + let usyncQuery = new USyncQuery() let contactEnabled = false - for (const jid of jids) { + for (const jid of phoneNumber) { if (isLidUser(jid)) { - // usyncQuery.withUser(new USyncUser().withLid(jid)) // intentional but needs JID too + logger?.warn('LIDs are not supported with onWhatsApp') continue } else { if (!contactEnabled) { @@ -284,27 +284,41 @@ export const makeSocket = (config: SocketConfig) => { const results = await executeUSyncQuery(usyncQuery) if (results) { - const lidOnly = results.list.filter(a => !!a.lid) - if (lidOnly.length > 0) { - const pairs = lidOnly.map(a => ({ pn: a.id, lid: a.lid as string })) - await signalRepository.lidMapping.storeLIDPNMappings(pairs) - for (const { pn, lid } of pairs) { - await signalRepository.migrateSession(pn, lid) - } - } - - return results.list - .filter(a => !!a.contact) - .map(({ contact, id, lid }) => ({ jid: id, exists: contact as boolean, lid: lid as string })) + return results.list.filter(a => !!a.contact).map(({ contact, id }) => ({ jid: id, exists: contact as boolean })) } } + const pnFromLIDUSync = async (jids: string[]): Promise => { + let usyncQuery = new USyncQuery().withLIDProtocol().withContext('background') + + for (const jid of jids) { + if (isLidUser(jid)) { + logger?.warn('LID user found in LID fetch call') + continue + } else { + usyncQuery.withUser(new USyncUser().withId(jid)) + } + } + + if (usyncQuery.users.length === 0) { + return [] // return early without forcing an empty query + } + + const results = await executeUSyncQuery(usyncQuery) + + if (results) { + return results.list.filter(a => !!a.lid).map(({ lid, id }) => ({ pn: id, lid: lid as string })) + } + + return [] + } + const ev = makeEventBuffer(logger) const { creds } = authState // add transaction capability const keys = addTransactionCapability(authState.keys, logger, transactionOpts) - const signalRepository = makeSignalRepository({ creds, keys }, logger, onWhatsApp) + const signalRepository = makeSignalRepository({ creds, keys }, logger, pnFromLIDUSync) let lastDateRecv: Date let epoch = 1 diff --git a/src/Types/Auth.ts b/src/Types/Auth.ts index f47c691f..0838c3fa 100644 --- a/src/Types/Auth.ts +++ b/src/Types/Auth.ts @@ -19,6 +19,11 @@ export type SignalIdentity = { identifierKey: Uint8Array } +export type LIDMapping = { + pn: string + lid: string +} + export type LTHashState = { version: number hash: Buffer diff --git a/src/Types/Socket.ts b/src/Types/Socket.ts index 985560c0..da68db01 100644 --- a/src/Types/Socket.ts +++ b/src/Types/Socket.ts @@ -2,7 +2,7 @@ import type { Agent } from 'https' import type { URL } from 'url' import { proto } from '../../WAProto/index.js' import type { ILogger } from '../Utils/logger' -import type { AuthenticationState, SignalAuthState, TransactionCapabilityOptions } from './Auth' +import type { AuthenticationState, LIDMapping, SignalAuthState, TransactionCapabilityOptions } from './Auth' import type { GroupMetadata } from './GroupMetadata' import { type MediaConnInfo } from './Message' import type { SignalRepositoryWithLIDStore } from './Signal' @@ -145,13 +145,6 @@ export type SocketConfig = { makeSignalRepository: ( auth: SignalAuthState, logger: ILogger, - onWhatsAppFunc?: (...jids: string[]) => Promise< - | { - jid: string - exists: boolean - lid: string - }[] - | undefined - > + pnToLIDFunc?: (jids: string[]) => Promise ) => SignalRepositoryWithLIDStore }