fix: detect identity key changes and reset sessions (align with WA Web) (#2307)

* feat(signal): add RetryReason enum and MAC error-based session recreation * feat(signal): add identity change detection with automatic session clearing * fix(signal): integrate identity change detection with pkmsg decryption This completes the identity change detection implementation by actually calling saveIdentity() during pkmsg decryption, which is CRITICAL for the feature to work. Changes: - Add extractIdentityFromPkmsg() function that parses PreKeyWhisperMessage protobuf to extract sender's identity key (33 bytes) - Call saveIdentity() BEFORE decryption in decryptMessage() for pkmsg type - Log when identity change is detected Flow: 1. Receive pkmsg from sender 2. Extract identity key from PreKeyWhisperMessage protobuf 3. Call storage.saveIdentity() which compares with stored key 4. If key changed → session is cleared atomically 5. Decryption proceeds with re-established session This matches WhatsApp Web's behavior where extractIdentityKey is called before handleNewSession (GysEGRAXCvh.js:40917, 48815). Ref: WhatsApp Web's extractIdentityKey (GysEGRAXCvh.js:48976-48998)
2026-02-05 16:31:26 -03:00
parent b5c174111f
commit b02390123a
3 changed files with 151 additions and 5 deletions
@@ -1,5 +1,7 @@
-/* @ts-ignore */
+// @ts-ignore
 import * as libsignal from 'libsignal'
+// @ts-ignore
+import { PreKeyWhisperMessage } from 'libsignal/src/protobufs'
 import { LRUCache } from 'lru-cache'
 import type { LIDMapping, SignalAuthState, SignalKeyStoreWithTransaction } from '../Types'
 import type { SignalRepositoryWithLIDStore } from '../Types/Signal'
@@ -20,6 +22,31 @@ import { SenderKeyRecord } from './Group/sender-key-record'
 import { GroupCipher, GroupSessionBuilder, SenderKeyDistributionMessage } from './Group'
 import { LIDMappingStore } from './lid-mapping'

+/** Extract identity key from PreKeyWhisperMessage for identity change detection */
+function extractIdentityFromPkmsg(ciphertext: Uint8Array): Uint8Array | undefined {
+	try {
+		if (!ciphertext || ciphertext.length < 2) {
+			return undefined
+		}
+
+		// Version byte check (version 3)
+		const version = ciphertext[0]!
+		if ((version & 0xf) !== 3) {
+			return undefined
+		}
+
+		// Parse protobuf (skip version byte)
+		const preKeyProto = PreKeyWhisperMessage.decode(ciphertext.slice(1))
+		if (preKeyProto.identityKey?.length === 33) {
+			return new Uint8Array(preKeyProto.identityKey)
+		}
+
+		return undefined
+	} catch {
+		return undefined
+	}
+}
+
 export function makeLibSignalRepository(
 	auth: SignalAuthState,
 	logger: ILogger,
@@ -79,6 +106,18 @@ export function makeLibSignalRepository(
 			const addr = jidToSignalProtocolAddress(jid)
 			const session = new libsignal.SessionCipher(storage, addr)

+			// Extract and save sender's identity key before decryption for identity change detection
+			if (type === 'pkmsg') {
+				const identityKey = extractIdentityFromPkmsg(ciphertext)
+				if (identityKey) {
+					const addrStr = addr.toString()
+					const identityChanged = await storage.saveIdentity(addrStr, identityKey)
+					if (identityChanged) {
+						logger.info({ jid, addr: addrStr }, 'identity key changed or new contact, session will be re-established')
+					}
+				}
+			}
+
 			async function doDecrypt() {
 				let result: Buffer
 				switch (type) {
@@ -359,7 +398,11 @@ const jidToSignalSenderKeyName = (group: string, user: string): SenderKeyName =>
 function signalStorage(
 	{ creds, keys }: SignalAuthState,
 	lidMapping: LIDMappingStore
-): SenderKeyStore & libsignal.SignalStorage {
+): SenderKeyStore &
+	libsignal.SignalStorage & {
+		loadIdentityKey(id: string): Promise<Uint8Array | undefined>
+		saveIdentity(id: string, identityKey: Uint8Array): Promise<boolean>
+	} {
 	// Shared function to resolve PN signal address to LID if mapping exists
 	const resolveLIDSignalAddress = async (id: string): Promise<string> => {
 		if (id.includes('.')) {
@@ -401,7 +444,38 @@ function signalStorage(
 			await keys.set({ session: { [wireJid]: session.serialize() } })
 		},
 		isTrustedIdentity: () => {
-			return true // todo: implement
+			return true // TOFU - Trust on First Use (same as WhatsApp Web)
+		},
+		loadIdentityKey: async (id: string) => {
+			const wireJid = await resolveLIDSignalAddress(id)
+			const { [wireJid]: key } = await keys.get('identity-key', [wireJid])
+			return key || undefined
+		},
+		saveIdentity: async (id: string, identityKey: Uint8Array): Promise<boolean> => {
+			const wireJid = await resolveLIDSignalAddress(id)
+			const { [wireJid]: existingKey } = await keys.get('identity-key', [wireJid])
+
+			const keysMatch =
+				existingKey &&
+				existingKey.length === identityKey.length &&
+				existingKey.every((byte, i) => byte === identityKey[i])
+
+			if (existingKey && !keysMatch) {
+				// Identity changed - clear session and update key
+				await keys.set({
+					session: { [wireJid]: null },
+					'identity-key': { [wireJid]: identityKey }
+				})
+				return true
+			}
+
+			if (!existingKey) {
+				// New contact - Trust on First Use (TOFU)
+				await keys.set({ 'identity-key': { [wireJid]: identityKey } })
+				return true
+			}
+
+			return false
 		},
 		loadPreKey: async (id: number | string) => {
 			const keyId = id.toString()
@@ -81,6 +81,7 @@ export type SignalDataTypeMap = {
 	'lid-mapping': string
 	'device-list': string[]
 	tctoken: { token: Buffer; timestamp?: string }
+	'identity-key': Uint8Array
 }

 export type SignalDataSet = { [T in keyof SignalDataTypeMap]?: { [id: string]: SignalDataTypeMap[T] | null } }
@@ -39,6 +39,29 @@ export interface RetryStatistics {
 	phoneRequests: number
 }

+// Retry reason codes matching WhatsApp Web's Signal error codes.
+export enum RetryReason {
+	UnknownError = 0,
+	SignalErrorNoSession = 1,
+	SignalErrorInvalidKey = 2,
+	SignalErrorInvalidKeyId = 3,
+	/** MAC verification failed - most common cause of decryption failures */
+	SignalErrorInvalidMessage = 4,
+	SignalErrorInvalidSignature = 5,
+	SignalErrorFutureMessage = 6,
+	/** Explicit MAC failure - session is definitely out of sync */
+	SignalErrorBadMac = 7,
+	SignalErrorInvalidSession = 8,
+	SignalErrorInvalidMsgKey = 9,
+	BadBroadcastEphemeralSetting = 10,
+	UnknownCompanionNoPrekey = 11,
+	AdvFailure = 12,
+	StatusRevokeDelay = 13
+}
+
+/** Error codes that indicate a MAC failure and require immediate session recreation */
+const MAC_ERROR_CODES = new Set([RetryReason.SignalErrorInvalidMessage, RetryReason.SignalErrorBadMac])
+
 export class MessageRetryManager {
 	private recentMessagesMap = new LRUCache<string, RecentMessage>({
 		max: RECENT_MESSAGES_SIZE,
@@ -107,9 +130,14 @@ export class MessageRetryManager {
 	}

 	/**
-	 * Check if a session should be recreated based on retry count and history
+	 * Check if a session should be recreated based on retry count, history, and error code.
+	 * MAC errors (codes 4 and 7) trigger immediate session recreation regardless of timeout.
 	 */
-	shouldRecreateSession(jid: string, hasSession: boolean): { reason: string; recreate: boolean } {
+	shouldRecreateSession(
+		jid: string,
+		hasSession: boolean,
+		errorCode?: RetryReason
+	): { reason: string; recreate: boolean } {
 		// If we don't have a session, always recreate
 		if (!hasSession) {
 			this.sessionRecreateHistory.set(jid, Date.now())
@@ -120,6 +148,20 @@ export class MessageRetryManager {
 			}
 		}

+		// IMMEDIATE recreation for MAC errors - session is definitely out of sync
+		if (errorCode !== undefined && MAC_ERROR_CODES.has(errorCode)) {
+			this.sessionRecreateHistory.set(jid, Date.now())
+			this.statistics.sessionRecreations++
+			this.logger.warn(
+				{ jid, errorCode: RetryReason[errorCode] },
+				'MAC error detected, forcing immediate session recreation'
+			)
+			return {
+				reason: `MAC error (code ${errorCode}: ${RetryReason[errorCode]}), immediate session recreation`,
+				recreate: true
+			}
+		}
+
 		const now = Date.now()
 		const prevTime = this.sessionRecreateHistory.get(jid)

@@ -136,6 +178,35 @@ export class MessageRetryManager {
 		return { reason: '', recreate: false }
 	}

+	/**
+	 * Parse error code from retry receipt's retry node.
+	 * Returns undefined if no error code is present.
+	 */
+	parseRetryErrorCode(errorAttr: string | undefined): RetryReason | undefined {
+		if (errorAttr === undefined || errorAttr === '') {
+			return undefined
+		}
+
+		const code = parseInt(errorAttr, 10)
+		if (Number.isNaN(code)) {
+			return undefined
+		}
+
+		// Validate it's a known RetryReason
+		if (code >= RetryReason.UnknownError && code <= RetryReason.StatusRevokeDelay) {
+			return code as RetryReason
+		}
+
+		return RetryReason.UnknownError
+	}
+
+	/**
+	 * Check if an error code indicates a MAC failure
+	 */
+	isMacError(errorCode: RetryReason | undefined): boolean {
+		return errorCode !== undefined && MAC_ERROR_CODES.has(errorCode)
+	}
+
 	/**
 	 * Increment retry counter for a message
 	 */