refactor: make MessageProcessor generator

Author: adiwajshing

package.json

@@ -23,7 +23,7 @@
 		"run:tsc": "node --experimental-strip-types",
 		"test:pure-js": "npm run run:tsc -- --import=./src/tests/load-purejs-crypto.ts --test",
 		"test:webcrypto": "npm run run:tsc -- --import=./src/tests/load-webcrypto.ts --test",
-		"handshake": "npm run run:tsc -- --import=./src/tests/load-purejs-crypto.ts ./src/scripts/handshake.ts",
+		"handshake": "npm run run:tsc -- --import=./src/tests/load-webcrypto.ts ./src/scripts/handshake.ts",
 		"lint": "eslint . --ext .js,.ts,.jsx,.tsx",
 		"lint:fix": "eslint . --fix --ext .js,.ts,.jsx,.tsx",
 		"prepare": "npm run build",

src/make-tls-client.ts

@@ -1,5 +1,5 @@
 import { crypto } from './crypto/index.ts'
-import type { CipherSuite, Key, KeyPair, ProcessPacket, TLSClientOptions, TLSHandshakeOptions, TLSKeyType, TLSPacket, TLSPacketContext, TLSProtocolVersion, TLSSessionTicket, X509Certificate } from './types/index.ts'
+import type { CipherSuite, Key, KeyPair, TLSClientOptions, TLSHandshakeOptions, TLSKeyType, TLSPacket, TLSPacketContext, TLSPacketWithType, TLSProtocolVersion, TLSSessionTicket, X509Certificate } from './types/index.ts'
 import { packClientHello } from './utils/client-hello.ts'
 import { CONTENT_TYPE_MAP, MAX_ENC_PACKET_SIZE, PACKET_TYPE, SUPPORTED_CIPHER_SUITE_MAP, SUPPORTED_NAMED_CURVE_MAP, SUPPORTED_NAMED_CURVES, SUPPORTED_RECORD_TYPE_MAP } from './utils/constants.ts'
 import type { SharedKeyData } from './utils/decryption-utils.ts'
@@ -65,7 +65,9 @@ export function makeTLSClient({
 	let clientCertificateRequested = false
 	let certificatesVerified = false
 
-	const processPacketUnsafe: ProcessPacket = async(type, { header, content }) => {
+	const processPacketUnsafe = async(
+		{ type, packet: { content, header } }: TLSPacketWithType
+	) => {
 		if(ended) {
 			logger.warn('connection closed, ignoring packet')
 			return
@@ -173,8 +175,8 @@ export function makeTLSClient({
 		}
 	}
 
-	const processPacket: ProcessPacket = (...args) => (
-		enqueueServerPacket(processPacketUnsafe, ...args)
+	const processPacket = (pkt: TLSPacketWithType) => (
+		enqueueServerPacket(processPacketUnsafe, pkt)
 	)
 
 	async function processRecord(
@@ -799,8 +801,10 @@ export function makeTLSClient({
 		 * Handle bytes received from the server.
 		 * Could be a complete or partial TLS packet
 		 */
-		handleReceivedBytes(data: Uint8Array) {
-			return processor.onData(data, processPacket)
+		async handleReceivedBytes(data: Uint8Array) {
+			for(const pkt of processor.onData(data)) {
+				await processPacket(pkt)
+			}
 		},
 		/**
 		 * Handle a complete TLS packet received

src/scripts/handshake.ts

@@ -20,7 +20,7 @@ await new Promise<void>((resolve, reject) => {
 		verifyServerCertificate: true,
 		logger: LOGGER,
 		// write raw bytes to the socket
-		async write({ header, content }) {
+		write({ header, content }) {
 			socket.write(header)
 			socket.write(content)
 		},
@@ -46,8 +46,14 @@ await new Promise<void>((resolve, reject) => {
 
 	socket.on('data', tls.handleReceivedBytes)
 	// start handshake as soon as the socket connects
-	socket.on('connect', () => tls.startHandshake())
+	socket.once('connect', () => tls.startHandshake())
 	socket.once('error', (err) => reject(err))
+	socket.once('close', () => {
+		console.log('Socket closed')
+		if(!tls.isHandshakeDone()) {
+			reject(new Error('Handshake was not completed before socket closed'))
+		}
+	})
 	// use the TCP socket to connect to the server
 	socket.connect({ host, port })
 })

src/tests/message-processor.test.ts

@@ -1,80 +1,61 @@
 import assert from 'node:assert'
-import { describe, it } from 'node:test'
-import type { TLSPacket } from '../types/index.ts'
+import { beforeEach, describe, it } from 'node:test'
+import type { PacketProcessor, TLSPacket, TLSPacketWithType } from '../types/index.ts'
 import { logger } from '../utils/index.ts'
 import { makeMessageProcessor } from '../utils/index.ts'
 import { expectBuffsEq } from './utils.ts'
 
 describe('TLS Message Processor', () => {
 
-	it('should process a complete message', async() => {
-		const processor = makeTestMsgProcessor()
-		const pkts = await processor.onData(
-			Buffer.from('15030300050101010101', 'hex')
+	let processor: PacketProcessor
+	beforeEach(() => {
+		processor = makeMessageProcessor(logger)
+	})
+
+	it('should process a complete message', () => {
+		const pkts = Array.from(
+			processor.onData(Buffer.from('15030300050101010101', 'hex'))
 		)
 		assert.equal(pkts.length, 1)
-		expectBuffsEq(pkts[0].header, Buffer.from('1503030005', 'hex'))
-		expectBuffsEq(pkts[0].content, Buffer.from('0101010101', 'hex'))
+		expectBuffsEq(pkts[0].packet.header, Buffer.from('1503030005', 'hex'))
+		expectBuffsEq(pkts[0].packet.content, Buffer.from('0101010101', 'hex'))
 	})
 
-	it('should process a message byte-by-byte', async() => {
-		const processor = makeTestMsgProcessor()
+	it('should process a message byte-by-byte', () => {
 		const buffer = Buffer.from('15030300050101010101', 'hex')
 		for(let i = 0;i < buffer.length;i++) {
-			const pkts = await processor.onData(buffer.subarray(i, i + 1))
+			const pkts = Array.from(processor.onData(buffer.subarray(i, i + 1)))
 			if(i < buffer.length - 1) {
 				assert.equal(pkts.length, 0)
 			} else {
 				assert.equal(pkts.length, 1)
-				expectBuffsEq(pkts[0].content, Buffer.from('0101010101', 'hex'))
+				expectBuffsEq(pkts[0].packet.content, Buffer.from('0101010101', 'hex'))
 			}
 		}
 	})
 
 	it('should process multiple messages', async() => {
-		const processor = makeTestMsgProcessor()
 		const buffers = [
 			Buffer.from('15030300050101010101', 'hex'),
 			Buffer.from('1503030006010101010101', 'hex')
 		]
-		const pkts = await processor.onData(Buffer.concat(buffers))
+		const pkts = Array.from(processor.onData(Buffer.concat(buffers)))
 		assert.equal(pkts.length, 2)
-		expectBuffsEq(pkts[0].content, Buffer.from('0101010101', 'hex'))
-		expectBuffsEq(pkts[1].content, Buffer.from('010101010101', 'hex'))
+		expectBuffsEq(pkts[0].packet.content, Buffer.from('0101010101', 'hex'))
+		expectBuffsEq(pkts[1].packet.content, Buffer.from('010101010101', 'hex'))
 	})
 
 	it('should process a message and a half', async() => {
-		const processor = makeTestMsgProcessor()
 		const msgAndHalfBuffer = Buffer.concat(
 			[
 				Buffer.from('15030300050101010101', 'hex'),
 				Buffer.from('1503030006', 'hex')
 			]
 		)
 		const finalBuffer = Buffer.from('010101010101', 'hex')
-		const pkts = await processor.onData(msgAndHalfBuffer)
+		const pkts = Array.from(processor.onData(msgAndHalfBuffer))
 		assert.equal(pkts.length, 1)
-		const pkts2 = await processor.onData(finalBuffer)
+		const pkts2 = Array.from(processor.onData(finalBuffer))
 		assert.equal(pkts2.length, 1)
 	})
-
-	// eslint-disable-next-line unicorn/consistent-function-scoping
-	function makeTestMsgProcessor() {
-		const processor = makeMessageProcessor(logger)
-
-		return {
-			...processor,
-			async onData(packet: Buffer) {
-				const packets: TLSPacket[] = []
-				await processor.onData(
-					packet,
-					(_, pkt) => {
-						packets.push(pkt)
-					}
-				)
-
-				return packets
-			}
-		}
-	}
 })

src/types/tls.ts

@@ -140,8 +140,11 @@ export type TLSEventHandlers = {
 	onSessionTicket?(ticket: TLSSessionTicket): void
 }
 
-export type ProcessPacket = (type: number, packet: TLSPacket) => void
+export type TLSPacketWithType = {
+	type: number
+	packet: TLSPacket
+}
 
 export type PacketProcessor = {
-	onData(data: Uint8Array, onChunk: ProcessPacket): void
+	onData(data: Uint8Array): Generator<TLSPacketWithType>
 }

src/utils/packets.ts

@@ -1,4 +1,4 @@
-import type { Logger, ProcessPacket, TLSProtocolVersion } from '../types/index.ts'
+import type { Logger, TLSPacketWithType, TLSProtocolVersion } from '../types/index.ts'
 import { PACKET_TYPE, TLS_PROTOCOL_VERSION_MAP } from './constants.ts'
 import { concatenateUint8Arrays, uint8ArrayToDataView } from './generics.ts'
 
@@ -116,7 +116,7 @@ export function makeMessageProcessor(logger: Logger) {
 		 * or a single packet
 		 * @param onChunk handle a complete packet
 		 */
-		async onData(packet: Uint8Array, onChunk: ProcessPacket) {
+		*onData(packet: Uint8Array) {
 			buffer = concatenateUint8Arrays([ buffer, packet ])
 			while(buffer.length) {
 				// if we already aren't processing a packet
@@ -151,7 +151,7 @@ export function makeMessageProcessor(logger: Logger) {
 					// remove the packet header
 					buffer = buffer.slice(5)
 					logger.trace(
-						{ bytesLeft, type: currentMessageType },
+						{ bytesLeft, recvBuffer: buffer.length, type: currentMessageType },
 						'starting processing packet'
 					)
 				}
@@ -165,17 +165,22 @@ export function makeMessageProcessor(logger: Logger) {
 				const body = buffer.slice(0, bytesLeft)
 
 				logger.trace({ type: currentMessageType }, 'got complete packet')
-				await onChunk(currentMessageType, {
-					header: currentMessageHeader!,
-					content: body
-				})
+				const pktWithType: TLSPacketWithType = {
+					type: currentMessageType,
+					packet: {
+						header: currentMessageHeader!,
+						content: body
+					}
+				}
 
 				currentMessageType = undefined
 
 				// if the current chunk we have still has bytes left
 				// then that means we have another packet in the chunk
 				// this will be processed in the next iteration of the loop
 				buffer = buffer.slice(body.length)
+
+				yield pktWithType
 			}
 		},
 		reset() {