feat: cleaner auth state management + store SK keys

!BREAKING_CHANGE

src/Utils/auth-utils.ts

@@ -1,86 +1,8 @@
 import { randomBytes } from 'crypto'
-import { proto } from '../../WAProto'
-import type { SignalKeyStore, AuthenticationCreds, KeyPair, LTHashState, AuthenticationState } from "../Types"
+import type { AuthenticationCreds, AuthenticationState, SignalDataTypeMap } from "../Types"
 import { Curve, signedKeyPair } from './crypto'
 import { generateRegistrationId, BufferJSON } from './generics'
 
-export const initInMemoryKeyStore = (
-	{ preKeys, sessions, senderKeys, appStateSyncKeys, appStateVersions }: { 
-		preKeys?: { [k: number]: KeyPair },
-		sessions?: { [k: string]: any },
-		senderKeys?: { [k: string]: any }
-		appStateSyncKeys?: { [k: string]: proto.IAppStateSyncKeyData },
-		appStateVersions?: { [k: string]: LTHashState },
-	} = { },
-	save: (data: any) => void
-) => {
-
-	preKeys = preKeys || { }
-	sessions = sessions || { }
-	senderKeys = senderKeys || { }
-	appStateSyncKeys = appStateSyncKeys || { }
-	appStateVersions = appStateVersions || { }
-
-	const keyData = {
-		preKeys,
-		sessions,
-		senderKeys,
-		appStateSyncKeys,
-		appStateVersions,
-	}
-
-	return {
-		...keyData,
-		getPreKey: keyId => preKeys[keyId],
-		setPreKey: (keyId, pair) => {
-			if(pair) preKeys[keyId] = pair
-			else delete preKeys[keyId]
-
-			save(keyData)
-		},
-		getSession: id => sessions[id],
-		setSession: (id, item) => {
-			if(item) sessions[id] = item
-			else delete sessions[id]
-
-			save(keyData)
-		},
-		getSenderKey: id => {
-			return senderKeys[id]
-		},
-		setSenderKey: (id, item) => {
-			if(item) senderKeys[id] = item
-			else delete senderKeys[id]
-
-			save(keyData)
-		},
-		getAppStateSyncKey: id => {
-			const obj = appStateSyncKeys[id]
-			if(obj) {
-				return proto.AppStateSyncKeyData.fromObject(obj)
-			}
-		},
-		setAppStateSyncKey: (id, item) => {
-			if(item) appStateSyncKeys[id] = item
-			else delete appStateSyncKeys[id]
-
-			save(keyData)
-		},
-		getAppStateSyncVersion: id => {
-			const obj = appStateVersions[id]
-			if(obj) {
-				return obj
-			}
-		},
-		setAppStateSyncVersion: (id, item) => {
-			if(item) appStateVersions[id] = item
-			else delete appStateVersions[id]
-
-			save(keyData)
-		}
-	} as SignalKeyStore
-}
-
 export const initAuthCreds = (): AuthenticationCreds => {
 	const identityKey = Curve.generateKeyPair()
 	return {
@@ -95,12 +17,22 @@ export const initAuthCreds = (): AuthenticationCreds => {
 		serverHasPreKeys: false
 	}
 }
+
+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) => {
+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
 	const { readFileSync, writeFileSync, existsSync } = require('fs')
-
-	let state: AuthenticationState = undefined
+	let creds: AuthenticationCreds
+	let keys: any = { }
 
 	// save the authentication state to a file
 	const saveState = () => {
@@ -108,26 +40,48 @@ export const useSingleFileAuthState = (filename: string) => {
 		writeFileSync(
 			filename,
 			// BufferJSON replacer utility saves buffers nicely
-			JSON.stringify(state, BufferJSON.replacer, 2)
+			JSON.stringify({ creds, keys }, BufferJSON.replacer, 2)
 		)
 	}
 	
     if(existsSync(filename)) {
-        const { creds, keys } = JSON.parse(
+        const result = JSON.parse(
             readFileSync(filename, { encoding: 'utf-8' }), 
             BufferJSON.reviver
         )
-        state = { 
-            creds: creds, 
-            // stores pre-keys, session & other keys in a JSON object
-            // we deserialize it here
-            keys: initInMemoryKeyStore(keys, saveState) 
-        }
+		creds = result.creds
+		keys = result.keys
     } else {
-        const creds = initAuthCreds()
-        const keys = initInMemoryKeyStore({ }, saveState)
-        state = { creds: creds, keys: keys }
+        creds = initAuthCreds()
+        keys = { }
     }
 
-	return { state, saveState }
+	return { 
+		state: { 
+			creds,
+			keys: {
+				get: (type, ids) => {
+					const key = KEY_MAP[type]
+					return ids.reduce(
+						(dict, id) => {
+							const value = keys[key]?.[id]
+							if(value) {
+								dict[id] = value
+							}
+							return dict
+						}, { }
+					)
+				},
+				set: (data) => {
+					for(const _key in data) {
+						const key = KEY_MAP[_key as keyof SignalDataTypeMap]
+						keys[key] = keys[key] || { }
+						Object.assign(keys[key], data[_key])
+					}
+					saveState()
+				}
+			}
+		}, 
+		saveState 
+	}
 }

src/Utils/chat-utils.ts

@@ -1,12 +1,14 @@
 import { Boom } from '@hapi/boom'
 import { aesDecrypt, hmacSign, aesEncrypt, hkdf } from "./crypto"
-import { AuthenticationState, WAPatchCreate, ChatMutation, WAPatchName, LTHashState, ChatModification, SignalKeyStore } from "../Types"
+import { WAPatchCreate, ChatMutation, WAPatchName, LTHashState, ChatModification, SignalKeyStore } from "../Types"
 import { proto } from '../../WAProto'
 import { LT_HASH_ANTI_TAMPERING } from './lt-hash'
 import { BinaryNode, getBinaryNodeChild, getBinaryNodeChildren } from '../WABinary'
 import { toNumber } from './generics'
 import { downloadContentFromMessage,  } from './messages-media'
 
+type FetchAppStateSyncKey = (keyId: string) => Promise<proto.IAppStateSyncKeyData> | proto.IAppStateSyncKeyData
+
 const mutationKeys = (keydata: Uint8Array) => {
     const expanded = hkdf(keydata, 160, { info: 'WhatsApp Mutation Keys' })
     return {
@@ -112,9 +114,9 @@ export const encodeSyncdPatch = async(
     { type, index, syncAction, apiVersion, operation }: WAPatchCreate,
     myAppStateKeyId: string,
     state: LTHashState,
-    keys: SignalKeyStore
+    getAppStateSyncKey: FetchAppStateSyncKey
 ) => {
-    const key = !!myAppStateKeyId ? await keys.getAppStateSyncKey(myAppStateKeyId) : undefined
+    const key = !!myAppStateKeyId ? await getAppStateSyncKey(myAppStateKeyId) : undefined
     if(!key) {
         throw new Boom(`myAppStateKey ("${myAppStateKeyId}") not present`, { statusCode: 404 })
     }
@@ -175,7 +177,7 @@ export const encodeSyncdPatch = async(
 export const decodeSyncdMutations = async(
     msgMutations: (proto.ISyncdMutation | proto.ISyncdRecord)[], 
     initialState: LTHashState,
-    getAppStateSyncKey: SignalKeyStore['getAppStateSyncKey'],
+    getAppStateSyncKey: FetchAppStateSyncKey,
     validateMacs: boolean
 ) => {
     const keyCache: { [_: string]: ReturnType<typeof mutationKeys> } = { }
@@ -247,7 +249,7 @@ export const decodeSyncdPatch = async(
     msg: proto.ISyncdPatch, 
     name: WAPatchName,
     initialState: LTHashState,
-    getAppStateSyncKey: SignalKeyStore['getAppStateSyncKey'],
+    getAppStateSyncKey: FetchAppStateSyncKey,
     validateMacs: boolean
 ) => {
     if(validateMacs) {
@@ -334,7 +336,7 @@ export const downloadExternalPatch = async(blob: proto.IExternalBlobReference) =
 export const decodeSyncdSnapshot = async(
     name: WAPatchName,
     snapshot: proto.ISyncdSnapshot,
-    getAppStateSyncKey: SignalKeyStore['getAppStateSyncKey'],
+    getAppStateSyncKey: FetchAppStateSyncKey,
     validateMacs: boolean = true
 ) => {
     const newState = newLTHashState()
@@ -370,7 +372,7 @@ export const decodePatches = async(
     name: WAPatchName,
     syncds: proto.ISyncdPatch[],
     initial: LTHashState,
-    getAppStateSyncKey: SignalKeyStore['getAppStateSyncKey'],
+    getAppStateSyncKey: FetchAppStateSyncKey,
     validateMacs: boolean = true
 ) => {
     const successfulMutations: ChatMutation[] = []

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, SignalAuthState, AuthenticationCreds } from "../Types/Auth"
-import { assertNodeErrorFree, BinaryNode, getBinaryNodeChild, getBinaryNodeChildBuffer, getBinaryNodeChildUInt, jidDecode, JidWithDevice } from "../WABinary"
+import { assertNodeErrorFree, BinaryNode, getBinaryNodeChild, getBinaryNodeChildBuffer, getBinaryNodeChildUInt, jidDecode, JidWithDevice, getBinaryNodeChildren } from "../WABinary"
 import { proto } from "../../WAProto"
 
 export const generateSignalPubKey = (pubKey: Uint8Array | Buffer) => {
@@ -33,13 +33,12 @@ export const createSignalIdentity = (
 	}
 }
 
-export const getPreKeys = async({ getPreKey }: SignalKeyStore, min: number, limit: number) => {
-	const dict: { [id: number]: KeyPair } = { }
+export const getPreKeys = async({ get }: SignalKeyStore, min: number, limit: number) => {
+	const idList: string[] = []
 	for(let id = min; id < limit;id++) {
-		const key = await getPreKey(id)
-		if(key) dict[+id] = key
+		idList.push(id.toString())
 	}
-	return dict
+	return get('pre-key', idList)
 }
 
 export const generateOrGetPreKeys = (creds: AuthenticationCreds, range: number) => {
@@ -84,20 +83,21 @@ export const xmppPreKey = (pair: KeyPair, id: number): BinaryNode => (
 )
 
 export const signalStorage = ({ creds, keys }: SignalAuthState) => ({
-	loadSession: async id => {
-		const sess = await keys.getSession(id)
+	loadSession: async (id: string) => {
+		const { [id]: sess } = await keys.get('session', [id])
 		if(sess) {
 			return libsignal.SessionRecord.deserialize(sess)
 		}
 	},
 	storeSession: async(id, session) => {
-		await keys.setSession(id, session.serialize())
+		await keys.set({ 'session': { [id]: session.serialize() } })
 	},
 	isTrustedIdentity: () => {
 		return true
 	},
-	loadPreKey: async(id: number) => {
-		const key = await keys.getPreKey(id)
+	loadPreKey: async(id: number | string) => {
+		const keyId = id.toString()
+		const { [keyId]: key } = await keys.get('pre-key', [keyId])
 		if(key) {
 			return {
 				privKey: Buffer.from(key.private),
@@ -105,7 +105,7 @@ export const signalStorage = ({ creds, keys }: SignalAuthState) => ({
 			}
 		}
 	},
-	removePreKey: (id: number) => keys.setPreKey(id, null),
+	removePreKey: (id: number) => keys.set({ 'pre-key': { [id]: null } }),
 	loadSignedPreKey: (keyId: number) => {
 		const key = creds.signedPreKey
 		return {
@@ -113,12 +113,12 @@ export const signalStorage = ({ creds, keys }: SignalAuthState) => ({
 			pubKey: Buffer.from(key.keyPair.public)
 		}
 	},
-	loadSenderKey: async(keyId) => {
-		const key = await keys.getSenderKey(keyId)
+	loadSenderKey: async(keyId: string) => {
+		const { [keyId]: key } = await keys.get('sender-key', [keyId])
 		if(key) return new SenderKeyRecord(key)
 	},
 	storeSenderKey: async(keyId, key) => {
-		await keys.setSenderKey(keyId, key.serialize())
+		await keys.set({ 'sender-key': { [keyId]: key.serialize() } })
 	},
 	getOurRegistrationId: () => (
 		creds.registrationId
@@ -148,10 +148,10 @@ export const processSenderKeyMessage = async(
 	const senderName = jidToSignalSenderKeyName(item.groupId, authorJid)
 
 	const senderMsg = new SenderKeyDistributionMessage(null, null, null, null, item.axolotlSenderKeyDistributionMessage)
-	const senderKey = await auth.keys.getSenderKey(senderName)
+	const { [senderName]: senderKey } = await auth.keys.get('sender-key', [senderName])
 	if(!senderKey) {
 		const record = new SenderKeyRecord()
-		await auth.keys.setSenderKey(senderName, record)
+		await auth.keys.set({ 'sender-key': { [senderName]: record } })
 	}
 	await builder.process(senderName, senderMsg)
 }
@@ -188,10 +188,10 @@ export const encryptSenderKeyMsgSignalProto = async(group: string, data: Uint8Ar
 	const senderName = jidToSignalSenderKeyName(group, meId)
 	const builder = new GroupSessionBuilder(storage)
 
-	const senderKey = await auth.keys.getSenderKey(senderName)
+	const { [senderName]: senderKey } = await auth.keys.get('sender-key', [senderName])
 	if(!senderKey) {
 		const record = new SenderKeyRecord()
-		await auth.keys.setSenderKey(senderName, record)
+		await auth.keys.set({ 'sender-key': { [senderName]: record } })
 	}
 
 	const senderKeyDistributionMessage = await builder.create(senderName)
@@ -202,7 +202,7 @@ export const encryptSenderKeyMsgSignalProto = async(group: string, data: Uint8Ar
 	}
 }
 
-export const parseAndInjectE2ESession = async(node: BinaryNode, auth: SignalAuthState) => {
+export const parseAndInjectE2ESessions = async(node: BinaryNode, auth: SignalAuthState) => {
 	const extractKey = (key: BinaryNode) => (
 		key ? ({
             keyId: getBinaryNodeChildUInt(key, 'id', 3),
@@ -212,23 +212,30 @@ export const parseAndInjectE2ESession = async(node: BinaryNode, auth: SignalAuth
             signature: getBinaryNodeChildBuffer(key, 'signature'),
         }) : undefined
 	)
-	node = getBinaryNodeChild(getBinaryNodeChild(node, 'list'), 'user')
-	assertNodeErrorFree(node)
-	
-	const signedKey = getBinaryNodeChild(node, 'skey')
-	const key = getBinaryNodeChild(node, 'key')
-	const identity = getBinaryNodeChildBuffer(node, 'identity')
-	const jid = node.attrs.jid
-	const registrationId = getBinaryNodeChildUInt(node, 'registration', 4)
-
-	const device = {
-		registrationId,
-		identityKey: generateSignalPubKey(identity),
-		signedPreKey: extractKey(signedKey),
-		preKey: extractKey(key)
+	const nodes = getBinaryNodeChildren(getBinaryNodeChild(node, 'list'), 'user')
+	for(const node of nodes) {
+		assertNodeErrorFree(node)
 	}
-	const cipher = new libsignal.SessionBuilder(signalStorage(auth), jidToSignalProtocolAddress(jid))
-	await cipher.initOutgoing(device)
+	await Promise.all(
+		nodes.map(
+			async node => {
+				const signedKey = getBinaryNodeChild(node, 'skey')
+				const key = getBinaryNodeChild(node, 'key')
+				const identity = getBinaryNodeChildBuffer(node, 'identity')
+				const jid = node.attrs.jid
+				const registrationId = getBinaryNodeChildUInt(node, 'registration', 4)
+			
+				const device = {
+					registrationId,
+					identityKey: generateSignalPubKey(identity),
+					signedPreKey: extractKey(signedKey),
+					preKey: extractKey(key)
+				}
+				const cipher = new libsignal.SessionBuilder(signalStorage(auth), jidToSignalProtocolAddress(jid))
+				await cipher.initOutgoing(device)
+			}
+		)
+	)
 }
 
 export const extractDeviceJids = (result: BinaryNode, myJid: string, excludeZeroDevices: boolean) => {