feat: implement transactions on auth state

src/Socket/messages-send.ts

@@ -234,6 +234,14 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
     const createParticipantNodes = async(jids: string[], bytes: Buffer) => {
         await assertSessions(jids, false)
+
+        if(authState.keys.isInTransaction()) {
+            await authState.keys.prefetch(
+                'session', 
+                jids.map(jid => jidToSignalProtocolAddress(jid).toString())
+            )
+        }
+        
         const nodes = await Promise.all(
             jids.map(
                 async jid => {
@@ -278,6 +286,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
             devices.push({ user, device })
         }
 
+        await authState.keys.transaction(
+            async() => {
                 if(isGroup) {
                     const { ciphertext, senderKeyDistributionMessageKey } = await encryptSenderKeyMsgSignalProto(destinationJid, encodedMsg, meId, authState)
                     
@@ -404,6 +414,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
                 logger.debug({ msgId }, `sending message to ${participants.length} devices`)
         
                 await sendNode(stanza)
+            }
+        )
 
         return msgId
     } 

src/Socket/socket.ts

@@ -5,7 +5,7 @@ import WebSocket from "ws"
 import { randomBytes } from 'crypto'
 import { proto } from '../../WAProto'
 import { DisconnectReason, SocketConfig, BaileysEventEmitter, ConnectionState, AuthenticationCreds } from "../Types"
-import { Curve, generateRegistrationNode, configureSuccessfulPairing, generateLoginNode, encodeBigEndian, promiseTimeout, generateOrGetPreKeys, xmppSignedPreKey, xmppPreKey, getPreKeys, makeNoiseHandler, useSingleFileAuthState } from "../Utils"
+import { Curve, generateRegistrationNode, configureSuccessfulPairing, generateLoginNode, encodeBigEndian, promiseTimeout, generateOrGetPreKeys, xmppSignedPreKey, xmppPreKey, getPreKeys, makeNoiseHandler, useSingleFileAuthState, addTransactionCapability } from "../Utils"
 import { DEFAULT_ORIGIN, DEF_TAG_PREFIX, DEF_CALLBACK_PREFIX, KEY_BUNDLE_TYPE } from "../Defaults"
 import { assertNodeErrorFree, BinaryNode, encodeBinaryNode, S_WHATSAPP_NET, getBinaryNodeChild } from '../WABinary'
 
@@ -539,7 +539,11 @@ export const makeSocket = ({
 	return {
         ws,
         ev,
-        authState,
+        authState: {
+            creds,
+            // add capability
+            keys: addTransactionCapability(authState.keys, logger)
+        },
         get user () {
             return authState.creds.me
         },

src/Types/Auth.ts

@@ -52,27 +52,19 @@ export type SignalDataTypeMap = {
     'app-state-sync-version': LTHashState
 }
 
-type SignalDataSet = { [T in keyof SignalDataTypeMap]?: { [id: string]: SignalDataTypeMap[T] | null } }
+export type SignalDataSet = { [T in keyof SignalDataTypeMap]?: { [id: string]: SignalDataTypeMap[T] | null } }
 
 type Awaitable<T> = T | Promise<T>
+
 export type SignalKeyStore = {
     get<T extends keyof SignalDataTypeMap>(type: T, ids: string[]): Awaitable<{ [id: string]: SignalDataTypeMap[T] }>
     set(data: SignalDataSet): Awaitable<void>
+}
 
-    /*getPreKey: (keyId: number) => Awaitable<KeyPair>
-    setPreKey: (keyId: number, pair: KeyPair | null) => Awaitable<void>
-
-    getSession: (sessionId: string) => Awaitable<any>
-    setSession: (sessionId: string, item: any | null) => Awaitable<void>
-
-    getSenderKey: (id: string) => Awaitable<any>
-    setSenderKey: (id: string, item: any | null) => Awaitable<void>
-
-    getAppStateSyncKey: (id: string) => Awaitable<proto.IAppStateSyncKeyData>
-    setAppStateSyncKey: (id: string, item: proto.IAppStateSyncKeyData | null) => Awaitable<void>
-
-    getAppStateSyncVersion: (name: WAPatchName) => Awaitable<LTHashState>
-    setAppStateSyncVersion: (id: WAPatchName, item: LTHashState) => Awaitable<void>*/
+export type SignalKeyStoreWithTransaction = SignalKeyStore & {
+    isInTransaction: () => boolean
+    transaction(exec: () => Promise<void>): Promise<void>
+    prefetch<T extends keyof SignalDataTypeMap>(type: T, ids: string[]): Promise<void>
 }
 
 export type SignalAuthState = {

src/Utils/auth-utils.ts

@@ -1,8 +1,97 @@
+import { Boom } from '@hapi/boom'
 import { randomBytes } from 'crypto'
-import type { AuthenticationCreds, AuthenticationState, SignalDataTypeMap } from "../Types"
+import type { Logger } from 'pino'
+import type { AuthenticationCreds, AuthenticationState, SignalDataTypeMap, SignalDataSet, SignalKeyStore, SignalKeyStoreWithTransaction } from "../Types"
 import { Curve, signedKeyPair } from './crypto'
 import { generateRegistrationId, BufferJSON } from './generics'
 
+const KEY_MAP: { [T in keyof SignalDataTypeMap]: string } = {
+	'pre-key': 'preKeys',
+	'session': 'sessions',
+	'sender-key': 'senderKeys',
+	'app-state-sync-key': 'appStateSyncKeys',
+	'app-state-sync-version': 'appStateVersions',
+	'sender-key-memory': 'senderKeyMemory'
+}
+
+export const addTransactionCapability = (state: SignalKeyStore, logger: Logger): SignalKeyStoreWithTransaction => {
+	let inTransaction = false
+	let transactionCache: SignalDataSet = { }
+	let mutations: SignalDataSet = { }
+
+	const prefetch = async(type: keyof SignalDataTypeMap, ids: string[]) => {
+		if(!inTransaction) {
+			throw new Boom('Cannot prefetch without transaction')
+		}
+
+		const dict = transactionCache[type]
+		const idsRequiringFetch = ids.filter(item => !dict?.[item])
+		const result = await state.get(type, idsRequiringFetch)
+
+		transactionCache[type] = transactionCache[type] || { }
+		Object.assign(transactionCache[type], result)
+	}
+
+	return {
+		get: async(type, ids) => {
+			if(inTransaction) {
+				await prefetch(type, ids)
+				return ids.reduce(
+					(dict, id) => {
+						const value = transactionCache[type]?.[id]
+						if(value) {
+							dict[id] = value
+						}
+						return dict
+					}, { }
+				)
+			} else {
+				return state.get(type, ids)
+			}
+		},
+		set: data => {
+			if(inTransaction) {
+				logger.trace({ types: Object.keys(data) }, `caching in transaction`)
+				for(const key in data) {
+					transactionCache[key] = transactionCache[key] || { }
+					Object.assign(transactionCache[key], data[key])
+
+					mutations[key] = mutations[key] || { }
+					Object.assign(mutations[key], data[key])
+				}
+			} else {
+				return state.set(data)
+			}
+		},
+		isInTransaction: () => inTransaction,
+		prefetch: (type, ids) => {
+			logger.trace({ type, ids }, `prefetching`)
+			return prefetch(type, ids)
+		},
+		transaction: async(work) => {
+			if(inTransaction) {
+				await work()
+			} else {
+				logger.debug('entering transaction')
+				inTransaction = true
+				try {
+					await work()
+					if(Object.keys(mutations).length) {
+						logger.debug('committing transaction')
+						await state.set(mutations)
+					} else {
+						logger.debug('no mutations in transaction')
+					}
+				} finally {
+					inTransaction = false
+					transactionCache = { }
+					mutations = { }
+				}
+			}
+		}
+	}
+}
+
 export const initAuthCreds = (): AuthenticationCreds => {
 	const identityKey = Curve.generateKeyPair()
 	return {
@@ -18,15 +107,6 @@ export const initAuthCreds = (): AuthenticationCreds => {
 	}
 }
 
-const KEY_MAP: { [T in keyof SignalDataTypeMap]: string } = {
-	'pre-key': 'preKeys',
-	'session': 'sessions',
-	'sender-key': 'senderKeys',
-	'app-state-sync-key': 'appStateSyncKeys',
-	'app-state-sync-version': 'appStateVersions',
-	'sender-key-memory': 'senderKeyMemory'
-}
-
 /** stores the full authentication state in a single JSON file */
 export const useSingleFileAuthState = (filename: string): { state: AuthenticationState, saveState: () => void } => {
 	// require fs here so that in case "fs" is not available -- the app does not crash