UNPKG

@baileys-md/baileys

Version:

Baileys WhatsApp API

220 lines (215 loc) 6.88 kB
//========================================// import { delay, generateRegistrationId } from "./generics.js"; import { DEFAULT_CACHE_TTLS } from "../Defaults/index.js"; import { PreKeyManager } from "./pre-key-manager.js"; import { Curve, signedKeyPair } from "./crypto.js"; import { AsyncLocalStorage } from "async_hooks"; import NodeCache from "@cacheable/node-cache"; import { randomBytes } from "crypto"; import { Mutex } from "async-mutex"; import PQueue from "p-queue"; export function makeCacheableSignalKeyStore(store, logger, _cache) { const cache = _cache || new NodeCache({ stdTTL: DEFAULT_CACHE_TTLS.SIGNAL_STORE, useClones: false, deleteOnExpire: true }); const cacheMutex = new Mutex(); function getUniqueId(type, id) { return `${type}.${id}`; } return { async get(type, ids) { return cacheMutex.runExclusive(async () => { const data = {}; const idsToFetch = []; for (const id of ids) { const item = (await cache.get(getUniqueId(type, id))); if (typeof item !== "undefined") { data[id] = item; } else { idsToFetch.push(id); } } if (idsToFetch.length) { logger?.trace({ items: idsToFetch.length }, "loading from store"); const fetched = await store.get(type, idsToFetch); for (const id of idsToFetch) { const item = fetched[id]; if (item) { data[id] = item; cache.set(getUniqueId(type, id), item); } } } return data; }); }, async set(data) { return cacheMutex.runExclusive(async () => { let keys = 0; for (const type in data) { for (const id in data[type]) { await cache.set(getUniqueId(type, id), data[type][id]); keys += 1; } } logger?.trace({ keys }, "updated cache"); await store.set(data); }); }, async clear() { await cache.flushAll(); await store.clear?.(); } }; } export const addTransactionCapability = (state, logger, { maxCommitRetries, delayBetweenTriesMs }) => { const txStorage = new AsyncLocalStorage(); const keyQueues = new Map(); const txMutexes = new Map(); const preKeyManager = new PreKeyManager(state, logger); function getQueue(key) { if (!keyQueues.has(key)) { keyQueues.set(key, new PQueue({ concurrency: 1 })); } return keyQueues.get(key); } function getTxMutex(key) { if (!txMutexes.has(key)) { txMutexes.set(key, new Mutex()); } return txMutexes.get(key); } function isInTransaction() { return !!txStorage.getStore(); } async function commitWithRetry(mutations) { if (Object.keys(mutations).length === 0) { logger.trace("no mutations in transaction"); return; } logger.trace("committing transaction"); for (let attempt = 0; attempt < maxCommitRetries; attempt++) { try { await state.set(mutations); logger.trace({ mutationCount: Object.keys(mutations).length }, "committed transaction"); return; } catch (error) { const retriesLeft = maxCommitRetries - attempt - 1; logger.warn(`failed to commit mutations, retries left=${retriesLeft}`); if (retriesLeft === 0) { throw error; } await delay(delayBetweenTriesMs); } } } return { get: async (type, ids) => { const ctx = txStorage.getStore(); if (!ctx) { return state.get(type, ids); } const cached = ctx.cache[type] || {}; const missing = ids.filter(id => !(id in cached)); if (missing.length > 0) { ctx.dbQueries++; logger.trace({ type, count: missing.length }, "fetching missing keys in transaction"); const fetched = await getTxMutex(type).runExclusive(() => state.get(type, missing)); ctx.cache[type] = ctx.cache[type] || {}; Object.assign(ctx.cache[type], fetched); } const result = {}; for (const id of ids) { const value = ctx.cache[type]?.[id]; if (value !== undefined && value !== null) { result[id] = value; } } return result; }, set: async (data) => { const ctx = txStorage.getStore(); if (!ctx) { const types = Object.keys(data); for (const type_ of types) { const type = type_; if (type === "pre-key") { await preKeyManager.validateDeletions(data, type); } } await Promise.all(types.map(type => getQueue(type).add(async () => { const typeData = { [type]: data[type] }; await state.set(typeData); }))); return; } logger.trace({ types: Object.keys(data) }, "caching in transaction"); for (const key_ in data) { const key = key_; ctx.cache[key] = ctx.cache[key] || {}; ctx.mutations[key] = ctx.mutations[key] || {}; if (key === "pre-key") { await preKeyManager.processOperations(data, key, ctx.cache, ctx.mutations, true); } else { Object.assign(ctx.cache[key], data[key]); Object.assign(ctx.mutations[key], data[key]); } } }, isInTransaction, transaction: async (work, key) => { const existing = txStorage.getStore(); if (existing) { logger.trace("reusing existing transaction context"); return work(); } return getTxMutex(key).runExclusive(async () => { const ctx = { cache: {}, mutations: {}, dbQueries: 0 }; logger.trace("entering transaction"); try { const result = await txStorage.run(ctx, work); await commitWithRetry(ctx.mutations); logger.trace({ dbQueries: ctx.dbQueries }, "transaction completed"); return result; } catch (error) { logger.error({ error }, "transaction failed, rolling back"); throw error; } }); } }; }; export const initAuthCreds = () => { const identityKey = Curve.generateKeyPair(); return { "noiseKey": Curve.generateKeyPair(), "pairingEphemeralKeyPair": Curve.generateKeyPair(), "signedIdentityKey": identityKey, "signedPreKey": signedKeyPair(identityKey, 1), "registrationId": generateRegistrationId(), "advSecretKey": randomBytes(32).toString("base64"), "processedHistoryMessages": [], "nextPreKeyId": 1, "firstUnuploadedPreKeyId": 1, "accountSyncCounter": 0, "accountSettings": { "unarchiveChats": false }, "registered": false, "pairingCode": undefined, "lastPropHash": undefined, "routingInfo": undefined, "additionalData": undefined }; };