fix(signal): align Signal Protocol handling with WABA Android behavior (#257)

7 corrections based on reverse-engineering of WhatsApp Business Android:

1. resolveCanonicalJid: PN→LID resolution for transaction lock keys (prevent race conditions)
2. Retain PN session during LID migration (avoid No Session errors)
3. Delayed PreKey deletion with 5-min grace period (prevent Invalid PreKey ID races)
4. Surgical session cleanup: only delete the specific corrupted device, not all devices
5. Identity dual storage: save identity key in both LID and PN addresses
6. MAC error cooldown reduced 10s→1s (faster recovery, aligned with WABA)
7. Allow session recreation on first retry (retryCount >= 1 instead of > 1)

* fix: remove unused jidDecode import from decode-wa-message.ts
* fix: add try/catch to delayed PreKey deletion to prevent unhandled rejection

The setTimeout async callback in removePreKey could cause unhandled promise
rejection if the keystore was destroyed (connection closed) before the
5-min grace period expired.

WAProto/WAProto.proto

@@ -1,7 +1,7 @@
 syntax = "proto3";
 package proto;
 
-/// WhatsApp Version: 2.3000.1034293476
+/// WhatsApp Version: 2.3000.1034427372
 
 message ADVDeviceIdentity {
     optional uint32 rawId = 1;

src/Defaults/baileys-version.json

@@ -1 +1 @@
-{"version":[2,3000,1034300341]}
+{"version": [2, 3000, 1034427372]}

src/Signal/libsignal.ts

@@ -302,6 +302,25 @@ export function makeLibSignalRepository(
 	// Promise instead of each spawning their own DB transactions.
 	const migrationInFlight = new Map<string, Promise<{ migrated: number; skipped: number; total: number }>>()
 
+	// Resolve PN JID to its canonical LID JID for transaction locking.
+	// This prevents PN/LID race conditions where concurrent operations for the
+	// same logical contact acquire different mutex locks because one uses PN
+	// and the other uses LID. (Aligned with WABA behavior — all operations use LID internally.)
+	const resolveCanonicalJid = async(jid: string): Promise<string> => {
+		if (isAnyLidUser(jid)) {
+			return jid
+		}
+
+		if (isAnyPnUser(jid)) {
+			const lid = await lidMapping.getLIDForPN(jid)
+			if (lid) {
+				return lid
+			}
+		}
+
+		return jid
+	}
+
 	const repository: SignalRepositoryWithLIDStore = {
 		decryptGroupMessage({ group, authorJid, msg }) {
 			const senderName = jidToSignalSenderKeyName(group, authorJid)
@@ -396,23 +415,24 @@ export function makeLibSignalRepository(
 				return result
 			}
 
-			// If it's not a sync message, we need to ensure atomicity
+			// Use canonical JID (PN→LID resolved) as transaction key to prevent
-			// For regular messages, we use a transaction to ensure atomicity
+			// PN/LID race conditions on the same logical session.
+			const canonicalJid = await resolveCanonicalJid(jid)
 			return parsedKeys.transaction(async () => {
 				return await doDecrypt()
-			}, jid)
+			}, canonicalJid)
 		},
 
 		async encryptMessage({ jid, data }) {
 			const addr = jidToSignalProtocolAddress(jid)
 			const cipher = new libsignal.SessionCipher(storage, addr)
 
-			// Use transaction to ensure atomicity
+			const canonicalJid = await resolveCanonicalJid(jid)
 			return parsedKeys.transaction(async () => {
 				const { type: sigType, body } = await cipher.encrypt(data)
 				const type = sigType === 3 ? 'pkmsg' : 'msg'
 				return { type, ciphertext: Buffer.from(body, 'binary') }
-			}, jid)
+			}, canonicalJid)
 		},
 
 		async encryptGroupMessage({ group, meId, data }) {
@@ -654,9 +674,10 @@ export function makeLibSignalRepository(
 								// Session exists (guaranteed from device discovery)
 								const fromSession = libsignal.SessionRecord.deserialize(pnSession)
 								if (fromSession.haveOpenSession()) {
-									// Queue for bulk update: copy to LID, delete from PN
+									// Queue for bulk update: copy to LID, retain PN session.
+									// WABA retains both PN and LID sessions during migration to avoid
+									// No Session errors if messages arrive via PN before migration completes.
 									sessionUpdates[lidAddrStr] = fromSession.serialize()
-									sessionUpdates[pnAddrStr] = null
 
 									migratedCount++
 								}
@@ -776,6 +797,13 @@ function signalStorage(
 		return id
 	}
 
+	// Delayed PreKey deletion: grace period to handle race conditions
+	// where two pkmsg with the same preKeyId arrive nearly simultaneously.
+	// WABA deletes immediately (33ms), but we add a 5-min grace period
+	// because we can't handle "Invalid PreKey ID" errors at the native level.
+	const PREKEY_GRACE_PERIOD_MS = 5 * 60 * 1000 // 5 minutes
+	const pendingPreKeyDeletions = new Map<string, ReturnType<typeof setTimeout>>()
+
 	return {
 		loadSession: async (id: string) => {
 			try {
@@ -808,7 +836,26 @@ function signalStorage(
 				}
 			}
 		},
-		removePreKey: (id: number) => keys.set({ 'pre-key': { [id]: null } }),
+		removePreKey: (id: number) => {
+			const keyId = id.toString()
+			// Clear any existing timer for this key
+			const existing = pendingPreKeyDeletions.get(keyId)
+			if (existing) {
+				clearTimeout(existing)
+			}
+
+			// Schedule deletion after grace period
+			const timer = setTimeout(async () => {
+				pendingPreKeyDeletions.delete(keyId)
+				try {
+					await keys.set({ 'pre-key': { [id]: null } })
+				} catch {
+					// Keystore may be destroyed if connection closed — safe to ignore
+				}
+			}, PREKEY_GRACE_PERIOD_MS)
+
+			pendingPreKeyDeletions.set(keyId, timer)
+		},
 		loadSignedPreKey: () => {
 			const key = creds.signedPreKey
 			return {
@@ -898,14 +945,24 @@ function signalStorage(
 					// IDENTITY KEY CHANGED - contact reinstalled WhatsApp or switched devices
 					const previousFingerprint = generateKeyFingerprint(existingKey)
 
-					// Delete old session and save new identity key atomically
+					// Delete old session and save new identity key atomically.
+					// Store identity in BOTH LID and PN addresses (WABA stores in both
+					// recipient_account_type=0 and type=1 with CONFLICT_REPLACE).
+					const identityUpdates: Record<string, Uint8Array> = { [wireJid]: identityKey }
+					if (wireJid !== id) {
+						identityUpdates[id] = identityKey
+					}
+
 					await keys.set({
 						session: { [wireJid]: null },
-						'identity-key': { [wireJid]: identityKey }
+						'identity-key': identityUpdates
 					})
 
-					// Update cache
+					// Update cache for both addresses
 					identityKeyCache.set(wireJid, identityKey)
+					if (wireJid !== id) {
+						identityKeyCache.set(id, identityKey)
+					}
 
 					// Record metrics
 					metrics.signalIdentityChanges?.inc({ type: 'changed' })
@@ -941,10 +998,19 @@ function signalStorage(
 
 				if (!existingKey) {
 					// NEW CONTACT - Trust On First Use (TOFU)
-					await keys.set({ 'identity-key': { [wireJid]: identityKey } })
+					// Store in both LID and PN addresses (aligned with WABA dual identity storage)
+					const identityUpdates: Record<string, Uint8Array> = { [wireJid]: identityKey }
+					if (wireJid !== id) {
+						identityUpdates[id] = identityKey
+					}
 
-					// Update cache
+					await keys.set({ 'identity-key': identityUpdates })
+
+					// Update cache for both addresses
 					identityKeyCache.set(wireJid, identityKey)
+					if (wireJid !== id) {
+						identityKeyCache.set(id, identityKey)
+					}
 
 					// Record metrics
 					metrics.signalIdentityChanges?.inc({ type: 'new' })

src/Socket/messages-recv.ts

@@ -69,7 +69,7 @@ import {
 	recordMessageReceived,
 	recordMessageRetry
 } from '../Utils/prometheus-metrics.js'
-import { isTcTokenExpired, resolveTcTokenJid } from '../Utils/tc-token-utils'
+import { isTcTokenExpired, resolveTcTokenJid, storeTcTokensFromIqResult } from '../Utils/tc-token-utils'
 import {
 	areJidsSameUser,
 	type BinaryNode,
@@ -1314,7 +1314,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 		let shouldRecreateSession = false
 		let recreateReason = ''
 
-		if (enableAutoSessionRecreation && messageRetryManager && retryCount > 1) {
+		if (enableAutoSessionRecreation && messageRetryManager && retryCount >= 1) {
 			try {
 				// Check if we have a session with this JID
 				const sessionId = signalRepository.jidToSignalProtocolAddress(fromJid)
@@ -1453,6 +1453,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 			})
 
 			// When a session is refreshed (identity change), re-issue tctoken fire-and-forget
+			// WABA Android: reissue stores senderTimestamp + realIssueTimestamp after IQ success
 			if (result.action === 'session_refreshed') {
 				const normalizedJid = jidNormalizedUser(from)
 				resolveTcTokenJid(normalizedJid, getLIDForPN)
@@ -1462,9 +1463,36 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 						if (entry?.token?.length && !isTcTokenExpired(entry.timestamp)) {
 							const senderTs = unixTimestampSeconds()
 							logTcToken('reissue', { jid: normalizedJid, reason: 'session_refreshed' })
-							getPrivacyTokens([normalizedJid], senderTs).catch(err => {
+							getPrivacyTokens([normalizedJid], senderTs)
-								logTcToken('reissue_fail', { jid: normalizedJid, error: err?.message })
+								.then(async (iqResult) => {
-							})
+									await storeTcTokensFromIqResult({
+										result: iqResult,
+										fallbackJid: normalizedJid,
+										keys: authState.keys,
+										getLIDForPN,
+										onNewJidStored: (storedJid) => {
+											tcTokenKnownJids.add(storedJid)
+											scheduleTcTokenIndexSave()
+										}
+									})
+									// Persist senderTimestamp + realIssueTimestamp after IQ success
+									const currentData = await authState.keys.get('tctoken', [tcJid])
+									const currentEntry = currentData[tcJid]
+									await authState.keys.set({
+										tctoken: {
+											[tcJid]: {
+												...currentEntry,
+												token: currentEntry?.token ?? Buffer.alloc(0),
+												senderTimestamp: senderTs,
+												realIssueTimestamp: 0
+											}
+										}
+									})
+									logTcToken('reissue_ok', { jid: normalizedJid, reason: 'session_refreshed' })
+								})
+								.catch(err => {
+									logTcToken('reissue_fail', { jid: normalizedJid, error: err?.message })
+								})
 						}
 					})
 					.catch(() => {
@@ -1912,7 +1940,10 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 						[storageJid]: {
 							...existing,
 							token: Buffer.from(content),
-							timestamp
+							timestamp,
+							// WABA Android: resets real_issue_timestamp when a new incoming token arrives
+							// (UPDATE wa_trusted_contacts_send SET real_issue_timestamp=null)
+							realIssueTimestamp: null
 						}
 					}
 				})
@@ -2002,7 +2033,7 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 		let shouldRecreateSession = false
 		let recreateReason = ''
 
-		if (enableAutoSessionRecreation && messageRetryManager && retryCount > 1) {
+		if (enableAutoSessionRecreation && messageRetryManager && retryCount >= 1) {
 			try {
 				const sessionId = signalRepository.jidToSignalProtocolAddress(participant)
 
@@ -2823,6 +2854,24 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 				const jid = jidNormalizedUser(attrs.from)
 				logTcToken('error_463', { jid, msgId })
 
+				// WABA Android: error 463 triggers getPrivacyTokens() fire-and-forget
+				// to ensure token is available for the retry below
+				getPrivacyTokens([jid])
+					.then(async (result) => {
+						await storeTcTokensFromIqResult({
+							result,
+							fallbackJid: jid,
+							keys: authState.keys,
+							getLIDForPN,
+							onNewJidStored: (storedJid) => {
+								tcTokenKnownJids.add(storedJid)
+								scheduleTcTokenIndexSave()
+							}
+						})
+						logTcToken('fetched', { jid, reason: 'error_463' })
+					})
+					.catch(() => { /* fire-and-forget */ })
+
 				// Single-retry: wait 1.5s for the server's tctoken notification to arrive,
 				// then resend. A Set prevents infinite retry loops.
 				// Composite key (jid:msgId) ensures retries are isolated per destination.
@@ -2858,7 +2907,24 @@ export const makeMessagesRecvSocket = (config: SocketConfig) => {
 					return
 				}
 			} else if (attrs.error === SERVER_ERROR_CODES.SmaxInvalid) {
-				logTcToken('error_479', { jid: attrs.from, msgId: attrs.id })
+				const jid479 = jidNormalizedUser(attrs.from)
+				logTcToken('error_479', { jid: jid479, msgId: attrs.id })
+				// WABA Android: error 479 (SmaxInvalid) also triggers token re-fetch
+				getPrivacyTokens([jid479])
+					.then(async (result) => {
+						await storeTcTokensFromIqResult({
+							result,
+							fallbackJid: jid479,
+							keys: authState.keys,
+							getLIDForPN,
+							onNewJidStored: (storedJid) => {
+								tcTokenKnownJids.add(storedJid)
+								scheduleTcTokenIndexSave()
+							}
+						})
+						logTcToken('fetched', { jid: jid479, reason: 'error_479' })
+					})
+					.catch(() => { /* fire-and-forget */ })
 			} else {
 				logger.warn({ attrs }, 'received error in ack')
 			}

src/Socket/messages-send.ts

@@ -1611,6 +1611,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 
 						// Persist senderTimestamp unconditionally — WA Web stores it in the chat table
 						// regardless of whether a token exists. Spread preserves token+timestamp if present.
+						// WABA Android: INSERT INTO wa_trusted_contacts_send (jid, sent_tc_token_timestamp, real_issue_timestamp)
+						// VALUES (?, ?, 0) — realIssueTimestamp=0 means issued but not yet confirmed by server
 						const currentData = await authState.keys.get('tctoken', [tcTokenJid])
 						const currentEntry = currentData[tcTokenJid]
 						await authState.keys.set({
@@ -1618,7 +1620,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
 								[tcTokenJid]: {
 									...currentEntry,
 									token: currentEntry?.token ?? Buffer.alloc(0),
-									senderTimestamp: issueTimestamp
+									senderTimestamp: issueTimestamp,
+									realIssueTimestamp: 0
 								}
 							}
 						})

src/Types/Auth.ts

@@ -80,7 +80,7 @@ export type SignalDataTypeMap = {
 	'app-state-sync-version': LTHashState
 	'lid-mapping': string
 	'device-list': string[]
-	tctoken: { token: Buffer; timestamp?: string; senderTimestamp?: number }
+	tctoken: { token: Buffer; timestamp?: string; senderTimestamp?: number; realIssueTimestamp?: number | null }
 	/** Identity key for Signal Protocol - used for detecting contact reinstalls */
 	'identity-key': Uint8Array
 }

src/Utils/decode-wa-message.ts

@@ -14,7 +14,6 @@ import {
 	isJidStatusBroadcast,
 	isLidUser,
 	isPnUser,
-	jidDecode
 	//	transferDevice
 } from '../WABinary'
 import { unpadRandomMax16 } from './generics'
@@ -488,8 +487,9 @@ export function isCorruptedSessionError(error: any): boolean {
 }
 
 /**
- * Clean up corrupted session by deleting all device sessions for a JID.
+ * Clean up corrupted session for a specific device JID.
- * Signal Protocol will automatically recreate the session on next message.
+ * WABA behavior: DELETE sessions WHERE recipient_id=? AND device_id=?
+ * Only deletes the exact device that was corrupted, not all devices.
  *
  * NOTE: This should NOT be called on every Bad MAC error (hot path).
  * Instead, let the retry+pkmsg flow handle recovery naturally (like WhatsApp does).
@@ -500,45 +500,7 @@ export async function cleanupCorruptedSession(
 	repository: SignalRepositoryWithLIDStore,
 	logger: ILogger
 ): Promise<number> {
-	const { user, device } = jidDecode(jid) || {}
+	await repository.deleteSession([jid])
-	if (!user) {
+	logger.info({ jid }, 'Cleaned up corrupted session for specific device')
-		logger.warn({ jid }, 'Cannot cleanup session - invalid JID')
+	return 1
-		return 0
-	}
-
-	// Build list of JIDs to delete (primary + secondary devices)
-	const jidsToDelete: string[] = []
-
-	// Determine domain type correctly (handle hosted JIDs)
-	// JID formats:
-	//   - PN: user@s.whatsapp.net
-	//   - LID: user@lid
-	//   - Hosted PN: user@hosted
-	//   - Hosted LID: user@hosted.lid
-	const isLID = jid.endsWith('@lid') || jid.endsWith('@hosted.lid')
-	const isHosted = jid.includes('@hosted')
-
-	let domain: string
-	if (isLID) {
-		domain = isHosted ? 'hosted.lid' : 'lid'
-	} else {
-		domain = isHosted ? 'hosted' : 's.whatsapp.net'
-	}
-
-	// Primary device (0)
-	jidsToDelete.push(`${user}@${domain}`)
-
-	// Secondary devices (1-5 common range for Web/Desktop/etc)
-	for (let i = 1; i <= 5; i++) {
-		jidsToDelete.push(`${user}:${i}@${domain}`)
-	}
-
-	// If specific device was identified and > 5, ensure it's included
-	if (device !== undefined && device > 5) {
-		jidsToDelete.push(`${user}:${device}@${domain}`)
-	}
-
-	await repository.deleteSession(jidsToDelete)
-
-	return jidsToDelete.length
 }

src/Utils/message-retry-manager.ts

@@ -218,7 +218,7 @@ export class MessageRetryManager {
 		if (errorCode !== undefined && MAC_ERROR_CODES.has(errorCode)) {
 			const now = Date.now()
 			const prevTime = this.sessionRecreateHistory.get(jid)
-			const MAC_ERROR_COOLDOWN_MS = 10_000 // 10 seconds
+			const MAC_ERROR_COOLDOWN_MS = 1_000 // 1 second — WABA recovers faster
 
 			if (prevTime && now - prevTime < MAC_ERROR_COOLDOWN_MS) {
 				const reasonName = RetryReason[errorCode] || `code_${errorCode}`

src/Utils/tc-token-utils.ts

@@ -159,7 +159,10 @@ export async function storeTcTokensFromIqResult({
 				[storageJid]: {
 					...existingEntry,
 					token: Buffer.from(tokenNode.content),
-					timestamp: tokenNode.attrs.t
+					timestamp: tokenNode.attrs.t,
+					// WABA Android: resets real_issue_timestamp to null when storing a new token
+					// (UPDATE wa_trusted_contacts_send SET real_issue_timestamp=null)
+					realIssueTimestamp: null
 				}
 			}
 		})