Add WebSocket support (ws)

Makefile

@@ -54,6 +54,11 @@ hono-sqlite3:
 hapi-postgres:
 	cd sample-apps/hapi-postgres && AIKIDO_DEBUG=true AIKIDO_BLOCKING=true node app.js
 
+
+.PHONY: ws-postgres
+ws-postgres:
+	cd sample-apps/ws-postgres && AIKIDO_DEBUG=true AIKIDO_BLOCKING=true node app.js
+
 .PHONY: install
 install:
 	mkdir -p build && cp library/package.json build/package.json

README.md

@@ -59,6 +59,10 @@ Aikido Firewall for Node.js 16+ is compatible with:
 * ✅ Google Cloud Functions
 * ✅ AWS Lambda
 
+### Real time communication
+
+* ✅ [`ws`](https://www.npmjs.com/package/ws) 8.x, 7.x
+
 ### ORMs and query builders
 
 See list above for supported database drivers.

end2end/package.json

@@ -3,7 +3,8 @@
   "private": true,
   "dependencies": {
     "@supercharge/promise-pool": "^3.1.1",
-    "tap": "^18.7.0"
+    "tap": "^18.7.0",
+    "ws": "^8.18.0"
   },
   "scripts": {
     "test": "tap tests/*.js --allow-empty-coverage -j 1"

end2end/tests/ws-postgres.test.js

@@ -0,0 +1,181 @@
+const t = require("tap");
+const { spawn } = require("child_process");
+const { resolve } = require("path");
+const timeout = require("../timeout");
+const { WebSocket } = require("ws");
+
+const pathToApp = resolve(__dirname, "../../sample-apps/ws-postgres", "app.js");
+
+t.test("it blocks in blocking mode", (t) => {
+  const server = spawn(`node`, [pathToApp, "4000"], {
+    env: { ...process.env, AIKIDO_DEBUG: "true", AIKIDO_BLOCKING: "true" },
+  });
+
+  server.on("close", () => {
+    t.end();
+  });
+
+  server.on("error", (err) => {
+    t.fail(err.message);
+  });
+
+  let stdout = "";
+  server.stdout.on("data", (data) => {
+    stdout += data.toString();
+  });
+
+  let stderr = "";
+  server.stderr.on("data", (data) => {
+    stderr += data.toString();
+  });
+
+  // Wait for the server to start
+  timeout(2000)
+    .then(async () => {
+      // Does not block normal messages
+      return await new Promise((resolve, reject) => {
+        const ws = new WebSocket(`ws://localhost:4000`);
+
+        ws.on("error", (err) => {
+          reject(err);
+        });
+
+        ws.on("open", () => {
+          ws.send("Hello world!");
+        });
+
+        ws.on("message", (data) => {
+          const str = data.toString();
+          if (str.includes("Welcome")) {
+            return;
+          }
+          t.match(str, /Hello world!/);
+          ws.close();
+          resolve();
+        });
+      });
+    })
+    .then(async () => {
+      // Does block sql injection
+      return await new Promise((resolve, reject) => {
+        const ws = new WebSocket(`ws://localhost:4000`);
+
+        ws.on("error", (err) => {
+          reject(err);
+        });
+
+        ws.on("open", () => {
+          ws.send("Bye'); DELETE FROM messages;--");
+        });
+
+        ws.on("message", (data) => {
+          const str = data.toString();
+          if (
+            str.includes("Welcome") ||
+            str.includes("Hello") ||
+            str.includes("Bye")
+          ) {
+            return;
+          }
+          t.match(str, /An error occurred/);
+          ws.close();
+          resolve();
+        });
+      });
+    })
+    .then(() => {
+      t.match(stdout, /Starting agent/);
+      t.match(stderr, /Aikido firewall has blocked an SQL injection/);
+    })
+    .catch((error) => {
+      t.fail(error.message);
+    })
+    .finally(() => {
+      server.kill();
+    });
+});
+
+t.test("it does not block in non-blocking mode", (t) => {
+  const server = spawn(`node`, [pathToApp, "4001"], {
+    env: { ...process.env, AIKIDO_DEBUG: "true" },
+  });
+
+  server.on("close", () => {
+    t.end();
+  });
+
+  server.on("error", (err) => {
+    t.fail(err.message);
+  });
+
+  let stdout = "";
+  server.stdout.on("data", (data) => {
+    stdout += data.toString();
+  });
+
+  let stderr = "";
+  server.stderr.on("data", (data) => {
+    stderr += data.toString();
+  });
+
+  // Wait for the server to start
+  timeout(2000)
+    .then(async () => {
+      // Does not block normal messages
+      return await new Promise((resolve, reject) => {
+        const ws = new WebSocket(`ws://localhost:4001`);
+
+        ws.on("error", (err) => {
+          reject(err);
+        });
+
+        ws.on("open", () => {
+          ws.send("Hello world!");
+        });
+
+        ws.on("message", (data) => {
+          const str = data.toString();
+          if (str.includes("Welcome")) {
+            return;
+          }
+          t.match(str, /Hello world!/);
+          ws.close();
+          resolve();
+        });
+      });
+    })
+    .then(async () => {
+      // Does block sql injection
+      return await new Promise((resolve, reject) => {
+        const ws = new WebSocket(`ws://localhost:4001`);
+
+        ws.on("error", (err) => {
+          reject(err);
+        });
+
+        ws.on("open", () => {
+          ws.send("Bye'); DELETE FROM messages;--");
+        });
+
+        ws.on("message", (data) => {
+          const str = data.toString();
+          if (str.includes("Welcome") || str.includes("Hello")) {
+            return;
+          }
+          t.match(str, /Bye/);
+          ws.close();
+          resolve();
+        });
+      });
+    })
+    .then(() => {
+      t.match(stdout, /Starting agent/);
+      t.notMatch(stderr, /Aikido firewall has blocked an SQL injection/);
+    })
+    .catch((error) => {
+      t.fail(error.message);
+    })
+    .finally(() => {
+      server.kill();
+    });
+});

library/agent/Context.ts

@@ -21,6 +21,7 @@ export type Context = {
   graphql?: string[];
   xml?: unknown;
   subdomains?: string[]; // https://expressjs.com/en/5x/api.html#req.subdomains
+  ws?: unknown; // Additional data related to WebSocket connections, like the last message received
 };
 
 /**
@@ -56,6 +57,7 @@ export function runWithContext<T>(context: Context, fn: () => T) {
     current.graphql = context.graphql;
     current.xml = context.xml;
     current.subdomains = context.subdomains;
+    current.ws = context.ws;
 
     return fn();
   }

library/agent/Source.ts

@@ -7,6 +7,7 @@ export const SOURCES = [
   "graphql",
   "xml",
   "subdomains",
+  "ws",
 ] as const;
 
 export type Source = (typeof SOURCES)[number];

library/agent/protect.ts

@@ -38,6 +38,7 @@
 import { XmlMinusJs } from "../sources/XmlMinusJs";
 import { Hapi } from "../sources/Hapi";
 import { Shelljs } from "../sinks/Shelljs";
+import { Ws } from "../sources/Ws";
 
 function isDebugging() {
   return (
@@ -135,10 +136,11 @@
     new XmlMinusJs(),
     new Shelljs(),
     new Hapi(),
+    new Ws(),
   ];
 }
 
 export function protect() {
   const agent = getAgent({
     serverless: undefined,
   });

library/helpers/getMaxWsMsgSize.test.ts

@@ -0,0 +1,46 @@
+import * as t from "tap";
+import { getMaxWsMsgSize } from "./getMaxWsMsgSize";
+
+t.test(
+  "returns default max body size when AIKIDO_MAX_WS_MSG_SIZE_MB is not set",
+  async (t) => {
+    delete process.env.AIKIDO_MAX_WS_MSG_SIZE_MB;
+    t.equal(
+      getMaxWsMsgSize(),
+      20 * 1024 * 1024,
+      "should return 20 MB as default"
+    );
+  }
+);
+
+t.test(
+  "returns parsed value from AIKIDO_MAX_WS_MSG_SIZE_MB without suffix",
+  async (t) => {
+    process.env.AIKIDO_MAX_WS_MSG_SIZE_MB = "10";
+    t.equal(getMaxWsMsgSize(), 10 * 1024 * 1024, "should return 10 MB");
+  }
+);
+
+t.test(
+  "returns default max body size for non-numeric AIKIDO_MAX_WS_MSG_SIZE_MB",
+  async (t) => {
+    process.env.AIKIDO_MAX_WS_MSG_SIZE_MB = "invalid";
+    t.equal(
+      getMaxWsMsgSize(),
+      20 * 1024 * 1024,
+      "should return 20 MB as default"
+    );
+  }
+);
+
+t.test(
+  "returns default max body size for negative AIKIDO_MAX_WS_MSG_SIZE_MB",
+  async (t) => {
+    process.env.AIKIDO_MAX_WS_MSG_SIZE_MB = "-5";
+    t.equal(
+      getMaxWsMsgSize(),
+      20 * 1024 * 1024,
+      "should return 20 MB as default"
+    );
+  }
+);

library/helpers/getMaxWsMsgSize.ts

@@ -0,0 +1,12 @@
+const MAX_WS_MSG_SIZE_MB = 20;
+
+export function getMaxWsMsgSize() {
+  if (process.env.AIKIDO_MAX_WS_MSG_SIZE_MB) {
+    const parsed = parseInt(process.env.AIKIDO_MAX_WS_MSG_SIZE_MB, 10);
+    if (!isNaN(parsed) && parsed > 0) {
+      return parsed * 1024 * 1024;
+    }
+  }
+
+  return MAX_WS_MSG_SIZE_MB * 1024 * 1024;
+}

library/package.json

@@ -47,6 +47,7 @@
     "@types/shell-quote": "^1.7.5",
     "@types/sinonjs__fake-timers": "^8.1.5",
     "@types/supertest": "^6.0.2",
+    "@types/ws": "^8.5.10",
     "@typescript-eslint/eslint-plugin": "^6.19.0",
     "@typescript-eslint/parser": "^6.19.0",
     "aws-sdk": "^2.1595.0",
@@ -72,6 +73,7 @@
     "tap": "^18.6.1",
     "typescript": "^5.3.3",
     "undici": "^6.12.0",
+    "ws": "^8.18.0",
     "xml-js": "^1.6.11",
     "xml2js": "^0.6.2"
   },

library/sources/HTTPServer.ts

@@ -2,8 +2,8 @@ import { Agent } from "../agent/Agent";
 import { Hooks } from "../agent/hooks/Hooks";
 import { Wrapper } from "../agent/Wrapper";
 import { isPackageInstalled } from "../helpers/isPackageInstalled";
-import { isPlainObject } from "../helpers/isPlainObject";
 import { createRequestListener } from "./http-server/createRequestListener";
+import { createUpgradeListener } from "./http-server/createUpgradeListener";
 
 export class HTTPServer implements Wrapper {
   private wrapRequestListener(args: unknown[], module: string, agent: Agent) {
@@ -40,10 +40,13 @@ export class HTTPServer implements Wrapper {
     ) {
       return args;
     }
-    if (args[0] !== "request") {
-      return args;
+    if (args[0] === "request") {
+      return this.wrapRequestListener(args, module, agent);
+    }
+    if (args[0] === "upgrade") {
+      return [args[0], createUpgradeListener(args[1], module, agent)];
     }
-    return this.wrapRequestListener(args, module, agent);
+    return args;
   }
 
   wrap(hooks: Hooks) {