feat: add websocket relay support to http tunnels

Author: dylanplayer

packages/cli/src/commands/http.ts

@@ -6,6 +6,9 @@ import type {
   HttpResponseEndMessage,
   HttpResponseMessage,
   HttpResponseStartMessage,
+  WsCloseMessage,
+  WsFrameMessage,
+  WsOpenMessage,
 } from '@auger/shared';
 import { buildWsUrl } from '../utils';
 
@@ -17,6 +20,71 @@ export type HttpCommandOptions = {
   subdomain?: string;
 };
 
+type LocalSocketEntry = {
+  socket: WebSocket;
+  open: boolean;
+  closingFromServer: boolean;
+  queuedFrames: Array<{ payload: Uint8Array; isBinary: boolean }>;
+};
+
+const textEncoder = new TextEncoder();
+const textDecoder = new TextDecoder();
+
+function normalizeCloseCode(code: number | undefined): number {
+  if (code === undefined) return 1000;
+  if (!Number.isInteger(code)) return 1000;
+  if (code < 1000 || code > 4999) return 1000;
+  return code;
+}
+
+async function messageDataToText(data: unknown): Promise<string> {
+  if (typeof data === 'string') return data;
+  if (data instanceof ArrayBuffer) return Buffer.from(data).toString('utf8');
+  if (ArrayBuffer.isView(data)) {
+    return Buffer.from(data.buffer, data.byteOffset, data.byteLength).toString('utf8');
+  }
+  if (data instanceof Blob) {
+    return Buffer.from(await data.arrayBuffer()).toString('utf8');
+  }
+  return String(data);
+}
+
+async function messageDataToPayload(
+  data: unknown
+): Promise<{ payload: Uint8Array; isBinary: boolean } | null> {
+  if (typeof data === 'string') {
+    return {
+      payload: textEncoder.encode(data),
+      isBinary: false,
+    };
+  }
+
+  if (data instanceof ArrayBuffer) {
+    return {
+      payload: new Uint8Array(data),
+      isBinary: true,
+    };
+  }
+
+  if (ArrayBuffer.isView(data)) {
+    return {
+      payload: new Uint8Array(
+        data.buffer.slice(data.byteOffset, data.byteOffset + data.byteLength)
+      ),
+      isBinary: true,
+    };
+  }
+
+  if (data instanceof Blob) {
+    return {
+      payload: new Uint8Array(await data.arrayBuffer()),
+      isBinary: true,
+    };
+  }
+
+  return null;
+}
+
 async function handleHttpRequest(
   ws: WebSocket,
   localPort: number,
@@ -111,11 +179,121 @@ export async function runHttpCommand(options: HttpCommandOptions): Promise<void>
     ? `${options.subdomain}:${options.localPort}`
     : `${options.localPort}`;
   let ws: WebSocket | null = null;
+  const localSockets = new Map<string, LocalSocketEntry>();
   let consecutiveFailures = 0;
   let reconnectScheduled = false;
   let fatalError = false;
 
+  const sendControl = (
+    payload:
+      | HelloMessage
+      | HttpResponseMessage
+      | HttpResponseStartMessage
+      | HttpResponseChunkMessage
+      | HttpResponseEndMessage
+      | WsFrameMessage
+      | WsCloseMessage
+  ) => {
+    if (ws?.readyState === WebSocket.OPEN) {
+      ws.send(toMessage(payload));
+    }
+  };
+
+  const closeAllLocalSockets = (code: number, reason: string) => {
+    for (const [id, entry] of localSockets.entries()) {
+      entry.closingFromServer = true;
+      localSockets.delete(id);
+      entry.socket.close(code, reason);
+    }
+  };
+
+  const flushQueuedFrames = (entry: LocalSocketEntry) => {
+    while (entry.queuedFrames.length > 0) {
+      const frame = entry.queuedFrames.shift();
+      if (!frame) continue;
+      if (frame.isBinary) {
+        entry.socket.send(frame.payload);
+      } else {
+        entry.socket.send(textDecoder.decode(frame.payload));
+      }
+    }
+  };
+
+  const handleWsOpen = (message: WsOpenMessage) => {
+    const localUrl = new URL(message.path, `ws://127.0.0.1:${options.localPort}`).toString();
+    const localSocket =
+      message.protocols.length > 0
+        ? new WebSocket(localUrl, message.protocols)
+        : new WebSocket(localUrl);
+    localSocket.binaryType = 'arraybuffer';
+
+    const entry: LocalSocketEntry = {
+      socket: localSocket,
+      open: false,
+      closingFromServer: false,
+      queuedFrames: [],
+    };
+    localSockets.set(message.id, entry);
+
+    localSocket.addEventListener('open', () => {
+      entry.open = true;
+      flushQueuedFrames(entry);
+    });
+
+    localSocket.addEventListener('message', async (event) => {
+      const payload = await messageDataToPayload(event.data);
+      if (!payload) return;
+      sendControl({
+        type: 'ws_frame',
+        id: message.id,
+        dataBase64: encodeBase64(payload.payload),
+        isBinary: payload.isBinary,
+      });
+    });
+
+    localSocket.addEventListener('close', (event) => {
+      localSockets.delete(message.id);
+      if (entry.closingFromServer) return;
+      sendControl({
+        type: 'ws_close',
+        id: message.id,
+        code: event.code,
+        reason: event.reason,
+      });
+    });
+  };
+
+  const handleWsFrame = (message: WsFrameMessage) => {
+    const entry = localSockets.get(message.id);
+    if (!entry) return;
+
+    const payload = decodeBase64(message.dataBase64);
+    if (!entry.open) {
+      entry.queuedFrames.push({
+        payload,
+        isBinary: message.isBinary,
+      });
+      return;
+    }
+
+    if (message.isBinary) {
+      entry.socket.send(payload);
+    } else {
+      entry.socket.send(textDecoder.decode(payload));
+    }
+  };
+
+  const handleWsClose = (message: WsCloseMessage) => {
+    const entry = localSockets.get(message.id);
+    if (!entry) return;
+
+    entry.closingFromServer = true;
+    localSockets.delete(message.id);
+    entry.socket.close(normalizeCloseCode(message.code), message.reason ?? '');
+  };
+
   const scheduleReconnect = (reason: string) => {
+    closeAllLocalSockets(1011, 'Tunnel connection lost');
     if (fatalError) {
       process.exit(1);
     }
@@ -136,6 +314,7 @@ export async function runHttpCommand(options: HttpCommandOptions): Promise<void>
 
   const connect = () => {
     ws = new WebSocket(wsUrl);
+    ws.binaryType = 'arraybuffer';
 
     ws.addEventListener('open', () => {
       consecutiveFailures = 0;
@@ -147,14 +326,11 @@ export async function runHttpCommand(options: HttpCommandOptions): Promise<void>
         requestedSubdomain: options.subdomain,
       };
 
-      ws?.send(toMessage(hello));
+      sendControl(hello);
     });
 
     ws.addEventListener('message', async (event) => {
-      const data =
-        typeof event.data === 'string'
-          ? event.data
-          : Buffer.from(event.data as ArrayBuffer).toString('utf8');
+      const data = await messageDataToText(event.data);
       const message = parseMessage(data);
 
       if (message.type === 'error') {
@@ -175,6 +351,21 @@ export async function runHttpCommand(options: HttpCommandOptions): Promise<void>
 
       if (message.type === 'http_request') {
         await handleHttpRequest(ws as WebSocket, options.localPort, message, label);
+        return;
+      }
+
+      if (message.type === 'ws_open') {
+        handleWsOpen(message);
+        return;
+      }
+
+      if (message.type === 'ws_frame') {
+        handleWsFrame(message);
+        return;
+      }
+
+      if (message.type === 'ws_close') {
+        handleWsClose(message);
       }
     });