Add WebSocket support

src/connection.ts

@@ -4,4 +4,5 @@ export default class Connection {
   hostname: string;
   clientId: string;
   websocket: HostipWebSocket;
+  sockets?: Map<string, HostipWebSocket>;
 }

src/handlers/handle-request.ts

@@ -43,15 +43,17 @@ const handleRequest = async function (request: Request, response: Response) {
     try {
       const forwardedResponseMessage: ForwardedResponseMessage =
         JSON.parse(text);
-      const body = Buffer.from(forwardedResponseMessage.body, "base64");
-      forwardedResponseMessage.headers["x-forwarded-for"] =
-        connection.websocket.ipAddress;
-
       // Bail if this handler is not for the request that created it
       if (forwardedResponseMessage.requestId !== requestId) {
         return;
       }
 
+      const body = forwardedResponseMessage.body
+        ? Buffer.from(forwardedResponseMessage.body, "base64")
+        : null;
+      forwardedResponseMessage.headers["x-forwarded-for"] =
+        connection.websocket.ipAddress;
+
       if (forwardedResponseMessage.type === "forwardedResponse") {
         response.status(forwardedResponseMessage.statusCode);
 

src/messages/websocket-client-message.ts

@@ -0,0 +1,5 @@
+export default interface WebSocketClientMessage {
+  type: "WebSocketClientMessage";
+  socketId: string;
+  data: string; // Not base64 encoded
+}

src/messages/websocket-close-message.ts

@@ -0,0 +1,6 @@
+export default interface WebSocketCloseMessage {
+  type: "WebSocketCloseMessage";
+  socketId: string;
+  code?: number;
+  data?: string;
+}

src/messages/websocket-host-message.ts

@@ -0,0 +1,5 @@
+export default interface WebSocketHostMessage {
+  type: "WebSocketHostMessage";
+  socketId: string;
+  data: string; // Not base64 encodedfer
+}

src/messages/websocket-open-message.ts

@@ -0,0 +1,6 @@
+export default interface WebSocketOpenMessage {
+  type: "WebSocketOpenMessage";
+  socketId: string;
+  url: string;
+  headers: any;
+}

src/proxy.ts

@@ -43,6 +43,8 @@ export default class Proxy {
 
       if (connection.clientId === clientId) {
         this.connections.splice(i, 1);
+
+        connection.sockets?.forEach((socket) => socket.terminate());
       }
     }
   }

src/websocket/host-ip-websocket.ts

@@ -4,6 +4,8 @@ export default class HostipWebSocket extends WebSocket {
   tunnelmoleClientId: string;
   connectionStart: number;
   isAlive: boolean;
+  /** Indicates that a websocket connection is used to tunnel a websocket */
+  dataTunnel: boolean;
   ipAddress: string;
   // stub for websocket
   sendMessage(object: any) {}

websocket.ts

@@ -1,10 +1,15 @@
+import { nanoid } from "nanoid";
 import { messageHandlers } from "./message-handlers";
 import HostipWebSocket from "./src/websocket/host-ip-websocket";
 import log from "./src/logging/log";
 import moment from "moment";
 import Proxy from "./src/proxy";
 import { IncomingMessage } from "http";
 import config from "./config";
+import type Connection from "./src/connection";
+import WebSocketOpenMessage from "./src/messages/websocket-open-message";
+import WebSocketCloseMessage from "./src/messages/websocket-close-message";
+import WebSocketHostMessage from "./src/messages/websocket-host-message";
 const inArray = require("in_array");
 
 /**
@@ -39,7 +44,11 @@ export default function websocket(
 
     // Skip any messages that are handled dynamically using other explicitly defined 'message' callbacks
     // Example: forwardedResponse handler in handleRequest that is set dynamically for every request
-    const dynamicallyHandledMessageTypes = ["forwardedResponse"];
+    const dynamicallyHandledMessageTypes = [
+      "forwardedResponse",
+      "WebSocketClientMessage",
+      "WebSocketCloseMessage",
+    ];
 
     if (inArray(message.type, dynamicallyHandledMessageTypes)) {
       return;
@@ -62,6 +71,8 @@ export default function websocket(
       log(message, "info");
     });
 
+  socketOverSocket(websocket, request);
+
   websocket.on("error", (code: number, reason: string) => {
     console.info("Caught an error. Error code: " + code + " Reason: " + reason);
   });
@@ -70,9 +81,98 @@ export default function websocket(
     websocket.terminate();
 
     const proxy = Proxy.getInstance();
-
     proxy.deleteConnection(websocket.tunnelmoleClientId);
 
     console.info("Connection Closed. Code: " + code + " Reason: " + reason);
   });
 }
+
+function socketOverSocket(
+  websocket: HostipWebSocket,
+  request: IncomingMessage,
+) {
+  if (request.url !== "/") openDatatunnel();
+  else {
+    // Must init within 2 seconds
+    setTimeout(() => {
+      openDatatunnel();
+    }, 2000);
+
+    // First message must be initialise
+    websocket.once("message", (text: string) => {
+      try {
+        const message = JSON.parse(text);
+        if (message.type === "initialise") websocket.dataTunnel = false;
+        else openDatatunnel();
+      } catch (error) {
+        openDatatunnel();
+      }
+    });
+  }
+
+  function openDatatunnel() {
+    // Only open one data tunnel per websocket
+    if (typeof websocket.dataTunnel !== "undefined") return;
+
+    websocket.dataTunnel = true;
+
+    const proxy = Proxy.getInstance();
+    const url = new URL("https://" + request.headers.host);
+    const hostname = url.hostname;
+    const connection: Connection = proxy.findConnectionByHostname(hostname);
+    if (!connection) return; // console.error("openDatatunnel.nope", hostname);
+
+    // console.log("openDatatunnel", connection?.clientId);
+
+    // Register this data tunnel
+    const socketId = nanoid();
+    if (!connection.sockets) connection.sockets = new Map();
+    connection.sockets.set(socketId, websocket);
+
+    // Let the client know the socket is open
+    const open: WebSocketOpenMessage = {
+      socketId,
+      type: "WebSocketOpenMessage",
+      url: request.url,
+      headers: request.headers,
+    };
+    connection.websocket.sendMessage(open);
+
+    // Forget socket on close
+    websocket.on("close", (code, data) => {
+      connection.sockets?.delete(socketId);
+      const closed: WebSocketCloseMessage = {
+        socketId,
+        type: "WebSocketCloseMessage",
+        code,
+        data,
+      };
+      console.log("websocket.on.close", closed);
+      connection.websocket.sendMessage(closed);
+    });
+
+    websocket.removeAllListeners("message");
+
+    // Send messages to the client
+    websocket.on("message", (text: string) => {
+      const forward: WebSocketHostMessage = {
+        socketId,
+        type: "WebSocketHostMessage",
+        data: text,
+      };
+      connection.websocket.sendMessage(forward);
+    });
+
+    // Send messages from the client
+    connection.websocket.on("message", (text: string) => {
+      const message = JSON.parse(text);
+      if (message.socketId !== socketId) return;
+
+      if (message.type === "WebSocketClientMessage") {
+        websocket.send(message.data);
+      } else if (message.type === "WebSocketCloseMessage") {
+        websocket.close(message.code, message.data);
+      }
+    });
+  }
+}