From 8f21a67e4ad3c027a9c83ef64417bf63462e308e Mon Sep 17 00:00:00 2001 From: Gustavo Quadri <87215048+gusquadri@users.noreply.github.com> Date: Sun, 7 Sep 2025 08:11:50 -0300 Subject: [PATCH] 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 --- src/Defaults/index.ts | 3 + src/Socket/messages-recv.ts | 48 +++++++++++++-- src/Socket/socket.ts | 116 ++++++++++++++++++++++++++++++++---- 3 files changed, 148 insertions(+), 19 deletions(-) diff --git a/src/Defaults/index.ts b/src/Defaults/index.ts index 4f7885aa..ee68a18d 100644 --- a/src/Defaults/index.ts +++ b/src/Defaults/index.ts @@ -111,6 +111,9 @@ export const MIN_PREKEY_COUNT = 5 export const INITIAL_PREKEY_COUNT = 30 +export const UPLOAD_TIMEOUT = 30000 // 30 seconds +export const MIN_UPLOAD_INTERVAL = 5000 // 5 seconds minimum between uploads + export const DEFAULT_CACHE_TTLS = { SIGNAL_STORE: 5 * 60, // 5 minutes MSG_RETRY: 60 * 60, // 1 hour diff --git a/src/Socket/messages-recv.ts b/src/Socket/messages-recv.ts index e2a90d2f..c39fb870 100644 --- a/src/Socket/messages-recv.ts +++ b/src/Socket/messages-recv.ts @@ -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 { diff --git a/src/Socket/socket.ts b/src/Socket/socket.ts index 80fe5858..4c238a66 100644 --- a/src/Socket/socket.ts +++ b/src/Socket/socket.ts @@ -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 | 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((_, 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 } }