feat: add ws upgrade handler (midwayjs#4360)

czy88840616 · web-flow · commit 20a1d945826e · 2025-07-13T20:31:17.000+08:00
* feat: add websocket upgrade handler

* fix: lint
diff --git a/packages/ws/src/framework.ts b/packages/ws/src/framework.ts
@@ -27,6 +27,7 @@ import {
   IMidwayWSConfigurationOptions,
   IMidwayWSContext,
   NextFunction,
+  UpgradeAuthHandler,
 } from './interface';
 import * as WebSocket from 'ws';
 
@@ -39,6 +40,7 @@ export class MidwayWSFramework extends BaseFramework<
   server: http.Server;
   protected heartBeatInterval: NodeJS.Timeout;
   protected connectionMiddlewareManager = this.createMiddlewareManager();
+  protected upgradeAuthHandler: UpgradeAuthHandler | null = null;
 
   configure(): IMidwayWSConfigurationOptions {
     return this.configService.getConfiguration('webSocket');
@@ -61,6 +63,9 @@ export class MidwayWSFramework extends BaseFramework<
       > => {
         return this.getConnectionMiddleware();
       },
+      onWebSocketUpgrade: (handler: UpgradeAuthHandler) => {
+        return this.onWebSocketUpgrade(handler);
+      },
     });
   }
   public app: IMidwayWSApplication;
@@ -82,7 +87,35 @@ export class MidwayWSFramework extends BaseFramework<
       server = this.configurationOptions.server ?? http.createServer();
     }
 
-    server.on('upgrade', (request, socket: any, head: Buffer) => {
+    server.on('upgrade', async (request, socket: any, head: Buffer) => {
+      // check if the upgrade auth handler is set
+      if (this.upgradeAuthHandler) {
+        try {
+          const authResult = await this.upgradeAuthHandler(
+            request,
+            socket,
+            head
+          );
+          if (!authResult) {
+            this.logger.warn(
+              '[midway:ws] WebSocket upgrade authentication failed'
+            );
+            socket.destroy();
+            return;
+          }
+          this.logger.debug(
+            '[midway:ws] WebSocket upgrade authentication passed'
+          );
+        } catch (error) {
+          this.logger.error(
+            '[midway:ws] WebSocket upgrade authentication error:',
+            error
+          );
+          socket.destroy();
+          return;
+        }
+      }
+
       this.app.handleUpgrade(request, socket, head, ws => {
         this.app.emit('connection', ws, request);
       });
@@ -120,6 +153,23 @@ export class MidwayWSFramework extends BaseFramework<
     return MidwayFrameworkType.WS;
   }
 
+  /**
+   * 设置升级前鉴权处理函数
+   * @param handler 鉴权处理函数，传入 null 可以禁用鉴权
+   */
+  public onWebSocketUpgrade(handler: UpgradeAuthHandler | null): void {
+    this.upgradeAuthHandler = handler;
+    if (handler) {
+      this.logger.info(
+        '[midway:ws] WebSocket upgrade authentication handler set'
+      );
+    } else {
+      this.logger.info(
+        '[midway:ws] WebSocket upgrade authentication handler removed'
+      );
+    }
+  }
+
   private async loadMidwayController() {
     // create room
     const controllerModules = listModule(WS_CONTROLLER_KEY);
diff --git a/packages/ws/src/interface.ts b/packages/ws/src/interface.ts
@@ -14,6 +14,7 @@ export type IMidwayWSApplication = IMidwayApplication<IMidwayWSContext, {
     middleware: CommonMiddlewareUnion<Context, NextFunction, undefined>
   ) => void;
   getConnectionMiddleware: ContextMiddlewareManager<Context, NextFunction, undefined>;
+  onWebSocketUpgrade: (handler: UpgradeAuthHandler | null) => void;
 }> & WebSocket.Server;
 
 export type IMidwayWSConfigurationOptions = {
@@ -38,3 +39,12 @@ export type IMidwayWSContext = IMidwayContext<WebSocket & {
 export type Application = IMidwayWSApplication;
 export type NextFunction = BaseNextFunction;
 export interface Context extends IMidwayWSContext {}
+
+/**
+ * WebSocket 升级前鉴权处理函数类型
+ */
+export type UpgradeAuthHandler = (
+  request: IncomingMessage,
+  socket: any,
+  head: Buffer
+) => Promise<boolean>;
diff --git a/packages/ws/test/fixtures/base-app-upgrade-auth/package.json b/packages/ws/test/fixtures/base-app-upgrade-auth/package.json
@@ -0,0 +1,3 @@
+{
+  "name": "base-app-upgrade-auth"
+} 
diff --git a/packages/ws/test/fixtures/base-app-upgrade-auth/src/configuration.ts b/packages/ws/test/fixtures/base-app-upgrade-auth/src/configuration.ts
@@ -0,0 +1,51 @@
+import { Configuration, App } from '@midwayjs/core';
+import { ILifeCycle } from '@midwayjs/core';
+import { Application } from '../../../../src';
+import { UpgradeAuthHandler } from '../../../../src';
+import * as http from 'http';
+
+@Configuration({
+  importConfigs: [
+    {
+      default: {
+        webSocket: {
+          port: 3000,
+        }
+      }
+    }
+  ]
+})
+export class AutoConfiguration implements ILifeCycle {
+
+  @App()
+  app: Application;
+
+  async onReady() {
+    // 设置升级鉴权处理函数
+    this.app.onWebSocketUpgrade(this.authHandler);
+  }
+
+  private authHandler: UpgradeAuthHandler = async (
+    request: http.IncomingMessage,
+    socket: any,
+    head: Buffer
+  ): Promise<boolean> => {
+    try {
+      // 从 URL 参数获取 token
+      const url = new URL(request.url || '', `http://${request.headers.host}`);
+      const token = url.searchParams.get('token');
+      
+      // 简单的 token 验证
+      if (token === 'valid-token') {
+        console.log('[Test Auth] Valid token, connection allowed');
+        return true;
+      }
+      
+      console.log('[Test Auth] Invalid or missing token, connection denied');
+      return false;
+    } catch (error) {
+      console.error('[Test Auth] Authentication error:', error);
+      return false;
+    }
+  };
+} 
diff --git a/packages/ws/test/fixtures/base-app-upgrade-auth/src/socket.ts b/packages/ws/test/fixtures/base-app-upgrade-auth/src/socket.ts
@@ -0,0 +1,26 @@
+import {
+  OnWSConnection,
+  OnWSMessage,
+  Provide,
+  WSController,
+} from '@midwayjs/core';
+
+@Provide()
+@WSController()
+export class HelloSocketController {
+  @OnWSConnection()
+  async onConnectionMethod() {
+    console.log('on connection');
+  }
+
+  @OnWSMessage('message')
+  async onMessage(data: any) {
+    // 处理 Buffer 数据
+    let messageData = data;
+    if (Buffer.isBuffer(data)) {
+      messageData = data.toString();
+    }
+    
+    return { echo: messageData, timestamp: Date.now() };
+  }
+} 
diff --git a/packages/ws/test/index.test.ts b/packages/ws/test/index.test.ts
@@ -1,4 +1,4 @@
-import { closeApp, createServer } from './utils';
+import { closeApp, createServer, testConnectionRejected } from './utils';
 import { sleep } from '@midwayjs/core';
 import { once } from 'events';
 import { createWebSocketClient } from '@midwayjs/mock';
@@ -109,4 +109,31 @@ describe('/test/index.test.ts', () => {
 
     await closeApp(app);
   });
+
+  it('should test onWebSocketUpgrade authentication', async () => {
+    const app = await createServer('base-app-upgrade-auth');
+
+    // 测试1: 没有 token 的连接应该被拒绝
+    const rejected1 = await testConnectionRejected('ws://localhost:3000');
+    expect(rejected1).toBe(true);
+
+    // 测试2: 无效 token 的连接应该被拒绝
+    const rejected2 = await testConnectionRejected('ws://localhost:3000?token=invalid-token');
+    expect(rejected2).toBe(true);
+
+    // 测试3: 有效 token 的连接应该成功
+    const client3 = await createWebSocketClient(`ws://localhost:3000?token=valid-token`);
+
+    // 发送消息测试连接是否正常工作
+    client3.send('test-message');
+    const gotEvent = once(client3, 'message');
+    const [data] = await gotEvent;
+    const response = JSON.parse(data);
+
+    expect(response.echo).toEqual('test-message');
+    expect(response.timestamp).toBeDefined();
+
+    await client3.close();
+    await closeApp(app);
+  });
 });
diff --git a/packages/ws/test/utils.ts b/packages/ws/test/utils.ts
@@ -15,3 +15,50 @@ export async function createServer(name: string, options: IMidwayWSConfiguration
 export async function closeApp(app) {
   return close(app);
 }
+
+/**
+ * 测试 WebSocket 连接是否被拒绝
+ * @param url WebSocket 连接 URL
+ * @param timeout 超时时间（毫秒）
+ * @returns Promise<boolean> true 表示连接被拒绝，false 表示连接成功
+ */
+export async function testConnectionRejected(url: string, timeout: number = 2000): Promise<boolean> {
+  return new Promise((resolve) => {
+    const WebSocket = require('ws');
+    const client = new WebSocket(url);
+    
+    let resolved = false;
+    const timer = setTimeout(() => {
+      if (!resolved) {
+        resolved = true;
+        client.terminate();
+        resolve(true); // 超时认为连接被拒绝
+      }
+    }, timeout);
+
+    client.on('open', () => {
+      if (!resolved) {
+        resolved = true;
+        clearTimeout(timer);
+        client.close();
+        resolve(false); // 连接成功
+      }
+    });
+
+    client.on('error', (error) => {
+      if (!resolved) {
+        resolved = true;
+        clearTimeout(timer);
+        resolve(true); // 连接被拒绝
+      }
+    });
+
+    client.on('close', () => {
+      if (!resolved) {
+        resolved = true;
+        clearTimeout(timer);
+        resolve(true); // 连接被拒绝
+      }
+    });
+  });
+}
diff --git a/site/docs/extensions/ws.md b/site/docs/extensions/ws.md
@@ -276,6 +276,101 @@ const client = new WebSocket('wss://websocket-echo.com/');
 client.on('ping', heartbeat);
 ```
 
+## 鉴权
+
+在 WebSocket 连接建立之前，您可能需要对客户端进行身份验证。从 v3.20.9 开始 Midway 提供了 `onWebSocketUpgrade` 方法来在 WebSocket 握手前进行鉴权。
+
+### 设置鉴权处理器
+
+您可以在应用启动时设置鉴权处理器：
+
+```typescript
+import { Configuration, Inject } from '@midwayjs/core';
+import { MidwayWSFramework } from '@midwayjs/ws';
+
+@Configuration()
+export class WSConfiguration {
+  @Inject()
+  wsFramework: MidwayWSFramework;
+
+  async onReady() {
+    // 设置升级前鉴权处理器
+    this.wsFramework.onWebSocketUpgrade(async (request, socket, head) => {
+      // 从 URL 参数中获取 token
+      const url = new URL(request.url, `http://${request.headers.host}`);
+      const token = url.searchParams.get('token');
+      
+      // 验证 token
+      if (token === 'valid-token') {
+        return true; // 允许连接
+      }
+      
+      return false; // 拒绝连接
+    });
+  }
+}
+```
+
+### 鉴权处理器参数
+
+鉴权处理器接收三个参数：
+
+- `request`: HTTP 请求对象 (`http.IncomingMessage`)
+- `socket`: 原始 socket 对象
+- `head`: WebSocket 握手的头部数据 (`Buffer`)
+
+处理器需要返回一个 `Promise<boolean>`：
+- `true`: 允许 WebSocket 连接
+- `false`: 拒绝 WebSocket 连接
+
+### 获取鉴权信息
+
+您可以从多个来源获取鉴权信息：
+
+**URL 参数**
+
+```typescript
+this.wsFramework.onWebSocketUpgrade(async (request, socket, head) => {
+  const url = new URL(request.url, `http://${request.headers.host}`);
+  const token = url.searchParams.get('token');
+  const userId = url.searchParams.get('userId');
+  
+  // 验证逻辑
+  return await this.validateToken(token, userId);
+});
+```
+
+**请求头**
+
+```typescript
+this.wsFramework.onWebSocketUpgrade(async (request, socket, head) => {
+  const authorization = request.headers.authorization;
+  
+  if (!authorization) {
+    return false;
+  }
+  
+  const token = authorization.replace('Bearer ', '');
+  return await this.validateToken(token);
+});
+```
+
+**Cookie**
+
+```typescript
+this.wsFramework.onWebSocketUpgrade(async (request, socket, head) => {
+  const cookie = request.headers.cookie;
+  
+  if (!cookie) {
+    return false;
+  }
+  
+  // 解析 cookie 获取 session 信息
+  const sessionId = this.parseCookie(cookie).sessionId;
+  return await this.validateSession(sessionId);
+});
+```
+
 
 
 ## 本地测试
diff --git a/site/i18n/en/docusaurus-plugin-content-docs/current/extensions/ws.md b/site/i18n/en/docusaurus-plugin-content-docs/current/extensions/ws.md

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+{`
	`2`	`+ "name": "base-app-upgrade-auth"`
	`3`	`+}`