feat: cleaner auth state management + store SK keys

!BREAKING_CHANGE

src/Socket/chats.ts

@@ -20,6 +20,11 @@ export const makeChatsSocket = (config: SocketConfig) => {
 
     const mutationMutex = makeMutex()
 
+    const getAppStateSyncKey = async(keyId: string) => {
+        const { [keyId]: key } = await authState.keys.get('app-state-sync-key', [keyId])
+        return key
+    }
+
     const interactiveQuery = async(userNodes: BinaryNode[], queryNode: BinaryNode) => {
         const result = await query({
             tag: 'iq',
@@ -175,7 +180,11 @@ export const makeChatsSocket = (config: SocketConfig) => {
 
         const states = { } as { [T in WAPatchName]: LTHashState }
         for(const name of collections) {
-            let state: LTHashState = fromScratch ? undefined : await authState.keys.getAppStateSyncVersion(name)
+            let state: LTHashState 
+            if(!fromScratch) {
+                const result = await authState.keys.get('app-state-sync-version', [name])
+                state = result[name]
+            }
             if(!state) state = newLTHashState()
 
             states[name] = state
@@ -213,16 +222,16 @@ export const makeChatsSocket = (config: SocketConfig) => {
             const name = key as WAPatchName
             const { patches, snapshot } = decoded[name]
             if(snapshot) {
-                const newState = await decodeSyncdSnapshot(name, snapshot, authState.keys.getAppStateSyncKey)
+                const newState = await decodeSyncdSnapshot(name, snapshot, getAppStateSyncKey)
                 states[name] = newState
 
                 logger.info(`restored state of ${name} from snapshot to v${newState.version}`)
             }
             // only process if there are syncd patches
             if(patches.length) {
-                const { newMutations, state: newState } = await decodePatches(name, patches, states[name], authState.keys.getAppStateSyncKey, true)
+                const { newMutations, state: newState } = await decodePatches(name, patches, states[name], getAppStateSyncKey, true)
 
-                await authState.keys.setAppStateSyncVersion(name, newState)
+                await authState.keys.set({ 'app-state-sync-version': { [name]: newState } })
     
                 logger.info(`synced ${name} to v${newState.version}`)
                 if(newMutations.length) {
@@ -415,12 +424,12 @@ export const makeChatsSocket = (config: SocketConfig) => {
                 logger.debug({ patch: patchCreate }, 'applying app patch')
 
                 await resyncAppState([name])
-                const initial = await authState.keys.getAppStateSyncVersion(name)
+                const { [name]: initial } = await authState.keys.get('app-state-sync-version', [name])
                 const { patch, state } = await encodeSyncdPatch(
                     patchCreate,
                     authState.creds.myAppStateKeyId!,
                     initial,
-                    authState.keys,
+                    getAppStateSyncKey,
                 )
 
                 const node: BinaryNode = {
@@ -456,10 +465,10 @@ export const makeChatsSocket = (config: SocketConfig) => {
                 }
                 await query(node)
         
-                await authState.keys.setAppStateSyncVersion(name, state)
+                await authState.keys.set({ 'app-state-sync-version': { [name]: state } })
                 
                 if(config.emitOwnEvents) {
-                    const result = await decodePatches(name, [{ ...patch, version: { version: state.version }, }], initial, authState.keys.getAppStateSyncKey)
+                    const result = await decodePatches(name, [{ ...patch, version: { version: state.version }, }], initial, getAppStateSyncKey)
                     processSyncActions(result.newMutations)
                 }
             }

src/Socket/messages-recv.ts

@@ -1,12 +1,11 @@
 
 import { SocketConfig, WAMessageStubType, ParticipantAction, Chat, GroupMetadata, WAMessageKey } from "../Types"
 import { decodeMessageStanza, encodeBigEndian, toNumber, downloadHistory, generateSignalPubKey, xmppPreKey, xmppSignedPreKey } from "../Utils"
-import { BinaryNode, jidDecode, jidEncode, isJidStatusBroadcast, areJidsSameUser, getBinaryNodeChildren, jidNormalizedUser, getAllBinaryNodeChildren, BinaryNodeAttributes } from '../WABinary'
+import { BinaryNode, jidDecode, jidEncode, isJidStatusBroadcast, areJidsSameUser, getBinaryNodeChildren, jidNormalizedUser, getAllBinaryNodeChildren, BinaryNodeAttributes, isJidGroup } from '../WABinary'
 import { proto } from "../../WAProto"
 import { KEY_BUNDLE_TYPE } from "../Defaults"
 import { makeChatsSocket } from "./chats"
 import { extractGroupMetadata } from "./groups"
-import { Boom } from "@hapi/boom"
 
 const getStatusFromReceiptType = (type: string | undefined) => {
     if(type === 'read' || type === 'read-self') {
@@ -25,7 +24,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 		ev,
         authState,
 		ws,
-        assertSession,
+        assertSessions,
         assertingPreKeys,
 		sendNode,
         relayMessage,
@@ -146,12 +145,12 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
                     if(keys?.length) {
                         let newAppStateSyncKeyId = ''
                         for(const { keyData, keyId } of keys) {
-                            const str = Buffer.from(keyId.keyId!).toString('base64')
+                            const strKeyId = Buffer.from(keyId.keyId!).toString('base64')
                             
-                            logger.info({ str }, 'injecting new app state sync key')
-                            await authState.keys.setAppStateSyncKey(str, keyData)
+                            logger.info({ strKeyId }, 'injecting new app state sync key')
+                            await authState.keys.set({ 'app-state-sync-key': { [strKeyId]: keyData } })
     
-                            newAppStateSyncKeyId = str
+                            newAppStateSyncKeyId = strKeyId
                         }
                         
                         ev.emit('creds.update', { myAppStateKeyId: newAppStateSyncKeyId })
@@ -473,7 +472,11 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
         )
 
         const participant = key.participant || key.remoteJid
-        await assertSession(participant, true)
+        await assertSessions([participant], true)
+
+        if(isJidGroup(key.remoteJid)) {
+            await authState.keys.set({ 'sender-key-memory': { [key.remoteJid]: null } })
+        }
 
         logger.debug({ participant }, 'forced new session for retry recp')
 

src/Socket/messages-send.ts

@@ -2,7 +2,7 @@
 import got from "got"
 import { Boom } from "@hapi/boom"
 import { SocketConfig, MediaConnInfo, AnyMessageContent, MiscMessageGenerationOptions, WAMediaUploadFunction, MessageRelayOptions } from "../Types"
-import { encodeWAMessage, generateMessageID, generateWAMessage, encryptSenderKeyMsgSignalProto, encryptSignalProto, extractDeviceJids, jidToSignalProtocolAddress, parseAndInjectE2ESession } from "../Utils"
+import { encodeWAMessage, generateMessageID, generateWAMessage, encryptSenderKeyMsgSignalProto, encryptSignalProto, extractDeviceJids, jidToSignalProtocolAddress, parseAndInjectE2ESessions } from "../Utils"
 import { BinaryNode, getBinaryNodeChild, getBinaryNodeChildren, isJidGroup, jidDecode, jidEncode, jidNormalizedUser, S_WHATSAPP_NET, BinaryNodeAttributes, JidWithDevice, reduceBinaryNodeToDictionary } from '../WABinary'
 import { proto } from "../../WAProto"
 import { WA_DEFAULT_EPHEMERAL, DEFAULT_ORIGIN, MEDIA_PATH_MAP } from "../Defaults"
@@ -189,15 +189,23 @@ export const makeMessagesSocket = (config: SocketConfig) => {
         return deviceResults
     }
 
-    const assertSession = async(jid: string, force: boolean) => {
-        const addr = jidToSignalProtocolAddress(jid).toString()
-        const session = await authState.keys.getSession(addr)
-        if(!session || force) {
-            logger.debug({ jid }, `fetching session`)
-            const identity: BinaryNode = {
-                tag: 'user',
-                attrs: { jid, reason: 'identity' },
+    const assertSessions = async(jids: string[], force: boolean) => {
+        let jidsRequiringFetch: string[] = []
+        if(force) {
+            jidsRequiringFetch = jids 
+        } else {
+            const addrs = jids.map(jid => jidToSignalProtocolAddress(jid).toString())
+            const sessions = await authState.keys.get('session', addrs)
+            for(const jid of jids) {
+                const signalId = jidToSignalProtocolAddress(jid).toString()
+                if(!sessions[signalId]) {
+                    jidsRequiringFetch.push(jid)
+                }
             }
+        }
+
+        if(jidsRequiringFetch.length) {
+            logger.debug({ jidsRequiringFetch }, `fetching sessions`)
             const result = await query({
                 tag: 'iq',
                 attrs: {
@@ -209,30 +217,41 @@ export const makeMessagesSocket = (config: SocketConfig) => {
                     {
                         tag: 'key',
                         attrs: { },
-                        content: [ identity ]
+                        content: jidsRequiringFetch.map(
+                            jid => ({
+                                tag: 'user',
+                                attrs: { jid, reason: 'identity' },
+                            })
+                        )
                     }
                 ]
             })
-            await parseAndInjectE2ESession(result, authState)
+            await parseAndInjectE2ESessions(result, authState)
             return true
         }
         return false
     }
 
-    const createParticipantNode = async(jid: string, bytes: Buffer) => {
-        await assertSession(jid, false)
-
-        const { type, ciphertext } = await encryptSignalProto(jid, bytes, authState)
-        const node: BinaryNode = {
-            tag: 'to',
-            attrs: { jid },
-            content: [{
-                tag: 'enc',
-                attrs: { v: '2', type },
-                content: ciphertext
-            }]
-        }
-        return node
+    const createParticipantNodes = async(jids: string[], bytes: Buffer) => {
+        await assertSessions(jids, false)
+        const nodes = await Promise.all(
+            jids.map(
+                async jid => {
+                    const { type, ciphertext } = await encryptSignalProto(jid, bytes, authState)
+                    const node: BinaryNode = {
+                        tag: 'to',
+                        attrs: { jid },
+                        content: [{
+                            tag: 'enc',
+                            attrs: { v: '2', type },
+                            content: ciphertext
+                        }]
+                    }
+                    return node
+                }
+            )
+        )
+        return nodes
     }
 
     const relayMessage = async(
@@ -248,10 +267,11 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
         const encodedMsg = encodeWAMessage(message)
         const participants: BinaryNode[] = []
-        let stanza: BinaryNode
 
         const destinationJid = jidEncode(user, isGroup ? 'g.us' : 's.whatsapp.net')
 
+        const binaryNodeContent: BinaryNode[] = []
+
         const devices: JidWithDevice[] = []
         if(participant) {
             const { user, device } = jidDecode(participant)
@@ -261,8 +281,17 @@ export const makeMessagesSocket = (config: SocketConfig) => {
         if(isGroup) {
             const { ciphertext, senderKeyDistributionMessageKey } = await encryptSenderKeyMsgSignalProto(destinationJid, encodedMsg, meId, authState)
             
-            let groupData = cachedGroupMetadata ? await cachedGroupMetadata(jid) : undefined 
-            if(!groupData) groupData = await groupMetadata(jid)
+            const [groupData, senderKeyMap] = await Promise.all([
+                (async() => {
+                    let groupData = cachedGroupMetadata ? await cachedGroupMetadata(jid) : undefined 
+                    if(!groupData) groupData = await groupMetadata(jid)
+                    return groupData
+                })(),
+                (async() => {
+                    const result = await authState.keys.get('sender-key-memory', [jid])
+                    return result[jid] || { }
+                })()
+            ])
 
             if(!participant) {
                 const participantsList = groupData.participants.map(p => p.id)
@@ -270,31 +299,31 @@ export const makeMessagesSocket = (config: SocketConfig) => {
                 devices.push(...additionalDevices)
             }
 
-            const encSenderKeyMsg = encodeWAMessage({
-                senderKeyDistributionMessage: {
-                    axolotlSenderKeyDistributionMessage: senderKeyDistributionMessageKey,
-                    groupId: destinationJid
-                }
-            })
-
+            const senderKeyJids: string[] = []
+            // ensure a connection is established with every device
             for(const {user, device} of devices) {
                 const jid = jidEncode(user, 's.whatsapp.net', device)
-                const participant = await createParticipantNode(jid, encSenderKeyMsg)
-                participants.push(participant)
+                if(!senderKeyMap[jid]) {
+                    senderKeyJids.push(jid)
+                    // store that this person has had the sender keys sent to them
+                    senderKeyMap[jid] = true
+                }
             }
+            // if there are some participants with whom the session has not been established
+            // if there are, we re-send the senderkey
+            if(senderKeyJids.length) {
+                logger.debug({ senderKeyJids }, 'sending new sender key')
 
-            const binaryNodeContent: BinaryNode[] = []
-            if( // if there are some participants with whom the session has not been established
-                // if there are, we overwrite the senderkey
-                !!participants.find((p) => (
-                    !!(p.content as BinaryNode[]).find(({ attrs }) => attrs.type == 'pkmsg')
-                ))
-            ) {
-                binaryNodeContent.push({
-                    tag: 'participants',
-                    attrs: { },
-                    content: participants
+                const encSenderKeyMsg = encodeWAMessage({
+                    senderKeyDistributionMessage: {
+                        axolotlSenderKeyDistributionMessage: senderKeyDistributionMessageKey,
+                        groupId: destinationJid
+                    }
                 })
+
+                participants.push(
+                    ...(await createParticipantNodes(senderKeyJids, encSenderKeyMsg))
+                )
             }
 
             binaryNodeContent.push({
@@ -303,25 +332,16 @@ export const makeMessagesSocket = (config: SocketConfig) => {
                 content: ciphertext
             })
 
-            stanza = {
-                tag: 'message',
-                attrs: {
-                    id: msgId,
-                    type: 'text',
-                    to: destinationJid
-                },
-                content: binaryNodeContent
-            }
+            await authState.keys.set({ 'sender-key-memory': { [jid]: senderKeyMap } })
         } else {
             const { user: meUser } = jidDecode(meId)
             
-            const messageToMyself: proto.IMessage = {
+            const encodedMeMsg = encodeWAMessage({
                 deviceSentMessage: {
                     destinationJid,
                     message
                 }
-            }
-            const encodedMeMsg = encodeWAMessage(messageToMyself)
+            })
 
             if(!participant) {
                 devices.push({ user })
@@ -331,47 +351,57 @@ export const makeMessagesSocket = (config: SocketConfig) => {
                 devices.push(...additionalDevices)
             }
 
+            const meJids: string[] = []
+            const otherJids: string[] = []
             for(const { user, device } of devices) {
+                const jid = jidEncode(user, 's.whatsapp.net', device)
                 const isMe = user === meUser
-                participants.push(
-                    await createParticipantNode(
-                        jidEncode(user, 's.whatsapp.net', device),
-                        isMe ? encodedMeMsg : encodedMsg
-                    )
-                )
+                if(isMe) meJids.push(jid)
+                else otherJids.push(jid)
             }
 
-            stanza = {
-                tag: 'message',
-                attrs: {
-                    id: msgId,
-                    type: 'text',
-                    to: destinationJid,
-                    ...(additionalAttributes || {})
-                },
-                content: [
-                    {
-                        tag: 'participants',
-                        attrs: { },
-                        content: participants
-                    },
-                ]
-            }
+            const [meNodes, otherNodes] = await Promise.all([
+                createParticipantNodes(meJids, encodedMeMsg),
+                createParticipantNodes(otherJids, encodedMsg)
+            ])
+            participants.push(...meNodes)
+            participants.push(...otherNodes)
         }
 
-        const shouldHaveIdentity = !!participants.find((p) => (
-            !!(p.content as BinaryNode[]).find(({ attrs }) => attrs.type == 'pkmsg')
-        ))
+        if(participants.length) {
+            binaryNodeContent.push({
+                tag: 'participants',
+                attrs: { },
+                content: participants
+            })
+        }
+
+        const stanza: BinaryNode = {
+            tag: 'message',
+            attrs: {
+                id: msgId,
+                type: 'text',
+                to: destinationJid,
+                ...(additionalAttributes || {})
+            },
+            content: binaryNodeContent
+        }
         
+        const shouldHaveIdentity = !!participants.find(
+            participant => (participant.content! as BinaryNode[]).find(n => n.attrs.type === 'pkmsg')
+        )
+
         if(shouldHaveIdentity) {
             (stanza.content as BinaryNode[]).push({
                 tag: 'device-identity',
                 attrs: { },
                 content: proto.ADVSignedDeviceIdentity.encode(authState.creds.account).finish()
             })
+
+            logger.debug({ jid }, 'adding device identity')
         }
 
-        logger.debug({ msgId }, `sending message to ${devices.length} devices`)
+        logger.debug({ msgId }, `sending message to ${participants.length} devices`)
 
         await sendNode(stanza)
 
@@ -427,7 +457,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 	return {
 		...sock,
-        assertSession,
+        assertSessions,
         relayMessage,
         sendDeliveryReceipt,
         sendReadReceipt,

src/Socket/socket.ts

@@ -196,10 +196,8 @@ export const makeSocket = ({
         if(!creds.serverHasPreKeys) {
             update.serverHasPreKeys = true
         }
-       
-        await Promise.all(
-            Object.keys(newPreKeys).map(k => authState.keys.setPreKey(+k, newPreKeys[+k]))
-        )
+
+        await authState.keys.set({ 'pre-key': newPreKeys })
 
         const preKeys = await getPreKeys(authState.keys, preKeysRange[0], preKeysRange[0] + preKeysRange[1])
         await execute(preKeys)
@@ -449,7 +447,7 @@ export const makeSocket = ({
         const genPairQR = () => {
             const ref = refs.shift()
             if(!ref) {
-                end(new Boom('QR refs attempts ended', { statusCode: DisconnectReason.restartRequired }))
+                end(new Boom('QR refs attempts ended', { statusCode: DisconnectReason.timedOut }))
                 return
             }