feat: add ws-proxy

src/module/ws-proxy/manager.ts

@@ -0,0 +1,84 @@
+import { nanoid } from 'nanoid';
+import { WebSocket } from 'ws';
+import { logger } from '../logger.ts';
+class WsMessage {
+  ws: WebSocket;
+  user?: string;
+  constructor({ ws, user }: WssMessageOptions) {
+    this.ws = ws;
+    this.user = user;
+  }
+  async sendData(data: any, opts?: { timeout?: number }) {
+    if (this.ws.readyState !== WebSocket.OPEN) {
+      return { code: 500, message: 'WebSocket is not open' };
+    }
+    const timeout = opts?.timeout || 10 * 6 * 1000; // 10 minutes
+    const id = nanoid();
+    const message = JSON.stringify({
+      id,
+      type: 'proxy',
+      data,
+    });
+    logger.info('ws-proxy sendData', message);
+    this.ws.send(message);
+    return new Promise((resolve) => {
+      const timer = setTimeout(() => {
+        resolve({
+          code: 500,
+          message: 'timeout',
+        });
+      }, timeout);
+      this.ws.once('message', (event: Buffer) => {
+        const eventData = event.toString();
+        if (!eventData) {
+          return;
+        }
+        const data = JSON.parse(eventData);
+        if (data.id === id) {
+          resolve(data.data);
+          clearTimeout(timer);
+        }
+      });
+    });
+  }
+}
+type WssMessageOptions = {
+  ws: WebSocket;
+  user?: string;
+};
+export class WsProxyManager {
+  wssMap: Map<string, WsMessage> = new Map();
+  constructor() {}
+  getId(id: string, user?: string) {
+    return id + '/' + user;
+  }
+  register(id: string, opts?: { ws: WebSocket; user: string }) {
+    const _id = this.getId(id, opts?.user || '');
+    if (this.wssMap.has(_id)) {
+      const value = this.wssMap.get(_id);
+      if (value) {
+        value.ws.close();
+      }
+    }
+    const value = new WsMessage({ ws: opts?.ws, user: opts?.user });
+    this.wssMap.set(_id, value);
+  }
+  unregister(id: string, user?: string) {
+    const _id = this.getId(id, user || '');
+    const value = this.wssMap.get(_id);
+    if (value) {
+      value.ws.close();
+    }
+    this.wssMap.delete(_id);
+  }
+  getIds() {
+    return Array.from(this.wssMap.keys());
+  }
+  get(id: string, user?: string) {
+    if (user) {
+      const _id = this.getId(id, user);
+      return this.wssMap.get(_id);
+    }
+    return this.wssMap.get(id);
+  }
+}