Merge pull request #25 from LJNeon/fortify-messages

Fortify message handler

src/index.ts

@@ -2,6 +2,7 @@ import { setupWatch } from "./fileWatch";
 import { config, loadConfig } from "./config";
 import { setupSocket } from "./networking/webSocket";
 import signal from "signal-js";
+import { RawData } from "ws";
 import {
   fileChangeEventToMsg,
   fileRemovalEventToMsg,
@@ -10,7 +11,7 @@ import {
 } from "./networking/messageGenerators";
 import { EventType } from "./eventTypes";
 import { messageHandler } from "./networking/messageHandler";
-import { FileEvent, Message } from "./interfaces";
+import { FileEvent } from "./interfaces";
 
 export async function start() {
   loadConfig();
@@ -18,7 +19,7 @@ export async function start() {
   const socket = setupSocket(signal);
 
   // Add a handler for received messages.
-  signal.on(EventType.MessageReceived, (msg: Message) => messageHandler(signal, msg, watch.paths));
+  signal.on(EventType.MessageReceived, (msg: RawData) => messageHandler(signal, msg, watch.paths));
 
   // Add a handler for when a connection to a game is made.
   signal.on(EventType.ConnectionMade, () => {

src/interfaces.ts

@@ -1,10 +1,35 @@
 import type { Stats } from "fs";
 
 export interface Message {
-  id: string;
+  jsonrpc: "2.0";
   method?: string;
-  jsonrpc: string;
-  params?: object;
+  result?: ResultType;
+  params?: FileMetadata;
+  error?: string;
+  id?: number;
+}
+
+type ResultType = string | number | string[] | FileContent[];
+type FileMetadata = FileData | FileContent | FileLocation | FileServer;
+
+export interface FileData {
+  filename: string;
+  content: string;
+  server: string;
+}
+
+export interface FileContent {
+  filename: string;
+  content: string;
+}
+
+export interface FileLocation {
+  filename: string;
+  server: string;
+}
+
+export interface FileServer {
+  server: string;
 }
 
 export interface FileEvent {

src/networking/messageGenerators.ts

@@ -14,7 +14,7 @@ export function fileChangeEventToMsg({ path }: FileEvent): Message {
       filename: addLeadingSlash(path),
       content: readFileSync(join(config.get("scriptsFolder"), path)).toString(),
     },
-    id: (messageCounter++).toString(),
+    id: messageCounter++,
   };
 }
 
@@ -26,7 +26,7 @@ export function fileRemovalEventToMsg({ path }: FileEvent): Message {
       server: "home",
       filename: addLeadingSlash(path),
     },
-    id: (messageCounter++).toString(),
+    id: messageCounter++,
   };
 }
 
@@ -34,7 +34,7 @@ export function requestDefinitionFile(): Message {
   return {
     jsonrpc: "2.0",
     method: "getDefinitionFile",
-    id: (messageCounter++).toString(),
+    id: messageCounter++,
   };
 }
 
@@ -45,7 +45,7 @@ export function requestFilenames(): Message {
     params: {
       server: "home",
     },
-    id: (messageCounter++).toString(),
+    id: messageCounter++,
   };
 }
 

src/networking/messageHandler.ts

@@ -1,31 +1,60 @@
 import { messageTracker } from "./messageTracker";
 import { Stats, writeFile } from "fs";
+import { RawData } from "ws";
 import { config } from "../config";
 import { EventType } from "../eventTypes";
 import { fileChangeEventToMsg } from "./messageGenerators";
 import type { Signal } from "signal-js";
 import { Message } from "../interfaces";
 
-export function messageHandler(signaller: Signal, msg: Message, paths: Map<string, Stats>) {
+function deserialize(data: RawData): Message {
+  const msg = JSON.parse(data.toString());
+
+  if (typeof msg.jsonrpc !== "string" || msg.jsonrpc !== "2.0" || typeof msg.id !== "number") {
+    throw Error("Malformed data received.");
+  }
+
+  const id: number = msg.id;
+  const request = messageTracker.get(id);
+
+  if (typeof request?.method !== "string") {
+    throw Error("Malformed JSON received.");
+  } else if (msg.error != null) {
+    throw Error(msg.error);
+  } else if (msg.result == null) {
+    throw Error("Malformed JSON received.");
+  }
+
+  return { jsonrpc: "2.0", method: request.method, result: msg.result, id };
+}
+
+export function isStringArray(s: Array<unknown>): s is string[] {
+  return s.every((s) => typeof s === "string");
+}
+
+export function messageHandler(signaller: Signal, data: RawData, paths: Map<string, Stats>) {
   let incoming;
 
   try {
-    incoming = JSON.parse(msg.toString());
+    incoming = deserialize(data);
   } catch (err) {
-    return console.log(err);
+    if (err instanceof Error) return console.log(err.message);
+    else throw err;
   }
-  console.log(incoming);
-  if (incoming.id == undefined) return;
 
-  if (incoming.result) {
-    const request = messageTracker.get(incoming.id);
-    if (request?.method && request.method == "getDefinitionFile" && incoming.result) {
+  switch (incoming.method) {
+    case "getDefinitionFile":
+      if (typeof incoming.result !== "string") return console.log("Malformed data received.");
+
       writeFile(config.get("definitionFile").location, incoming.result, (err) => {
         if (err) return console.log(err);
       });
-    }
 
-    if (request?.method && request.method == "getFileNames" && incoming.result) {
+      break;
+    case "getFileNames": {
+      if (!Array.isArray(incoming.result) || !isStringArray(incoming.result))
+        return console.log("Malformed data received.");
+
       const gameFiles = incoming.result.map((file: string) => removeLeadingSlash(file));
 
       paths.forEach((stats, fileName) => {

src/networking/messageTracker.ts

@@ -1,10 +1,12 @@
 import type { Message } from "../interfaces";
 
 class MessageTracker {
-  data = new Map<string, Message>();
+  data = new Map<number, Message>();
   #maxLength = 200;
 
   push(msg: Message) {
+    if (typeof msg.id !== "number") return;
+
     this.data.set(msg.id, msg);
 
     if (this.data.size > this.#maxLength) {
@@ -13,7 +15,7 @@ class MessageTracker {
     }
   }
 
-  get(index: string) {
+  get(index: number) {
     return this.data.get(index);
   }
 }

src/networking/webSocket.ts

@@ -3,7 +3,6 @@ import { WebSocketServer } from "ws";
 import { config } from "../config";
 import { EventType } from "../eventTypes";
 import { Message } from "../interfaces";
-import { requestDefinitionFile } from "./messageGenerators";
 import { messageTracker } from "./messageTracker";
 
 export function setupSocket(signaller: Signal) {