From a0d271d5575bbcd73588d8c01e45958263c720b4 Mon Sep 17 00:00:00 2001
From: Francesco Guardiani <me@slinkydeveloper.com>
Date: Thu, 8 Feb 2024 15:32:58 +0100
Subject: [PATCH] Check protocol version (#250)

---
 src/io/decoder.ts           | 16 +++++++++++++++-
 test/message_coders.test.ts | 29 ++++++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/io/decoder.ts b/src/io/decoder.ts
index f6af7102..e1cab043 100644
--- a/src/io/decoder.ts
+++ b/src/io/decoder.ts
@@ -32,6 +32,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 +59,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 5c6cbe9a..53150a76 100644
--- a/test/message_coders.test.ts
+++ b/test/message_coders.test.ts
@@ -10,10 +10,15 @@
  */
 
 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 { 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 +49,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"
+    );
+  });
 });