fix: sync pre-keys on socket start and add specific pre-key error approach (#1663)

* fix: sync prekeys when storage mismatch * fix: verify if current pre-key exists instead of matching * add: pre-keys error specific approach * fix: lint * fix: lint
2025-09-07 08:11:50 -03:00
parent 0fbce52fd3
commit 8f21a67e4a
3 changed files with 148 additions and 19 deletions
@@ -711,8 +711,11 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 									updateSendMessageAgainCount(ids[0], key.participant)
 									logger.debug({ attrs, key }, 'recv retry request')
 									await sendMessagesAgain(key, ids, retryNode!)
-								} catch (error: any) {
-									logger.error({ key, ids, trace: error.stack }, 'error in sending message again')
+								} catch (error: unknown) {
+									logger.error(
+										{ key, ids, trace: error instanceof Error ? error.stack : 'Unknown error' },
+										'error in sending message again'
+									)
 								}
 							} else {
 								logger.info({ attrs, key }, 'recv retry for not fromMe message')
@@ -823,19 +826,52 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 							return sendMessageAck(node, NACK_REASONS.ParsingError)
 						}

+						const errorMessage = msg?.messageStubParameters?.[0] || ''
+						const isPreKeyError = errorMessage.includes('PreKey')
+
+						console.debug(`[handleMessage] Attempting retry request for failed decryption`)
+
+						// Handle both pre-key and normal retries in single mutex
 						retryMutex.mutex(async () => {
-							if (ws.isOpen) {
-								if (getBinaryNodeChild(node, 'unavailable')) {
+							try {
+								if (!ws.isOpen) {
+									logger.debug({ node }, 'Connection closed, skipping retry')
 									return
 								}

+								if (getBinaryNodeChild(node, 'unavailable')) {
+									logger.debug('Message unavailable, skipping retry')
+									return
+								}
+
+								// Handle pre-key errors with upload and delay
+								if (isPreKeyError) {
+									logger.info({ error: errorMessage }, 'PreKey error detected, uploading and retrying')
+
+									try {
+										logger.debug('Uploading pre-keys for error recovery')
+										await uploadPreKeys(5)
+										logger.debug('Waiting for server to process new pre-keys')
+										await delay(1000)
+									} catch (uploadErr) {
+										logger.error({ uploadErr }, 'Pre-key upload failed, proceeding with retry anyway')
+									}
+								}
+
 								const encNode = getBinaryNodeChild(node, 'enc')
 								await sendRetryRequest(node, !encNode)
 								if (retryRequestDelayMs) {
 									await delay(retryRequestDelayMs)
 								}
-							} else {
-								logger.debug({ node }, 'connection closed, ignoring retry req')
+							} catch (err) {
+								logger.error({ err, isPreKeyError }, 'Failed to handle retry, attempting basic retry')
+								// Still attempt retry even if pre-key upload failed
+								try {
+									const encNode = getBinaryNodeChild(node, 'enc')
+									await sendRetryRequest(node, !encNode)
+								} catch (retryErr) {
+									logger.error({ retryErr }, 'Failed to send retry after error handling')
+								}
 							}
 						})
 					} else {
@@ -8,7 +8,9 @@ import {
 	DEF_TAG_PREFIX,
 	INITIAL_PREKEY_COUNT,
 	MIN_PREKEY_COUNT,
-	NOISE_WA_HEADER
+	MIN_UPLOAD_INTERVAL,
+	NOISE_WA_HEADER,
+	UPLOAD_TIMEOUT
 } from '../Defaults'
 import { cleanupQueues } from '../Signal/Group/queue-job'
 import type { SocketConfig } from '../Types'
@@ -298,24 +300,112 @@ export const makeSocket = (config: SocketConfig) => {
 		return +countChild.attrs.value!
 	}

+	// Pre-key upload state management
+	let uploadPreKeysPromise: Promise<void> | null = null
+	let lastUploadTime = 0
+
 	/** generates and uploads a set of pre-keys to the server */
-	const uploadPreKeys = async (count = INITIAL_PREKEY_COUNT) => {
-		await keys.transaction(async () => {
-			logger.info({ count }, 'uploading pre-keys')
-			const { update, node } = await getNextPreKeysNode({ creds, keys }, count)
+	const uploadPreKeys = async (count = INITIAL_PREKEY_COUNT, retryCount = 0) => {
+		// Check minimum interval (except for retries)
+		if (retryCount === 0) {
+			const timeSinceLastUpload = Date.now() - lastUploadTime
+			if (timeSinceLastUpload < MIN_UPLOAD_INTERVAL) {
+				logger.debug(`Skipping upload, only ${timeSinceLastUpload}ms since last upload`)
+				return
+			}
+		}

-			await query(node)
-			ev.emit('creds.update', update)
+		// Prevent multiple concurrent uploads
+		if (uploadPreKeysPromise) {
+			logger.debug('Pre-key upload already in progress, waiting for completion')
+			return uploadPreKeysPromise
+		}

-			logger.info({ count }, 'uploaded pre-keys')
-		})
+		const uploadLogic = async () => {
+			logger.info({ count, retryCount }, 'uploading pre-keys')
+
+			// Generate and save pre-keys atomically (prevents ID collisions on retry)
+			const node = await keys.transaction(async () => {
+				logger.debug({ requestedCount: count }, 'generating pre-keys with requested count')
+				const { update, node } = await getNextPreKeysNode({ creds, keys }, count)
+				// Update credentials immediately to prevent duplicate IDs on retry
+				ev.emit('creds.update', update)
+				return node // Only return node since update is already used
+			})
+
+			// Upload to server (outside transaction, can fail without affecting local keys)
+			try {
+				await query(node)
+				logger.info({ count }, 'uploaded pre-keys successfully')
+				lastUploadTime = Date.now()
+			} catch (uploadError) {
+				logger.error({ uploadError, count }, 'Failed to upload pre-keys to server')
+
+				// Exponential backoff retry (max 3 retries)
+				if (retryCount < 3) {
+					const backoffDelay = Math.min(1000 * Math.pow(2, retryCount), 10000)
+					logger.info(`Retrying pre-key upload in ${backoffDelay}ms`)
+					await new Promise(resolve => setTimeout(resolve, backoffDelay))
+					return uploadPreKeys(count, retryCount + 1)
+				}
+
+				throw uploadError
+			}
+		}
+
+		// Add timeout protection
+		uploadPreKeysPromise = Promise.race([
+			uploadLogic(),
+			new Promise<void>((_, reject) =>
+				setTimeout(() => reject(new Boom('Pre-key upload timeout', { statusCode: 408 })), UPLOAD_TIMEOUT)
+			)
+		])
+
+		try {
+			await uploadPreKeysPromise
+		} finally {
+			uploadPreKeysPromise = null
+		}
+	}
+
+	const verifyCurrentPreKeyExists = async () => {
+		const currentPreKeyId = creds.nextPreKeyId - 1
+		if (currentPreKeyId <= 0) {
+			return { exists: false, currentPreKeyId: 0 }
+		}
+
+		const preKeys = await keys.get('pre-key', [currentPreKeyId.toString()])
+		const exists = !!preKeys[currentPreKeyId.toString()]
+
+		return { exists, currentPreKeyId }
 	}

 	const uploadPreKeysToServerIfRequired = async () => {
-		const preKeyCount = await getAvailablePreKeysOnServer()
-		logger.info(`${preKeyCount} pre-keys found on server`)
-		if (preKeyCount <= MIN_PREKEY_COUNT) {
-			await uploadPreKeys()
+		try {
+			const preKeyCount = await getAvailablePreKeysOnServer()
+			const { exists: currentPreKeyExists, currentPreKeyId } = await verifyCurrentPreKeyExists()
+
+			logger.info(`${preKeyCount} pre-keys found on server`)
+			logger.info(`Current prekey ID: ${currentPreKeyId}, exists in storage: ${currentPreKeyExists}`)
+
+			const lowServerCount = preKeyCount <= MIN_PREKEY_COUNT
+			const missingCurrentPreKey = !currentPreKeyExists && currentPreKeyId > 0
+
+			const shouldUpload = lowServerCount || missingCurrentPreKey
+
+			if (shouldUpload) {
+				const reasons = []
+				if (lowServerCount) reasons.push(`server count low (${preKeyCount})`)
+				if (missingCurrentPreKey) reasons.push(`current prekey ${currentPreKeyId} missing from storage`)
+
+				logger.info(`Uploading PreKeys due to: ${reasons.join(', ')}`)
+				await uploadPreKeys()
+			} else {
+				logger.info(`PreKey validation passed - Server: ${preKeyCount}, Current prekey ${currentPreKeyId} exists`)
+			}
+		} catch (error) {
+			logger.error({ error }, 'Failed to check/upload pre-keys during initialization')
+			// Don't throw - allow connection to continue even if pre-key check fails
 		}
 	}