feat: add ws-proxy

src/index.ts

@@ -3,6 +3,8 @@ import { config } from './module/config.ts';
 import { app } from './app.ts';
 import './route/route.ts';
 import net from 'net';
+import { WssApp } from './module/ws-proxy/index.ts';
+
 const port = config?.proxy?.port || 3005;
 
 app
@@ -21,10 +23,16 @@ app.listen(port, () => {
 
 app.server.on(handleRequest);
 
+const wssApp = new WssApp();
 const main = () => {
   console.log('Upgrade initialization started');
 
-  app.server.server.on('upgrade', (req, socket, head) => {
+  app.server.server.on('upgrade', async (req, socket, head) => {
+    const isUpgrade = wssApp.upgrade(req, socket, head);
+    if (isUpgrade) {
+      console.log('WebSocket upgrade successful for path:', req.url);
+      return;
+    }
     const proxyApiList = config?.apiList || [];
     const proxyApi = proxyApiList.find((item) => req.url.startsWith(item.path));
 
@@ -40,12 +48,12 @@ const main = () => {
       const proxySocket = net.connect(options.port, options.hostname, () => {
         proxySocket.write(
           `GET ${options.path} HTTP/1.1\r\n` +
-          `Host: ${options.hostname}\r\n` +
-          `Connection: Upgrade\r\n` +
-          `Upgrade: websocket\r\n` +
-          `Sec-WebSocket-Key: ${req.headers['sec-websocket-key']}\r\n` +
-          `Sec-WebSocket-Version: ${req.headers['sec-websocket-version']}\r\n` +
-          `\r\n`
+            `Host: ${options.hostname}\r\n` +
+            `Connection: Upgrade\r\n` +
+            `Upgrade: websocket\r\n` +
+            `Sec-WebSocket-Key: ${req.headers['sec-websocket-key']}\r\n` +
+            `Sec-WebSocket-Version: ${req.headers['sec-websocket-version']}\r\n` +
+            `\r\n`,
         );
         proxySocket.pipe(socket);
         socket.pipe(proxySocket);

src/module/index.ts

@@ -13,6 +13,7 @@ import { getLoginUser } from '@/middleware/auth.ts';
 import { rediretHome } from './user-home/index.ts';
 import { aiProxy } from './proxy/ai-proxy.ts';
 import { logger } from './logger.ts';
+import { UserV1Proxy } from './ws-proxy/proxy.ts';
 const domain = config?.proxy?.domain;
 const allowedOrigins = config?.proxy?.allowedOrigin || [];
 
@@ -221,8 +222,8 @@ export const handleRequest = async (req: http.IncomingMessage, res: http.ServerR
     res.write('Server Error\n');
     res.end();
   };
-  const createNotFoundPage = async (msg?: string) => {
-    res.writeHead(404, { 'Content-Type': 'text/html; charset=utf-8' });
+  const createNotFoundPage = async (msg?: string, code = 404) => {
+    res.writeHead(code, { 'Content-Type': 'text/html; charset=utf-8' });
     res.write(msg || 'Not Found App\n');
     res.end();
   };
@@ -231,6 +232,11 @@ export const handleRequest = async (req: http.IncomingMessage, res: http.ServerR
       createNotFoundPage,
     });
   }
+  if (user !== 'api' && app === 'v1') {
+    return UserV1Proxy(req, res, {
+      createNotFoundPage,
+    });
+  }
 
   const userApp = new UserApp({ user, app });
   let isExist = await userApp.getExist();

src/module/ws-proxy/index.ts

@@ -0,0 +1,72 @@
+import { WebSocketServer } from 'ws';
+import { nanoid } from 'nanoid';
+import { WsProxyManager } from './manager.ts';
+import { getLoginUser } from '@/middleware/auth.ts';
+export const wsProxyManager = new WsProxyManager();
+
+export const upgrade = async (request: any, socket: any, head: any) => {
+  const req = request as any;
+  const url = new URL(req.url, 'http://localhost');
+  const id = url.searchParams.get('id');
+  if (url.pathname === '/ws/proxy') {
+    console.log('upgrade', request.url, id);
+    wss.handleUpgrade(req, socket, head, (ws) => {
+      // 这里手动触发 connection 事件
+      // @ts-ignore
+      wss.emit('connection', ws, req);
+    });
+    return true;
+  }
+  return false;
+};
+export const wss = new WebSocketServer({
+  noServer: true,
+  path: '/ws/proxy',
+});
+
+wss.on('connection', async (ws, req) => {
+  console.log('connected', req.url);
+  // const user = await getLoginUser(req);
+  // if (!user) {
+  //   ws.send(
+  //     JSON.stringify({
+  //       type: 'error',
+  //       message: 'Invalid authorization',
+  //     }),
+  //   );
+  //   ws.close();
+  //   return;
+  // }
+  const url = new URL(req.url, 'http://localhost');
+  const id = url?.searchParams?.get('id') || nanoid();
+  const user = 'root';
+  wsProxyManager.register(id, { user, ws });
+  ws.send(
+    JSON.stringify({
+      type: 'connected',
+      id,
+    }),
+  );
+  ws.on('message', async (event: Buffer) => {
+    const eventData = event.toString();
+    if (!eventData) {
+      return;
+    }
+    const data = JSON.parse(eventData);
+    console.log('message', data);
+  });
+  ws.on('close', () => {
+    console.log('ws closed');
+    wsProxyManager.unregister(id, user);
+  });
+});
+
+export class WssApp {
+  wss: WebSocketServer;
+  constructor() {
+    this.wss = wss;
+  }
+  upgrade(request: any, socket: any, head: any) {
+    return upgrade(request, socket, head);
+  }
+}

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);
+  }
+}

src/module/ws-proxy/proxy.ts

@@ -0,0 +1,35 @@
+import { IncomingMessage, ServerResponse } from 'http';
+import { wsProxyManager } from './index.ts';
+
+import { App } from '@kevisual/router';
+import { log } from 'console';
+import { logger } from '../logger.ts';
+
+type ProxyOptions = {
+  createNotFoundPage: (msg?: string) => any;
+};
+export const UserV1Proxy = async (req: IncomingMessage, res: ServerResponse, opts?: ProxyOptions) => {
+  const { url } = req;
+  const { pathname } = new URL(url || '', `http://localhost`);
+  const [user, app, userAppKey] = pathname.split('/').slice(1);
+  if (!user || !app || !userAppKey) {
+    opts?.createNotFoundPage?.('应用未启动');
+    return false;
+  }
+  const data = await App.handleRequest(req, res);
+  logger.debug('data', data);
+  const client = wsProxyManager.get(userAppKey, user);
+  const ids = wsProxyManager.getIds();
+  if (!client) {
+    opts?.createNotFoundPage?.(`应用未启动, 未找到应用, ${userAppKey}, ${ids.join(',')}`);
+    return false;
+  }
+  const value = await client.sendData(data);
+  if (value) {
+    res.writeHead(200, { 'Content-Type': 'application/json' });
+    res.end(JSON.stringify(value));
+    return true;
+  }
+  opts?.createNotFoundPage?.('应用未启动');
+  return true;
+};