From 7863c0e4c2baa1c19aaf7ca71ba67f05a7299b19 Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Fri, 17 Jun 2022 13:16:41 +0530
Subject: [PATCH] perf: nicer "shouldIncludeDeviceIdentity" flag computation

---
 src/Socket/messages-send.ts | 34 ++++++++++++++++++++++------------
 src/Utils/signal.ts         |  8 +++-----
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/src/Socket/messages-send.ts b/src/Socket/messages-send.ts
index 524ecb0..e357e84 100644
--- a/src/Socket/messages-send.ts
+++ b/src/Socket/messages-send.ts
@@ -266,11 +266,16 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 		return didFetchNewSession
 	}
 
-	const createParticipantNodes = (jids: string[], bytes: Buffer) => (
-		Promise.all(
+	const createParticipantNodes = async(jids: string[], bytes: Buffer) => {
+		let shouldIncludeDeviceIdentity = false
+		const nodes = await Promise.all(
 			jids.map(
 				async jid => {
 					const { type, ciphertext } = await encryptSignalProto(jid, bytes, authState)
+					if(type === 'pkmsg') {
+						shouldIncludeDeviceIdentity = true
+					}
+
 					const node: BinaryNode = {
 						tag: 'to',
 						attrs: { jid },
@@ -284,7 +289,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 				}
 			)
 		)
-	)
+		return { nodes, shouldIncludeDeviceIdentity }
+	}
 
 	const relayMessage = async(
 		jid: string,
@@ -293,6 +299,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 	) => {
 		const meId = authState.creds.me!.id
 
+		let shouldIncludeDeviceIdentity = false
+
 		const { user, server } = jidDecode(jid)
 		const isGroup = server === 'g.us'
 		msgId = msgId || generateMessageID()
@@ -376,9 +384,10 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 						await assertSessions(senderKeyJids, false)
 
-						participants.push(
-							...(await createParticipantNodes(senderKeyJids, encSenderKeyMsg))
-						)
+						const result = await createParticipantNodes(senderKeyJids, encSenderKeyMsg)
+						shouldIncludeDeviceIdentity = shouldIncludeDeviceIdentity || result.shouldIncludeDeviceIdentity
+
+						participants.push(...result.nodes)
 					}
 
 					binaryNodeContent.push({
@@ -423,12 +432,17 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 					await assertSessions(allJids, false)
 
-					const [meNodes, otherNodes] = await Promise.all([
+					const [
+						{ nodes: meNodes, shouldIncludeDeviceIdentity: s1 },
+						{ nodes: otherNodes, shouldIncludeDeviceIdentity: s2 }
+					] = await Promise.all([
 						createParticipantNodes(meJids, encodedMeMsg),
 						createParticipantNodes(otherJids, encodedMsg)
 					])
 					participants.push(...meNodes)
 					participants.push(...otherNodes)
+
+					shouldIncludeDeviceIdentity = shouldIncludeDeviceIdentity || s1 || s2
 				}
 
 				if(participants.length) {
@@ -465,11 +479,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 					stanza.attrs.to = destinationJid
 				}
 
-				const shouldHaveIdentity = !!participants.find(
-					participant => (participant.content! as BinaryNode[]).find(n => n.attrs.type === 'pkmsg')
-				)
-
-				if(shouldHaveIdentity) {
+				if(shouldIncludeDeviceIdentity) {
 					(stanza.content as BinaryNode[]).push({
 						tag: 'device-identity',
 						attrs: { },
diff --git a/src/Utils/signal.ts b/src/Utils/signal.ts
index 10ca261..4a76034 100644
--- a/src/Utils/signal.ts
+++ b/src/Utils/signal.ts
@@ -175,11 +175,9 @@ export const encryptSignalProto = async(user: string, buffer: Buffer, auth: Sign
 	const addr = jidToSignalProtocolAddress(user)
 	const cipher = new libsignal.SessionCipher(signalStorage(auth), addr)
 
-	const { type, body } = await cipher.encrypt(buffer)
-	return {
-		type: type === 3 ? 'pkmsg' : 'msg',
-		ciphertext: Buffer.from(body, 'binary')
-	}
+	const { type: sigType, body } = await cipher.encrypt(buffer)
+	const type = sigType === 3 ? 'pkmsg' : 'msg'
+	return { type, ciphertext: Buffer.from(body, 'binary') }
 }
 
 export const encryptSenderKeyMsgSignalProto = async(group: string, data: Uint8Array | Buffer, meId: string, auth: SignalAuthState) => {