@midwayjs/ws 3.20.8 → 3.20.10

package/dist/framework.d.ts

@@ -2,11 +2,12 @@
 /// <reference types="node" />
 import { BaseFramework, CommonMiddlewareUnion, ContextMiddlewareManager, IMidwayBootstrapOptions, MidwayFrameworkType } from '@midwayjs/core';
 import * as http from 'http';
-import { Application, Context, IMidwayWSApplication, IMidwayWSConfigurationOptions, NextFunction } from './interface';
+import { Application, Context, IMidwayWSApplication, IMidwayWSConfigurationOptions, NextFunction, UpgradeAuthHandler } from './interface';
 export declare class MidwayWSFramework extends BaseFramework<Application, Context, IMidwayWSConfigurationOptions> {
     server: http.Server;
     protected heartBeatInterval: NodeJS.Timeout;
     protected connectionMiddlewareManager: ContextMiddlewareManager<Context, unknown, unknown>;
+    protected upgradeAuthHandler: UpgradeAuthHandler | null;
     configure(): IMidwayWSConfigurationOptions;
     applicationInitialize(options: IMidwayBootstrapOptions): void;
     app: IMidwayWSApplication;
@@ -14,6 +15,11 @@ export declare class MidwayWSFramework extends BaseFramework<Application, Contex
     run(): Promise<void>;
     protected beforeStop(): Promise<void>;
     getFrameworkType(): MidwayFrameworkType;
+    /**
+     * 设置升级前鉴权处理函数
+     * @param handler 鉴权处理函数，传入 null 可以禁用鉴权
+     */
+    onWebSocketUpgrade(handler: UpgradeAuthHandler | null): void;
     private loadMidwayController;
     private addNamespace;
     private bindSocketResponse;

package/dist/framework.js

@@ -16,6 +16,7 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
     constructor() {
         super(...arguments);
         this.connectionMiddlewareManager = this.createMiddlewareManager();
+        this.upgradeAuthHandler = null;
     }
     configure() {
         return this.configService.getConfiguration('webSocket');
@@ -31,6 +32,9 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
             getConnectionMiddleware: () => {
                 return this.getConnectionMiddleware();
             },
+            onWebSocketUpgrade: (handler) => {
+                return this.onWebSocketUpgrade(handler);
+            },
         });
     }
     async afterContainerReady(options) {
@@ -46,7 +50,34 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
         else {
             server = (_a = this.configurationOptions.server) !== null && _a !== void 0 ? _a : http.createServer();
         }
-        server.on('upgrade', (request, socket, head) => {
+        if (this.configurationOptions.enableServerHeartbeatCheck) {
+            if (server.listening) {
+                this.startHeartBeat();
+            }
+            else {
+                server.on('listening', () => {
+                    this.startHeartBeat();
+                });
+            }
+        }
+        server.on('upgrade', async (request, socket, head) => {
+            // 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);
             });
@@ -56,13 +87,19 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
             await new Promise(resolve => {
                 server.listen(this.configurationOptions.port, () => {
                     this.logger.info(`[midway:ws] WebSocket server port = ${this.configurationOptions.port} start success.`);
-                    if (this.configurationOptions.enableServerHeartbeatCheck) {
-                        this.startHeartBeat();
-                    }
                     resolve();
                 });
             });
         }
+        this.app.on('error', err => {
+            this.logger.error('socket server got error', err);
+        });
+        this.app.on('close', () => {
+            if (this.heartBeatInterval) {
+                clearInterval(this.heartBeatInterval);
+            }
+            this.logger.info('socket server close');
+        });
     }
     async beforeStop() {
         return new Promise(resolve => {
@@ -77,6 +114,19 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
     getFrameworkType() {
         return core_1.MidwayFrameworkType.WS;
     }
+    /**
+     * 设置升级前鉴权处理函数
+     * @param handler 鉴权处理函数，传入 null 可以禁用鉴权
+     */
+    onWebSocketUpgrade(handler) {
+        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');
+        }
+    }
     async loadMidwayController() {
         // create room
         const controllerModules = (0, core_1.listModule)(core_1.WS_CONTROLLER_KEY);
@@ -191,15 +241,6 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
                 }
             }
         });
-        this.app.on('error', err => {
-            this.logger.error('socket server got error', err);
-        });
-        this.app.on('close', () => {
-            if (this.heartBeatInterval) {
-                clearInterval(this.heartBeatInterval);
-            }
-            this.logger.info('socket server close');
-        });
     }
     async bindSocketResponse(result, socket, propertyName, methodMap) {
         if (!result)

package/dist/interface.d.ts

@@ -1,10 +1,13 @@
 /// <reference types="node" />
+/// <reference types="node" />
+/// <reference types="node" />
 import * as WebSocket from 'ws';
 import { CommonMiddlewareUnion, ContextMiddlewareManager, IConfigurationOptions, IMidwayApplication, IMidwayContext, NextFunction as BaseNextFunction } from '@midwayjs/core';
 import type { IncomingMessage } from 'http';
 export type IMidwayWSApplication = IMidwayApplication<IMidwayWSContext, {
     useConnectionMiddleware: (middleware: CommonMiddlewareUnion<Context, NextFunction, undefined>) => void;
     getConnectionMiddleware: ContextMiddlewareManager<Context, NextFunction, undefined>;
+    onWebSocketUpgrade: (handler: UpgradeAuthHandler | null) => void;
 }> & WebSocket.Server;
 export type IMidwayWSConfigurationOptions = {
     pubClient?: any;
@@ -27,4 +30,8 @@ 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>;
 //# sourceMappingURL=interface.d.ts.map

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@midwayjs/ws",
-  "version": "3.20.8",
+  "version": "3.20.10",
   "description": "Midway Web Framework for ws",
   "main": "dist/index.js",
   "typings": "index.d.ts",
@@ -24,6 +24,7 @@
   "license": "MIT",
   "devDependencies": {
     "@midwayjs/core": "^3.20.4",
+    "@midwayjs/koa": "^3.20.4",
     "@midwayjs/mock": "^3.20.4",
     "fs-extra": "11.3.0"
   },
@@ -39,5 +40,5 @@
   "engines": {
     "node": ">=12"
   },
-  "gitHead": "561d0e178875f28fcb98efaf6570e981b17b3870"
+  "gitHead": "60ef5c3dd5ebb2ee79388a4b942f07a9ef70b854"
 }