From 0eaa5af909d90df935f79a4cc96740b2e89ff518 Mon Sep 17 00:00:00 2001
From: Rajeh Taher <rajeh@reforward.dev>
Date: Wed, 15 May 2024 18:34:37 +0300
Subject: [PATCH] feat(add-routing-info): initial commit (#773)

---
 src/Socket/socket.ts       | 15 ++++++++++++++-
 src/Types/Auth.ts          |  1 +
 src/Utils/auth-utils.ts    |  1 +
 src/Utils/noise-handler.ts | 20 ++++++++++++++++++--
 4 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/Socket/socket.ts b/src/Socket/socket.ts
index 78af0e5d2db..823b7636470 100644
--- a/src/Socket/socket.ts
+++ b/src/Socket/socket.ts
@@ -76,6 +76,10 @@ export const makeSocket = (config: SocketConfig) => {
 		url = new URL(`tcp://${MOBILE_ENDPOINT}:${MOBILE_PORT}`)
 	}
 
+	if(!config.mobile && url.protocol === 'wss' && authState?.creds?.routingInfo) {
+		url.searchParams.append('ED', authState.creds.routingInfo.toString('base64url'))
+	}
+
 	const ws = config.socket ? config.socket : config.mobile ? new MobileSocketClient(url, config) : new WebSocketClient(url, config)
 
 	ws.connect()
@@ -88,7 +92,8 @@ export const makeSocket = (config: SocketConfig) => {
 		keyPair: ephemeralKeyPair,
 		NOISE_HEADER: config.mobile ? MOBILE_NOISE_HEADER : NOISE_WA_HEADER,
 		mobile: config.mobile,
-		logger
+		logger,
+		routingInfo: authState?.creds?.routingInfo
 	})
 
 	const { creds } = authState
@@ -671,6 +676,14 @@ export const makeSocket = (config: SocketConfig) => {
 		end(new Boom('Multi-device beta not joined', { statusCode: DisconnectReason.multideviceMismatch }))
 	})
 
+	ws.on('CB:ib,,edge_routing', (node: BinaryNode) => {
+		const edgeRoutingNode = getBinaryNodeChild(node, 'edge_routing')
+		const routingInfo = getBinaryNodeChild(edgeRoutingNode, 'routing_info')
+		if(routingInfo?.content) {
+			authState.creds.routingInfo = Buffer.from(routingInfo?.content as Uint8Array)
+		}
+	})
+
 	let didStartBuffer = false
 	process.nextTick(() => {
 		if(creds.me?.id) {
diff --git a/src/Types/Auth.ts b/src/Types/Auth.ts
index a353521ef22..4338286f2e5 100644
--- a/src/Types/Auth.ts
+++ b/src/Types/Auth.ts
@@ -69,6 +69,7 @@ export type AuthenticationCreds = SignalCreds & {
 	registration: RegistrationOptions
     pairingCode: string | undefined
     lastPropHash: string | undefined
+    routingInfo: Buffer | undefined
 }
 
 export type SignalDataTypeMap = {
diff --git a/src/Utils/auth-utils.ts b/src/Utils/auth-utils.ts
index cc8f5e4e57a..82f3dc85136 100644
--- a/src/Utils/auth-utils.ts
+++ b/src/Utils/auth-utils.ts
@@ -217,5 +217,6 @@ export const initAuthCreds = (): AuthenticationCreds => {
 		registration: {} as never,
 		pairingCode: undefined,
 		lastPropHash: undefined,
+		routingInfo: undefined,
 	}
 }
\ No newline at end of file
diff --git a/src/Utils/noise-handler.ts b/src/Utils/noise-handler.ts
index a6928560084..5e37ba3c728 100644
--- a/src/Utils/noise-handler.ts
+++ b/src/Utils/noise-handler.ts
@@ -18,11 +18,13 @@ export const makeNoiseHandler = ({
 	NOISE_HEADER,
 	mobile,
 	logger,
+	routingInfo
 }: {
 	keyPair: KeyPair
 	NOISE_HEADER: Uint8Array
 	mobile: boolean
 	logger: Logger
+	routingInfo?: Buffer | undefined
 }) => {
 	logger = logger.child({ class: 'ns' })
 
@@ -133,11 +135,25 @@ export const makeNoiseHandler = ({
 				data = encrypt(data)
 			}
 
-			const introSize = sentIntro ? 0 : NOISE_HEADER.length
+			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)
 
 			if(!sentIntro) {
-				frame.set(NOISE_HEADER)
+				frame.set(header)
 				sentIntro = true
 			}