feat: add cleanup handlers to prevent memory leaks on socket close

src/Signal/libsignal.ts

@@ -184,6 +184,11 @@ export function makeLibSignalRepository(
 			}, `delete-${jids.length}-sessions`)
 		},
 
+		close() {
+			migratedSessionCache.clear()
+			lidMapping.close()
+		},
+
 		async migrateSession(
 			fromJid: string,
 			toJid: string

src/Signal/lid-mapping.ts

@@ -206,4 +206,11 @@ export class LIDMappingStore {
 		this.logger.trace(`Found reverse mapping: ${lid} → ${pnJid}`)
 		return pnJid
 	}
+
+	/**
+	 * Close the cache and release resources
+	 */
+	close(): void {
+		this.mappingCache.clear()
+	}
 }

src/Socket/chats.ts

@@ -67,7 +67,17 @@
 		getMessage
 	} = config
 	const sock = makeSocket(config)
-	const { ev, ws, authState, generateMessageTag, sendNode, query, signalRepository, onUnexpectedError } = sock
+	const {
+		ev,
+		ws,
+		authState,
+		generateMessageTag,
+		sendNode,
+		query,
+		signalRepository,
+		onUnexpectedError,
+		registerSocketEndHandler
+	} = sock
 
 	let privacySettings: { [_: string]: string } | undefined
 
@@ -95,10 +105,6 @@
 			useClones: false
 		}) as CacheStore)
 
-	if (!config.placeholderResendCache) {
-		config.placeholderResendCache = placeholderResendCache
-	}
-
 	/** helper function to fetch the given app state sync key */
 	const getAppStateSyncKey = async (keyId: string) => {
 		const { [keyId]: key } = await authState.keys.get('app-state-sync-key', [keyId])
@@ -575,7 +581,7 @@
 								// collection is done with sync
 								collectionsToHandle.delete(name)
 							}
 						} catch (error: any) {
 							// if retry attempts overshoot
 							// or key not found
 							const isIrrecoverableError =
@@ -1186,6 +1192,20 @@
 		}, 20_000)
 	})
 
+	registerSocketEndHandler(() => {
+		if (awaitingSyncTimeout) {
+			clearTimeout(awaitingSyncTimeout)
+			awaitingSyncTimeout = undefined
+		}
+
+		if (!config.placeholderResendCache && placeholderResendCache.close) {
+			placeholderResendCache.close()
+		}
+
+		syncState = SyncState.Connecting
+		privacySettings = undefined
+	})
+
 	return {
 		...sock,
 		createCallLink,

src/Socket/messages-recv.ts

@@ -85,7 +85,8 @@
 		sendReceipt,
 		uploadPreKeys,
 		sendPeerDataOperationMessage,
-		messageRetryManager
+		messageRetryManager,
+		registerSocketEndHandler
 	} = sock
 
 	/** this mutex ensures that each retryRequest will wait for the previous one to finish */
@@ -185,7 +186,7 @@
 			return
 		}
 
 		let data: any
 		try {
 			data = JSON.parse(mexNode.content.toString())
 		} catch (error) {
@@ -281,7 +282,7 @@
 			case 'update':
 				const settingsNode = getBinaryNodeChild(child, 'settings')
 				if (settingsNode) {
 					const update: Record<string, any> = {}
 					const nameNode = getBinaryNodeChild(settingsNode, 'name')
 					if (nameNode?.content) update.name = nameNode.content.toString()
 
@@ -1575,6 +1576,23 @@
 		}
 	})
 
+	registerSocketEndHandler(() => {
+		if (!config.msgRetryCounterCache && msgRetryCache.close) {
+			msgRetryCache.close()
+		}
+
+		if (!config.callOfferCache && callOfferCache.close) {
+			callOfferCache.close()
+		}
+
+		if (!config.placeholderResendCache && placeholderResendCache.close) {
+			placeholderResendCache.close()
+		}
+
+		identityAssertDebounce.close()
+		sendActiveReceipts = false
+	})
+
 	return {
 		...sock,
 		sendMessageAck,

src/Socket/messages-send.ts

@@ -78,7 +78,8 @@
 		fetchPrivacySettings,
 		sendNode,
 		groupMetadata,
-		groupToggleEphemeral
+		groupToggleEphemeral,
+		registerSocketEndHandler
 	} = sock
 
 	const userDevicesCache =
@@ -508,8 +509,8 @@
 		const meLid = authState.creds.me?.lid
 		const meLidUser = meLid ? jidDecode(meLid)?.user : null
 
 		const encryptionPromises = (patchedMessages as any).map(
 			async ({ recipientJid: jid, message: patchedMessage }: any) => {
 				if (!jid) return null
 				let msgToEncrypt = patchedMessage
 				if (dsmMessage) {
@@ -1054,6 +1055,21 @@
 
 	const waitForMsgMediaUpdate = bindWaitForEvent(ev, 'messages.media-update')
 
+	registerSocketEndHandler(() => {
+		if (!config.userDevicesCache && userDevicesCache.close) {
+			userDevicesCache.close()
+		}
+
+		if (peerSessionsCache.close) {
+			peerSessionsCache.close()
+		}
+
+		mediaConn = undefined as any
+		if (messageRetryManager) {
+			messageRetryManager.clear()
+		}
+	})
+
 	return {
 		...sock,
 		getPrivacyTokens,
@@ -1098,7 +1114,7 @@
 								content.url = getUrlFromDirectPath(content.directPath!)
 
 								logger.debug({ directPath: media.directPath, key: result.key }, 'media update successful')
 							} catch (err: any) {
 								error = err
 							}
 						}

src/Socket/socket.ts

@@ -192,7 +192,7 @@
 
 		const msgId = node.attrs.id
 
 		const result = await promiseTimeout<any>(timeoutMs, async (resolve, reject) => {
 			const result = waitForMessage(msgId, timeoutMs).catch(reject)
 			sendNode(node)
 				.then(async () => resolve(await result))
@@ -364,6 +364,8 @@
 	let qrTimer: NodeJS.Timeout
 	let closed = false
 
+	const socketEndHandlers: Array<(error: Error | undefined) => void> = []
+
 	/** log & process any unexpected errors */
 	const onUnexpectedError = (err: Error | Boom, msg: string) => {
 		logger.error({ err }, `unexpected error in '${msg}'`)
@@ -621,12 +623,22 @@
 		ws.removeAllListeners('open')
 		ws.removeAllListeners('message')
 
+		signalRepository.close?.()
+
 		if (!ws.isClosed && !ws.isClosing) {
 			try {
 				ws.close()
 			} catch {}
 		}
 
+		for (const handler of socketEndHandlers) {
+			try {
+				handler(error)
+			} catch (err) {
+				logger.error({ err }, 'error in socket end handler')
+			}
+		}
+
 		ev.emit('connection.update', {
 			connection: 'close',
 			lastDisconnect: {
@@ -635,6 +647,7 @@
 			}
 		})
 		ev.removeAllListeners('connection.update')
+		ev.destroy()
 	}
 
 	const waitForSocketOpen = async () => {
@@ -826,7 +839,7 @@
 	ws.on('open', async () => {
 		try {
 			await validateConnection()
 		} catch (err: any) {
 			logger.error({ err }, 'error in validating connection')
 			end(err)
 		}
@@ -894,7 +907,7 @@
 			ev.emit('connection.update', { isNewLogin: true, qr: undefined })
 
 			await sendNode(reply)
 		} catch (error: any) {
 			logger.info({ trace: error.stack }, 'error in pairing')
 			end(error)
 		}
@@ -1029,6 +1042,10 @@
 		Object.assign(creds, update)
 	})
 
+	const registerSocketEndHandler = (handler: (error: Error | undefined) => void) => {
+		socketEndHandlers.push(handler)
+	}
+
 	return {
 		type: 'md' as 'md',
 		ws,
@@ -1046,6 +1063,7 @@
 		sendNode,
 		logout,
 		end,
+		registerSocketEndHandler,
 		onUnexpectedError,
 		uploadPreKeys,
 		uploadPreKeysToServerIfRequired,

src/Types/Signal.ts

@@ -73,4 +73,5 @@ export type SignalRepository = {
 // Optimized repository with pre-loaded LID mapping store
 export interface SignalRepositoryWithLIDStore extends SignalRepository {
 	lidMapping: LIDMappingStore
+	close?: () => void
 }

src/Types/Socket.ts

@@ -19,6 +19,7 @@ export type CacheStore = {
 	del(key: string): void | Promise<void> | number | boolean
 	/** flush all data */
 	flushAll(): void | Promise<void>
+	close?: () => void
 }
 
 export type PossiblyExtendedCacheStore = CacheStore & {

src/Utils/event-buffer.ts

@@ -61,6 +61,8 @@ type BaileysBufferableEventEmitter = BaileysEventEmitter & {
 	flush(): boolean
 	/** is there an ongoing buffer */
 	isBuffering(): boolean
+	/** destroy the event buffer, clearing all resources */
+	destroy(): void
 }
 
 /**
@@ -225,7 +227,24 @@ export const makeEventBuffer = (logger: ILogger): BaileysBufferableEventEmitter
 		},
 		on: (...args) => ev.on(...args),
 		off: (...args) => ev.off(...args),
-		removeAllListeners: (...args) => ev.removeAllListeners(...args)
+		removeAllListeners: (...args) => ev.removeAllListeners(...args),
+		destroy() {
+			// Clear buffer timeout
+			if (bufferTimeout) {
+				clearTimeout(bufferTimeout)
+				bufferTimeout = null
+			}
+
+			// Clear history cache
+			historyCache.clear()
+			// Reset buffer data
+			data = makeBufferData()
+			isBuffering = false
+			bufferCount = 0
+			// Remove all listeners
+			ev.removeAllListeners()
+			logger.debug('Event buffer destroyed')
+		}
 	}
 }
 

src/Utils/message-retry-manager.ts

@@ -213,6 +213,25 @@ export class MessageRetryManager {
 		}
 	}
 
+	clear(): void {
+		this.recentMessagesMap.clear()
+		this.messageKeyIndex.clear()
+		this.sessionRecreateHistory.clear()
+		this.retryCounters.clear()
+		for (const messageId of Object.keys(this.pendingPhoneRequests)) {
+			this.cancelPendingPhoneRequest(messageId)
+		}
+
+		this.statistics = {
+			totalRetries: 0,
+			successfulRetries: 0,
+			failedRetries: 0,
+			mediaRetries: 0,
+			sessionRecreations: 0,
+			phoneRequests: 0
+		}
+	}
+
 	private keyToString(key: RecentMessageKey): string {
 		return `${key.to}${MESSAGE_KEY_SEPARATOR}${key.id}`
 	}