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

This commit is contained in:
Rajeh Taher
2025-10-03 00:04:39 +03:00
parent 1060f44bff
commit 592f70b81d
6 changed files with 119 additions and 77 deletions
+10 -21
View File
@@ -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
+65 -29
View File
@@ -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)
}
/**
+5
View File
@@ -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[] } = {}
+32 -18
View File
@@ -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<LIDMapping[] | undefined> => {
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
+5
View File
@@ -19,6 +19,11 @@ export type SignalIdentity = {
identifierKey: Uint8Array
}
export type LIDMapping = {
pn: string
lid: string
}
export type LTHashState = {
version: number
hash: Buffer
+2 -9
View File
@@ -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<LIDMapping[] | undefined>
) => SignalRepositoryWithLIDStore
}