@midwayjs/ws 4.0.0-beta.1 → 4.0.0-beta.2

package/dist/framework.d.ts

@@ -2,17 +2,23 @@
 /// <reference types="node" />
 import { BaseFramework, CommonMiddlewareUnion, ContextMiddlewareManager, IMidwayBootstrapOptions } 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;
     protected frameworkLoggerName: string;
     configure(): IMidwayWSConfigurationOptions;
     applicationInitialize(options: IMidwayBootstrapOptions): Promise<void>;
     app: IMidwayWSApplication;
     run(): Promise<void>;
     protected beforeStop(): Promise<void>;
+    /**
+     * 设置升级前鉴权处理函数
+     * @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;
         this.frameworkLoggerName = 'wsLogger';
     }
     configure() {
@@ -32,6 +33,9 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
             getConnectionMiddleware: () => {
                 return this.getConnectionMiddleware();
             },
+            onWebSocketUpgrade: (handler) => {
+                return this.onWebSocketUpgrade(handler);
+            },
         });
         await this.loadMidwayController();
     }
@@ -44,7 +48,34 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
         else {
             server = this.configurationOptions.server ?? 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);
             });
@@ -54,13 +85,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 => {
@@ -72,6 +109,19 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
             this.server.close();
         });
     }
+    /**
+     * 设置升级前鉴权处理函数
+     * @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 = core_1.DecoratorManager.listModule(core_1.WS_CONTROLLER_KEY);
@@ -95,6 +145,7 @@ let MidwayWSFramework = class MidwayWSFramework extends core_1.BaseFramework {
             // create request context
             this.app.createAnonymousContext(socket);
             socket.requestContext.registerObject('socket', socket);
+            socket.request = request;
             socket.app = this.app;
             // run connection middleware
             const connectFn = await this.middlewareService.compose([
@@ -182,15 +233,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,8 +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;
@@ -19,9 +24,14 @@ export type IMidwayWSConfigurationOptions = {
 export type IMidwayWSContext = IMidwayContext<WebSocket & {
     app: IMidwayWSApplication;
     isAlive: boolean;
+    request: IncomingMessage;
 }>;
 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": "4.0.0-beta.1",
+  "version": "4.0.0-beta.2",
   "description": "Midway Web Framework for ws",
   "main": "dist/index.js",
   "typings": "index.d.ts",
@@ -23,8 +23,9 @@
   ],
   "license": "MIT",
   "devDependencies": {
-    "@midwayjs/core": "^4.0.0-beta.1",
-    "@midwayjs/mock": "^4.0.0-beta.1",
+    "@midwayjs/core": "^4.0.0-beta.2",
+    "@midwayjs/koa": "^4.0.0-beta.2",
+    "@midwayjs/mock": "^4.0.0-beta.2",
     "fs-extra": "11.3.0"
   },
   "dependencies": {
@@ -37,7 +38,7 @@
     "url": "https://github.com/midwayjs/midway.git"
   },
   "engines": {
-    "node": ">=12"
+    "node": ">=20"
   },
-  "gitHead": "832961ec3aff123c033197d8c00cb2bc9bad7ff8"
+  "gitHead": "53bfef4c5279da5f09025e4610bdbf64f94f60bd"
 }