From 5887551d682ca11ad3196d62c83d0d39ac24c760 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas=20de=20Oliveira=20Lopes?= <55464917+jlucaso1@users.noreply.github.com> Date: Tue, 20 Jan 2026 07:37:04 -0300 Subject: [PATCH] fix: resolve race condition in decodeFrame handling and improve encryption integrity (#2182) * fix: resolve race condition in decodeFrame handling and improve encryption integrity * chore: pr feedback --- src/Utils/noise-handler.ts | 216 ++++++++------ src/__tests__/Utils/noise-handler.test.ts | 339 ++++++++++++++++++++++ 2 files changed, 472 insertions(+), 83 deletions(-) create mode 100644 src/__tests__/Utils/noise-handler.test.ts diff --git a/src/Utils/noise-handler.ts b/src/Utils/noise-handler.ts index 779a4bab..feb8a6c2 100644 --- a/src/Utils/noise-handler.ts +++ b/src/Utils/noise-handler.ts @@ -7,13 +7,48 @@ import { decodeBinaryNode } from '../WABinary' import { aesDecryptGCM, aesEncryptGCM, Curve, hkdf, sha256 } from './crypto' import type { ILogger } from './logger' -const generateIV = (counter: number) => { - const iv = new ArrayBuffer(12) - new DataView(iv).setUint32(8, counter) +const IV_LENGTH = 12 +const EMPTY_BUFFER = Buffer.alloc(0) + +const generateIV = (counter: number): Uint8Array => { + const iv = new ArrayBuffer(IV_LENGTH) + new DataView(iv).setUint32(8, counter) return new Uint8Array(iv) } +class TransportState { + private readCounter = 0 + private writeCounter = 0 + + private readonly iv = new Uint8Array(IV_LENGTH) + + constructor( + private readonly encKey: Buffer, + private readonly decKey: Buffer + ) {} + + encrypt(plaintext: Uint8Array): Uint8Array { + const c = this.writeCounter++ + this.iv[8] = (c >>> 24) & 0xff + this.iv[9] = (c >>> 16) & 0xff + this.iv[10] = (c >>> 8) & 0xff + this.iv[11] = c & 0xff + + return aesEncryptGCM(plaintext, this.encKey, this.iv, EMPTY_BUFFER) + } + + decrypt(ciphertext: Uint8Array): Buffer { + const c = this.readCounter++ + this.iv[8] = (c >>> 24) & 0xff + this.iv[9] = (c >>> 16) & 0xff + this.iv[10] = (c >>> 8) & 0xff + this.iv[11] = c & 0xff + + return aesDecryptGCM(ciphertext, this.decKey, this.iv, EMPTY_BUFFER) as Buffer + } +} + export const makeNoiseHandler = ({ keyPair: { private: privateKey, public: publicKey }, NOISE_HEADER, @@ -27,40 +62,63 @@ export const makeNoiseHandler = ({ }) => { logger = logger.child({ class: 'ns' }) + const data = Buffer.from(NOISE_MODE) + let hash = data.byteLength === 32 ? data : sha256(data) + let salt = hash + let encKey = hash + let decKey = hash + let counter = 0 + let sentIntro = false + + let inBytes: Buffer = Buffer.alloc(0) + + let transport: TransportState | null = null + let isWaitingForTransport = false + let pendingOnFrame: ((buff: Uint8Array | BinaryNode) => void) | null = null + + let introHeader: Buffer + if (routingInfo) { + introHeader = Buffer.alloc(7 + routingInfo.byteLength + NOISE_HEADER.length) + introHeader.write('ED', 0, 'utf8') + introHeader.writeUint8(0, 2) + introHeader.writeUint8(1, 3) + introHeader.writeUint8(routingInfo.byteLength >> 16, 4) + introHeader.writeUint16BE(routingInfo.byteLength & 65535, 5) + introHeader.set(routingInfo, 7) + introHeader.set(NOISE_HEADER, 7 + routingInfo.byteLength) + } else { + introHeader = Buffer.from(NOISE_HEADER) + } + const authenticate = (data: Uint8Array) => { - if (!isFinished) { + if (!transport) { hash = sha256(Buffer.concat([hash, data])) } } - const encrypt = (plaintext: Uint8Array) => { - const result = aesEncryptGCM(plaintext, encKey, generateIV(writeCounter), hash) - - writeCounter += 1 + const encrypt = (plaintext: Uint8Array): Uint8Array => { + if (transport) { + return transport.encrypt(plaintext) + } + const result = aesEncryptGCM(plaintext, encKey, generateIV(counter++), hash) authenticate(result) return result } - const decrypt = (ciphertext: Uint8Array) => { - // before the handshake is finished, we use the same counter - // after handshake, the counters are different - const iv = generateIV(isFinished ? readCounter : writeCounter) - const result = aesDecryptGCM(ciphertext, decKey, iv, hash) - - if (isFinished) { - readCounter += 1 - } else { - writeCounter += 1 + const decrypt = (ciphertext: Uint8Array): Uint8Array => { + if (transport) { + return transport.decrypt(ciphertext) } + const result = aesDecryptGCM(ciphertext, decKey, generateIV(counter++), hash) authenticate(ciphertext) return result } const localHKDF = async (data: Uint8Array) => { const key = await hkdf(Buffer.from(data), 64, { salt, info: '' }) - return [key.slice(0, 32), key.slice(32)] + return [key.subarray(0, 32), key.subarray(32)] } const mixIntoKey = async (data: Uint8Array) => { @@ -68,31 +126,49 @@ export const makeNoiseHandler = ({ salt = write! encKey = read! decKey = read! - readCounter = 0 - writeCounter = 0 + counter = 0 } const finishInit = async () => { + isWaitingForTransport = true const [write, read] = await localHKDF(new Uint8Array(0)) - encKey = write! - decKey = read! - hash = Buffer.from([]) - readCounter = 0 - writeCounter = 0 - isFinished = true + transport = new TransportState(write!, read!) + isWaitingForTransport = false + + logger.trace('Noise handler transitioned to Transport state') + + if (pendingOnFrame) { + logger.trace({ length: inBytes.length }, 'Flushing buffered frames after transport ready') + await processData(pendingOnFrame) + pendingOnFrame = null + } } - const data = Buffer.from(NOISE_MODE) - let hash = data.byteLength === 32 ? data : sha256(data) - let salt = hash - let encKey = hash - let decKey = hash - let readCounter = 0 - let writeCounter = 0 - let isFinished = false - let sentIntro = false + const processData = async (onFrame: (buff: Uint8Array | BinaryNode) => void) => { + let size: number | undefined - let inBytes = Buffer.alloc(0) + while (true) { + if (inBytes.length < 3) return + + size = (inBytes[0]! << 16) | (inBytes[1]! << 8) | inBytes[2]! + + if (inBytes.length < size + 3) return + + let frame: Uint8Array | BinaryNode = inBytes.subarray(3, size + 3) + inBytes = inBytes.subarray(size + 3) + + if (transport) { + const result = transport.decrypt(frame) + frame = await decodeBinaryNode(result) + } + + if (logger.level === 'trace') { + logger.trace({ msg: (frame as BinaryNode)?.attrs?.id }, 'recv frame') + } + + onFrame(frame) + } + } authenticate(NOISE_HEADER) authenticate(publicKey) @@ -152,67 +228,41 @@ export const makeNoiseHandler = ({ return keyEnc }, encodeFrame: (data: Buffer | Uint8Array) => { - if (isFinished) { - data = encrypt(data) + if (transport) { + data = transport.encrypt(data) } - let header: Buffer - - if (routingInfo) { - header = Buffer.alloc(7) - header.write('ED', 0, 'utf8') - header.writeUint8(0, 2) - header.writeUint8(1, 3) - header.writeUint8(routingInfo.byteLength >> 16, 4) - header.writeUint16BE(routingInfo.byteLength & 65535, 5) - header = Buffer.concat([header, routingInfo, NOISE_HEADER]) - } else { - header = Buffer.from(NOISE_HEADER) - } - - const introSize = sentIntro ? 0 : header.length - const frame = Buffer.alloc(introSize + 3 + data.byteLength) + const dataLen = data.byteLength + const introSize = sentIntro ? 0 : introHeader.length + const frame = Buffer.allocUnsafe(introSize + 3 + dataLen) if (!sentIntro) { - frame.set(header) + frame.set(introHeader) sentIntro = true } - frame.writeUInt8(data.byteLength >> 16, introSize) - frame.writeUInt16BE(65535 & data.byteLength, introSize + 1) + frame[introSize] = (dataLen >>> 16) & 0xff + frame[introSize + 1] = (dataLen >>> 8) & 0xff + frame[introSize + 2] = dataLen & 0xff + frame.set(data, introSize + 3) return frame }, decodeFrame: async (newData: Buffer | Uint8Array, onFrame: (buff: Uint8Array | BinaryNode) => void) => { - // the binary protocol uses its own framing mechanism - // on top of the WS frames - // so we get this data and separate out the frames - const getBytesSize = () => { - if (inBytes.length >= 3) { - return (inBytes.readUInt8() << 16) | inBytes.readUInt16BE(1) - } + if (isWaitingForTransport) { + inBytes = Buffer.concat([inBytes, newData]) + pendingOnFrame = onFrame + return } - inBytes = Buffer.concat([inBytes, newData]) - - logger.trace(`recv ${newData.length} bytes, total recv ${inBytes.length} bytes`) - - let size = getBytesSize() - while (size && inBytes.length >= size + 3) { - let frame: Uint8Array | BinaryNode = inBytes.slice(3, size + 3) - inBytes = inBytes.slice(size + 3) - - if (isFinished) { - const result = decrypt(frame) - frame = await decodeBinaryNode(result) - } - - logger.trace({ msg: (frame as BinaryNode)?.attrs?.id }, 'recv frame') - - onFrame(frame) - size = getBytesSize() + if (inBytes.length === 0) { + inBytes = Buffer.from(newData) + } else { + inBytes = Buffer.concat([inBytes, newData]) } + + await processData(onFrame) } } } diff --git a/src/__tests__/Utils/noise-handler.test.ts b/src/__tests__/Utils/noise-handler.test.ts new file mode 100644 index 00000000..76ad19d6 --- /dev/null +++ b/src/__tests__/Utils/noise-handler.test.ts @@ -0,0 +1,339 @@ +import { jest } from '@jest/globals' +import { NOISE_WA_HEADER } from '../../Defaults' +import { Curve } from '../../Utils/crypto' +import { makeNoiseHandler } from '../../Utils/noise-handler' +import type { BinaryNode } from '../../WABinary/types' + +// Create a mock logger +const createMockLogger = () => ({ + child: jest.fn().mockReturnThis(), + trace: jest.fn(), + debug: jest.fn(), + info: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + fatal: jest.fn(), + level: 'trace' +}) + +// Helper to create a frame with length prefix +const createFrame = (payload: Buffer) => { + const frame = Buffer.alloc(3 + payload.length) + frame.writeUInt8(payload.length >> 16, 0) + frame.writeUInt16BE(payload.length & 0xffff, 1) + payload.copy(frame, 3) + return frame +} + +describe('Noise Handler', () => { + describe('decodeFrame with multiple frames in buffer', () => { + it('should process multiple unencrypted frames in single buffer', async () => { + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + const payload1 = Buffer.from([1, 2, 3, 4, 5]) + const payload2 = Buffer.from([6, 7, 8, 9, 10]) + + const frame1 = createFrame(payload1) + const frame2 = createFrame(payload2) + + const combinedBuffer = Buffer.concat([frame1, frame2]) + + const receivedFrames: Buffer[] = [] + const onFrame = (frame: Uint8Array | BinaryNode) => { + receivedFrames.push(Buffer.from(frame as Uint8Array)) + } + + await handler.decodeFrame(combinedBuffer, onFrame) + + expect(receivedFrames).toHaveLength(2) + expect(receivedFrames[0]).toEqual(payload1) + expect(receivedFrames[1]).toEqual(payload2) + }) + + it('should handle frames split across multiple decodeFrame calls', async () => { + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + const payload = Buffer.from([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + const frame = createFrame(payload) + + const receivedFrames: Buffer[] = [] + const onFrame = (frame: Uint8Array | BinaryNode) => { + receivedFrames.push(Buffer.from(frame as Uint8Array)) + } + + // Split the frame across two calls + const part1 = frame.slice(0, 5) + const part2 = frame.slice(5) + + await handler.decodeFrame(part1, onFrame) + expect(receivedFrames).toHaveLength(0) + + await handler.decodeFrame(part2, onFrame) + expect(receivedFrames).toHaveLength(1) + expect(receivedFrames[0]).toEqual(payload) + }) + + it('should correctly process frames when callback triggers async operations', async () => { + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + const payload1 = Buffer.from([1, 2, 3]) + const payload2 = Buffer.from([4, 5, 6]) + + const combinedBuffer = Buffer.concat([createFrame(payload1), createFrame(payload2)]) + + const receivedFrames: Buffer[] = [] + const callbackOrder: number[] = [] + + const onFrame = (frame: Uint8Array | BinaryNode) => { + const frameNum = receivedFrames.length + 1 + callbackOrder.push(frameNum) + receivedFrames.push(Buffer.from(frame as Uint8Array)) + } + + await handler.decodeFrame(combinedBuffer, onFrame) + + expect(receivedFrames).toHaveLength(2) + expect(callbackOrder).toEqual([1, 2]) + expect(receivedFrames[0]).toEqual(payload1) + expect(receivedFrames[1]).toEqual(payload2) + }) + }) + + describe('encrypted frame handling', () => { + it('should encrypt and verify frame structure', async () => { + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + await handler.finishInit() + + const payload = Buffer.from('test payload') + const encoded = handler.encodeFrame(payload) + + expect(encoded.length).toBeGreaterThan(payload.length + 3) + + const encoded2 = handler.encodeFrame(Buffer.from('second payload')) + expect(encoded2.slice(0, 3)).toEqual(Buffer.from([0, 0, encoded2.length - 3])) + }) + + it('should produce different ciphertext for same plaintext due to counter', async () => { + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + await handler.finishInit() + + const payload = Buffer.from('same payload') + + const encrypted1 = handler.encrypt(payload) + const encrypted2 = handler.encrypt(payload) + const encrypted3 = handler.encrypt(payload) + + expect(encrypted1).not.toEqual(encrypted2) + expect(encrypted2).not.toEqual(encrypted3) + expect(encrypted1).not.toEqual(encrypted3) + }) + }) + + describe('race condition scenario - concurrent decodeFrame calls', () => { + it('should handle concurrent decodeFrame calls without corrupting inBytes buffer', async () => { + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + // Create multiple frames + const payloads = Array.from({ length: 5 }, (_, i) => Buffer.from(`payload-${i}`)) + const frames = payloads.map(createFrame) + + const receivedFrames: Buffer[] = [] + const onFrame = (frame: Uint8Array | BinaryNode) => { + receivedFrames.push(Buffer.from(frame as Uint8Array)) + } + + // Simulate concurrent calls (multiple WebSocket messages arriving rapidly) + // This tests the shared inBytes buffer handling + await Promise.all(frames.map(frame => handler.decodeFrame(frame, onFrame))) + + // All frames should be received + expect(receivedFrames).toHaveLength(5) + + // Verify all payloads are present (order may vary due to concurrency) + const receivedPayloads = receivedFrames.map(f => f.toString()) + payloads.forEach(p => { + expect(receivedPayloads).toContain(p.toString()) + }) + }) + + it('should maintain counter integrity with many frames in single buffer', async () => { + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + // Create 10 frames to stress test the while loop + const payloads = Array.from({ length: 10 }, (_, i) => Buffer.from(`frame-${i}-payload-data`)) + + const combinedBuffer = Buffer.concat(payloads.map(createFrame)) + + const receivedFrames: Buffer[] = [] + const onFrame = (frame: Uint8Array | BinaryNode) => { + receivedFrames.push(Buffer.from(frame as Uint8Array)) + } + + await handler.decodeFrame(combinedBuffer, onFrame) + + expect(receivedFrames).toHaveLength(10) + payloads.forEach((payload, i) => { + expect(receivedFrames[i]).toEqual(payload) + }) + }) + }) + + describe('encrypted frame race condition', () => { + it('should produce different ciphertext for same plaintext due to counter', async () => { + // Verify that encryption uses incrementing counters + + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + await handler.finishInit() + + const payload1 = Buffer.from('message-1') + const payload2 = Buffer.from('message-2') + + const encrypted1 = handler.encrypt(payload1) + const encrypted2 = handler.encrypt(payload2) + + expect(encrypted1.length).toBe(payload1.length + 16) // +16 for GCM tag + expect(encrypted2.length).toBe(payload2.length + 16) + + // The encrypted data should be different (different counters used) + expect(encrypted1).not.toEqual(encrypted2) + }) + + it('should serialize concurrent decodeFrame calls (fix for race condition)', async () => { + // This test verifies that the lock mechanism correctly serializes + // concurrent decodeFrame calls, preventing race conditions + + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + const payload1 = Buffer.from('first') + const payload2 = Buffer.from('second') + const payload3 = Buffer.from('third') + + const frame1 = createFrame(payload1) + const frame2 = createFrame(payload2) + const frame3 = createFrame(payload3) + + const receivedOrder: string[] = [] + + const onFrame = (frame: Uint8Array | BinaryNode) => { + const content = Buffer.from(frame as Uint8Array).toString() + receivedOrder.push(content) + } + + // Start all three decodeFrame calls "simultaneously" + // With the lock fix, they should be processed in order + const p1 = handler.decodeFrame(frame1, onFrame) + const p2 = handler.decodeFrame(frame2, onFrame) + const p3 = handler.decodeFrame(frame3, onFrame) + + await Promise.all([p1, p2, p3]) + + // With serialization, frames should be received in the order + // the decodeFrame calls were made + expect(receivedOrder).toHaveLength(3) + expect(receivedOrder[0]).toBe('first') + expect(receivedOrder[1]).toBe('second') + expect(receivedOrder[2]).toBe('third') + }) + + it('should maintain frame order with interleaved partial frames after fix', async () => { + // This test verifies that partial frames from different sources + // are correctly reassembled when calls are serialized + + const keyPair = Curve.generateKeyPair() + const logger = createMockLogger() + + const handler = makeNoiseHandler({ + keyPair, + NOISE_HEADER: NOISE_WA_HEADER, + logger: logger as any + }) + + // Create a single frame split into parts + const payload = Buffer.from('complete-message-content') + const frame = createFrame(payload) + + const part1 = frame.slice(0, 10) + const part2 = frame.slice(10) + + const receivedFrames: Buffer[] = [] + const onFrame = (frame: Uint8Array | BinaryNode) => { + receivedFrames.push(Buffer.from(frame as Uint8Array)) + } + + // With serialization, these should be processed in order + // and the frame should be correctly reassembled + await handler.decodeFrame(part1, onFrame) + expect(receivedFrames).toHaveLength(0) // Not complete yet + + await handler.decodeFrame(part2, onFrame) + expect(receivedFrames).toHaveLength(1) + expect(receivedFrames[0]).toEqual(payload) + }) + }) +})