@midwayjs/ws 3.20.4 → 3.20.9

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,24 @@ 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) => {
+        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);
             });
@@ -77,6 +98,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);
@@ -102,6 +136,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([

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": "3.20.4",
+  "version": "3.20.9",
   "description": "Midway Web Framework for ws",
   "main": "dist/index.js",
   "typings": "index.d.ts",
@@ -39,5 +39,5 @@
   "engines": {
     "node": ">=12"
   },
-  "gitHead": "c3fb65a7ada8829635f3c6af5ef83c65c3a43d79"
+  "gitHead": "d9032e8d591cb4e4997e31d757d7f055078536d6"
 }