From 355b0664df977df1e0402ea40f0661cd0b9386d7 Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Sat, 2 Apr 2022 13:37:10 +0530
Subject: [PATCH] feat: handle encrypt notification

---
 src/Socket/messages-recv.ts | 37 +++++++++++++++++++++++++++++++++----
 src/Socket/socket.ts        |  7 ++++---
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts
index 55b24c6..1e627c3 100644
--- a/src/Socket/messages-recv.ts
+++ b/src/Socket/messages-recv.ts
@@ -5,10 +5,12 @@ import { BaileysEventMap, MessageUserReceipt, SocketConfig, WAMessageStubType }
 import { debouncedTimeout, decodeMessageStanza, delay, encodeBigEndian, generateSignalPubKey, getStatusFromReceiptType, normalizeMessageContent, xmppPreKey, xmppSignedPreKey } from '../Utils'
 import { makeKeyedMutex, makeMutex } from '../Utils/make-mutex'
 import processMessage from '../Utils/process-message'
-import { areJidsSameUser, BinaryNode, BinaryNodeAttributes, getAllBinaryNodeChildren, getBinaryNodeChildren, isJidGroup, jidDecode, jidEncode, jidNormalizedUser } from '../WABinary'
+import { areJidsSameUser, BinaryNode, BinaryNodeAttributes, getAllBinaryNodeChildren, getBinaryNodeChild, getBinaryNodeChildren, isJidGroup, jidDecode, jidEncode, jidNormalizedUser, S_WHATSAPP_NET } from '../WABinary'
 import { makeChatsSocket } from './chats'
 import { extractGroupMetadata } from './groups'
 
+const MIN_PREKEY_COUNT = 5
+
 export const makeMessagesRecvSocket = (config: SocketConfig) => {
 	const {
 		logger,
@@ -28,6 +30,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 		sendReceipt,
 		resyncMainAppState,
 		emitEventsFromMap,
+		uploadPreKeys
 	} = sock
 
 	/** this mutex ensures that the notifications (receipts, messages etc.) are processed in order */
@@ -167,7 +170,30 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 		return newEvents
 	}
 
-	const processNotification = (node: BinaryNode): Partial<proto.IWebMessageInfo> => {
+	const handleEncryptNotification = async(node: BinaryNode) => {
+		const from = node.attrs.from
+		if(from === S_WHATSAPP_NET) {
+			const countChild = getBinaryNodeChild(node, 'count')
+			const count = +countChild.attrs.value
+			const shouldUploadMorePreKeys = count < MIN_PREKEY_COUNT
+
+			logger.debug({ count, shouldUploadMorePreKeys }, 'recv pre-key count')
+			if(shouldUploadMorePreKeys) {
+				await uploadPreKeys()
+			}
+		} else {
+			const identityNode = getBinaryNodeChild(node, 'identity')
+			if(identityNode) {
+				logger.info({ jid: from }, 'identity changed')
+				// not handling right now
+				// signal will override new identity anyway
+			} else {
+				logger.info({ node }, 'unknown encrypt notification')
+			}
+		}
+	}
+
+	const processNotification = async(node: BinaryNode): Promise<Partial<proto.IWebMessageInfo>> => {
 		const result: Partial<proto.IWebMessageInfo> = { }
 		const [child] = getAllBinaryNodeChildren(node)
 
@@ -242,6 +268,9 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 					logger.info({ deviceJids }, 'got my own devices')
 				}
 
+				break
+			case 'encrypt':
+				handleEncryptNotification(node)
 				break
 			}
 		}
@@ -431,8 +460,8 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 		await sendMessageAck(node, { class: 'notification', type: node.attrs.type })
 		await processingMutex.mutex(
 			remoteJid,
-			() => {
-				const msg = processNotification(node)
+			async() => {
+				const msg = await processNotification(node)
 				if(msg) {
 					const fromMe = areJidsSameUser(node.attrs.participant || node.attrs.from, authState.creds.me!.id)
 					msg.key = {
diff --git a/src/Socket/socket.ts b/src/Socket/socket.ts
index 63e671d..45b3914 100644
--- a/src/Socket/socket.ts
+++ b/src/Socket/socket.ts
@@ -244,9 +244,9 @@ export const makeSocket = ({
 		ev.emit('creds.update', update)
 	}
 
-	/** generates and uploads a set of pre-keys */
-	const uploadPreKeys = async() => {
-		await assertingPreKeys(INITIAL_PREKEY_COUNT, async preKeys => {
+	/** generates and uploads a set of pre-keys to the server */
+	const uploadPreKeys = async(count = INITIAL_PREKEY_COUNT) => {
+		await assertingPreKeys(count, async preKeys => {
 			const node: BinaryNode = {
 				tag: 'iq',
 				attrs: {
@@ -602,6 +602,7 @@ export const makeSocket = ({
 		logout,
 		end,
 		onUnexpectedError,
+		uploadPreKeys,
 		/** Waits for the connection to WA to reach a state */
 		waitForConnectionUpdate: bindWaitForConnectionUpdate(ev)
 	}