feat: update server & userscript

src/server/main.ts

@@ -11,7 +11,8 @@ const wss = new WebSocketServer({
   port: PORT,
 });
 
-let whatsapp: WebSocket | undefined;
+const bridgeSockets = new Set<WebSocket>();
+const clientSockets = new Set<WebSocket>();
 
 const pending = new Map<
   string,
@@ -26,11 +27,24 @@ function requestId() {
   return crypto.randomUUID();
 }
 
+function getWhatsApp(): WebSocket | undefined {
+  for (const ws of bridgeSockets) {
+    if (ws.readyState === WebSocket.OPEN) {
+      return ws;
+    }
+
+    bridgeSockets.delete(ws);
+  }
+
+  return undefined;
+}
+
 function call<T>(
   method: BridgeRequest["method"],
   params?: unknown,
 ): Promise<T> {
-  if (!whatsapp || whatsapp.readyState !== WebSocket.OPEN) {
+  const wa = getWhatsApp();
+  if (!wa || wa.readyState !== WebSocket.OPEN) {
     throw new Error("WhatsApp Web is not connected");
   }
 
@@ -49,7 +63,7 @@ function call<T>(
       timer,
     });
 
-    whatsapp!.send(
+    wa!.send(
       JSON.stringify({
         type: "request",
         id,
@@ -94,6 +108,11 @@ function handleWhatsAppMessage(data: string) {
     // console.log("[WhatsApp event]", message.event, message.data);
 
     if (message.event === "message") {
+      // forward to clients
+      clientSockets.forEach((ws) => {
+        ws.send(JSON.stringify(message));
+      });
+
       const m = message.data as MessageDTO;
 
       insertMessage({
@@ -109,15 +128,14 @@ function handleWhatsAppMessage(data: string) {
 }
 
 wss.on("connection", (ws) => {
-  console.log("[WS] client connected");
+  console.log("[WS] socket connected");
 
-  /*
-   * First connection from Tampermonkey is the
-   * WhatsApp bridge.
-   */
-  if (!whatsapp && !ws.url.includes("?client=1")) {
-    whatsapp = ws;
+  const url = new URL(ws.url ?? "", "http://localhost");
+  const isBridge =
+    !url.searchParams.get("client") || url.searchParams.get("bridge") === "1";
 
+  if (isBridge) {
+    bridgeSockets.add(ws);
     console.log("[WS] WhatsApp bridge attached");
 
     ws.on("message", (data) => {
@@ -125,32 +143,23 @@ wss.on("connection", (ws) => {
     });
 
     ws.on("close", () => {
-      if (whatsapp === ws) {
-        whatsapp = undefined;
-      }
-
-      console.log("[WS] WhatsApp disconnected");
+      bridgeSockets.delete(ws);
+      console.log("[WS] WhatsApp bridge disconnected");
     });
 
+    ws.on("error", () => {
+      bridgeSockets.delete(ws);
+    });
     return;
   }
 
-  /*
-   * This is an external application/client.
-   *
-   * Example:
-   *   my backend
-   *   CLI
-   *   React application
-   *   another service
-   */
+  clientSockets.add(ws);
 
   ws.on("message", async (data) => {
     let request: any;
 
     try {
       request = JSON.parse(data.toString()) as BridgeRequest;
-
       if (request.type !== "request") {
         return;
       }
@@ -214,8 +223,13 @@ wss.on("connection", (ws) => {
   });
 
   ws.on("close", () => {
+    clientSockets.delete(ws);
     console.log("[WS] external client disconnected");
   });
+
+  ws.on("error", () => {
+    clientSockets.delete(ws);
+  });
 });
 
 initDb();

src/userscript/index.ts

@@ -7,7 +7,7 @@ import {
   serializeMessage,
 } from "./serializer";
 
-const WS_URL = "wss://wa-bridge.rul.sh/ws";
+const WS_URL = "wss://wa-bridge.rul.sh/ws?bridge=1";
 
 let socket: WebSocket | undefined;