From b42a152e4654a2fc13c7d98adaddc737eb6346f0 Mon Sep 17 00:00:00 2001
From: Jerry <JerryTech.Y@outlook.com>
Date: Sun, 13 Apr 2025 01:55:12 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8E=A5=E5=85=A5ws=E6=A8=A1=E5=9D=97?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .env                               |  2 +
 .idea/codeStyles/Project.xml       |  7 +--
 package.json                       |  4 +-
 pnpm-lock.yaml                     | 27 +++++++++++
 src/app.ts                         |  1 +
 src/services/ws/handler.ts         | 30 ++++++++++++
 src/services/ws/wsClientManager.ts | 44 ++++++++++++++++++
 src/services/ws/wsServer.ts        | 73 ++++++++++++++++++++++++++++++
 src/types/wsMessage.ts             |  6 +++
 src/utils/redis/serializer.ts      | 34 --------------
 10 files changed, 187 insertions(+), 41 deletions(-)
 create mode 100644 src/services/ws/handler.ts
 create mode 100644 src/services/ws/wsClientManager.ts
 create mode 100644 src/services/ws/wsServer.ts
 create mode 100644 src/types/wsMessage.ts
 delete mode 100644 src/utils/redis/serializer.ts

diff --git a/.env b/.env
index 2e9a2c1..7c035c3 100644
--- a/.env
+++ b/.env
@@ -2,3 +2,5 @@ PORT=3000
 DEBUG=true
 RD_PORT=6379
 RD_ADD=127.0.0.1
+WS_SECRET=114514
+WS_PORT=3001
diff --git a/.idea/codeStyles/Project.xml b/.idea/codeStyles/Project.xml
index 79e151b..196d4e1 100644
--- a/.idea/codeStyles/Project.xml
+++ b/.idea/codeStyles/Project.xml
@@ -14,14 +14,9 @@
     </JSCodeStyleSettings>
     <TypeScriptCodeStyleSettings version="0">
       <option name="FORCE_SEMICOLON_STYLE" value="true" />
-      <option name="FILE_NAME_STYLE" value="CAMEL_CASE" />
       <option name="SPACE_BEFORE_FUNCTION_LEFT_PARENTH" value="false" />
-      <option name="USE_PUBLIC_MODIFIER" value="true" />
       <option name="USE_DOUBLE_QUOTES" value="false" />
       <option name="FORCE_QUOTE_STYlE" value="true" />
-      <option name="PREFER_EXPLICIT_TYPES_VARS_FIELDS" value="true" />
-      <option name="PREFER_EXPLICIT_TYPES_FUNCTION_RETURNS" value="true" />
-      <option name="PREFER_EXPLICIT_TYPES_FUNCTION_EXPRESSION_RETURNS" value="true" />
       <option name="ENFORCE_TRAILING_COMMA" value="WhenMultiline" />
       <option name="SPACES_WITHIN_OBJECT_LITERAL_BRACES" value="true" />
       <option name="SPACES_WITHIN_IMPORTS" value="true" />
@@ -61,4 +56,4 @@
       </indentOptions>
     </codeStyleSettings>
   </code_scheme>
-</component>
\ No newline at end of file
+</component>
diff --git a/package.json b/package.json
index 7b5d7f9..abcaa9f 100644
--- a/package.json
+++ b/package.json
@@ -11,12 +11,14 @@
     "dotenv": "^16.0.0",
     "express": "^4.18.0",
     "ioredis": "^5.6.0",
-    "mkdirp": "^3.0.1"
+    "mkdirp": "^3.0.1",
+    "ws": "^8.18.1"
   },
   "devDependencies": {
     "@types/express": "^4.17.0",
     "@types/mkdirp": "^2.0.0",
     "@types/node": "^18.0.0",
+    "@types/ws": "^8.18.1",
     "prettier": "^3.5.3",
     "ts-node-dev": "^2.0.0",
     "typescript": "^5.0.0"
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 3c2a0da..78e7f01 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -23,6 +23,9 @@ importers:
       mkdirp:
         specifier: ^3.0.1
         version: 3.0.1
+      ws:
+        specifier: ^8.18.1
+        version: 8.18.1
     devDependencies:
       '@types/express':
         specifier: ^4.17.0
@@ -33,6 +36,9 @@ importers:
       '@types/node':
         specifier: ^18.0.0
         version: 18.19.86
+      '@types/ws':
+        specifier: ^8.18.1
+        version: 8.18.1
       prettier:
         specifier: ^3.5.3
         version: 3.5.3
@@ -117,6 +123,9 @@ packages:
   '@types/strip-json-comments@0.0.30':
     resolution: {integrity: sha512-7NQmHra/JILCd1QqpSzl8+mJRc8ZHz3uDm8YV1Ks9IhK0epEiTw8aIErbvH9PI+6XbqhyIQy3462nEsn7UVzjQ==}
 
+  '@types/ws@8.18.1':
+    resolution: {integrity: sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==}
+
   accepts@1.3.8:
     resolution: {integrity: sha512-PYAthTa2m2VKxuvSD3DPC/Gy+U+sOA1LAuT8mkmRuvw+NACSaeXEQ+NHcVF7rONl6qcaxV3Uuemwawk+7+SJLw==}
     engines: {node: '>= 0.6'}
@@ -669,6 +678,18 @@ packages:
   wrappy@1.0.2:
     resolution: {integrity: sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==}
 
+  ws@8.18.1:
+    resolution: {integrity: sha512-RKW2aJZMXeMxVpnZ6bck+RswznaxmzdULiBr6KY7XkTnW8uvt0iT9H5DkHUChXrc+uurzwa0rVI16n/Xzjdz1w==}
+    engines: {node: '>=10.0.0'}
+    peerDependencies:
+      bufferutil: ^4.0.1
+      utf-8-validate: '>=5.0.2'
+    peerDependenciesMeta:
+      bufferutil:
+        optional: true
+      utf-8-validate:
+        optional: true
+
   xtend@4.0.2:
     resolution: {integrity: sha512-LKYU1iAXJXUgAXn9URjiu+MWhyUXHsvfp7mcuYm9dSUKK0/CjtrUwFAxD82/mCWbtLsGjFIad0wIsod4zrTAEQ==}
     engines: {node: '>=0.4'}
@@ -756,6 +777,10 @@ snapshots:
 
   '@types/strip-json-comments@0.0.30': {}
 
+  '@types/ws@8.18.1':
+    dependencies:
+      '@types/node': 18.19.86
+
   accepts@1.3.8:
     dependencies:
       mime-types: 2.1.35
@@ -1316,6 +1341,8 @@ snapshots:
 
   wrappy@1.0.2: {}
 
+  ws@8.18.1: {}
+
   xtend@4.0.2: {}
 
   yn@3.1.1: {}
diff --git a/src/app.ts b/src/app.ts
index 6c5bacc..80b6eb5 100644
--- a/src/app.ts
+++ b/src/app.ts
@@ -4,6 +4,7 @@ import paths from './utils/core/path';
 import sampleController from './modules/sample/sample.controller';
 import imageController from './modules/image/image.controller';
 import config from './utils/core/config';
+import './services/ws/wsServer';
 
 const apps = {
   async createApp() {
diff --git a/src/services/ws/handler.ts b/src/services/ws/handler.ts
new file mode 100644
index 0000000..ffe5831
--- /dev/null
+++ b/src/services/ws/handler.ts
@@ -0,0 +1,30 @@
+import WebSocket from 'ws';
+import WsMessage from '../../types/wsMessage';
+
+class WSMessageHandler {
+  public async handle(socket: WebSocket, clientID: string, msg: WsMessage) {
+    switch (msg.type) {
+      case 'test':
+        await this.reply(socket, { type: 'test', data: 'hi' });
+        break;
+      case 'ping':
+        await this.reply(socket, { type: 'pong' });
+        break;
+      default:
+        await this.reply(socket, { type: 'error', message: 'Unknown message' });
+        break;
+    }
+  }
+
+  private async reply(socket: WebSocket, data: any): Promise<void> {
+    return new Promise((resolve, reject) => {
+      socket.send(JSON.stringify(data), (err) => {
+        if (err) reject(err);
+        else resolve();
+      });
+    });
+  }
+}
+
+const wsHandler = new WSMessageHandler();
+export default wsHandler;
diff --git a/src/services/ws/wsClientManager.ts b/src/services/ws/wsClientManager.ts
new file mode 100644
index 0000000..2d89e46
--- /dev/null
+++ b/src/services/ws/wsClientManager.ts
@@ -0,0 +1,44 @@
+import WebSocket from 'ws';
+
+type ClientID = string;
+
+class WSClientManager {
+  private clients = new Map<ClientID, WebSocket>();
+
+  public add(id: ClientID, socket: WebSocket) {
+    this.clients.set(id, socket);
+  }
+
+  public remove(id: ClientID) {
+    this.clients.delete(id);
+  }
+
+  public get(id: ClientID): WebSocket | undefined {
+    return this.clients.get(id);
+  }
+
+  public async send(id: ClientID, payload: any): Promise<boolean> {
+    const socket = this.clients.get(id);
+    if (!socket || socket.readyState !== WebSocket.OPEN) return false;
+    return this.safeSend(socket, payload);
+  }
+
+  public async broadcast(payload: any): Promise<void> {
+    const tasks = Array.from(this.clients.values()).map((socket) => {
+      socket.readyState === WebSocket.OPEN ? this.safeSend(socket, payload) : Promise.resolve();
+    });
+    await Promise.all(tasks);
+  }
+
+  private async safeSend(socket: WebSocket, data: any): Promise<boolean> {
+    return new Promise((resolve, reject) => {
+      socket.send(JSON.stringify(data), (err) => {
+        if (err) reject(false);
+        else resolve(true);
+      });
+    });
+  }
+}
+
+const wsClientManager = new WSClientManager();
+export default wsClientManager;
diff --git a/src/services/ws/wsServer.ts b/src/services/ws/wsServer.ts
new file mode 100644
index 0000000..268a049
--- /dev/null
+++ b/src/services/ws/wsServer.ts
@@ -0,0 +1,73 @@
+import WebSocket, { WebSocketServer } from 'ws';
+import config from '../../utils/core/config';
+import wsClientManager from './wsClientManager';
+import wsHandler from './handler';
+import logger from '../../utils/core/logger';
+
+interface AuthenticatedSocket extends WebSocket {
+  isAuthed?: boolean;
+  clientId?: string;
+}
+
+class WSServer {
+  private wss: WebSocketServer;
+  private PORT = config.get('WS_PORT');
+  private WS_SECRET = config.get('WS_SECRET');
+
+  constructor() {
+    this.wss = new WebSocketServer({ port: Number(this.PORT) });
+    this.init();
+    logger.info(`WebSocket Server started at ws://localhost:${this.PORT}`);
+  }
+
+  private init() {
+    this.wss.on('connection', (socket: AuthenticatedSocket) => {
+      socket.on('message', async (raw) => {
+        let msg: any;
+        try {
+          msg = JSON.parse(raw.toString());
+        } catch {
+          return this.send(socket, { type: 'error', message: 'JSON 解析失败' });
+        }
+
+        // 鉴权
+        if (!socket.isAuthed) {
+          if (msg.type === 'auth' && msg.secret === this.WS_SECRET && msg.clientId) {
+            socket.isAuthed = true;
+            socket.clientId = msg.clientId;
+            wsClientManager.add(msg.clientId, socket);
+            return this.send(socket, { type: 'auth', success: true });
+          }
+          return this.send(socket, { type: 'auth', success: false });
+        }
+
+        // 业务处理
+        if (socket.clientId) {
+          try {
+            await wsHandler.handle(socket, socket.clientId, msg);
+          } catch (e) {
+            await this.send(socket, { type: 'error', message: '处理出错' });
+          }
+        }
+      });
+
+      socket.on('close', () => {
+        if (socket.clientId) {
+          wsClientManager.remove(socket.clientId);
+        }
+      });
+    });
+  }
+
+  private async send(socket: WebSocket, data: any): Promise<void> {
+    return new Promise((resolve, reject) => {
+      socket.send(JSON.stringify(data), (err) => {
+        if (err) reject(err);
+        else resolve();
+      });
+    });
+  }
+}
+
+const wsServer = new WSServer();
+export default wsServer;
diff --git a/src/types/wsMessage.ts b/src/types/wsMessage.ts
new file mode 100644
index 0000000..fb2af4a
--- /dev/null
+++ b/src/types/wsMessage.ts
@@ -0,0 +1,6 @@
+interface wsMessage {
+  type: string;
+  [key: string]: any;
+}
+
+export default wsMessage;
diff --git a/src/utils/redis/serializer.ts b/src/utils/redis/serializer.ts
deleted file mode 100644
index b80288b..0000000
--- a/src/utils/redis/serializer.ts
+++ /dev/null
@@ -1,34 +0,0 @@
-import logger from '../core/logger';
-
-class RedisSerializer {
-  public serialize<T>(data: T): string {
-    return JSON.stringify(data);
-  }
-
-  public deserialize<T>(jsonString: string): T | undefined {
-    try {
-      return JSON.parse(jsonString) as T;
-    } catch (err) {
-      logger.error(`Redis反序列化失败: ${err}`);
-    }
-  }
-
-  /**
-   * 处理Date类型
-   */
-  public reviveDates<T>(obj: T): T {
-    const dateRegex = /^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}.\d{3}Z$/;
-
-    const reviver = (_: string, value: any) => {
-      if (typeof value === 'string' && dateRegex.test(value)) {
-        return new Date(value);
-      }
-      return value;
-    };
-
-    return JSON.parse(JSON.stringify(obj), reviver);
-  }
-}
-
-const serializer = new RedisSerializer();
-export default serializer;