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 
+	}
 }