feat: cleaner auth state management + store SK keys

!BREAKING_CHANGE

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,