naruyaizumi
Version:
A WebSockets library for interacting with WhatsApp Web
275 lines (274 loc) • 10 kB
JavaScript
import NodeCache from "@cacheable/node-cache";
import { AsyncLocalStorage } from "async_hooks";
import { Mutex } from "async-mutex";
import { randomBytes } from "crypto";
import PQueue from "p-queue";
import { DEFAULT_CACHE_TTLS } from "../Defaults/index.js";
import { Curve, signedKeyPair } from "./crypto.js";
import { delay, generateRegistrationId } from "./generics.js";
import { PreKeyManager } from "./pre-key-manager.js";
/**
* Adds caching capability to a SignalKeyStore
* @param store the store to add caching to
* @param logger to log trace events
* @param _cache cache store to use
*/
export function makeCacheableSignalKeyStore(store, logger, _cache) {
const cache =
_cache ||
new NodeCache({
stdTTL: DEFAULT_CACHE_TTLS.SIGNAL_STORE, // 5 minutes
useClones: false,
deleteOnExpire: true,
});
// Mutex for protecting cache operations
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;
await 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?.();
},
};
}
/**
* Adds DB-like transaction capability to the SignalKeyStore
* Uses AsyncLocalStorage for automatic context management
* @param state the key store to apply this capability to
* @param logger logger to log events
* @returns SignalKeyStore with transaction capability
*/
export const addTransactionCapability = (
state,
logger,
{ maxCommitRetries, delayBetweenTriesMs }
) => {
const txStorage = new AsyncLocalStorage();
// Queues for concurrency control
const keyQueues = new Map();
const txMutexes = new Map();
// Pre-key manager for specialized operations
const preKeyManager = new PreKeyManager(state, logger);
/**
* Get or create a queue for a specific key type
*/
function getQueue(key) {
if (!keyQueues.has(key)) {
keyQueues.set(key, new PQueue({ concurrency: 1 }));
}
return keyQueues.get(key);
}
/**
* Get or create a transaction mutex
*/
function getTxMutex(key) {
if (!txMutexes.has(key)) {
txMutexes.set(key, new Mutex());
}
return txMutexes.get(key);
}
/**
* Check if currently in a transaction
*/
function isInTransaction() {
return !!txStorage.getStore();
}
/**
* Commit transaction with retries
*/
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) {
// No transaction - direct read without exclusive lock for concurrency
return state.get(type, ids);
}
// In transaction - check cache first
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));
// Update cache
ctx.cache[type] = ctx.cache[type] || {};
Object.assign(ctx.cache[type], fetched);
}
// Return requested ids from cache
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) {
// No transaction - direct write with queue protection
const types = Object.keys(data);
// Process pre-keys with validation
for (const type_ of types) {
const type = type_;
if (type === "pre-key") {
await preKeyManager.validateDeletions(data, type);
}
}
// Write all data in parallel
await Promise.all(
types.map((type) =>
getQueue(type).add(async () => {
const typeData = { [type]: data[type] };
await state.set(typeData);
})
)
);
return;
}
// In transaction - update cache and mutations
logger.trace({ types: Object.keys(data) }, "caching in transaction");
for (const key_ in data) {
const key = key_;
// Ensure structures exist
ctx.cache[key] = ctx.cache[key] || {};
ctx.mutations[key] = ctx.mutations[key] || {};
// Special handling for pre-keys
if (key === "pre-key") {
await preKeyManager.processOperations(
data,
key,
ctx.cache,
ctx.mutations,
true
);
} else {
// Normal key types
Object.assign(ctx.cache[key], data[key]);
Object.assign(ctx.mutations[key], data[key]);
}
}
},
isInTransaction,
transaction: async (work, key) => {
const existing = txStorage.getStore();
// Nested transaction - reuse existing context
if (existing) {
logger.trace("reusing existing transaction context");
return work();
}
// New transaction - acquire mutex and create context
return getTxMutex(key).runExclusive(async () => {
const ctx = {
cache: {},
mutations: {},
dbQueries: 0,
};
logger.trace("entering transaction");
try {
const result = await txStorage.run(ctx, work);
// Commit mutations
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,
};
};
//# sourceMappingURL=auth-utils.js.map