From dee815448b4d7fc9ec946b06533186fa487eb02b Mon Sep 17 00:00:00 2001 From: Adhiraj Singh Date: Sun, 7 Nov 2021 19:51:14 +0530 Subject: [PATCH] feat: cache user devices --- package.json | 1 + src/Socket/messages-send.ts | 51 +++++++++++++++++++++++++++++-------- src/Types/index.ts | 3 +++ src/Utils/signal.ts | 4 +-- src/WABinary/jid-utils.ts | 5 ++++ yarn.lock | 12 +++++++++ 6 files changed, 64 insertions(+), 12 deletions(-) diff --git a/package.json b/package.json index 5832a0e..2ed692d 100644 --- a/package.json +++ b/package.json @@ -39,6 +39,7 @@ "jimp": "^0.16.1", "libsignal": "git+https://github.com/adiwajshing/libsignal-node", "music-metadata": "^7.4.1", + "node-cache": "^5.1.2", "pino": "^6.7.0", "protobufjs": "^6.10.1", "ws": "^7.3.1" diff --git a/src/Socket/messages-send.ts b/src/Socket/messages-send.ts index bb933d4..cd97fcc 100644 --- a/src/Socket/messages-send.ts +++ b/src/Socket/messages-send.ts @@ -3,10 +3,11 @@ import got from "got" import { Boom } from "@hapi/boom" import { SocketConfig, MediaConnInfo, AnyMessageContent, MiscMessageGenerationOptions, WAMediaUploadFunction, MessageRelayOptions, WAMessageKey } from "../Types" import { encodeWAMessage, generateMessageID, generateWAMessage, encryptSenderKeyMsgSignalProto, encryptSignalProto, extractDeviceJids, jidToSignalProtocolAddress, parseAndInjectE2ESession } from "../Utils" -import { BinaryNode, getBinaryNodeChild, getBinaryNodeChildren, isJidGroup, jidDecode, jidEncode, jidNormalizedUser, S_WHATSAPP_NET, BinaryNodeAttributes } from '../WABinary' +import { BinaryNode, getBinaryNodeChild, getBinaryNodeChildren, isJidGroup, jidDecode, jidEncode, jidNormalizedUser, S_WHATSAPP_NET, BinaryNodeAttributes, JidWithDevice } from '../WABinary' import { proto } from "../../WAProto" import { WA_DEFAULT_EPHEMERAL, DEFAULT_ORIGIN, MEDIA_PATH_MAP } from "../Defaults" import { makeGroupsSocket } from "./groups" +import NodeCache from "node-cache" export const makeMessagesSocket = (config: SocketConfig) => { const { logger } = config @@ -21,6 +22,10 @@ export const makeMessagesSocket = (config: SocketConfig) => { groupToggleEphemeral } = sock + const userDevicesCache = config.userDevicesCache || new NodeCache({ + stdTTL: 300, // 5 minutes + useClones: false + }) let privacySettings: { [_: string]: string } | undefined const fetchPrivacySettings = async(force: boolean = false) => { @@ -122,11 +127,22 @@ export const makeMessagesSocket = (config: SocketConfig) => { } const getUSyncDevices = async(jids: string[], ignoreZeroDevices: boolean) => { + const deviceResults: JidWithDevice[] = [] + + const users: BinaryNode[] = [] jids = Array.from(new Set(jids)) - const users = jids.map(jid => ({ - tag: 'user', - attrs: { jid: jidNormalizedUser(jid) } - })) + for(let jid of jids) { + const user = jidDecode(jid).user + jid = jidNormalizedUser(jid) + if(userDevicesCache.has(user)) { + const devices: JidWithDevice[] = userDevicesCache.get(user) + deviceResults.push(...devices) + + logger.trace({ user }, 'using cache for devices') + } else { + users.push({ tag: 'user', attrs: { jid } }) + } + } const iq: BinaryNode = { tag: 'iq', @@ -163,7 +179,22 @@ export const makeMessagesSocket = (config: SocketConfig) => { } const result = await query(iq) const { device } = jidDecode(authState.creds.me!.id) - return extractDeviceJids(result, device, ignoreZeroDevices) + + const extracted = extractDeviceJids(result, device, ignoreZeroDevices) + const deviceMap: { [_: string]: JidWithDevice[] } = {} + + for(const item of extracted) { + deviceMap[item.user] = deviceMap[item.user] || [] + deviceMap[item.user].push(item) + + deviceResults.push(item) + } + + for(const key in deviceMap) { + userDevicesCache.set(key, deviceMap[key]) + } + + return deviceResults } const assertSession = async(jid: string, force: boolean) => { @@ -245,8 +276,8 @@ export const makeMessagesSocket = (config: SocketConfig) => { } }) - for(const {user, device, agent} of devices) { - const jid = jidEncode(user, 's.whatsapp.net', device, agent) + for(const {user, device} of devices) { + const jid = jidEncode(user, 's.whatsapp.net', device) const participant = await createParticipantNode(jid, encSenderKeyMsg) participants.push(participant) } @@ -301,11 +332,11 @@ export const makeMessagesSocket = (config: SocketConfig) => { logger.debug(`got ${devices.length} additional devices`) - for(const { user, device, agent } of devices) { + for(const { user, device } of devices) { const isMe = user === meUser participants.push( await createParticipantNode( - jidEncode(user, 's.whatsapp.net', device, agent), + jidEncode(user, 's.whatsapp.net', device), isMe ? encodedMeMsg : encodedMsg ) ) diff --git a/src/Types/index.ts b/src/Types/index.ts index 470dcc7..bcc7b1f 100644 --- a/src/Types/index.ts +++ b/src/Types/index.ts @@ -9,6 +9,7 @@ import type EventEmitter from "events" import type { Agent } from "https" import type { Logger } from "pino" import type { URL } from "url" +import type NodeCache from 'node-cache' import { AuthenticationState } from './Auth' import { Chat, PresenceData } from './Chat' @@ -45,6 +46,8 @@ export type SocketConfig = { printQRInTerminal: boolean /** should events be emitted for actions done by this socket connection */ emitOwnEvents: boolean + /** provide a cache to store a user's device list */ + userDevicesCache?: NodeCache } export enum DisconnectReason { diff --git a/src/Utils/signal.ts b/src/Utils/signal.ts index d479491..c7c90c8 100644 --- a/src/Utils/signal.ts +++ b/src/Utils/signal.ts @@ -3,7 +3,7 @@ import { encodeBigEndian } from "./generics" import { Curve } from "./crypto" import { SenderKeyDistributionMessage, GroupSessionBuilder, SenderKeyRecord, SenderKeyName, GroupCipher } from '../../WASignalGroup' import { SignalIdentity, SignalKeyStore, SignedKeyPair, KeyPair, AuthenticationState } from "../Types/Auth" -import { assertNodeErrorFree, BinaryNode, getBinaryNodeChild, getBinaryNodeChildBuffer, getBinaryNodeChildUInt, jidDecode } from "../WABinary" +import { assertNodeErrorFree, BinaryNode, getBinaryNodeChild, getBinaryNodeChildBuffer, getBinaryNodeChildUInt, jidDecode, JidWithDevice } from "../WABinary" import { proto } from "../../WAProto" export const generateSignalPubKey = (pubKey: Uint8Array | Buffer) => { @@ -232,7 +232,7 @@ export const parseAndInjectE2ESession = async(node: BinaryNode, auth: Authentica } export const extractDeviceJids = (result: BinaryNode, myDeviceId: number, excludeZeroDevices: boolean) => { - const extracted: { user: string, device?: number, agent?: number }[] = [] + const extracted: JidWithDevice[] = [] for(const node of result.content as BinaryNode[]) { const list = getBinaryNodeChild(node, 'list')?.content if(list && Array.isArray(list)) { diff --git a/src/WABinary/jid-utils.ts b/src/WABinary/jid-utils.ts index 492e050..5b09c00 100644 --- a/src/WABinary/jid-utils.ts +++ b/src/WABinary/jid-utils.ts @@ -6,6 +6,11 @@ export const STORIES_JID = 'status@broadcast' export type JidServer = 'c.us' | 'g.us' | 'broadcast' | 's.whatsapp.net' | 'call' +export type JidWithDevice = { + user: string + device?: number +} + export const jidEncode = (user: string | number | null, server: JidServer, device?: number, agent?: number) => { return `${user || ''}${!!agent ? `_${agent}` : ''}${!!device ? `:${device}` : ''}@${server}` } diff --git a/yarn.lock b/yarn.lock index c81de94..902a756 100644 --- a/yarn.lock +++ b/yarn.lock @@ -1471,6 +1471,11 @@ clone-response@^1.0.2: dependencies: mimic-response "^1.0.0" +clone@2.x: + version "2.1.2" + resolved "https://registry.yarnpkg.com/clone/-/clone-2.1.2.tgz#1b7f4b9f591f1e8f83670401600345a02887435f" + integrity sha1-G39Ln1kfHo+DZwQBYANFoCiHQ18= + co@^4.6.0: version "4.6.0" resolved "https://registry.yarnpkg.com/co/-/co-4.6.0.tgz#6ea6bdf3d853ae54ccb8e47bfa0bf3f9031fb184" @@ -2897,6 +2902,13 @@ neo-async@^2.6.0: resolved "https://registry.yarnpkg.com/neo-async/-/neo-async-2.6.2.tgz#b4aafb93e3aeb2d8174ca53cf163ab7d7308305f" integrity sha512-Yd3UES5mWCSqR+qNT93S3UoYUkqAZ9lLg8a7g9rimsWmYGK8cVToA4/sF3RrshdyV3sAGMXVUmpMYOw+dLpOuw== +node-cache@^5.1.2: + version "5.1.2" + resolved "https://registry.yarnpkg.com/node-cache/-/node-cache-5.1.2.tgz#f264dc2ccad0a780e76253a694e9fd0ed19c398d" + integrity sha512-t1QzWwnk4sjLWaQAS8CHgOJ+RAfmHpxFWmc36IWTiWHQfs0w5JDMBS1b1ZxQteo0vVVuWJvIUKHDkkeK7vIGCg== + dependencies: + clone "2.x" + node-int64@^0.4.0: version "0.4.0" resolved "https://registry.yarnpkg.com/node-int64/-/node-int64-0.4.0.tgz#87a9065cdb355d3182d8f94ce11188b825c68a3b"