general: Bring back USYNC calls for LID getting, and improve

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<LIDMapping[] | undefined>
+): 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

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<LIDMapping[] | undefined>
+
+	constructor(
+		keys: SignalKeyStoreWithTransaction,
+		logger: ILogger,
+		pnToLIDFunc?: (jids: string[]) => Promise<LIDMapping[] | undefined>
+	) {
 		this.keys = keys
+		this.pnToLIDFunc = pnToLIDFunc
 		this.logger = logger
 	}
 
 	/**
 	 * Store LID-PN mapping - USER LEVEL
 	 */
-	async storeLIDPNMappings(pairs: { lid: string; pn: string }[]): Promise<void> {
+	async storeLIDPNMappings(pairs: LIDMapping[]): Promise<void> {
 		// 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<string | null> {
-		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<LIDMapping[] | null> {
+		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)
 	}
 
 	/**