From 5f6174a2e51644c16414ba2efb61e3178435a249 Mon Sep 17 00:00:00 2001
From: slinkydeveloper <francescoguard@gmail.com>
Date: Tue, 6 Feb 2024 12:53:47 +0100
Subject: [PATCH] Check protocol version

---
 src/io/decoder.ts           | 17 ++++++++++++++++-
 test/message_coders.test.ts | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/src/io/decoder.ts b/src/io/decoder.ts
index 1129f93b..0cca11e9 100644
--- a/src/io/decoder.ts
+++ b/src/io/decoder.ts
@@ -25,6 +25,7 @@ import { PROTOBUF_MESSAGE_BY_TYPE } from "../types/protocol";
 import { Header, Message } from "../types/types";
 import assert from "assert";
 import { ensureError } from "../types/errors";
+import { serialize } from "node:v8";
 
 type Output = { push(msg: Message): void };
 type DecoderState = { state: number; header: Header | undefined; buf: Buffer };
@@ -32,6 +33,8 @@ type DecoderState = { state: number; header: Header | undefined; buf: Buffer };
 const WAITING_FOR_HEADER = 0;
 const WAITING_FOR_BODY = 1;
 
+export const SUPPORTED_PROTOCOL_VERSION = 0;
+
 function initalDecoderState(buf: Buffer): DecoderState {
   return {
     state: WAITING_FOR_HEADER,
@@ -57,8 +60,20 @@ function decodeMessages(decoderState: DecoderState, out: Output): DecoderState {
         }
         const h = buf.readBigUInt64BE();
         buf = buf.subarray(8);
-        decoderState.header = Header.fromU64be(h);
+        const materializedHeader = Header.fromU64be(h);
+        decoderState.header = materializedHeader;
         decoderState.state = WAITING_FOR_BODY;
+
+        // Check protocol version
+        if (
+          materializedHeader.protocolVersion !== undefined &&
+          materializedHeader.protocolVersion !== SUPPORTED_PROTOCOL_VERSION
+        ) {
+          throw new Error(
+            `Unsupported protocol version ${materializedHeader.protocolVersion}, only version ${SUPPORTED_PROTOCOL_VERSION} is supported`
+          );
+        }
+
         break;
       }
       case WAITING_FOR_BODY: {
diff --git a/test/message_coders.test.ts b/test/message_coders.test.ts
index 9ebd2159..db0e8179 100644
--- a/test/message_coders.test.ts
+++ b/test/message_coders.test.ts
@@ -10,10 +10,19 @@
  */
 
 import { describe, expect } from "@jest/globals";
-import { streamDecoder } from "../src/io/decoder";
+import {
+  decodeMessagesBuffer,
+  streamDecoder,
+  SUPPORTED_PROTOCOL_VERSION,
+} from "../src/io/decoder";
 import { Message } from "../src/types/types";
-import { backgroundInvokeMessage } from "./protoutils";
+import {
+  backgroundInvokeMessage,
+  startMessage,
+  toStateEntries,
+} from "./protoutils";
 import { encodeMessage } from "../src/io/encoder";
+import { START_MESSAGE_TYPE, StartMessage } from "../src/types/protocol";
 
 describe("The stream decoder", () => {
   it("should handle decoding of messages across chunks", () => {
@@ -44,4 +53,26 @@ describe("The stream decoder", () => {
     expect(result[0]).toStrictEqual(largeMessage);
     expect(callbackCounter).toStrictEqual(2);
   });
+
+  it("should fail when unsupported protocol version", () => {
+    const startMessage = encodeMessage(
+      new Message(
+        START_MESSAGE_TYPE,
+        StartMessage.create({
+          id: Buffer.from(
+            "f311f1fdcb9863f0018bd3400ecd7d69b547204e776218b2",
+            "hex"
+          ),
+          debugId: "8xHx_cuYY_AAYvTQA7NfWm1RyBOd2IYsg",
+        }),
+        undefined,
+        SUPPORTED_PROTOCOL_VERSION + 1,
+        undefined
+      )
+    );
+
+    expect(() => decodeMessagesBuffer(Buffer.from(startMessage))).toThrow(
+      "Unsupported protocol version"
+    );
+  });
 });