From 9ed4c28b8ebd4651541fc084ce052318efef76ed Mon Sep 17 00:00:00 2001
From: Adhiraj Singh <adhirajsingh1001@gmail.com>
Date: Mon, 15 Nov 2021 09:09:55 +0530
Subject: [PATCH] fix: minor retry logic cleanup

---
 src/Socket/messages-recv.ts | 21 +++++++++++----------
 src/Types/index.ts          |  2 ++
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts
index 49dcfe5..3040fcf 100644
--- a/src/Socket/messages-recv.ts
+++ b/src/Socket/messages-recv.ts
@@ -1,5 +1,5 @@
 
-import { SocketConfig, WAMessageStubType, ParticipantAction, Chat, GroupMetadata } from "../Types"
+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 { proto } from "../../WAProto"
@@ -22,6 +22,8 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
         sendDeliveryReceipt,
 	} = sock
 
+    const msgRetryMap = config.msgRetryCounterMap || { }
+
     const sendMessageAck = async({ tag, attrs }: BinaryNode, extraAttrs: BinaryNodeAttributes) => {
         const stanza: BinaryNode = {
             tag: 'ack',
@@ -38,14 +40,15 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
         await sendNode(stanza)
     }
 
-    const retries = new Map<string, number>()
     const sendRetryRequest = async(node: BinaryNode) => {
-        if (retries.has(node.attrs.id) && retries.get(node.attrs.id)! >= 5) {
-            retries.delete(node.attrs.id)
+        const msgId = node.attrs.id
+        const retryCount = msgRetryMap[msgId] || 1
+        if(retryCount >= 5) {
+            logger.debug({ retryCount, msgId }, 'reached retry limit, clearing')
+            delete msgRetryMap[msgId]
             return
         }
-        const retryCount = retries.get(node.attrs.id) || 1
-        retries.set(node.attrs.id, retryCount + 1)
+        msgRetryMap[msgId] = retryCount+1
 
         const isGroup = !!node.attrs.participant
         const { account, signedPreKey, signedIdentityKey: identityKey } = authState.creds
@@ -59,7 +62,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
             const receipt: BinaryNode = {
                 tag: 'receipt',
                 attrs: {
-                    id: node.attrs.id,
+                    id: msgId,
                     type: 'retry',
                     to: isGroup ? node.attrs.from : jidEncode(decFrom!.user, 's.whatsapp.net', decFrom!.device, 0)
                 },
@@ -102,9 +105,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
             }
             await sendNode(receipt)
 
-            logger.info({ msgId: node.attrs.id, retryCount }, 'sent retry receipt')
-
-            ev.emit('auth-state.update', authState)
+            logger.info({ msgAttrs: node.attrs, retryCount }, 'sent retry receipt')
         })
     }
     const processMessage = async(message: proto.IWebMessageInfo, chatUpdate: Partial<Chat>) => {
diff --git a/src/Types/index.ts b/src/Types/index.ts
index a860819..a8e3f77 100644
--- a/src/Types/index.ts
+++ b/src/Types/index.ts
@@ -50,6 +50,8 @@ export type SocketConfig = {
     userDevicesCache?: NodeCache
     /** provide a cache to store media, so does not have to be re-uploaded */
     mediaCache?: NodeCache
+    /** map to store the retry counts for failed messages */
+    msgRetryCounterMap?: { [msgId: string]: number }
 }
 
 export enum DisconnectReason {