refactor: cleaner & faster app state sync management

1. Is orders of magnitude faster than the previous edition
2. Stores lesser data, so more memory efficient
3. This breaks the current app state in baileys, but baileys will auto-resync & update the state

src/Types/Auth.ts

@@ -1,6 +1,6 @@
 import type { Contact } from "./Contact"
 import type { proto } from "../../WAProto"
-import type { WAPatchName, ChatMutation } from "./Chat"
+import type { WAPatchName } from "./Chat"
 
 export type KeyPair = { public: Uint8Array, private: Uint8Array }
 export type SignedKeyPair = { keyPair: KeyPair, signature: Uint8Array, keyId: number }
@@ -14,7 +14,13 @@ export type SignalIdentity = {
 	identifierKey: Uint8Array
 }
 
-export type LTHashState = { version: number, hash: Buffer, mutations: ChatMutation[] }
+export type LTHashState = { 
+    version: number
+    hash: Buffer
+    indexValueMap: { 
+        [indexMacBase64: string]: { valueMac: Uint8Array | Buffer }
+    }
+}
 
 export type SignalCreds = {
     readonly signedIdentityKey: KeyPair

src/Utils/chat-utils.ts

@@ -51,29 +51,40 @@ const to64BitNetworkOrder = function(e) {
 
 type Mac = { indexMac: Uint8Array, valueMac: Uint8Array, operation: proto.SyncdMutation.SyncdMutationSyncdOperation }
 
-const computeLtHash = (initial: Uint8Array, macs: Mac[], getPrevSetValueMac: (index: Uint8Array, internalIndex: number) => { valueMac: Uint8Array, operation: number }) => {
+const makeLtHashGenerator = ({ indexValueMap, hash }: Pick<LTHashState, 'hash' | 'indexValueMap'>) => {
+    indexValueMap = { ...indexValueMap }
     const addBuffs: ArrayBuffer[] = []
     const subBuffs: ArrayBuffer[] = []
-    for(let i = 0; i < macs.length;i++) {
+
-        const { indexMac, valueMac, operation } = macs[i]
+    return {
-        const subOp = getPrevSetValueMac(indexMac, i)
+        mix: ({ indexMac, valueMac, operation }: Mac) => {
-        if(operation === proto.SyncdMutation.SyncdMutationSyncdOperation.REMOVE) {
+            const indexMacBase64 = Buffer.from(indexMac).toString('base64')
-            if(!subOp) {
+            const prevOp = indexValueMap[indexMacBase64]
-                throw new Boom('tried remove, but no buffer', { data: { indexMac, valueMac } })
+            if(operation === proto.SyncdMutation.SyncdMutationSyncdOperation.REMOVE) {
+                if(!prevOp) {
+                    throw new Boom('tried remove, but no previous op', { data: { indexMac, valueMac } })
+                }
+                // remove from index value mac, since this mutation is erased
+                delete indexValueMap[indexMacBase64]
+            } else {
+                addBuffs.push(new Uint8Array(valueMac).buffer)
+                // add this index into the history map
+                indexValueMap[indexMacBase64] = { valueMac }
             }
-        } else {
+            if(prevOp) {
-            addBuffs.push(new Uint8Array(valueMac).buffer)
+                subBuffs.push(new Uint8Array(prevOp.valueMac).buffer)
-        }
+            }
-        if(subOp) {
+        },
-            if(subOp.operation === proto.SyncdMutation.SyncdMutationSyncdOperation.SET) {
+        finish: () => {
-                subBuffs.push(new Uint8Array(subOp.valueMac).buffer)
+            const result = LT_HASH_ANTI_TAMPERING.subtractThenAdd(new Uint8Array(hash).buffer, addBuffs, subBuffs)
+            const buffer = Buffer.from(result)
+
+            return {
+                hash: buffer,
+                indexValueMap
             }
         }
     }
-
-    const result = LT_HASH_ANTI_TAMPERING.subtractThenAdd(new Uint8Array(initial).buffer, addBuffs, subBuffs)
-    const buff = Buffer.from(result)
-    return buff
 }
 
 const generateSnapshotMac = (lthash: Uint8Array, version: number, name: WAPatchName, key: Buffer) => {
@@ -125,11 +136,11 @@ export const encodeSyncdPatch = async(
     const valueMac = generateMac(operation, encValue, encKeyId, keyValue.valueMacKey)
     const indexMac = hmacSign(indexBuffer, keyValue.indexKey)
 
-    state.hash = computeLtHash(
+    // update LT hash
-        state.hash, 
+    const generator = makeLtHashGenerator(state)
-        [ { indexMac, valueMac, operation } ],
+    generator.mix({ indexMac, valueMac, operation })
-        (index) => [...state.mutations].reverse().find(m => Buffer.compare(m.indexMac, index) === 0)
+    Object.assign(state, generator.finish())
-    )
+
     state.version += 1
 
     const snapshotMac = generateSnapshotMac(state.hash, state.version, type, keyValue.snapshotMacKey)
@@ -154,21 +165,15 @@ export const encodeSyncdPatch = async(
         ]
     }
 
-    state.mutations = [
+    const base64Index = indexMac.toString('base64')
-        ...state.mutations,
+    state.indexValueMap[base64Index] = { valueMac }
-        {
+
-            action: syncAction,
-            index,
-            valueMac,
-            indexMac,
-            operation
-        }
-    ]
     return { patch, state }
 }
 
 export const decodeSyncdMutations = async(
-    msgMutations: proto.ISyncdMutation[], 
+    msgMutations: (proto.ISyncdMutation | proto.ISyncdRecord)[], 
+    initialState: LTHashState,
     getAppStateSyncKey: SignalKeyStore['getAppStateSyncKey'],
     validateMacs: boolean
 ) => {
@@ -188,11 +193,18 @@ export const decodeSyncdMutations = async(
         return key
     }
 
+    const ltGenerator = makeLtHashGenerator(initialState)
+
     const mutations: ChatMutation[] = []
     // indexKey used to HMAC sign record.index.blob
     // valueEncryptionKey used to AES-256-CBC encrypt record.value.blob[0:-32]
     // the remaining record.value.blob[0:-32] is the mac, it the HMAC sign of key.keyId + decoded proto data + length of bytes in keyId
-    for(const { operation, record } of msgMutations!) {
+    for(const msgMutation of msgMutations!) {
+        // if it's a syncdmutation, get the operation property
+        // otherwise, if it's only a record -- it'll be a SET mutation
+        const operation = 'operation' in msgMutation ? msgMutation.operation : proto.SyncdMutation.SyncdMutationSyncdOperation.SET
+        const record = ('record' in msgMutation && !!msgMutation.record) ? msgMutation.record : msgMutation as proto.ISyncdRecord
+        
         const key = await getKey(record.keyId!.id!)
         const content = Buffer.from(record.value!.blob!)
         const encContent = content.slice(0, -32)
@@ -215,21 +227,25 @@ export const decodeSyncdMutations = async(
         }
 
         const indexStr = Buffer.from(syncAction.index).toString()
-        mutations.push({ 
+        const mutation: ChatMutation = { 
             action: syncAction.value!, 
             index: JSON.parse(indexStr),
             indexMac: record.index!.blob!,
             valueMac: ogValueMac,
             operation: operation
-        })
+        }
+        mutations.push(mutation)
+
+        ltGenerator.mix(mutation)
     }
 
-    return { mutations }
+    return { mutations, ...ltGenerator.finish() }
 }
 
 export const decodeSyncdPatch = async(
     msg: proto.ISyncdPatch, 
     name: WAPatchName,
+    initialState: LTHashState,
     getAppStateSyncKey: SignalKeyStore['getAppStateSyncKey'],
     validateMacs: boolean
 ) => {
@@ -245,7 +261,7 @@ export const decodeSyncdPatch = async(
         } 
     }
 
-    const result = await decodeSyncdMutations(msg!.mutations!, getAppStateSyncKey, validateMacs)
+    const result = await decodeSyncdMutations(msg!.mutations!, initialState, getAppStateSyncKey, validateMacs)
     return result
 }
 
@@ -286,10 +302,6 @@ export const extractSyncdPatches = async(result: BinaryNode) => {
                         if(!syncd.version) {
                             syncd.version = { version: +collectionNode.attrs.version+1 }
                         }
-                        if(syncd.externalMutations) {
-                            const ref = await downloadExternalPatch(syncd.externalMutations)
-                            syncd.mutations.push(...ref.mutations)
-                        }
                         syncds.push(syncd)
                     }
                 }
@@ -324,21 +336,18 @@ export const decodeSyncdSnapshot = async(
     getAppStateSyncKey: SignalKeyStore['getAppStateSyncKey'],
     validateMacs: boolean = true
 ) => {
+    const newState = newLTHashState()
+    newState.version = toNumber(snapshot.version!.version!)
+    
+    const records = snapshot.records!
 
-    const version = toNumber(snapshot.version!.version!)
+    const ltGenerator = makeLtHashGenerator(newState)
+    for(const { index, value } of records) {
+        const valueMac = value.blob!.slice(-32)!
+        ltGenerator.mix({ indexMac: index.blob!, valueMac, operation: 0 })
+    }
 
-    const mappedRecords = snapshot.records!.map(
+    Object.assign(newState, ltGenerator.finish())
-        record => ({ record, operation: proto.SyncdMutation.SyncdMutationSyncdOperation.SET })
-    )
-    const macs = mappedRecords.map(m => ({
-        operation: m.operation!,
-        indexMac: m.record.index!.blob!,
-        valueMac: m.record.value!.blob!.slice(-32)
-    }))
-    const { mutations } = await decodeSyncdMutations(mappedRecords, getAppStateSyncKey, validateMacs)
-
-    let hash = Buffer.alloc(128)
-    hash = computeLtHash(hash, macs, () => undefined)
 
     if(validateMacs) {
         const base64Key = Buffer.from(snapshot.keyId!.id!).toString('base64')
@@ -347,14 +356,13 @@ export const decodeSyncdSnapshot = async(
             throw new Boom(`failed to find key "${base64Key}" to decode mutation`, { statusCode: 500 })
         }
         const result = mutationKeys(keyEnc.keyData!)
-        const computedSnapshotMac = generateSnapshotMac(hash, version, name, result.snapshotMacKey)
+        const computedSnapshotMac = generateSnapshotMac(newState.hash, newState.version, name, result.snapshotMacKey)
         if(Buffer.compare(snapshot.mac!, computedSnapshotMac) !== 0) {
-            throw new Boom(`failed to verify LTHash at ${version} of ${name} from snapshot`, { statusCode: 500 })
+            throw new Boom(`failed to verify LTHash at ${newState.version} of ${name} from snapshot`, { statusCode: 500 })
         }
     }
-    
+
-    const state: LTHashState = { version, mutations, hash }
+    return newState
-    return state
 }
 
 export const decodePatches = async(
@@ -366,72 +374,42 @@ export const decodePatches = async(
 ) => {
     const successfulMutations: ChatMutation[] = []
 
-    let current = initial.hash
+    const newState: LTHashState = {
-    let currentVersion = initial.version
+        ...initial,
+        indexValueMap: { ...initial.indexValueMap }
+    }
 
     for(const syncd of syncds) {
-        const { mutations, version, keyId, snapshotMac } = syncd
+        const { version, keyId, snapshotMac } = syncd
-        const macs = mutations.map(
+        if(syncd.externalMutations) {
-            m => ({
+            const ref = await downloadExternalPatch(syncd.externalMutations)
-                operation: m.operation!,
+            syncd.mutations.push(...ref.mutations)
-                indexMac: m.record.index!.blob!,
+        }
-                valueMac: m.record.value!.blob!.slice(-32)
-            })
-        )
         
-        currentVersion = toNumber(version.version!)
+        newState.version = toNumber(version.version!)
 
-        current = computeLtHash(current, macs, (index, maxIndex) => {
+        const decodeResult = await decodeSyncdPatch(syncd, name, newState, getAppStateSyncKey, validateMacs)
-            let result: { valueMac: Uint8Array, operation: number }
-            for(const item of initial.mutations) {
-                if(Buffer.compare(item.indexMac, index) === 0) {
-                    result = item
-                }
-            }
-            for(const { version, mutations } of syncds) {
-                const versionNum = toNumber(version.version!)
-                const mutationIdx = mutations.findIndex(m => {
-                    return Buffer.compare(m.record!.index!.blob, index) === 0
-                })
 
-                if(mutationIdx >= 0 && (versionNum < currentVersion || mutationIdx < maxIndex)) {
+        newState.hash = decodeResult.hash
-                    const mut = mutations[mutationIdx]
+        newState.indexValueMap = decodeResult.indexValueMap
-                    result = {
+        successfulMutations.push(...decodeResult.mutations)
-                        valueMac: mut.record!.value!.blob!.slice(-32),
-                        operation: mut.operation
-                    }
-                }
-
-                if(versionNum >= currentVersion) {
-                    break
-                }
-            }
-            return result
-        })
         
         if(validateMacs) {
             const base64Key = Buffer.from(keyId!.id!).toString('base64')
             const keyEnc = await getAppStateSyncKey(base64Key)
             if(!keyEnc) {
-                throw new Boom(`failed to find key "${base64Key}" to decode mutation`, { statusCode: 500 })
+                throw new Boom(`failed to find key "${base64Key}" to decode mutation`)
             }
             const result = mutationKeys(keyEnc.keyData!)
-            const computedSnapshotMac = generateSnapshotMac(current, currentVersion, name, result.snapshotMacKey)
+            const computedSnapshotMac = generateSnapshotMac(newState.hash, newState.version, name, result.snapshotMacKey)
             if(Buffer.compare(snapshotMac, computedSnapshotMac) !== 0) {
-                throw new Boom(`failed to verify LTHash at ${currentVersion} of ${name}`, { statusCode: 500 })
+                throw new Boom(`failed to verify LTHash at ${newState.version} of ${name}`)
             }
         }
-
-        const decodeResult = await decodeSyncdPatch(syncd, name, getAppStateSyncKey, validateMacs)
-        successfulMutations.push(...decodeResult.mutations)
     }
     return {
         newMutations: successfulMutations,
-        state: { 
+        state: newState
-            hash: current, 
-            version: currentVersion, 
-            mutations: [...initial.mutations, ...successfulMutations]
-        } as LTHashState
     }
 }