From 4e2f6a73bc9d68701869cde0ef48dbebb3b741c6 Mon Sep 17 00:00:00 2001 From: akitasummer Date: Thu, 9 Jul 2026 16:37:11 +0800 Subject: [PATCH 1/9] feat(controller): add websocket controller support --- core/controller-decorator/index.ts | 5 + .../src/decorator/http/HTTPParam.ts | 60 ++- .../websocket/WebSocketController.ts | 26 ++ .../decorator/websocket/WebSocketMethod.ts | 19 + .../src/decorator/websocket/WebSocketParam.ts | 84 ++++ .../WebSocketControllerMetaBuilder.ts | 65 +++ .../WebSocketControllerMethodMetaBuilder.ts | 98 +++++ .../src/model/WebSocketControllerMeta.ts | 61 +++ .../src/model/WebSocketMethodMeta.ts | 149 +++++++ core/controller-decorator/src/model/index.ts | 2 + .../src/util/WebSocketInfoUtil.ts | 78 ++++ .../types/controller-decorator/MetadataKey.ts | 6 + .../controller-decorator/WebSocketContext.ts | 20 + .../WebSocketController.ts | 5 + .../controller-decorator/WebSocketMethod.ts | 4 + .../controller-decorator/WebSocketParam.ts | 11 + core/types/controller-decorator/index.ts | 4 + .../types/controller-decorator/model/types.ts | 12 + plugin/controller/app.ts | 9 + .../lib/impl/websocket/WebSocketContext.ts | 35 ++ .../websocket/WebSocketControllerRegister.ts | 406 ++++++++++++++++++ plugin/controller/package.json | 2 + .../app/controller/WebSocketController.ts | 94 ++++ .../test/websocket/websocket.test.ts | 146 +++++++ .../test/websocket/websocketCluster.test.ts | 126 ++++++ 25 files changed, 1526 insertions(+), 1 deletion(-) create mode 100644 core/controller-decorator/src/decorator/websocket/WebSocketController.ts create mode 100644 core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts create mode 100644 core/controller-decorator/src/decorator/websocket/WebSocketParam.ts create mode 100644 core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts create mode 100644 core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts create mode 100644 core/controller-decorator/src/model/WebSocketControllerMeta.ts create mode 100644 core/controller-decorator/src/model/WebSocketMethodMeta.ts create mode 100644 core/controller-decorator/src/util/WebSocketInfoUtil.ts create mode 100644 core/types/controller-decorator/WebSocketContext.ts create mode 100644 core/types/controller-decorator/WebSocketController.ts create mode 100644 core/types/controller-decorator/WebSocketMethod.ts create mode 100644 core/types/controller-decorator/WebSocketParam.ts create mode 100644 plugin/controller/lib/impl/websocket/WebSocketContext.ts create mode 100644 plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts create mode 100644 plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts create mode 100644 plugin/controller/test/websocket/websocket.test.ts create mode 100644 plugin/controller/test/websocket/websocketCluster.test.ts diff --git a/core/controller-decorator/index.ts b/core/controller-decorator/index.ts index 03cd745ac..893b336d4 100644 --- a/core/controller-decorator/index.ts +++ b/core/controller-decorator/index.ts @@ -1,4 +1,5 @@ import './src/impl/http/HTTPControllerMetaBuilder'; +import './src/impl/websocket/WebSocketControllerMetaBuilder'; import './src/impl/mcp/MCPControllerMetaBuilder'; export * from '@eggjs/tegg-types/controller-decorator'; @@ -10,6 +11,9 @@ export * from './src/decorator/http/HTTPController'; export * from './src/decorator/http/HTTPMethod'; export * from './src/decorator/http/HTTPParam'; export * from './src/decorator/http/Host'; +export * from './src/decorator/websocket/WebSocketController'; +export * from './src/decorator/websocket/WebSocketMethod'; +export * from './src/decorator/websocket/WebSocketParam'; export * from './src/decorator/mcp/MCPController'; export * from './src/decorator/mcp/MCPPrompt'; export * from './src/decorator/mcp/MCPResource'; @@ -19,6 +23,7 @@ export * from './src/builder/ControllerMetaBuilderFactory'; export * from './src/util/ControllerMetadataUtil'; export * from './src/util/MCPInfoUtil'; export * from './src/util/HTTPPriorityUtil'; +export { default as WebSocketInfoUtil } from './src/util/WebSocketInfoUtil'; export { default as ControllerInfoUtil } from './src/util/ControllerInfoUtil'; export { default as MethodInfoUtil } from './src/util/MethodInfoUtil'; diff --git a/core/controller-decorator/src/decorator/http/HTTPParam.ts b/core/controller-decorator/src/decorator/http/HTTPParam.ts index 4161decf5..e74dbf30e 100644 --- a/core/controller-decorator/src/decorator/http/HTTPParam.ts +++ b/core/controller-decorator/src/decorator/http/HTTPParam.ts @@ -1,8 +1,9 @@ import assert from 'node:assert'; -import { HTTPParamType } from '@eggjs/tegg-types'; +import { HTTPParamType, WebSocketParamType } from '@eggjs/tegg-types'; import type { EggProtoImplClass, HTTPParamParams, HTTPQueriesParams, HTTPQueryParams } from '@eggjs/tegg-types'; import HTTPInfoUtil from '../../util/HTTPInfoUtil'; import { ObjectUtils } from '@eggjs/tegg-common-util'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; // TODO url params // /foo/:id @@ -28,6 +29,17 @@ export function HTTPHeaders() { }; } +export function Headers() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.HEADERS, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.HEADERS, parameterIndex, controllerClazz, methodName); + }; +} + export function HTTPQuery(param?: HTTPQueryParams) { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { assert(typeof propertyKey === 'string', @@ -41,6 +53,21 @@ export function HTTPQuery(param?: HTTPQueryParams) { }; } +export function Query(param?: HTTPQueryParams) { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); + const name = param?.name || argNames[parameterIndex]; + HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.QUERY, parameterIndex, controllerClazz, methodName); + HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.QUERY, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); + }; +} + export function HTTPQueries(param?: HTTPQueriesParams) { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { assert(typeof propertyKey === 'string', @@ -54,6 +81,21 @@ export function HTTPQueries(param?: HTTPQueriesParams) { }; } +export function Queries(param?: HTTPQueriesParams) { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); + const name = param?.name || argNames[parameterIndex]; + HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.QUERIES, parameterIndex, controllerClazz, methodName); + HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.QUERIES, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); + }; +} + export function HTTPParam(param?: HTTPParamParams) { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { assert(typeof propertyKey === 'string', @@ -67,6 +109,21 @@ export function HTTPParam(param?: HTTPParamParams) { }; } +export function Param(param?: HTTPParamParams) { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); + const name = param?.name || argNames[parameterIndex]; + HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.PARAM, parameterIndex, controllerClazz, methodName); + HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.PARAM, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); + }; +} + export function Request() { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { const [ nodeMajor ] = process.versions.node.split('.').map(v => Number(v)); @@ -77,6 +134,7 @@ export function Request() { const methodName = propertyKey as string; const controllerClazz = target.constructor as EggProtoImplClass; HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.REQUEST, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.REQUEST, parameterIndex, controllerClazz, methodName); }; } diff --git a/core/controller-decorator/src/decorator/websocket/WebSocketController.ts b/core/controller-decorator/src/decorator/websocket/WebSocketController.ts new file mode 100644 index 000000000..1672a66e7 --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket/WebSocketController.ts @@ -0,0 +1,26 @@ +import { PrototypeUtil, SingletonProto } from '@eggjs/core-decorator'; +import { StackUtil } from '@eggjs/tegg-common-util'; +import type { EggProtoImplClass, WebSocketControllerParams } from '@eggjs/tegg-types'; +import { AccessLevel, ControllerType } from '@eggjs/tegg-types'; +import ControllerInfoUtil from '../../util/ControllerInfoUtil'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +export function WebSocketController(param?: WebSocketControllerParams) { + return function(constructor: EggProtoImplClass) { + ControllerInfoUtil.setControllerType(constructor, ControllerType.WEBSOCKET); + if (param?.controllerName) { + ControllerInfoUtil.setControllerName(constructor, param.controllerName); + } + if (param?.path) { + WebSocketInfoUtil.setWebSocketPath(param.path, constructor); + } + + const func = SingletonProto({ + accessLevel: AccessLevel.PUBLIC, + name: param?.protoName, + }); + func(constructor); + + PrototypeUtil.setFilePath(constructor, StackUtil.getCalleeFromStack(false, 5)); + }; +} diff --git a/core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts b/core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts new file mode 100644 index 000000000..b2fbc6d99 --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts @@ -0,0 +1,19 @@ +import assert from 'node:assert'; +import { ControllerType } from '@eggjs/tegg-types'; +import type { EggProtoImplClass, WebSocketMethodParams } from '@eggjs/tegg-types'; +import MethodInfoUtil from '../../util/MethodInfoUtil'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +export function WebSocketMethod(param: WebSocketMethodParams) { + return function(target: any, propertyKey: PropertyKey) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const controllerClazz = target.constructor as EggProtoImplClass; + const methodName = propertyKey as string; + MethodInfoUtil.setMethodControllerType(controllerClazz, methodName, ControllerType.WEBSOCKET); + WebSocketInfoUtil.setWebSocketMethodPath(param.path, controllerClazz, methodName); + if (param.priority !== undefined) { + WebSocketInfoUtil.setWebSocketMethodPriority(param.priority, controllerClazz, methodName); + } + }; +} diff --git a/core/controller-decorator/src/decorator/websocket/WebSocketParam.ts b/core/controller-decorator/src/decorator/websocket/WebSocketParam.ts new file mode 100644 index 000000000..510785c1e --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket/WebSocketParam.ts @@ -0,0 +1,84 @@ +import assert from 'node:assert'; +import { WebSocketParamType } from '@eggjs/tegg-types'; +import type { EggProtoImplClass, WebSocketParamParams, WebSocketQueriesParams, WebSocketQueryParams } from '@eggjs/tegg-types'; +import { ObjectUtils } from '@eggjs/tegg-common-util'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +export function WebSocketHeaders() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.HEADERS, parameterIndex, controllerClazz, methodName); + }; +} + +export function WebSocketQuery(param?: WebSocketQueryParams) { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); + const name = param?.name || argNames[parameterIndex]; + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.QUERY, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); + }; +} + +export function WebSocketQueries(param?: WebSocketQueriesParams) { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); + const name = param?.name || argNames[parameterIndex]; + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.QUERIES, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); + }; +} + +export function WebSocketParam(param?: WebSocketParamParams) { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); + const name = param?.name || argNames[parameterIndex]; + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.PARAM, parameterIndex, controllerClazz, methodName); + WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); + }; +} + +export function WebSocketRequest() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.REQUEST, parameterIndex, controllerClazz, methodName); + }; +} + +export function WebSocketSocket() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.SOCKET, parameterIndex, controllerClazz, methodName); + }; +} + +export function WebSocketStream() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.STREAM, parameterIndex, controllerClazz, methodName); + }; +} diff --git a/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts b/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts new file mode 100644 index 000000000..5aef7a594 --- /dev/null +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts @@ -0,0 +1,65 @@ +import assert from 'node:assert'; +import { PrototypeUtil } from '@eggjs/core-decorator'; +import { ObjectUtils } from '@eggjs/tegg-common-util'; +import { ClassUtil } from '@eggjs/tegg-metadata'; +import type { EggProtoImplClass } from '@eggjs/tegg-types'; +import { ControllerType } from '@eggjs/tegg-types'; +import { ControllerMetaBuilderFactory } from '../../builder/ControllerMetaBuilderFactory'; +import { WebSocketControllerMeta, WebSocketMethodMeta } from '../../model'; +import ControllerInfoUtil from '../../util/ControllerInfoUtil'; +import { ControllerMetadataUtil } from '../../util/ControllerMetadataUtil'; +import { ControllerValidator } from '../../util/validator/ControllerValidator'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; +import { WebSocketControllerMethodMetaBuilder } from './WebSocketControllerMethodMetaBuilder'; + +export class WebSocketControllerMetaBuilder { + private readonly clazz: EggProtoImplClass; + + constructor(clazz: EggProtoImplClass) { + this.clazz = clazz; + } + + private buildMethod(): WebSocketMethodMeta[] { + const methodNames = ObjectUtils.getProperties(this.clazz.prototype); + const methods: WebSocketMethodMeta[] = []; + for (const methodName of methodNames) { + const builder = new WebSocketControllerMethodMetaBuilder(this.clazz, methodName); + const methodMeta = builder.build(); + if (methodMeta) { + methods.push(methodMeta); + } + } + return methods; + } + + build(): WebSocketControllerMeta { + ControllerValidator.validate(this.clazz); + const controllerType = ControllerInfoUtil.getControllerType(this.clazz); + assert(controllerType === ControllerType.WEBSOCKET, 'invalidate controller type'); + const webSocketPath = WebSocketInfoUtil.getWebSocketPath(this.clazz); + const middlewares = ControllerInfoUtil.getControllerMiddlewares(this.clazz); + const methods = this.buildMethod(); + const clazzName = this.clazz.name; + const controllerName = ControllerInfoUtil.getControllerName(this.clazz) || clazzName; + const property = PrototypeUtil.getProperty(this.clazz); + const protoName = property!.name as string; + const hosts = ControllerInfoUtil.getControllerHosts(this.clazz); + const metadata = new WebSocketControllerMeta( + clazzName, protoName, controllerName, webSocketPath, middlewares, methods, hosts); + ControllerMetadataUtil.setControllerMetadata(this.clazz, metadata); + for (const method of metadata.methods) { + const realPath = metadata.getMethodRealPath(method); + if (!realPath.startsWith('/')) { + const desc = ClassUtil.classDescription(this.clazz); + throw new Error(`class ${desc} method ${method.name} path ${realPath} not start with /`); + } + } + return metadata; + } + + static create(clazz: EggProtoImplClass) { + return new WebSocketControllerMetaBuilder(clazz); + } +} + +ControllerMetaBuilderFactory.registerControllerMetaBuilder(ControllerType.WEBSOCKET, WebSocketControllerMetaBuilder.create); diff --git a/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts new file mode 100644 index 000000000..e73d756d6 --- /dev/null +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts @@ -0,0 +1,98 @@ +import path from 'node:path'; +import { ClassUtil } from '@eggjs/tegg-metadata'; +import type { EggProtoImplClass } from '@eggjs/tegg-types'; +import { WebSocketMethodMeta, WebSocketParamMeta, WebSocketParamMetaUtil } from '../../model'; +import { MethodValidator } from '../../util/validator/MethodValidator'; +import MethodInfoUtil from '../../util/MethodInfoUtil'; +import { HTTPPriorityUtil } from '../../util/HTTPPriorityUtil'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +export class WebSocketControllerMethodMetaBuilder { + private readonly clazz: EggProtoImplClass; + private readonly methodName: string; + + constructor(clazz: EggProtoImplClass, methodName: string) { + this.clazz = clazz; + this.methodName = methodName; + } + + private checkParamDecorators() { + const method = this.clazz.prototype[this.methodName]; + const functionLength = method.length; + const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); + const contextIndex = MethodInfoUtil.getMethodContextIndex(this.clazz, this.methodName); + const hasAnnotationParamCount = typeof contextIndex === 'undefined' + ? paramIndexList.length + : paramIndexList.length + 1; + const maxParamCount = Math.max(functionLength, hasAnnotationParamCount); + + for (let i = 0; i < maxParamCount; ++i) { + if (i === contextIndex) { + continue; + } + const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(i, this.clazz, this.methodName); + if (!paramType) { + const classDesc = ClassUtil.classDescription(this.clazz); + throw new Error(`${classDesc}:${this.methodName} param ${i} has no websocket param type, Please add @WebSocketParam, @WebSocketQuery, @WebSocketQueries, @WebSocketHeaders, @WebSocketRequest, @WebSocketSocket, @WebSocketStream or @Context`); + } + } + } + + private buildParamType(webSocketPath: string): Map { + this.checkParamDecorators(); + + const paramTypeMap = new Map(); + const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); + for (const paramIndex of paramIndexList) { + const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(paramIndex, this.clazz, this.methodName)!; + const paramName = WebSocketInfoUtil.getWebSocketMethodParamName(paramIndex, this.clazz, this.methodName); + const paramMeta = WebSocketParamMetaUtil.createParam(paramType, paramName); + + try { + paramMeta.validate(webSocketPath); + } catch (e) { + const classDesc = ClassUtil.classDescription(this.clazz); + e.message = `build websocket controller ${classDesc} method ${this.methodName} param ${paramName} failed: ${e.message}`; + throw e; + } + + paramTypeMap.set(paramIndex, paramMeta); + } + return paramTypeMap; + } + + getPriority() { + const priority = WebSocketInfoUtil.getWebSocketMethodPriority(this.clazz, this.methodName); + if (priority !== undefined) { + return priority; + } + const controllerPath = WebSocketInfoUtil.getWebSocketPath(this.clazz); + const methodPath = WebSocketInfoUtil.getWebSocketMethodPath(this.clazz, this.methodName)!; + const realPath = controllerPath ? path.posix.join(controllerPath, methodPath) : methodPath; + const defaultPriority = HTTPPriorityUtil.calcPathPriority(realPath); + if (defaultPriority > HTTPPriorityUtil.DEFAULT_PRIORITY) { + throw new Error(`path ${realPath} is too long, should set priority manually`); + } + return defaultPriority; + } + + build(): WebSocketMethodMeta | undefined { + MethodValidator.validate(this.clazz, this.methodName); + const controllerType = MethodInfoUtil.getMethodControllerType(this.clazz, this.methodName); + if (!controllerType) { + return undefined; + } + const parentPath = WebSocketInfoUtil.getWebSocketPath(this.clazz); + const webSocketPath = WebSocketInfoUtil.getWebSocketMethodPath(this.clazz, this.methodName)!; + const contextIndex = MethodInfoUtil.getMethodContextIndex(this.clazz, this.methodName); + const middlewares = MethodInfoUtil.getMethodMiddlewares(this.clazz, this.methodName); + const hosts = MethodInfoUtil.getMethodHosts(this.clazz, this.methodName); + const realPath = parentPath + ? path.posix.join(parentPath, webSocketPath) + : webSocketPath; + const paramTypeMap = this.buildParamType(realPath); + const priority = this.getPriority(); + return new WebSocketMethodMeta( + this.methodName, webSocketPath, middlewares, contextIndex, paramTypeMap, priority, hosts); + } +} diff --git a/core/controller-decorator/src/model/WebSocketControllerMeta.ts b/core/controller-decorator/src/model/WebSocketControllerMeta.ts new file mode 100644 index 000000000..c54535b6a --- /dev/null +++ b/core/controller-decorator/src/model/WebSocketControllerMeta.ts @@ -0,0 +1,61 @@ +import path from 'node:path'; +import type { ControllerMetadata, EggPrototypeName, MiddlewareFunc } from '@eggjs/tegg-types'; +import { ControllerType } from '@eggjs/tegg-types'; +import { WebSocketMethodMeta } from './WebSocketMethodMeta'; + +export class WebSocketControllerMeta implements ControllerMetadata { + readonly protoName: EggPrototypeName; + readonly controllerName: string; + readonly className: string; + public readonly type = ControllerType.WEBSOCKET; + public readonly path?: string; + public readonly middlewares: readonly MiddlewareFunc[]; + public readonly methods: readonly WebSocketMethodMeta[]; + public readonly hosts?: string[]; + + constructor( + className: string, + protoName: EggPrototypeName, + controllerName: string, + path: string | undefined, + middlewares: MiddlewareFunc[], + methods: WebSocketMethodMeta[], + hosts: string[] | undefined, + ) { + this.protoName = protoName; + this.controllerName = controllerName; + this.className = className; + this.path = path; + this.middlewares = middlewares; + this.methods = methods; + this.hosts = hosts; + } + + getMethodRealPath(method: WebSocketMethodMeta) { + if (this.path) { + return path.posix.join(this.path, method.path); + } + return method.path; + } + + getMethodHosts(method: WebSocketMethodMeta): string[] | undefined { + if (this.hosts) { + return this.hosts; + } + return method.hosts; + } + + getMethodName(method: WebSocketMethodMeta) { + return `WEBSOCKET ${this.controllerName}.${method.name}`; + } + + getMethodMiddlewares(method: WebSocketMethodMeta) { + if (this.middlewares.length) { + return [ + ...this.middlewares, + ...method.middlewares, + ]; + } + return [ ...method.middlewares ]; + } +} diff --git a/core/controller-decorator/src/model/WebSocketMethodMeta.ts b/core/controller-decorator/src/model/WebSocketMethodMeta.ts new file mode 100644 index 000000000..8fe7d6827 --- /dev/null +++ b/core/controller-decorator/src/model/WebSocketMethodMeta.ts @@ -0,0 +1,149 @@ +import assert from 'node:assert'; +import pathToRegexp from 'path-to-regexp'; +import { WebSocketParamType } from '@eggjs/tegg-types'; +import type { MethodMeta, MiddlewareFunc } from '@eggjs/tegg-types'; + +export abstract class WebSocketParamMeta { + type: WebSocketParamType; + + abstract validate(webSocketPath: string); +} + +export class WebSocketPathParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.PARAM; + name: string; + + constructor(name: string) { + super(); + this.name = name; + } + + validate(webSocketPath: string) { + const names: pathToRegexp.Key[] = []; + pathToRegexp(webSocketPath, names); + if (!names.find(name => String(name.name) === this.name)) { + throw new Error(`can not find param ${this.name} in path ${webSocketPath}`); + } + } +} + +export class WebSocketQueryParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.QUERY; + name: string; + + constructor(name: string) { + super(); + this.name = name; + } + + validate() { + return; + } +} + +export class WebSocketQueriesParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.QUERIES; + name: string; + + constructor(name: string) { + super(); + this.name = name; + } + + validate() { + return; + } +} + +export class WebSocketHeadersParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.HEADERS; + + validate() { + return; + } +} + +export class WebSocketRequestParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.REQUEST; + + validate() { + return; + } +} + +export class WebSocketSocketParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.SOCKET; + + validate() { + return; + } +} + +export class WebSocketStreamParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.STREAM; + + validate() { + return; + } +} + +export class WebSocketMethodMeta implements MethodMeta { + public readonly name: string; + public readonly path: string; + public readonly middlewares: readonly MiddlewareFunc[]; + public readonly contextParamIndex: number | undefined; + public readonly paramMap: Map; + public readonly priority: number; + public readonly hosts: string[] | undefined; + + constructor( + name: string, + path: string, + middlewares: MiddlewareFunc[], + contextParamIndex: number | undefined, + paramMap: Map, + priority: number, + hosts: string[] | undefined, + ) { + this.name = name; + this.path = path; + this.middlewares = middlewares; + this.contextParamIndex = contextParamIndex; + this.paramMap = paramMap; + this.priority = priority; + this.hosts = hosts; + } +} + +export class WebSocketParamMetaUtil { + static createParam(type: WebSocketParamType, name?: string) { + switch (type) { + case WebSocketParamType.PARAM: { + assert(name, 'websocket path param must has name'); + return new WebSocketPathParamMeta(name!); + } + case WebSocketParamType.QUERY: { + assert(name, 'websocket query param must has name'); + return new WebSocketQueryParamMeta(name!); + } + case WebSocketParamType.QUERIES: { + assert(name, 'websocket queries param must has name'); + return new WebSocketQueriesParamMeta(name!); + } + case WebSocketParamType.HEADERS: { + return new WebSocketHeadersParamMeta(); + } + case WebSocketParamType.REQUEST: { + return new WebSocketRequestParamMeta(); + } + case WebSocketParamType.SOCKET: { + return new WebSocketSocketParamMeta(); + } + case WebSocketParamType.STREAM: { + return new WebSocketStreamParamMeta(); + } + default: + assert.fail('never arrive'); + } + } +} diff --git a/core/controller-decorator/src/model/index.ts b/core/controller-decorator/src/model/index.ts index 4d355ee7f..f64e67dc0 100644 --- a/core/controller-decorator/src/model/index.ts +++ b/core/controller-decorator/src/model/index.ts @@ -3,6 +3,8 @@ export * from './HTTPControllerMeta'; export * from './HTTPRequest'; export * from './HTTPResponse'; export * from './HTTPCookies'; +export * from './WebSocketMethodMeta'; +export * from './WebSocketControllerMeta'; export * from './MCPControllerMeta'; export * from './MCPPromptMeta'; export * from './MCPResourceMeta'; diff --git a/core/controller-decorator/src/util/WebSocketInfoUtil.ts b/core/controller-decorator/src/util/WebSocketInfoUtil.ts new file mode 100644 index 000000000..e8f9babdb --- /dev/null +++ b/core/controller-decorator/src/util/WebSocketInfoUtil.ts @@ -0,0 +1,78 @@ +import { MetadataUtil } from '@eggjs/core-decorator'; +import { + CONTROLLER_WEBSOCKET_METHOD_PARAM_NAME_MAP, + CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP, + CONTROLLER_WEBSOCKET_METHOD_PATH_MAP, + CONTROLLER_WEBSOCKET_METHOD_PRIORITY, + CONTROLLER_WEBSOCKET_PATH, +} from '@eggjs/tegg-types'; +import type { EggProtoImplClass, WebSocketParamType } from '@eggjs/tegg-types'; +import { MapUtil } from '@eggjs/tegg-common-util'; + +type WebSocketMethodPathMap = Map; +type WebSocketMethodParamTypeMap = Map>; +type WebSocketMethodParamNameMap = Map>; +type WebSocketMethodPriorityMap = Map; + +export default class WebSocketInfoUtil { + static setWebSocketPath(path: string, clazz: EggProtoImplClass) { + MetadataUtil.defineMetaData(CONTROLLER_WEBSOCKET_PATH, path, clazz); + } + + static getWebSocketPath(clazz: EggProtoImplClass): string | undefined { + return MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_PATH, clazz); + } + + static setWebSocketMethodPath(path: string, clazz: EggProtoImplClass, methodName: string) { + const methodPathMap = MetadataUtil.initOwnMapMetaData(CONTROLLER_WEBSOCKET_METHOD_PATH_MAP, clazz, new Map()); + methodPathMap.set(methodName, path); + } + + static getWebSocketMethodPath(clazz: EggProtoImplClass, methodName: string): string | undefined { + const methodPathMap: WebSocketMethodPathMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_METHOD_PATH_MAP, clazz); + return methodPathMap?.get(methodName); + } + + static setWebSocketMethodParamType(paramType: WebSocketParamType, parameterIndex: number, clazz: EggProtoImplClass, methodName: string) { + const methodParamMap: WebSocketMethodParamTypeMap = MetadataUtil.initOwnMapMetaData(CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP, clazz, new Map()); + const paramMap = MapUtil.getOrStore(methodParamMap, methodName, new Map()); + paramMap.set(parameterIndex, paramType); + } + + static getParamIndexList(clazz: EggProtoImplClass, methodName: string): number[] { + const methodParamMap: WebSocketMethodParamTypeMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP, clazz); + const paramMap = methodParamMap?.get(methodName); + if (!paramMap) { + return []; + } + return Array.from(paramMap.keys()); + } + + static getWebSocketMethodParamType(parameterIndex: number, clazz: EggProtoImplClass, methodName: string): WebSocketParamType | undefined { + const methodParamMap: WebSocketMethodParamTypeMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP, clazz); + const paramMap = methodParamMap?.get(methodName); + return paramMap?.get(parameterIndex); + } + + static setWebSocketMethodParamName(paramName: string, parameterIndex: number, clazz: EggProtoImplClass, methodName: string) { + const methodParamNameMap: WebSocketMethodParamNameMap = MetadataUtil.initOwnMapMetaData(CONTROLLER_WEBSOCKET_METHOD_PARAM_NAME_MAP, clazz, new Map()); + const paramMap = MapUtil.getOrStore(methodParamNameMap, methodName, new Map()); + paramMap.set(parameterIndex, paramName); + } + + static getWebSocketMethodParamName(parameterIndex: number, clazz: EggProtoImplClass, methodName: string): string | undefined { + const methodParamNameMap: WebSocketMethodParamNameMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_METHOD_PARAM_NAME_MAP, clazz); + const paramMap = methodParamNameMap?.get(methodName); + return paramMap?.get(parameterIndex); + } + + static getWebSocketMethodPriority(clazz: EggProtoImplClass, methodName: string): number | undefined { + const methodPriorityMap: WebSocketMethodPriorityMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_METHOD_PRIORITY, clazz); + return methodPriorityMap?.get(methodName); + } + + static setWebSocketMethodPriority(priority: number, clazz: EggProtoImplClass, methodName: string) { + const methodPriorityMap: WebSocketMethodPriorityMap = MetadataUtil.initOwnMapMetaData(CONTROLLER_WEBSOCKET_METHOD_PRIORITY, clazz, new Map()); + methodPriorityMap.set(methodName, priority); + } +} diff --git a/core/types/controller-decorator/MetadataKey.ts b/core/types/controller-decorator/MetadataKey.ts index 5eb77550e..7faecad30 100644 --- a/core/types/controller-decorator/MetadataKey.ts +++ b/core/types/controller-decorator/MetadataKey.ts @@ -15,6 +15,12 @@ export const CONTROLLER_METHOD_PARAM_TYPE_MAP = Symbol.for('EggPrototype#control export const CONTROLLER_METHOD_PARAM_NAME_MAP = Symbol.for('EggPrototype#controller#method#http#params#name'); export const CONTROLLER_METHOD_PRIORITY = Symbol.for('EggPrototype#controller#method#http#priority'); +export const CONTROLLER_WEBSOCKET_PATH = Symbol.for('EggPrototype#controller#websocket#path'); +export const CONTROLLER_WEBSOCKET_METHOD_PATH_MAP = Symbol.for('EggPrototype#controller#method#websocket#path'); +export const CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP = Symbol.for('EggPrototype#controller#method#websocket#params#type'); +export const CONTROLLER_WEBSOCKET_METHOD_PARAM_NAME_MAP = Symbol.for('EggPrototype#controller#method#websocket#params#name'); +export const CONTROLLER_WEBSOCKET_METHOD_PRIORITY = Symbol.for('EggPrototype#controller#method#websocket#priority'); + export const METHOD_CONTROLLER_TYPE_MAP = Symbol.for('EggPrototype#controller#mthods'); export const METHOD_CONTROLLER_HOST = Symbol.for('EggPrototype#controller#mthods#host'); export const METHOD_CONTEXT_INDEX = Symbol.for('EggPrototype#controller#method#context'); diff --git a/core/types/controller-decorator/WebSocketContext.ts b/core/types/controller-decorator/WebSocketContext.ts new file mode 100644 index 000000000..eb82f5d7e --- /dev/null +++ b/core/types/controller-decorator/WebSocketContext.ts @@ -0,0 +1,20 @@ +import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; + +export interface WebSocketLike { + on(event: string, listener: (...args: any[]) => void): this; + once(event: string, listener: (...args: any[]) => void): this; + send(data: any, cb?: (err?: Error) => void): void; + close(code?: number, reason?: string | Buffer): void; + readonly readyState?: number; +} + +export interface WebSocketContext { + readonly socket: Socket; + readonly request: IncomingMessage; + readonly path: string; + readonly host: string; + readonly params: Record; + readonly query: Record; + readonly queries: Record; + readonly headers: IncomingHttpHeaders; +} diff --git a/core/types/controller-decorator/WebSocketController.ts b/core/types/controller-decorator/WebSocketController.ts new file mode 100644 index 000000000..bec77b2e4 --- /dev/null +++ b/core/types/controller-decorator/WebSocketController.ts @@ -0,0 +1,5 @@ +export interface WebSocketControllerParams { + path?: string; + controllerName?: string; + protoName?: string; +} diff --git a/core/types/controller-decorator/WebSocketMethod.ts b/core/types/controller-decorator/WebSocketMethod.ts new file mode 100644 index 000000000..9019ce114 --- /dev/null +++ b/core/types/controller-decorator/WebSocketMethod.ts @@ -0,0 +1,4 @@ +export interface WebSocketMethodParams { + path: string; + priority?: number; +} diff --git a/core/types/controller-decorator/WebSocketParam.ts b/core/types/controller-decorator/WebSocketParam.ts new file mode 100644 index 000000000..25111609a --- /dev/null +++ b/core/types/controller-decorator/WebSocketParam.ts @@ -0,0 +1,11 @@ +export interface WebSocketQueryParams { + name?: string; +} + +export interface WebSocketQueriesParams { + name?: string; +} + +export interface WebSocketParamParams { + name?: string; +} diff --git a/core/types/controller-decorator/index.ts b/core/types/controller-decorator/index.ts index 828f96679..db3b985a5 100644 --- a/core/types/controller-decorator/index.ts +++ b/core/types/controller-decorator/index.ts @@ -5,6 +5,10 @@ export * from './builder'; export * from './HTTPController'; export * from './HTTPMethod'; export * from './HTTPParam'; +export * from './WebSocketController'; +export * from './WebSocketMethod'; +export * from './WebSocketParam'; +export * from './WebSocketContext'; export * from './MetadataKey'; export * from './MCPController'; export * from './MCPPromptParams'; diff --git a/core/types/controller-decorator/model/types.ts b/core/types/controller-decorator/model/types.ts index 4ace7fe56..c7bd158b5 100644 --- a/core/types/controller-decorator/model/types.ts +++ b/core/types/controller-decorator/model/types.ts @@ -7,6 +7,7 @@ export type { IncomingHttpHeaders } from 'node:http'; export enum ControllerType { HTTP = 'HTTP', + WEBSOCKET = 'WEBSOCKET', MCP = 'MCP', SOFA_RPC = 'SOFA_RPC', SOFA_RPC_STREAM = 'SOFA_RPC_STREAM', @@ -23,6 +24,7 @@ export type ControllerTypeLike = ControllerType | string; export enum MethodType { HTTP = 'HTTP', + WEBSOCKET = 'WEBSOCKET', SOFA_RPC = 'SOFA_RPC', SOFA_RPC_STREAM = 'SOFA_RPC_STREAM', MGW_RPC = 'MGW_RPC', @@ -53,6 +55,16 @@ export enum HTTPParamType { COOKIES = 'COOKIES', } +export enum WebSocketParamType { + PARAM = 'PARAM', + QUERY = 'QUERY', + QUERIES = 'QUERIES', + HEADERS = 'HEADERS', + REQUEST = 'REQUEST', + SOCKET = 'SOCKET', + STREAM = 'STREAM', +} + export enum MCPProtocols { STDIO = 'STDIO', SSE = 'SSE', diff --git a/plugin/controller/app.ts b/plugin/controller/app.ts index 5093705da..9eb7a63ae 100644 --- a/plugin/controller/app.ts +++ b/plugin/controller/app.ts @@ -4,6 +4,7 @@ import { AppLoadUnitControllerHook } from './lib/AppLoadUnitControllerHook'; import { GlobalGraph, LoadUnitLifecycleContext } from '@eggjs/tegg-metadata'; import { ControllerMetaBuilderFactory, ControllerType } from '@eggjs/tegg'; import { HTTPControllerRegister } from './lib/impl/http/HTTPControllerRegister'; +import { WebSocketControllerRegister } from './lib/impl/websocket/WebSocketControllerRegister'; import { ControllerRegisterFactory } from './lib/ControllerRegisterFactory'; import { ControllerLoadUnitHandler } from './lib/ControllerLoadUnitHandler'; import { LoadUnitInstanceLifecycleContext, ModuleLoadUnitInstance } from '@eggjs/tegg-runtime'; @@ -60,6 +61,7 @@ export default class ControllerAppBootHook { return new EggControllerLoader(unitPath); }); this.controllerRegisterFactory.registerControllerRegister(ControllerType.HTTP, HTTPControllerRegister.create); + this.controllerRegisterFactory.registerControllerRegister(ControllerType.WEBSOCKET, WebSocketControllerRegister.create); this.app.loadUnitFactory.registerLoadUnitCreator( CONTROLLER_LOAD_UNIT, (ctx: LoadUnitLifecycleContext): ControllerLoadUnit => { @@ -148,10 +150,16 @@ export default class ControllerAppBootHook { // The HTTPControllerRegister will collect all the methods // and register methods after collect is done. HTTPControllerRegister.instance?.doRegister(this.app.rootProtoManager); + WebSocketControllerRegister.instance?.doRegister(); + WebSocketControllerRegister.instance?.listen(); this.app.config.mcp.hooks = this.mcpControllerRegister?.hooks; } + serverDidReady() { + WebSocketControllerRegister.instance?.listen(); + } + configDidLoad() { GlobalGraph.instance!.registerBuildHook(middlewareGraphHook); } @@ -180,6 +188,7 @@ export default class ControllerAppBootHook { this.app.eggPrototypeLifecycleUtil.deleteLifecycle(this.controllerPrototypeHook); ControllerMetadataManager.instance.clear(); HTTPControllerRegister.clean(); + WebSocketControllerRegister.clean(); this.mcpControllerRegister?.clean(); } } diff --git a/plugin/controller/lib/impl/websocket/WebSocketContext.ts b/plugin/controller/lib/impl/websocket/WebSocketContext.ts new file mode 100644 index 000000000..937e52eb6 --- /dev/null +++ b/plugin/controller/lib/impl/websocket/WebSocketContext.ts @@ -0,0 +1,35 @@ +import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; +import type { WebSocket } from 'ws'; +import type { WebSocketContext } from '@eggjs/tegg'; + +export class WebSocketContextImpl implements WebSocketContext { + readonly socket: WebSocket; + readonly request: IncomingMessage; + readonly path: string; + readonly host: string; + readonly params: Record; + readonly query: Record; + readonly queries: Record; + readonly headers: IncomingHttpHeaders; + + constructor(options: { + socket: WebSocket; + request: IncomingMessage; + url: URL; + params: Record; + }) { + this.socket = options.socket; + this.request = options.request; + this.path = options.url.pathname; + this.host = options.url.host; + this.params = options.params; + this.headers = options.request.headers; + this.query = {}; + this.queries = {}; + + for (const key of options.url.searchParams.keys()) { + this.query[key] = options.url.searchParams.get(key) ?? undefined; + this.queries[key] = options.url.searchParams.getAll(key); + } + } +} diff --git a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts new file mode 100644 index 000000000..6be45b800 --- /dev/null +++ b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts @@ -0,0 +1,406 @@ +import assert from 'node:assert'; +import { IncomingMessage, ServerResponse } from 'node:http'; +import { Duplex, pipeline } from 'node:stream'; +import compose from 'koa-compose'; +import pathToRegexp from 'path-to-regexp'; +import { Application, Context } from 'egg'; +import { EggRouter } from '@eggjs/router'; +import { FrameworkErrorFormater } from 'egg-errors'; +import { createWebSocketStream, WebSocket, WebSocketServer } from 'ws'; +import { + CONTROLLER_META_DATA, + ControllerMetadata, + ControllerType, + Next, + WebSocketControllerMeta, + WebSocketMethodMeta, + WebSocketParamType, + WebSocketPathParamMeta, + WebSocketQueriesParamMeta, + WebSocketQueryParamMeta, +} from '@eggjs/tegg'; +import { EggContainerFactory } from '@eggjs/tegg-runtime'; +import { EggPrototype } from '@eggjs/tegg-metadata'; +import { ROOT_PROTO } from '@eggjs/egg-module-common'; +import { ControllerRegister } from '../../ControllerRegister'; +import { RouterConflictError } from '../../errors'; +import { WebSocketContextImpl } from './WebSocketContext'; + +const noop = () => { + // noop +}; + +interface WebSocketRoute { + controllerProto: EggPrototype; + controllerMeta: WebSocketControllerMeta; + methodMeta: WebSocketMethodMeta; + methodRealPath: string; + methodName: string; + host?: string; + keys: pathToRegexp.Key[]; + regexp: RegExp; +} + +export class WebSocketControllerRegister implements ControllerRegister { + static instance?: WebSocketControllerRegister; + + private readonly app: Application; + private readonly eggContainerFactory: typeof EggContainerFactory; + private readonly checkRouters: Map; + private controllerProtos: EggPrototype[] = []; + private routes: WebSocketRoute[] = []; + private webSocketServer?: WebSocketServer; + private upgradeHandler?: (req: IncomingMessage, socket: Duplex, head: Buffer) => void; + private listening = false; + + static create(proto: EggPrototype, controllerMeta: ControllerMetadata, app: Application) { + assert( + controllerMeta.type === ControllerType.WEBSOCKET, + 'controller meta type is not WEBSOCKET', + ); + if (!WebSocketControllerRegister.instance) { + WebSocketControllerRegister.instance = new WebSocketControllerRegister(app); + } + WebSocketControllerRegister.instance.controllerProtos.push(proto); + return WebSocketControllerRegister.instance; + } + + constructor(app: Application) { + this.app = app; + this.eggContainerFactory = app.eggContainerFactory; + this.checkRouters = new Map(); + this.checkRouters.set('default', new EggRouter({ sensitive: true }, {} as any)); + } + + register(): Promise { + return Promise.resolve(); + } + + static clean() { + this.instance?.close(); + this.instance = undefined; + } + + doRegister() { + const methodMap = new Map(); + for (const proto of this.controllerProtos) { + const metadata = proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta; + for (const method of metadata.methods) { + methodMap.set(method, proto); + } + } + const allMethods = Array.from(methodMap.keys()) + .sort((a, b) => b.priority - a.priority); + + for (const method of allMethods) { + const controllerProto = methodMap.get(method)!; + const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta; + this.checkDuplicate(controllerMeta, method); + } + + this.routes = allMethods.flatMap(method => { + const controllerProto = methodMap.get(method)!; + const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta; + const hosts = this.getMethodHosts(controllerMeta, method) || [ undefined ]; + return hosts.map(host => this.createRoute(controllerProto, controllerMeta, method, host)); + }); + } + + listen() { + if (this.listening || !this.routes.length) { + return; + } + const server = (this.app as any).server; + if (!server) { + return; + } + + this.webSocketServer = new WebSocketServer({ noServer: true }); + this.upgradeHandler = (req, socket, head) => { + this.handleUpgrade(req, socket, head).catch(error => { + this.app.logger.error('[tegg/websocket] handle upgrade failed: %s', error.stack || error.message); + this.rejectSocket(socket, 500, 'Internal Server Error'); + }); + }; + server.on('upgrade', this.upgradeHandler); + this.listening = true; + } + + close() { + const server = (this.app as any).server; + if (server && this.upgradeHandler) { + server.removeListener('upgrade', this.upgradeHandler); + } + this.webSocketServer?.close(); + this.webSocketServer = undefined; + this.upgradeHandler = undefined; + this.listening = false; + this.routes = []; + this.controllerProtos = []; + this.checkRouters.clear(); + } + + private createRoute( + controllerProto: EggPrototype, + controllerMeta: WebSocketControllerMeta, + methodMeta: WebSocketMethodMeta, + host: string | undefined, + ): WebSocketRoute { + const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); + const keys: pathToRegexp.Key[] = []; + const regexp = pathToRegexp(methodRealPath, keys, { sensitive: true }); + return { + controllerProto, + controllerMeta, + methodMeta, + methodRealPath, + methodName: this.getMethodName(controllerMeta, methodMeta), + host, + keys, + regexp, + }; + } + + private checkDuplicate(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { + let router = this.checkRouters.get('default')!; + const hosts = this.getMethodHosts(controllerMeta, methodMeta) || []; + if (!hosts.length) { + this.checkDuplicateInRouter(router, controllerMeta, methodMeta); + this.registerToRouter(router, controllerMeta, methodMeta); + return; + } + + hosts.forEach(host => { + router = this.checkRouters.get(host)!; + if (!router) { + router = new EggRouter({ sensitive: true }, {} as any); + this.checkRouters.set(host, router); + } + this.checkDuplicateInRouter(router, controllerMeta, methodMeta); + this.registerToRouter(router, controllerMeta, methodMeta); + }); + } + + private registerToRouter(router: EggRouter, controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { + const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); + const methodName = this.getMethodName(controllerMeta, methodMeta); + Reflect.apply(router.get, router, [ methodName, methodRealPath, noop ]); + } + + private checkDuplicateInRouter(router: EggRouter, controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { + const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); + const matched = router.match(methodRealPath, 'GET'); + const methodName = this.getMethodName(controllerMeta, methodMeta); + if (matched.route) { + const [ layer ] = matched.path; + const err = new RouterConflictError(`register websocket controller ${methodName} failed, ${controllerMeta.type} ${methodRealPath} is conflict with exists rule ${layer.path}`); + throw FrameworkErrorFormater.format(err); + } + } + + private getMethodRealPath(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { + return controllerMeta.getMethodRealPath(methodMeta); + } + + private getMethodHosts(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta): string[] | undefined { + return controllerMeta.getMethodHosts(methodMeta); + } + + private getMethodName(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta): string { + return controllerMeta.getMethodName(methodMeta); + } + + private getMethodMiddlewares(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { + return controllerMeta.getMethodMiddlewares(methodMeta); + } + + private async handleUpgrade(req: IncomingMessage, socket: Duplex, head: Buffer) { + const url = this.createURL(req); + const matched = this.matchRoute(req, url); + if (!matched) { + this.rejectSocket(socket, 404, 'Not Found'); + return; + } + + this.webSocketServer!.handleUpgrade(req, socket, head, webSocket => { + this.handleConnection(matched.route, matched.params, url, req, webSocket) + .catch(error => { + this.app.logger.error('[tegg/websocket] handle connection failed: %s', error.stack || error.message); + if (webSocket.readyState === WebSocket.OPEN || webSocket.readyState === WebSocket.CONNECTING) { + webSocket.close(1011, 'Internal Server Error'); + } + }); + }); + } + + private createURL(req: IncomingMessage): URL { + const host = req.headers.host || 'localhost'; + return new URL(req.url || '/', `http://${host}`); + } + + private matchRoute(req: IncomingMessage, url: URL): { route: WebSocketRoute; params: Record } | undefined { + for (const route of this.routes) { + if (route.host && !this.matchHost(route.host, req.headers.host)) { + continue; + } + const matched = route.regexp.exec(url.pathname); + if (!matched) { + continue; + } + const params: Record = {}; + route.keys.forEach((key, index) => { + params[String(key.name)] = decodeURIComponent(matched[index + 1]); + }); + return { route, params }; + } + } + + private matchHost(expectedHost: string, requestHost: string | undefined) { + if (!requestHost) { + return false; + } + return requestHost === expectedHost || requestHost.split(':')[0] === expectedHost; + } + + private async handleConnection( + route: WebSocketRoute, + params: Record, + url: URL, + req: IncomingMessage, + webSocket: WebSocket, + ) { + const res = new ServerResponse(req); + const eggCtx = this.app.createContext(req, res) as unknown as Context; + (eggCtx as any)[ROOT_PROTO] = route.controllerProto; + const webSocketCtx = new WebSocketContextImpl({ + socket: webSocket, + request: req, + url, + params, + }); + const lifecycleMiddleware = this.app.middleware.teggCtxLifecycleMiddleware(); + const methodMiddlewares = this.getMethodMiddlewares(route.controllerMeta, route.methodMeta); + let invoked = false; + const handler = async (_ctx: Context, next: Next) => { + invoked = true; + await this.invokeController(route, webSocketCtx); + await next(); + }; + const composed = compose([ ...methodMiddlewares, handler ]); + + await this.app.ctxStorage.run(eggCtx, async () => { + await lifecycleMiddleware(eggCtx, async () => { + await composed(eggCtx, async () => { + // final middleware + }); + if (!invoked) { + webSocket.close(1008, 'WebSocket middleware did not call next'); + return; + } + await this.waitWebSocketClose(webSocket); + }); + }); + } + + private async invokeController(route: WebSocketRoute, webSocketCtx: WebSocketContextImpl) { + const methodMeta = route.methodMeta as WebSocketMethodMeta; + const argsLength = methodMeta.paramMap.size; + const hasContext = methodMeta.contextParamIndex !== undefined; + const contextIndex = methodMeta.contextParamIndex; + const methodArgsLength = argsLength + (hasContext ? 1 : 0); + const eggObj = await this.eggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); + const realObj = eggObj.obj; + const realMethod = realObj[methodMeta.name]; + const args: Array = new Array(methodArgsLength); + let webSocketStream: Duplex | undefined; + const getWebSocketStream = () => { + if (!webSocketStream) { + webSocketStream = createWebSocketStream(webSocketCtx.socket); + } + return webSocketStream; + }; + if (hasContext) { + args[contextIndex!] = webSocketCtx; + } + for (const [ index, param ] of methodMeta.paramMap) { + switch (param.type) { + case WebSocketParamType.PARAM: { + const pathParam = param as WebSocketPathParamMeta; + args[index] = webSocketCtx.params[pathParam.name]; + break; + } + case WebSocketParamType.QUERY: { + const queryParam = param as WebSocketQueryParamMeta; + args[index] = webSocketCtx.query[queryParam.name]; + break; + } + case WebSocketParamType.QUERIES: { + const queryParam = param as WebSocketQueriesParamMeta; + args[index] = webSocketCtx.queries[queryParam.name] || []; + break; + } + case WebSocketParamType.HEADERS: { + args[index] = webSocketCtx.headers; + break; + } + case WebSocketParamType.REQUEST: { + args[index] = webSocketCtx.request; + break; + } + case WebSocketParamType.SOCKET: { + args[index] = webSocketCtx.socket; + break; + } + case WebSocketParamType.STREAM: { + args[index] = getWebSocketStream(); + break; + } + default: + assert.fail('never arrive'); + } + } + + const result = await Reflect.apply(realMethod, realObj, args); + this.pipeResponseStream(result, webSocketCtx.socket, getWebSocketStream); + } + + private pipeResponseStream(result: unknown, webSocket: WebSocket, getWebSocketStream: () => Duplex) { + if (!this.isReadableStream(result)) { + return; + } + + pipeline(result, getWebSocketStream(), error => { + if (!error) { + return; + } + if (webSocket.readyState === WebSocket.CLOSED || webSocket.readyState === WebSocket.CLOSING) { + return; + } + this.app.logger.error('[tegg/websocket] pipe response stream failed: %s', error.stack || error.message); + webSocket.close(1011, 'Internal Server Error'); + }); + } + + private isReadableStream(result: unknown): result is NodeJS.ReadableStream { + return !!result && typeof (result as NodeJS.ReadableStream).pipe === 'function'; + } + + private waitWebSocketClose(webSocket: WebSocket): Promise { + if (webSocket.readyState === WebSocket.CLOSED || webSocket.readyState === WebSocket.CLOSING) { + return Promise.resolve(); + } + return new Promise(resolve => { + webSocket.once('close', () => resolve()); + webSocket.once('error', () => resolve()); + }); + } + + private rejectSocket(socket: Duplex, status: number, message: string) { + if (socket.destroyed) { + return; + } + const body = message; + socket.write(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); + socket.destroy(); + } +} diff --git a/plugin/controller/package.json b/plugin/controller/package.json index 9019818a0..2025b00f4 100644 --- a/plugin/controller/package.json +++ b/plugin/controller/package.json @@ -65,6 +65,7 @@ "path-to-regexp": "^1.8.0", "raw-body": "^2.5.2", "sdk-base": "^4.2.0", + "ws": "^8.18.0", "zod": "^4.0.0" }, "devDependencies": { @@ -73,6 +74,7 @@ "@eggjs/tegg-plugin": "^3.84.2", "@types/mocha": "^10.0.1", "@types/node": "^20.2.4", + "@types/ws": "^8.5.12", "cross-env": "^7.0.3", "egg": "^3.9.1", "egg-mock": "^5.5.0", diff --git a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts new file mode 100644 index 000000000..a29793590 --- /dev/null +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts @@ -0,0 +1,94 @@ +import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; +import { PassThrough, type Readable } from 'node:stream'; +import type { WebSocket } from 'ws'; +import { + Context, + Headers, + Inject, + Param, + Queries, + Query, + Request, + WebSocketContext, + WebSocketController, + WebSocketMethod, + WebSocketSocket, + WebSocketStream, +} from '@eggjs/tegg'; +import AppService from '../../modules/multi-module-service/AppService'; + +@WebSocketController({ + path: '/ws', +}) +export class AppWebSocketController { + @Inject() + appService: AppService; + + @WebSocketMethod({ + path: '/echo/:id', + }) + async echo( + @WebSocketSocket() socket: WebSocket, + @Param() id: string, + @Query() name: string, + @Queries({ name: 'tag' }) tags: string[], + @Headers() headers: IncomingHttpHeaders, + @Request() request: IncomingMessage, + @Context() ctx: WebSocketContext, + ) { + await this.appService.save({ + name: id, + desc: `ws:${name}`, + }); + const app = await this.appService.findApp(id); + + socket.send(JSON.stringify({ + type: 'ready', + id, + name, + tags, + app, + header: headers['x-client-id'], + url: request.url, + path: ctx.path, + pid: process.pid, + })); + + socket.on('message', data => { + socket.send(JSON.stringify({ + type: 'echo', + data: data.toString(), + pid: process.pid, + })); + }); + } + + @WebSocketMethod({ + path: '/stream/:id', + }) + async stream( + @WebSocketStream() input: Readable, + @Param() id: string, + @Query() name: string, + ) { + const output = new PassThrough(); + output.write(JSON.stringify({ + type: 'ready', + id, + name, + pid: process.pid, + })); + + input.on('data', data => { + output.write(JSON.stringify({ + type: 'stream', + data: data.toString(), + pid: process.pid, + })); + }); + input.on('end', () => output.end()); + input.on('error', error => output.destroy(error)); + + return output; + } +} diff --git a/plugin/controller/test/websocket/websocket.test.ts b/plugin/controller/test/websocket/websocket.test.ts new file mode 100644 index 000000000..fd516a6ab --- /dev/null +++ b/plugin/controller/test/websocket/websocket.test.ts @@ -0,0 +1,146 @@ +import { strict as assert } from 'node:assert'; +import path from 'node:path'; +import mm from 'egg-mock'; +import WebSocket from 'ws'; +import type { RawData } from 'ws'; + +const pluginRoot = process.cwd(); +const eggFramework = path.join(pluginRoot, '../../node_modules/egg'); + +function requestUrl(app, pathname: string) { + return app.httpRequest().get(pathname).url.replace(/^http/, 'ws'); +} + +async function createClient(url: string, headers?: Record): Promise { + const socket = new WebSocket(url, { headers }); + await new Promise((resolve, reject) => { + socket.once('open', resolve); + socket.once('error', reject); + }); + return socket; +} + +async function receiveJSON(socket: WebSocket): Promise { + const data = await new Promise((resolve, reject) => { + socket.once('message', resolve); + socket.once('error', reject); + socket.once('close', () => reject(new Error('websocket closed before message'))); + }); + return JSON.parse(data.toString()); +} + +async function createClientWithFirstMessage(url: string, headers?: Record) { + const socket = new WebSocket(url, { headers }); + const firstMessage = receiveJSON(socket); + await new Promise((resolve, reject) => { + socket.once('open', resolve); + socket.once('error', reject); + }); + return { socket, firstMessage }; +} + +async function closeClient(socket: WebSocket) { + if (socket.readyState === WebSocket.CLOSED) { + return; + } + await new Promise(resolve => { + socket.once('close', () => resolve()); + socket.close(); + }); +} + +describe('plugin/controller/test/websocket/websocket.test.ts', () => { + let app; + + beforeEach(() => { + mm(process.env, 'EGG_TYPESCRIPT', true); + }); + + afterEach(() => { + mm.restore(); + }); + + before(async () => { + mm(process.env, 'EGG_TYPESCRIPT', true); + mm(process, 'cwd', () => { + return pluginRoot; + }); + app = mm.app({ + baseDir: path.join(pluginRoot, 'test/fixtures/apps/controller-app'), + framework: eggFramework, + }); + await app.ready(); + }); + + after(async () => { + await app?.close(); + }); + + it('should handle websocket controller', async () => { + const { socket, firstMessage } = await createClientWithFirstMessage( + requestUrl(app, '/ws/echo/foo?name=bar&tag=a&tag=b'), + { 'x-client-id': 'client-1' }, + ); + try { + const ready = await firstMessage; + assert.deepEqual(ready, { + type: 'ready', + id: 'foo', + name: 'bar', + tags: [ 'a', 'b' ], + app: { + name: 'foo', + desc: 'ws:bar', + }, + header: 'client-1', + url: '/ws/echo/foo?name=bar&tag=a&tag=b', + path: '/ws/echo/foo', + pid: ready.pid, + }); + assert.equal(typeof ready.pid, 'number'); + + socket.send('hello'); + const echo = await receiveJSON(socket); + assert.deepEqual(echo, { + type: 'echo', + data: 'hello', + pid: ready.pid, + }); + } finally { + await closeClient(socket); + } + }); + + it('should handle websocket stream param and returned stream', async () => { + const { socket, firstMessage } = await createClientWithFirstMessage( + requestUrl(app, '/ws/stream/foo?name=bar'), + ); + try { + const ready = await firstMessage; + assert.deepEqual(ready, { + type: 'ready', + id: 'foo', + name: 'bar', + pid: ready.pid, + }); + assert.equal(typeof ready.pid, 'number'); + + socket.send('hello-stream'); + const echo = await receiveJSON(socket); + assert.deepEqual(echo, { + type: 'stream', + data: 'hello-stream', + pid: ready.pid, + }); + } finally { + await closeClient(socket); + } + }); + + it('should reject unmatched websocket route', async () => { + await assert.rejects( + () => createClient(requestUrl(app, '/ws/not-found')), + /Unexpected server response: 404/, + ); + }); +}); diff --git a/plugin/controller/test/websocket/websocketCluster.test.ts b/plugin/controller/test/websocket/websocketCluster.test.ts new file mode 100644 index 000000000..516a3d0b2 --- /dev/null +++ b/plugin/controller/test/websocket/websocketCluster.test.ts @@ -0,0 +1,126 @@ +import { strict as assert } from 'node:assert'; +import path from 'node:path'; +import mm from 'egg-mock'; +import WebSocket from 'ws'; +import type { RawData } from 'ws'; + +const pluginRoot = process.cwd(); +const eggFramework = path.join(pluginRoot, '../../node_modules/egg'); + +function requestUrl(app, pathname: string) { + return app.httpRequest().get(pathname).url.replace(/^http/, 'ws'); +} + +async function receiveJSON(socket: WebSocket): Promise { + const data = await new Promise((resolve, reject) => { + socket.once('message', resolve); + socket.once('error', reject); + socket.once('close', () => reject(new Error('websocket closed before message'))); + }); + return JSON.parse(data.toString()); +} + +async function createClientWithFirstMessage(url: string) { + const socket = new WebSocket(url); + const firstMessage = receiveJSON(socket); + await new Promise((resolve, reject) => { + socket.once('open', resolve); + socket.once('error', reject); + }); + return { socket, firstMessage }; +} + +async function closeClient(socket: WebSocket) { + if (socket.readyState === WebSocket.CLOSED) { + return; + } + await new Promise(resolve => { + socket.once('close', () => resolve()); + socket.close(); + }); +} + +describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { + let app; + + before(async () => { + mm(process.env, 'EGG_TYPESCRIPT', true); + mm(process, 'cwd', () => { + return pluginRoot; + }); + app = mm.cluster({ + baseDir: path.join(pluginRoot, 'test/fixtures/apps/controller-app'), + framework: eggFramework, + workers: 2, + sticky: false, + opt: { + env: { + ...process.env, + NODE_OPTIONS: '--require ts-node/register tsconfig-paths/register', + }, + }, + }); + await app.ready(); + }); + + after(async () => { + await app?.close(); + }); + + it('should handle websocket controller in cluster', async () => { + const entries = await Promise.all([ 1, 2, 3, 4 ].map(index => { + return createClientWithFirstMessage(requestUrl(app, `/ws/echo/cluster-${index}?name=bar&tag=${index}`)); + })); + const clients = entries.map(entry => entry.socket); + try { + const readyMessages = await Promise.all(entries.map(entry => entry.firstMessage)); + assert.equal(readyMessages.length, 4); + for (const [ index, ready ] of readyMessages.entries()) { + assert.equal(ready.type, 'ready'); + assert.equal(ready.id, `cluster-${index + 1}`); + assert.equal(ready.name, 'bar'); + assert.deepEqual(ready.tags, [ String(index + 1) ]); + assert.equal(ready.path, `/ws/echo/cluster-${index + 1}`); + assert.equal(typeof ready.pid, 'number'); + } + + clients.forEach((socket, index) => socket.send(`hello-${index}`)); + const echoMessages = await Promise.all(clients.map(socket => receiveJSON(socket))); + for (const [ index, echo ] of echoMessages.entries()) { + assert.equal(echo.type, 'echo'); + assert.equal(echo.data, `hello-${index}`); + assert.equal(typeof echo.pid, 'number'); + } + } finally { + await Promise.all(clients.map(socket => closeClient(socket))); + } + }); + + it('should handle websocket stream params in cluster', async () => { + const entries = await Promise.all([ 1, 2, 3, 4 ].map(index => { + return createClientWithFirstMessage(requestUrl(app, `/ws/stream/stream-${index}?name=bar`)); + })); + const clients = entries.map(entry => entry.socket); + try { + const readyMessages = await Promise.all(entries.map(entry => entry.firstMessage)); + assert.equal(readyMessages.length, 4); + for (const [ index, ready ] of readyMessages.entries()) { + assert.equal(ready.type, 'ready'); + assert.equal(ready.id, `stream-${index + 1}`); + assert.equal(ready.name, 'bar'); + assert.equal(typeof ready.pid, 'number'); + } + + clients.forEach((socket, index) => socket.send(`stream-${index}`)); + const echoMessages = await Promise.all(clients.map(socket => receiveJSON(socket))); + for (const [ index, echo ] of echoMessages.entries()) { + assert.equal(echo.type, 'stream'); + assert.equal(echo.data, `stream-${index}`); + assert.equal(typeof echo.pid, 'number'); + } + } finally { + await Promise.all(clients.map(socket => closeClient(socket))); + } + }); + +}); From c003e95da838e0ebed247a6b574efadca6b7516c Mon Sep 17 00:00:00 2001 From: akitasummer Date: Thu, 9 Jul 2026 17:35:09 +0800 Subject: [PATCH 2/9] feat(controller): add websocket fetch controller --- core/controller-decorator/index.ts | 4 + .../WebSocketFetchController.ts | 26 ++ .../websocket-fetch/WebSocketFetchMethod.ts | 57 ++++ .../websocket-fetch/WebSocketFetchParam.ts | 47 +++ .../WebSocketFetchControllerMetaBuilder.ts | 108 +++++++ ...bSocketFetchControllerMethodMetaBuilder.ts | 150 +++++++++ .../WebSocketControllerMethodMetaBuilder.ts | 13 + .../src/model/WebSocketFetchControllerMeta.ts | 73 +++++ .../src/model/WebSocketFetchMethodMeta.ts | 34 ++ .../src/model/WebSocketMethodMeta.ts | 55 ++++ core/controller-decorator/src/model/index.ts | 2 + .../src/util/WebSocketInfoUtil.ts | 14 +- .../types/controller-decorator/MetadataKey.ts | 1 + .../WebSocketFetchController.ts | 5 + .../WebSocketFetchMethod.ts | 3 + .../WebSocketFetchParam.ts | 1 + core/types/controller-decorator/index.ts | 3 + .../types/controller-decorator/model/types.ts | 15 + plugin/controller/app.ts | 1 + .../websocket/WebSocketControllerRegister.ts | 230 ++++++++++++-- .../app/controller/WebSocketController.ts | 128 +++++++- .../test/websocket/websocket.test.ts | 300 ++++++++++++++++++ .../test/websocket/websocketCluster.test.ts | 114 +++++++ 23 files changed, 1351 insertions(+), 33 deletions(-) create mode 100644 core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts create mode 100644 core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts create mode 100644 core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchParam.ts create mode 100644 core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts create mode 100644 core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts create mode 100644 core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts create mode 100644 core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts create mode 100644 core/types/controller-decorator/WebSocketFetchController.ts create mode 100644 core/types/controller-decorator/WebSocketFetchMethod.ts create mode 100644 core/types/controller-decorator/WebSocketFetchParam.ts diff --git a/core/controller-decorator/index.ts b/core/controller-decorator/index.ts index 893b336d4..910e0217e 100644 --- a/core/controller-decorator/index.ts +++ b/core/controller-decorator/index.ts @@ -1,5 +1,6 @@ import './src/impl/http/HTTPControllerMetaBuilder'; import './src/impl/websocket/WebSocketControllerMetaBuilder'; +import './src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder'; import './src/impl/mcp/MCPControllerMetaBuilder'; export * from '@eggjs/tegg-types/controller-decorator'; @@ -14,6 +15,9 @@ export * from './src/decorator/http/Host'; export * from './src/decorator/websocket/WebSocketController'; export * from './src/decorator/websocket/WebSocketMethod'; export * from './src/decorator/websocket/WebSocketParam'; +export * from './src/decorator/websocket-fetch/WebSocketFetchController'; +export * from './src/decorator/websocket-fetch/WebSocketFetchMethod'; +export * from './src/decorator/websocket-fetch/WebSocketFetchParam'; export * from './src/decorator/mcp/MCPController'; export * from './src/decorator/mcp/MCPPrompt'; export * from './src/decorator/mcp/MCPResource'; diff --git a/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts new file mode 100644 index 000000000..b0a8a6214 --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts @@ -0,0 +1,26 @@ +import { PrototypeUtil, SingletonProto } from '@eggjs/core-decorator'; +import { StackUtil } from '@eggjs/tegg-common-util'; +import type { EggProtoImplClass, WebSocketFetchControllerParams } from '@eggjs/tegg-types'; +import { AccessLevel, ControllerType } from '@eggjs/tegg-types'; +import ControllerInfoUtil from '../../util/ControllerInfoUtil'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +export function WebSocketFetchController(param?: WebSocketFetchControllerParams) { + return function(constructor: EggProtoImplClass) { + ControllerInfoUtil.setControllerType(constructor, ControllerType.WEBSOCKET_FETCH); + if (param?.controllerName) { + ControllerInfoUtil.setControllerName(constructor, param.controllerName); + } + if (param?.path) { + WebSocketInfoUtil.setWebSocketPath(param.path, constructor); + } + + const func = SingletonProto({ + accessLevel: AccessLevel.PUBLIC, + name: param?.protoName, + }); + func(constructor); + + PrototypeUtil.setFilePath(constructor, StackUtil.getCalleeFromStack(false, 5)); + }; +} diff --git a/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts new file mode 100644 index 000000000..809c6a13a --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts @@ -0,0 +1,57 @@ +import assert from 'node:assert'; +import { ControllerType, WebSocketFetchMethodType } from '@eggjs/tegg-types'; +import type { EggProtoImplClass, WebSocketFetchMethodParams } from '@eggjs/tegg-types'; +import MethodInfoUtil from '../../util/MethodInfoUtil'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +function setWebSocketFetchMethodType( + target: any, + propertyKey: PropertyKey, + methodType: WebSocketFetchMethodType, +) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const controllerClazz = target.constructor as EggProtoImplClass; + const methodName = propertyKey as string; + MethodInfoUtil.setMethodControllerType(controllerClazz, methodName, ControllerType.WEBSOCKET_FETCH); + WebSocketInfoUtil.setWebSocketFetchMethodType(methodType, controllerClazz, methodName); + return { controllerClazz, methodName }; +} + +export function WebSocketFetchMethod(param?: WebSocketFetchMethodParams) { + return function(target: any, propertyKey: PropertyKey) { + const { controllerClazz, methodName } = setWebSocketFetchMethodType( + target, + propertyKey, + WebSocketFetchMethodType.DATA, + ); + WebSocketInfoUtil.setWebSocketMethodPath('', controllerClazz, methodName); + if (param?.priority !== undefined) { + WebSocketInfoUtil.setWebSocketMethodPriority(param.priority, controllerClazz, methodName); + } + }; +} + +export function WebSocketFetchOnConnection() { + return function(target: any, propertyKey: PropertyKey) { + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CONNECTION); + }; +} + +export function WebSocketFetchOnOpen() { + return function(target: any, propertyKey: PropertyKey) { + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.OPEN); + }; +} + +export function WebSocketFetchOnError() { + return function(target: any, propertyKey: PropertyKey) { + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.ERROR); + }; +} + +export function WebSocketFetchOnClose() { + return function(target: any, propertyKey: PropertyKey) { + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CLOSE); + }; +} diff --git a/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchParam.ts b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchParam.ts new file mode 100644 index 000000000..6c61b83e6 --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchParam.ts @@ -0,0 +1,47 @@ +import assert from 'node:assert'; +import { WebSocketParamType } from '@eggjs/tegg-types'; +import type { EggProtoImplClass } from '@eggjs/tegg-types'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +function setWebSocketFetchParamType( + target: any, + propertyKey: PropertyKey, + parameterIndex: number, + paramType: WebSocketParamType, +) { + assert(typeof propertyKey === 'string', + `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); + const methodName = propertyKey as string; + const controllerClazz = target.constructor as EggProtoImplClass; + WebSocketInfoUtil.setWebSocketMethodParamType(paramType, parameterIndex, controllerClazz, methodName); +} + +export function WebSocketData() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.DATA); + }; +} + +export function WebSocketClose() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.CLOSE); + }; +} + +export function WebSocketError() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.ERROR); + }; +} + +export function WebSocketCloseCode() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.CLOSE_CODE); + }; +} + +export function WebSocketCloseReason() { + return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { + setWebSocketFetchParamType(target, propertyKey, parameterIndex, WebSocketParamType.CLOSE_REASON); + }; +} diff --git a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts new file mode 100644 index 000000000..7fa07eb99 --- /dev/null +++ b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts @@ -0,0 +1,108 @@ +import assert from 'node:assert'; +import { PrototypeUtil } from '@eggjs/core-decorator'; +import { ObjectUtils } from '@eggjs/tegg-common-util'; +import { ClassUtil } from '@eggjs/tegg-metadata'; +import type { EggProtoImplClass } from '@eggjs/tegg-types'; +import { ControllerType, WebSocketFetchMethodType } from '@eggjs/tegg-types'; +import { ControllerMetaBuilderFactory } from '../../builder/ControllerMetaBuilderFactory'; +import { WebSocketFetchControllerMeta, WebSocketFetchMethodMeta } from '../../model'; +import ControllerInfoUtil from '../../util/ControllerInfoUtil'; +import { ControllerMetadataUtil } from '../../util/ControllerMetadataUtil'; +import { ControllerValidator } from '../../util/validator/ControllerValidator'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; +import { WebSocketFetchControllerMethodMetaBuilder } from './WebSocketFetchControllerMethodMetaBuilder'; + +export class WebSocketFetchControllerMetaBuilder { + private readonly clazz: EggProtoImplClass; + + constructor(clazz: EggProtoImplClass) { + this.clazz = clazz; + } + + private buildMethods() { + const methodNames = ObjectUtils.getProperties(this.clazz.prototype); + const methods: WebSocketFetchMethodMeta[] = []; + let connectionMethod: WebSocketFetchMethodMeta | undefined; + let openMethod: WebSocketFetchMethodMeta | undefined; + let errorMethod: WebSocketFetchMethodMeta | undefined; + let closeMethod: WebSocketFetchMethodMeta | undefined; + + for (const methodName of methodNames) { + const builder = new WebSocketFetchControllerMethodMetaBuilder(this.clazz, methodName); + const methodMeta = builder.build(); + if (!methodMeta) { + continue; + } + switch (methodMeta.methodType) { + case WebSocketFetchMethodType.DATA: { + methods.push(methodMeta); + break; + } + case WebSocketFetchMethodType.CONNECTION: { + this.assertSingleLifecycleMethod(connectionMethod, methodMeta.methodType); + connectionMethod = methodMeta; + break; + } + case WebSocketFetchMethodType.OPEN: { + this.assertSingleLifecycleMethod(openMethod, methodMeta.methodType); + openMethod = methodMeta; + break; + } + case WebSocketFetchMethodType.ERROR: { + this.assertSingleLifecycleMethod(errorMethod, methodMeta.methodType); + errorMethod = methodMeta; + break; + } + case WebSocketFetchMethodType.CLOSE: { + this.assertSingleLifecycleMethod(closeMethod, methodMeta.methodType); + closeMethod = methodMeta; + break; + } + default: + assert.fail('never arrive'); + } + } + return { methods, connectionMethod, openMethod, errorMethod, closeMethod }; + } + + private assertSingleLifecycleMethod(methodMeta: WebSocketFetchMethodMeta | undefined, methodType: WebSocketFetchMethodType) { + if (!methodMeta) { + return; + } + const classDesc = ClassUtil.classDescription(this.clazz); + throw new Error(`build websocket fetch controller ${classDesc} failed: duplicate ${methodType} lifecycle method`); + } + + build(): WebSocketFetchControllerMeta { + ControllerValidator.validate(this.clazz); + const controllerType = ControllerInfoUtil.getControllerType(this.clazz); + assert(controllerType === ControllerType.WEBSOCKET_FETCH, 'invalidate controller type'); + const webSocketPath = WebSocketInfoUtil.getWebSocketPath(this.clazz); + const middlewares = ControllerInfoUtil.getControllerMiddlewares(this.clazz); + const { methods, connectionMethod, openMethod, errorMethod, closeMethod } = this.buildMethods(); + assert(methods.length === 1, 'websocket fetch controller must has exactly one @WebSocketFetchMethod'); + const clazzName = this.clazz.name; + const controllerName = ControllerInfoUtil.getControllerName(this.clazz) || clazzName; + const property = PrototypeUtil.getProperty(this.clazz); + const protoName = property!.name as string; + const hosts = ControllerInfoUtil.getControllerHosts(this.clazz); + const metadata = new WebSocketFetchControllerMeta( + clazzName, protoName, controllerName, webSocketPath, middlewares, methods, hosts, + connectionMethod, openMethod, errorMethod, closeMethod); + ControllerMetadataUtil.setControllerMetadata(this.clazz, metadata); + for (const method of metadata.methods) { + const realPath = metadata.getMethodRealPath(method); + if (!realPath.startsWith('/')) { + const desc = ClassUtil.classDescription(this.clazz); + throw new Error(`class ${desc} method ${method.name} path ${realPath} not start with /`); + } + } + return metadata; + } + + static create(clazz: EggProtoImplClass) { + return new WebSocketFetchControllerMetaBuilder(clazz); + } +} + +ControllerMetaBuilderFactory.registerControllerMetaBuilder(ControllerType.WEBSOCKET_FETCH, WebSocketFetchControllerMetaBuilder.create); diff --git a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts new file mode 100644 index 000000000..3197e223d --- /dev/null +++ b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts @@ -0,0 +1,150 @@ +import path from 'node:path'; +import { ClassUtil } from '@eggjs/tegg-metadata'; +import { ControllerType, WebSocketFetchMethodType, WebSocketParamType } from '@eggjs/tegg-types'; +import type { EggProtoImplClass } from '@eggjs/tegg-types'; +import { WebSocketFetchMethodMeta, WebSocketParamMeta, WebSocketParamMetaUtil } from '../../model'; +import { MethodValidator } from '../../util/validator/MethodValidator'; +import MethodInfoUtil from '../../util/MethodInfoUtil'; +import { HTTPPriorityUtil } from '../../util/HTTPPriorityUtil'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +export class WebSocketFetchControllerMethodMetaBuilder { + private readonly clazz: EggProtoImplClass; + private readonly methodName: string; + + constructor(clazz: EggProtoImplClass, methodName: string) { + this.clazz = clazz; + this.methodName = methodName; + } + + private checkParamDecorators(methodType: WebSocketFetchMethodType) { + const method = this.clazz.prototype[this.methodName]; + const functionLength = method.length; + const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); + const contextIndex = MethodInfoUtil.getMethodContextIndex(this.clazz, this.methodName); + const hasAnnotationParamCount = typeof contextIndex === 'undefined' + ? paramIndexList.length + : paramIndexList.length + 1; + const maxParamCount = Math.max(functionLength, hasAnnotationParamCount); + + for (let i = 0; i < maxParamCount; ++i) { + if (i === contextIndex) { + continue; + } + const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(i, this.clazz, this.methodName); + if (!paramType) { + const classDesc = ClassUtil.classDescription(this.clazz); + throw new Error(`${classDesc}:${this.methodName} param ${i} has no websocket fetch param type, Please add @WebSocketData, @WebSocketClose, @WebSocketError, @WebSocketCloseCode, @WebSocketCloseReason, @WebSocketParam, @WebSocketQuery, @WebSocketQueries, @WebSocketHeaders, @WebSocketRequest, @WebSocketSocket or @Context`); + } + if (!this.isAllowedParamType(methodType, paramType)) { + const classDesc = ClassUtil.classDescription(this.clazz); + throw new Error(`${classDesc}:${this.methodName} param ${i} type ${paramType} is not allowed in websocket fetch ${methodType} method`); + } + } + } + + private isAllowedParamType(methodType: WebSocketFetchMethodType, paramType: WebSocketParamType) { + if (paramType === WebSocketParamType.STREAM) { + return false; + } + if (this.isCommonParamType(paramType)) { + return true; + } + if (paramType === WebSocketParamType.CLOSE) { + return methodType !== WebSocketFetchMethodType.CLOSE; + } + switch (methodType) { + case WebSocketFetchMethodType.DATA: { + return paramType === WebSocketParamType.DATA; + } + case WebSocketFetchMethodType.ERROR: { + return paramType === WebSocketParamType.ERROR; + } + case WebSocketFetchMethodType.CLOSE: { + return paramType === WebSocketParamType.CLOSE_CODE || + paramType === WebSocketParamType.CLOSE_REASON; + } + case WebSocketFetchMethodType.CONNECTION: + case WebSocketFetchMethodType.OPEN: { + return false; + } + default: + return false; + } + } + + private isCommonParamType(paramType: WebSocketParamType) { + return paramType === WebSocketParamType.PARAM || + paramType === WebSocketParamType.QUERY || + paramType === WebSocketParamType.QUERIES || + paramType === WebSocketParamType.HEADERS || + paramType === WebSocketParamType.REQUEST || + paramType === WebSocketParamType.SOCKET; + } + + private buildParamType(webSocketPath: string, methodType: WebSocketFetchMethodType): Map { + this.checkParamDecorators(methodType); + + const paramTypeMap = new Map(); + const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); + for (const paramIndex of paramIndexList) { + const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(paramIndex, this.clazz, this.methodName)!; + const paramName = WebSocketInfoUtil.getWebSocketMethodParamName(paramIndex, this.clazz, this.methodName); + const paramMeta = WebSocketParamMetaUtil.createParam(paramType, paramName); + + if (methodType === WebSocketFetchMethodType.DATA) { + try { + paramMeta.validate(webSocketPath); + } catch (e) { + const classDesc = ClassUtil.classDescription(this.clazz); + e.message = `build websocket fetch controller ${classDesc} method ${this.methodName} param ${paramName} failed: ${e.message}`; + throw e; + } + } + + paramTypeMap.set(paramIndex, paramMeta); + } + return paramTypeMap; + } + + getPriority() { + const priority = WebSocketInfoUtil.getWebSocketMethodPriority(this.clazz, this.methodName); + if (priority !== undefined) { + return priority; + } + const controllerPath = WebSocketInfoUtil.getWebSocketPath(this.clazz); + const methodPath = WebSocketInfoUtil.getWebSocketMethodPath(this.clazz, this.methodName) || ''; + const realPath = controllerPath ? path.posix.join(controllerPath, methodPath) : methodPath; + const defaultPriority = HTTPPriorityUtil.calcPathPriority(realPath); + if (defaultPriority > HTTPPriorityUtil.DEFAULT_PRIORITY) { + throw new Error(`path ${realPath} is too long, should set priority manually`); + } + return defaultPriority; + } + + build(): WebSocketFetchMethodMeta | undefined { + MethodValidator.validate(this.clazz, this.methodName); + const controllerType = MethodInfoUtil.getMethodControllerType(this.clazz, this.methodName); + if (!controllerType || controllerType !== ControllerType.WEBSOCKET_FETCH) { + return undefined; + } + const methodType = WebSocketInfoUtil.getWebSocketFetchMethodType(this.clazz, this.methodName); + if (!methodType) { + return undefined; + } + const parentPath = WebSocketInfoUtil.getWebSocketPath(this.clazz); + const webSocketPath = methodType === WebSocketFetchMethodType.DATA + ? WebSocketInfoUtil.getWebSocketMethodPath(this.clazz, this.methodName) || '' + : ''; + const contextIndex = MethodInfoUtil.getMethodContextIndex(this.clazz, this.methodName); + const middlewares = MethodInfoUtil.getMethodMiddlewares(this.clazz, this.methodName); + const hosts = MethodInfoUtil.getMethodHosts(this.clazz, this.methodName); + const realPath = parentPath + ? path.posix.join(parentPath, webSocketPath) + : webSocketPath; + const paramTypeMap = this.buildParamType(realPath, methodType); + const priority = methodType === WebSocketFetchMethodType.DATA ? this.getPriority() : HTTPPriorityUtil.DEFAULT_PRIORITY; + return new WebSocketFetchMethodMeta( + this.methodName, webSocketPath, middlewares, contextIndex, paramTypeMap, priority, hosts, methodType); + } +} diff --git a/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts index e73d756d6..213109efd 100644 --- a/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts @@ -1,5 +1,6 @@ import path from 'node:path'; import { ClassUtil } from '@eggjs/tegg-metadata'; +import { WebSocketParamType } from '@eggjs/tegg-types'; import type { EggProtoImplClass } from '@eggjs/tegg-types'; import { WebSocketMethodMeta, WebSocketParamMeta, WebSocketParamMetaUtil } from '../../model'; import { MethodValidator } from '../../util/validator/MethodValidator'; @@ -45,6 +46,10 @@ export class WebSocketControllerMethodMetaBuilder { const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); for (const paramIndex of paramIndexList) { const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(paramIndex, this.clazz, this.methodName)!; + if (this.isFetchOnlyParamType(paramType)) { + const classDesc = ClassUtil.classDescription(this.clazz); + throw new Error(`${classDesc}:${this.methodName} param ${paramIndex} is websocket fetch only`); + } const paramName = WebSocketInfoUtil.getWebSocketMethodParamName(paramIndex, this.clazz, this.methodName); const paramMeta = WebSocketParamMetaUtil.createParam(paramType, paramName); @@ -61,6 +66,14 @@ export class WebSocketControllerMethodMetaBuilder { return paramTypeMap; } + private isFetchOnlyParamType(paramType: WebSocketParamType) { + return paramType === WebSocketParamType.DATA || + paramType === WebSocketParamType.CLOSE || + paramType === WebSocketParamType.ERROR || + paramType === WebSocketParamType.CLOSE_CODE || + paramType === WebSocketParamType.CLOSE_REASON; + } + getPriority() { const priority = WebSocketInfoUtil.getWebSocketMethodPriority(this.clazz, this.methodName); if (priority !== undefined) { diff --git a/core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts b/core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts new file mode 100644 index 000000000..823f90fe3 --- /dev/null +++ b/core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts @@ -0,0 +1,73 @@ +import path from 'node:path'; +import type { ControllerMetadata, EggPrototypeName, MiddlewareFunc } from '@eggjs/tegg-types'; +import { ControllerType } from '@eggjs/tegg-types'; +import { WebSocketFetchMethodMeta } from './WebSocketFetchMethodMeta'; + +export class WebSocketFetchControllerMeta implements ControllerMetadata { + readonly protoName: EggPrototypeName; + readonly controllerName: string; + readonly className: string; + public readonly type = ControllerType.WEBSOCKET_FETCH; + public readonly path?: string; + public readonly middlewares: readonly MiddlewareFunc[]; + public readonly methods: readonly WebSocketFetchMethodMeta[]; + public readonly hosts?: string[]; + public readonly connectionMethod?: WebSocketFetchMethodMeta; + public readonly openMethod?: WebSocketFetchMethodMeta; + public readonly errorMethod?: WebSocketFetchMethodMeta; + public readonly closeMethod?: WebSocketFetchMethodMeta; + + constructor( + className: string, + protoName: EggPrototypeName, + controllerName: string, + path: string | undefined, + middlewares: MiddlewareFunc[], + methods: WebSocketFetchMethodMeta[], + hosts: string[] | undefined, + connectionMethod: WebSocketFetchMethodMeta | undefined, + openMethod: WebSocketFetchMethodMeta | undefined, + errorMethod: WebSocketFetchMethodMeta | undefined, + closeMethod: WebSocketFetchMethodMeta | undefined, + ) { + this.protoName = protoName; + this.controllerName = controllerName; + this.className = className; + this.path = path; + this.middlewares = middlewares; + this.methods = methods; + this.hosts = hosts; + this.connectionMethod = connectionMethod; + this.openMethod = openMethod; + this.errorMethod = errorMethod; + this.closeMethod = closeMethod; + } + + getMethodRealPath(method: WebSocketFetchMethodMeta) { + if (this.path) { + return path.posix.join(this.path, method.path); + } + return method.path; + } + + getMethodHosts(method: WebSocketFetchMethodMeta): string[] | undefined { + if (this.hosts) { + return this.hosts; + } + return method.hosts; + } + + getMethodName(method: WebSocketFetchMethodMeta) { + return `WEBSOCKET_FETCH ${this.controllerName}.${method.name}`; + } + + getMethodMiddlewares(method: WebSocketFetchMethodMeta) { + if (this.middlewares.length) { + return [ + ...this.middlewares, + ...method.middlewares, + ]; + } + return [ ...method.middlewares ]; + } +} diff --git a/core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts b/core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts new file mode 100644 index 000000000..ec0546a0b --- /dev/null +++ b/core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts @@ -0,0 +1,34 @@ +import { WebSocketFetchMethodType } from '@eggjs/tegg-types'; +import type { MethodMeta, MiddlewareFunc } from '@eggjs/tegg-types'; +import { WebSocketParamMeta } from './WebSocketMethodMeta'; + +export class WebSocketFetchMethodMeta implements MethodMeta { + public readonly name: string; + public readonly path: string; + public readonly middlewares: readonly MiddlewareFunc[]; + public readonly contextParamIndex: number | undefined; + public readonly paramMap: Map; + public readonly priority: number; + public readonly hosts: string[] | undefined; + public readonly methodType: WebSocketFetchMethodType; + + constructor( + name: string, + path: string, + middlewares: MiddlewareFunc[], + contextParamIndex: number | undefined, + paramMap: Map, + priority: number, + hosts: string[] | undefined, + methodType: WebSocketFetchMethodType, + ) { + this.name = name; + this.path = path; + this.middlewares = middlewares; + this.contextParamIndex = contextParamIndex; + this.paramMap = paramMap; + this.priority = priority; + this.hosts = hosts; + this.methodType = methodType; + } +} diff --git a/core/controller-decorator/src/model/WebSocketMethodMeta.ts b/core/controller-decorator/src/model/WebSocketMethodMeta.ts index 8fe7d6827..527783a9b 100644 --- a/core/controller-decorator/src/model/WebSocketMethodMeta.ts +++ b/core/controller-decorator/src/model/WebSocketMethodMeta.ts @@ -87,6 +87,46 @@ export class WebSocketStreamParamMeta extends WebSocketParamMeta { } } +export class WebSocketDataParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.DATA; + + validate() { + return; + } +} + +export class WebSocketCloseParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.CLOSE; + + validate() { + return; + } +} + +export class WebSocketErrorParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.ERROR; + + validate() { + return; + } +} + +export class WebSocketCloseCodeParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.CLOSE_CODE; + + validate() { + return; + } +} + +export class WebSocketCloseReasonParamMeta extends WebSocketParamMeta { + type = WebSocketParamType.CLOSE_REASON; + + validate() { + return; + } +} + export class WebSocketMethodMeta implements MethodMeta { public readonly name: string; public readonly path: string; @@ -142,6 +182,21 @@ export class WebSocketParamMetaUtil { case WebSocketParamType.STREAM: { return new WebSocketStreamParamMeta(); } + case WebSocketParamType.DATA: { + return new WebSocketDataParamMeta(); + } + case WebSocketParamType.CLOSE: { + return new WebSocketCloseParamMeta(); + } + case WebSocketParamType.ERROR: { + return new WebSocketErrorParamMeta(); + } + case WebSocketParamType.CLOSE_CODE: { + return new WebSocketCloseCodeParamMeta(); + } + case WebSocketParamType.CLOSE_REASON: { + return new WebSocketCloseReasonParamMeta(); + } default: assert.fail('never arrive'); } diff --git a/core/controller-decorator/src/model/index.ts b/core/controller-decorator/src/model/index.ts index f64e67dc0..80ed915ae 100644 --- a/core/controller-decorator/src/model/index.ts +++ b/core/controller-decorator/src/model/index.ts @@ -5,6 +5,8 @@ export * from './HTTPResponse'; export * from './HTTPCookies'; export * from './WebSocketMethodMeta'; export * from './WebSocketControllerMeta'; +export * from './WebSocketFetchMethodMeta'; +export * from './WebSocketFetchControllerMeta'; export * from './MCPControllerMeta'; export * from './MCPPromptMeta'; export * from './MCPResourceMeta'; diff --git a/core/controller-decorator/src/util/WebSocketInfoUtil.ts b/core/controller-decorator/src/util/WebSocketInfoUtil.ts index e8f9babdb..0a9fed55b 100644 --- a/core/controller-decorator/src/util/WebSocketInfoUtil.ts +++ b/core/controller-decorator/src/util/WebSocketInfoUtil.ts @@ -4,15 +4,17 @@ import { CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP, CONTROLLER_WEBSOCKET_METHOD_PATH_MAP, CONTROLLER_WEBSOCKET_METHOD_PRIORITY, + CONTROLLER_WEBSOCKET_FETCH_METHOD_TYPE_MAP, CONTROLLER_WEBSOCKET_PATH, } from '@eggjs/tegg-types'; -import type { EggProtoImplClass, WebSocketParamType } from '@eggjs/tegg-types'; +import type { EggProtoImplClass, WebSocketFetchMethodType, WebSocketParamType } from '@eggjs/tegg-types'; import { MapUtil } from '@eggjs/tegg-common-util'; type WebSocketMethodPathMap = Map; type WebSocketMethodParamTypeMap = Map>; type WebSocketMethodParamNameMap = Map>; type WebSocketMethodPriorityMap = Map; +type WebSocketFetchMethodTypeMap = Map; export default class WebSocketInfoUtil { static setWebSocketPath(path: string, clazz: EggProtoImplClass) { @@ -75,4 +77,14 @@ export default class WebSocketInfoUtil { const methodPriorityMap: WebSocketMethodPriorityMap = MetadataUtil.initOwnMapMetaData(CONTROLLER_WEBSOCKET_METHOD_PRIORITY, clazz, new Map()); methodPriorityMap.set(methodName, priority); } + + static setWebSocketFetchMethodType(type: WebSocketFetchMethodType, clazz: EggProtoImplClass, methodName: string) { + const methodTypeMap: WebSocketFetchMethodTypeMap = MetadataUtil.initOwnMapMetaData(CONTROLLER_WEBSOCKET_FETCH_METHOD_TYPE_MAP, clazz, new Map()); + methodTypeMap.set(methodName, type); + } + + static getWebSocketFetchMethodType(clazz: EggProtoImplClass, methodName: string): WebSocketFetchMethodType | undefined { + const methodTypeMap: WebSocketFetchMethodTypeMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_FETCH_METHOD_TYPE_MAP, clazz); + return methodTypeMap?.get(methodName); + } } diff --git a/core/types/controller-decorator/MetadataKey.ts b/core/types/controller-decorator/MetadataKey.ts index 7faecad30..901822b49 100644 --- a/core/types/controller-decorator/MetadataKey.ts +++ b/core/types/controller-decorator/MetadataKey.ts @@ -20,6 +20,7 @@ export const CONTROLLER_WEBSOCKET_METHOD_PATH_MAP = Symbol.for('EggPrototype#con export const CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP = Symbol.for('EggPrototype#controller#method#websocket#params#type'); export const CONTROLLER_WEBSOCKET_METHOD_PARAM_NAME_MAP = Symbol.for('EggPrototype#controller#method#websocket#params#name'); export const CONTROLLER_WEBSOCKET_METHOD_PRIORITY = Symbol.for('EggPrototype#controller#method#websocket#priority'); +export const CONTROLLER_WEBSOCKET_FETCH_METHOD_TYPE_MAP = Symbol.for('EggPrototype#controller#method#websocketFetch#type'); export const METHOD_CONTROLLER_TYPE_MAP = Symbol.for('EggPrototype#controller#mthods'); export const METHOD_CONTROLLER_HOST = Symbol.for('EggPrototype#controller#mthods#host'); diff --git a/core/types/controller-decorator/WebSocketFetchController.ts b/core/types/controller-decorator/WebSocketFetchController.ts new file mode 100644 index 000000000..5cb488dc9 --- /dev/null +++ b/core/types/controller-decorator/WebSocketFetchController.ts @@ -0,0 +1,5 @@ +export interface WebSocketFetchControllerParams { + path?: string; + controllerName?: string; + protoName?: string; +} diff --git a/core/types/controller-decorator/WebSocketFetchMethod.ts b/core/types/controller-decorator/WebSocketFetchMethod.ts new file mode 100644 index 000000000..df43f5ee7 --- /dev/null +++ b/core/types/controller-decorator/WebSocketFetchMethod.ts @@ -0,0 +1,3 @@ +export interface WebSocketFetchMethodParams { + priority?: number; +} diff --git a/core/types/controller-decorator/WebSocketFetchParam.ts b/core/types/controller-decorator/WebSocketFetchParam.ts new file mode 100644 index 000000000..d58aa9da5 --- /dev/null +++ b/core/types/controller-decorator/WebSocketFetchParam.ts @@ -0,0 +1 @@ +export type WebSocketFetchClose = (code?: number, reason?: string | Buffer) => void; diff --git a/core/types/controller-decorator/index.ts b/core/types/controller-decorator/index.ts index db3b985a5..98dda1a6b 100644 --- a/core/types/controller-decorator/index.ts +++ b/core/types/controller-decorator/index.ts @@ -9,6 +9,9 @@ export * from './WebSocketController'; export * from './WebSocketMethod'; export * from './WebSocketParam'; export * from './WebSocketContext'; +export * from './WebSocketFetchController'; +export * from './WebSocketFetchMethod'; +export * from './WebSocketFetchParam'; export * from './MetadataKey'; export * from './MCPController'; export * from './MCPPromptParams'; diff --git a/core/types/controller-decorator/model/types.ts b/core/types/controller-decorator/model/types.ts index c7bd158b5..bfdb0805f 100644 --- a/core/types/controller-decorator/model/types.ts +++ b/core/types/controller-decorator/model/types.ts @@ -8,6 +8,7 @@ export type { IncomingHttpHeaders } from 'node:http'; export enum ControllerType { HTTP = 'HTTP', WEBSOCKET = 'WEBSOCKET', + WEBSOCKET_FETCH = 'WEBSOCKET_FETCH', MCP = 'MCP', SOFA_RPC = 'SOFA_RPC', SOFA_RPC_STREAM = 'SOFA_RPC_STREAM', @@ -25,6 +26,7 @@ export type ControllerTypeLike = ControllerType | string; export enum MethodType { HTTP = 'HTTP', WEBSOCKET = 'WEBSOCKET', + WEBSOCKET_FETCH = 'WEBSOCKET_FETCH', SOFA_RPC = 'SOFA_RPC', SOFA_RPC_STREAM = 'SOFA_RPC_STREAM', MGW_RPC = 'MGW_RPC', @@ -63,6 +65,19 @@ export enum WebSocketParamType { REQUEST = 'REQUEST', SOCKET = 'SOCKET', STREAM = 'STREAM', + DATA = 'DATA', + CLOSE = 'CLOSE', + ERROR = 'ERROR', + CLOSE_CODE = 'CLOSE_CODE', + CLOSE_REASON = 'CLOSE_REASON', +} + +export enum WebSocketFetchMethodType { + DATA = 'DATA', + CONNECTION = 'CONNECTION', + OPEN = 'OPEN', + ERROR = 'ERROR', + CLOSE = 'CLOSE', } export enum MCPProtocols { diff --git a/plugin/controller/app.ts b/plugin/controller/app.ts index 9eb7a63ae..3bcca2760 100644 --- a/plugin/controller/app.ts +++ b/plugin/controller/app.ts @@ -62,6 +62,7 @@ export default class ControllerAppBootHook { }); this.controllerRegisterFactory.registerControllerRegister(ControllerType.HTTP, HTTPControllerRegister.create); this.controllerRegisterFactory.registerControllerRegister(ControllerType.WEBSOCKET, WebSocketControllerRegister.create); + this.controllerRegisterFactory.registerControllerRegister(ControllerType.WEBSOCKET_FETCH, WebSocketControllerRegister.create); this.app.loadUnitFactory.registerLoadUnitCreator( CONTROLLER_LOAD_UNIT, (ctx: LoadUnitLifecycleContext): ControllerLoadUnit => { diff --git a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts index 6be45b800..c3da3d308 100644 --- a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts +++ b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts @@ -6,13 +6,16 @@ import pathToRegexp from 'path-to-regexp'; import { Application, Context } from 'egg'; import { EggRouter } from '@eggjs/router'; import { FrameworkErrorFormater } from 'egg-errors'; -import { createWebSocketStream, WebSocket, WebSocketServer } from 'ws'; +import { createWebSocketStream, RawData, WebSocket, WebSocketServer } from 'ws'; import { CONTROLLER_META_DATA, ControllerMetadata, ControllerType, Next, WebSocketControllerMeta, + WebSocketFetchClose, + WebSocketFetchControllerMeta, + WebSocketFetchMethodMeta, WebSocketMethodMeta, WebSocketParamType, WebSocketPathParamMeta, @@ -32,8 +35,8 @@ const noop = () => { interface WebSocketRoute { controllerProto: EggPrototype; - controllerMeta: WebSocketControllerMeta; - methodMeta: WebSocketMethodMeta; + controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta; + methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta; methodRealPath: string; methodName: string; host?: string; @@ -41,6 +44,15 @@ interface WebSocketRoute { regexp: RegExp; } +interface WebSocketMethodPayload { + data?: RawData; + close?: WebSocketFetchClose; + error?: Error; + closeCode?: number; + closeReason?: Buffer; + getWebSocketStream?: () => Duplex; +} + export class WebSocketControllerRegister implements ControllerRegister { static instance?: WebSocketControllerRegister; @@ -55,8 +67,8 @@ export class WebSocketControllerRegister implements ControllerRegister { static create(proto: EggPrototype, controllerMeta: ControllerMetadata, app: Application) { assert( - controllerMeta.type === ControllerType.WEBSOCKET, - 'controller meta type is not WEBSOCKET', + controllerMeta.type === ControllerType.WEBSOCKET || controllerMeta.type === ControllerType.WEBSOCKET_FETCH, + 'controller meta type is not WEBSOCKET or WEBSOCKET_FETCH', ); if (!WebSocketControllerRegister.instance) { WebSocketControllerRegister.instance = new WebSocketControllerRegister(app); @@ -82,9 +94,9 @@ export class WebSocketControllerRegister implements ControllerRegister { } doRegister() { - const methodMap = new Map(); + const methodMap = new Map(); for (const proto of this.controllerProtos) { - const metadata = proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta; + const metadata = proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; for (const method of metadata.methods) { methodMap.set(method, proto); } @@ -94,13 +106,13 @@ export class WebSocketControllerRegister implements ControllerRegister { for (const method of allMethods) { const controllerProto = methodMap.get(method)!; - const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta; + const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; this.checkDuplicate(controllerMeta, method); } this.routes = allMethods.flatMap(method => { const controllerProto = methodMap.get(method)!; - const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta; + const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; const hosts = this.getMethodHosts(controllerMeta, method) || [ undefined ]; return hosts.map(host => this.createRoute(controllerProto, controllerMeta, method, host)); }); @@ -142,8 +154,8 @@ export class WebSocketControllerRegister implements ControllerRegister { private createRoute( controllerProto: EggPrototype, - controllerMeta: WebSocketControllerMeta, - methodMeta: WebSocketMethodMeta, + controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, + methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, host: string | undefined, ): WebSocketRoute { const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); @@ -161,7 +173,7 @@ export class WebSocketControllerRegister implements ControllerRegister { }; } - private checkDuplicate(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { + private checkDuplicate(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { let router = this.checkRouters.get('default')!; const hosts = this.getMethodHosts(controllerMeta, methodMeta) || []; if (!hosts.length) { @@ -181,13 +193,13 @@ export class WebSocketControllerRegister implements ControllerRegister { }); } - private registerToRouter(router: EggRouter, controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { + private registerToRouter(router: EggRouter, controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); const methodName = this.getMethodName(controllerMeta, methodMeta); Reflect.apply(router.get, router, [ methodName, methodRealPath, noop ]); } - private checkDuplicateInRouter(router: EggRouter, controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { + private checkDuplicateInRouter(router: EggRouter, controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); const matched = router.match(methodRealPath, 'GET'); const methodName = this.getMethodName(controllerMeta, methodMeta); @@ -198,20 +210,20 @@ export class WebSocketControllerRegister implements ControllerRegister { } } - private getMethodRealPath(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { - return controllerMeta.getMethodRealPath(methodMeta); + private getMethodRealPath(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + return (controllerMeta as any).getMethodRealPath(methodMeta); } - private getMethodHosts(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta): string[] | undefined { - return controllerMeta.getMethodHosts(methodMeta); + private getMethodHosts(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): string[] | undefined { + return (controllerMeta as any).getMethodHosts(methodMeta); } - private getMethodName(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta): string { - return controllerMeta.getMethodName(methodMeta); + private getMethodName(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): string { + return (controllerMeta as any).getMethodName(methodMeta); } - private getMethodMiddlewares(controllerMeta: WebSocketControllerMeta, methodMeta: WebSocketMethodMeta) { - return controllerMeta.getMethodMiddlewares(methodMeta); + private getMethodMiddlewares(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + return (controllerMeta as any).getMethodMiddlewares(methodMeta); } private async handleUpgrade(req: IncomingMessage, socket: Duplex, head: Buffer) { @@ -283,7 +295,11 @@ export class WebSocketControllerRegister implements ControllerRegister { let invoked = false; const handler = async (_ctx: Context, next: Next) => { invoked = true; - await this.invokeController(route, webSocketCtx); + if (route.controllerMeta.type === ControllerType.WEBSOCKET_FETCH) { + await this.invokeFetchController(route, webSocketCtx, eggCtx); + } else { + await this.invokeController(route, webSocketCtx); + } await next(); }; const composed = compose([ ...methodMiddlewares, handler ]); @@ -304,14 +320,9 @@ export class WebSocketControllerRegister implements ControllerRegister { private async invokeController(route: WebSocketRoute, webSocketCtx: WebSocketContextImpl) { const methodMeta = route.methodMeta as WebSocketMethodMeta; - const argsLength = methodMeta.paramMap.size; - const hasContext = methodMeta.contextParamIndex !== undefined; - const contextIndex = methodMeta.contextParamIndex; - const methodArgsLength = argsLength + (hasContext ? 1 : 0); const eggObj = await this.eggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); const realObj = eggObj.obj; const realMethod = realObj[methodMeta.name]; - const args: Array = new Array(methodArgsLength); let webSocketStream: Duplex | undefined; const getWebSocketStream = () => { if (!webSocketStream) { @@ -319,6 +330,21 @@ export class WebSocketControllerRegister implements ControllerRegister { } return webSocketStream; }; + const args = this.buildMethodArgs(methodMeta, webSocketCtx, { getWebSocketStream }); + const result = await Reflect.apply(realMethod, realObj, args); + this.pipeResponseStream(result, webSocketCtx.socket, getWebSocketStream); + } + + private buildMethodArgs( + methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, + webSocketCtx: WebSocketContextImpl, + payload: WebSocketMethodPayload, + ) { + const argsLength = methodMeta.paramMap.size; + const hasContext = methodMeta.contextParamIndex !== undefined; + const contextIndex = methodMeta.contextParamIndex; + const methodArgsLength = argsLength + (hasContext ? 1 : 0); + const args: unknown[] = new Array(methodArgsLength); if (hasContext) { args[contextIndex!] = webSocketCtx; } @@ -352,16 +378,158 @@ export class WebSocketControllerRegister implements ControllerRegister { break; } case WebSocketParamType.STREAM: { - args[index] = getWebSocketStream(); + assert(payload.getWebSocketStream, '@WebSocketStream can not be used here'); + args[index] = payload.getWebSocketStream(); + break; + } + case WebSocketParamType.DATA: { + args[index] = payload.data; + break; + } + case WebSocketParamType.CLOSE: { + args[index] = payload.close; + break; + } + case WebSocketParamType.ERROR: { + args[index] = payload.error; + break; + } + case WebSocketParamType.CLOSE_CODE: { + args[index] = payload.closeCode; + break; + } + case WebSocketParamType.CLOSE_REASON: { + args[index] = payload.closeReason; break; } default: assert.fail('never arrive'); } } + return args; + } - const result = await Reflect.apply(realMethod, realObj, args); - this.pipeResponseStream(result, webSocketCtx.socket, getWebSocketStream); + private async invokeFetchController(route: WebSocketRoute, webSocketCtx: WebSocketContextImpl, eggCtx: Context) { + const controllerMeta = route.controllerMeta as WebSocketFetchControllerMeta; + const methodMeta = route.methodMeta as WebSocketFetchMethodMeta; + const eggObj = await this.eggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); + const realObj = eggObj.obj; + const close = this.createFetchClose(webSocketCtx.socket); + const invokeMethod = async ( + targetMethodMeta: WebSocketFetchMethodMeta | undefined, + payload: WebSocketMethodPayload = {}, + ) => { + if (!targetMethodMeta) { + return; + } + const realMethod = realObj[targetMethodMeta.name]; + const args = this.buildMethodArgs(targetMethodMeta, webSocketCtx, { + close, + ...payload, + }); + return await Reflect.apply(realMethod, realObj, args); + }; + const handleError = async (error: Error) => { + if (!controllerMeta.errorMethod) { + this.app.logger.error('[tegg/websocket-fetch] handle error: %s', error.stack || error.message); + return; + } + await invokeMethod(controllerMeta.errorMethod, { error }); + }; + const closeHandled = new Promise(resolve => { + webSocketCtx.socket.once('close', (code, reason) => { + this.app.ctxStorage.run(eggCtx, async () => { + try { + await invokeMethod(controllerMeta.closeMethod, { + closeCode: code, + closeReason: reason, + }); + } catch (error) { + this.app.logger.error('[tegg/websocket-fetch] handle close failed: %s', error.stack || error.message); + } finally { + resolve(); + } + }).catch(error => { + this.app.logger.error('[tegg/websocket-fetch] handle close failed: %s', error.stack || error.message); + resolve(); + }); + }); + }); + + webSocketCtx.socket.on('message', data => { + this.app.ctxStorage.run(eggCtx, async () => { + try { + const result = await invokeMethod(methodMeta, { data }); + this.sendFetchResponseStream(result, webSocketCtx.socket, handleError); + } catch (error) { + await handleError(error); + } + }).catch(error => { + this.app.logger.error('[tegg/websocket-fetch] handle message failed: %s', error.stack || error.message); + }); + }); + webSocketCtx.socket.on('error', error => { + this.app.ctxStorage.run(eggCtx, async () => { + await handleError(error); + }).catch(error => { + this.app.logger.error('[tegg/websocket-fetch] handle socket error failed: %s', error.stack || error.message); + }); + }); + + await invokeMethod(controllerMeta.connectionMethod); + await invokeMethod(controllerMeta.openMethod); + await closeHandled; + } + + private createFetchClose(webSocket: WebSocket): WebSocketFetchClose { + return (code = 1000, reason?: string | Buffer) => { + if (webSocket.readyState !== WebSocket.OPEN && webSocket.readyState !== WebSocket.CONNECTING) { + return; + } + webSocket.close(code, reason); + }; + } + + private sendFetchResponseStream( + result: unknown, + webSocket: WebSocket, + handleError: (error: Error) => Promise, + ) { + if (result === undefined || result === null) { + return; + } + if (!this.isReadableStream(result)) { + handleError(new Error('WebSocketFetch method must return a readable stream or void')).catch(error => { + this.app.logger.error('[tegg/websocket-fetch] handle invalid response failed: %s', error.stack || error.message); + }); + return; + } + result.on('data', chunk => { + this.sendFetchChunk(webSocket, chunk, handleError); + }); + result.once('error', error => { + handleError(error).catch(err => { + this.app.logger.error('[tegg/websocket-fetch] handle stream error failed: %s', err.stack || err.message); + }); + }); + } + + private sendFetchChunk( + webSocket: WebSocket, + chunk: unknown, + handleError: (error: Error) => Promise, + ) { + if (webSocket.readyState !== WebSocket.OPEN) { + return; + } + webSocket.send(chunk as any, error => { + if (!error) { + return; + } + handleError(error).catch(err => { + this.app.logger.error('[tegg/websocket-fetch] handle send error failed: %s', err.stack || err.message); + }); + }); } private pipeResponseStream(result: unknown, webSocket: WebSocket, getWebSocketStream: () => Duplex) { diff --git a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts index a29793590..9198921c6 100644 --- a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts @@ -1,6 +1,6 @@ import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; import { PassThrough, type Readable } from 'node:stream'; -import type { WebSocket } from 'ws'; +import type { RawData, WebSocket } from 'ws'; import { Context, Headers, @@ -11,12 +11,31 @@ import { Request, WebSocketContext, WebSocketController, + WebSocketClose, + WebSocketCloseCode, + WebSocketCloseReason, + WebSocketData, + WebSocketError, + WebSocketFetchClose, + WebSocketFetchController, + WebSocketFetchMethod, + WebSocketFetchOnClose, + WebSocketFetchOnConnection, + WebSocketFetchOnError, + WebSocketFetchOnOpen, WebSocketMethod, WebSocketSocket, WebSocketStream, } from '@eggjs/tegg'; import AppService from '../../modules/multi-module-service/AppService'; +interface WebSocketFetchRequest { + content: string; + close?: boolean; + error?: boolean; + delayMs?: number; +} + @WebSocketController({ path: '/ws', }) @@ -92,3 +111,110 @@ export class AppWebSocketController { return output; } } + +@WebSocketFetchController({ + path: '/ws-fetch/:id', +}) +export class AppWebSocketFetchController { + @Inject() + appService: AppService; + + @WebSocketFetchOnConnection() + onConnection( + @Headers() headers: IncomingHttpHeaders, + @Context() ctx: WebSocketContext, + ) { + ctx.socket.send(JSON.stringify({ + type: 'connection', + header: headers['x-client-id'], + path: ctx.path, + pid: process.pid, + })); + } + + @WebSocketFetchOnOpen() + onOpen( + @Context() ctx: WebSocketContext, + ) { + ctx.socket.send(JSON.stringify({ + type: 'open', + path: ctx.path, + pid: process.pid, + })); + } + + @WebSocketFetchMethod() + onData( + @WebSocketData() data: RawData, + @WebSocketClose() close: WebSocketFetchClose, + @Param() id: string, + @Query() name: string, + @Queries({ name: 'tag' }) tags: string[], + @Headers() headers: IncomingHttpHeaders, + @Context() ctx: WebSocketContext, + ) { + const body = JSON.parse(data.toString()) as WebSocketFetchRequest; + const output = new PassThrough(); + output.write(JSON.stringify({ + type: 'data', + phase: 1, + id, + name, + tags, + header: headers['x-client-id'], + path: ctx.path, + content: body.content, + pid: process.pid, + })); + + setTimeout(() => { + if (output.destroyed) { + return; + } + if (body.error) { + output.destroy(new Error(`fetch error: ${body.content}`)); + return; + } + output.write(JSON.stringify({ + type: 'data', + phase: 2, + id, + content: body.content, + pid: process.pid, + })); + output.end(); + }, body.delayMs ?? 10); + + if (body.close) { + output.once('end', () => setImmediate(() => close(1000, 'server done'))); + } + return output; + } + + @WebSocketFetchOnError() + onError( + @WebSocketError() error: Error, + @Headers() headers: IncomingHttpHeaders, + @Context() ctx: WebSocketContext, + ) { + ctx.socket.send(JSON.stringify({ + type: 'error', + message: error.message, + header: headers['x-client-id'], + path: ctx.path, + pid: process.pid, + })); + } + + @WebSocketFetchOnClose() + async onClose( + @WebSocketCloseCode() code: number, + @WebSocketCloseReason() reason: Buffer, + @Param() id: string, + ) { + await this.appService.save({ + name: `ws-fetch-close-${id}`, + desc: `${code}:${reason.toString()}`, + }); + } +} diff --git a/plugin/controller/test/websocket/websocket.test.ts b/plugin/controller/test/websocket/websocket.test.ts index fd516a6ab..0145d7afa 100644 --- a/plugin/controller/test/websocket/websocket.test.ts +++ b/plugin/controller/test/websocket/websocket.test.ts @@ -29,6 +29,94 @@ async function receiveJSON(socket: WebSocket): Promise { return JSON.parse(data.toString()); } +function createJSONReceiver(socket: WebSocket) { + const messages: any[] = []; + const waiters: Array<{ resolve: (value: any) => void; reject: (error: Error) => void }> = []; + + socket.on('message', data => { + const message = JSON.parse(data.toString()); + const waiter = waiters.shift(); + if (waiter) { + waiter.resolve(message); + return; + } + messages.push(message); + }); + socket.once('error', error => { + while (waiters.length) { + waiters.shift()!.reject(error); + } + }); + socket.once('close', () => { + while (waiters.length) { + waiters.shift()!.reject(new Error('websocket closed before message')); + } + }); + + return { + next() { + if (messages.length) { + return Promise.resolve(messages.shift()); + } + return new Promise((resolve, reject) => { + waiters.push({ resolve, reject }); + }); + }, + }; +} + +function createTaggedJSONReceiver(entries: Array<{ name: string; socket: WebSocket }>) { + const messages: Array<{ name: string; message: any }> = []; + const waiters: Array<{ resolve: (value: { name: string; message: any }) => void; reject: (error: Error) => void }> = []; + + for (const entry of entries) { + entry.socket.on('message', data => { + const message = { + name: entry.name, + message: JSON.parse(data.toString()), + }; + const waiter = waiters.shift(); + if (waiter) { + waiter.resolve(message); + return; + } + messages.push(message); + }); + entry.socket.once('error', error => { + while (waiters.length) { + waiters.shift()!.reject(error); + } + }); + entry.socket.once('close', () => { + while (waiters.length) { + waiters.shift()!.reject(new Error(`websocket ${entry.name} closed before message`)); + } + }); + } + + return { + next() { + if (messages.length) { + return Promise.resolve(messages.shift()!); + } + return new Promise<{ name: string; message: any }>((resolve, reject) => { + waiters.push({ resolve, reject }); + }); + }, + }; +} + +async function receiveClose(socket: WebSocket): Promise<{ code: number; reason: string }> { + return await new Promise(resolve => { + socket.once('close', (code, reason) => { + resolve({ + code, + reason: reason.toString(), + }); + }); + }); +} + async function createClientWithFirstMessage(url: string, headers?: Record) { const socket = new WebSocket(url, { headers }); const firstMessage = receiveJSON(socket); @@ -137,6 +225,218 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { } }); + it('should handle websocket fetch controller', async () => { + const socket = new WebSocket(requestUrl(app, '/ws-fetch/fetch-1?name=bar&tag=a&tag=b'), { + headers: { 'x-client-id': 'client-fetch' }, + }); + const receiver = createJSONReceiver(socket); + await new Promise((resolve, reject) => { + socket.once('open', resolve); + socket.once('error', reject); + }); + try { + const connection = await receiver.next(); + assert.deepEqual(connection, { + type: 'connection', + header: 'client-fetch', + path: '/ws-fetch/fetch-1', + pid: connection.pid, + }); + assert.equal(typeof connection.pid, 'number'); + + const open = await receiver.next(); + assert.deepEqual(open, { + type: 'open', + path: '/ws-fetch/fetch-1', + pid: open.pid, + }); + assert.equal(typeof open.pid, 'number'); + + socket.send(JSON.stringify({ + content: 'first', + })); + const first = await receiver.next(); + assert.deepEqual(first, { + type: 'data', + phase: 1, + id: 'fetch-1', + name: 'bar', + tags: [ 'a', 'b' ], + header: 'client-fetch', + path: '/ws-fetch/fetch-1', + content: 'first', + pid: first.pid, + }); + assert.equal(typeof first.pid, 'number'); + assert.deepEqual(await receiver.next(), { + type: 'data', + phase: 2, + id: 'fetch-1', + content: 'first', + pid: first.pid, + }); + assert.equal(socket.readyState, WebSocket.OPEN); + + socket.send(JSON.stringify({ + content: 'broken', + error: true, + })); + const beforeError = await receiver.next(); + assert.equal(beforeError.type, 'data'); + assert.equal(beforeError.phase, 1); + assert.equal(beforeError.content, 'broken'); + const error = await receiver.next(); + assert.deepEqual(error, { + type: 'error', + message: 'fetch error: broken', + header: 'client-fetch', + path: '/ws-fetch/fetch-1', + pid: error.pid, + }); + assert.equal(socket.readyState, WebSocket.OPEN); + + const closePromise = receiveClose(socket); + socket.send(JSON.stringify({ + content: 'done', + close: true, + })); + const closeData = await receiver.next(); + assert.equal(closeData.type, 'data'); + assert.equal(closeData.phase, 1); + assert.equal(closeData.content, 'done'); + assert.deepEqual(await receiver.next(), { + type: 'data', + phase: 2, + id: 'fetch-1', + content: 'done', + pid: closeData.pid, + }); + assert.deepEqual(await closePromise, { + code: 1000, + reason: 'server done', + }); + + const res = await app.httpRequest() + .get('/apps/ws-fetch-close-fetch-1') + .expect(200); + assert.equal(res.body.app.name, 'ws-fetch-close-fetch-1'); + assert.equal(res.body.app.desc, '1000:server done'); + } finally { + await closeClient(socket); + } + }); + + it('should isolate two websocket fetch connections with interleaved data', async () => { + const firstSocket = new WebSocket(requestUrl(app, '/ws-fetch/cross-1?name=first&tag=a'), { + headers: { 'x-client-id': 'client-cross-1' }, + }); + const secondSocket = new WebSocket(requestUrl(app, '/ws-fetch/cross-2?name=second&tag=b'), { + headers: { 'x-client-id': 'client-cross-2' }, + }); + const receiver = createTaggedJSONReceiver([ + { name: 'first', socket: firstSocket }, + { name: 'second', socket: secondSocket }, + ]); + await Promise.all([ + new Promise((resolve, reject) => { + firstSocket.once('open', resolve); + firstSocket.once('error', reject); + }), + new Promise((resolve, reject) => { + secondSocket.once('open', resolve); + secondSocket.once('error', reject); + }), + ]); + + try { + const lifecycleMessages = [ + await receiver.next(), + await receiver.next(), + await receiver.next(), + await receiver.next(), + ]; + assert.deepEqual( + lifecycleMessages + .map(item => `${item.name}:${item.message.type}:${item.message.path}`) + .sort(), + [ + 'first:connection:/ws-fetch/cross-1', + 'first:open:/ws-fetch/cross-1', + 'second:connection:/ws-fetch/cross-2', + 'second:open:/ws-fetch/cross-2', + ], + ); + + firstSocket.send(JSON.stringify({ + content: 'from-first', + delayMs: 80, + })); + const firstPhase1 = await receiver.next(); + assert.equal(firstPhase1.name, 'first'); + assert.deepEqual(firstPhase1.message, { + type: 'data', + phase: 1, + id: 'cross-1', + name: 'first', + tags: [ 'a' ], + header: 'client-cross-1', + path: '/ws-fetch/cross-1', + content: 'from-first', + pid: firstPhase1.message.pid, + }); + + secondSocket.send(JSON.stringify({ + content: 'from-second', + delayMs: 10, + })); + const secondPhase1 = await receiver.next(); + assert.equal(secondPhase1.name, 'second'); + assert.deepEqual(secondPhase1.message, { + type: 'data', + phase: 1, + id: 'cross-2', + name: 'second', + tags: [ 'b' ], + header: 'client-cross-2', + path: '/ws-fetch/cross-2', + content: 'from-second', + pid: secondPhase1.message.pid, + }); + + const secondPhase2 = await receiver.next(); + assert.deepEqual(secondPhase2, { + name: 'second', + message: { + type: 'data', + phase: 2, + id: 'cross-2', + content: 'from-second', + pid: secondPhase1.message.pid, + }, + }); + + const firstPhase2 = await receiver.next(); + assert.deepEqual(firstPhase2, { + name: 'first', + message: { + type: 'data', + phase: 2, + id: 'cross-1', + content: 'from-first', + pid: firstPhase1.message.pid, + }, + }); + + assert.equal(firstSocket.readyState, WebSocket.OPEN); + assert.equal(secondSocket.readyState, WebSocket.OPEN); + } finally { + await Promise.all([ + closeClient(firstSocket), + closeClient(secondSocket), + ]); + } + }); + it('should reject unmatched websocket route', async () => { await assert.rejects( () => createClient(requestUrl(app, '/ws/not-found')), diff --git a/plugin/controller/test/websocket/websocketCluster.test.ts b/plugin/controller/test/websocket/websocketCluster.test.ts index 516a3d0b2..54b7c304c 100644 --- a/plugin/controller/test/websocket/websocketCluster.test.ts +++ b/plugin/controller/test/websocket/websocketCluster.test.ts @@ -20,6 +20,53 @@ async function receiveJSON(socket: WebSocket): Promise { return JSON.parse(data.toString()); } +function createJSONReceiver(socket: WebSocket) { + const messages: any[] = []; + const waiters: Array<{ resolve: (value: any) => void; reject: (error: Error) => void }> = []; + + socket.on('message', data => { + const message = JSON.parse(data.toString()); + const waiter = waiters.shift(); + if (waiter) { + waiter.resolve(message); + return; + } + messages.push(message); + }); + socket.once('error', error => { + while (waiters.length) { + waiters.shift()!.reject(error); + } + }); + socket.once('close', () => { + while (waiters.length) { + waiters.shift()!.reject(new Error('websocket closed before message')); + } + }); + + return { + next() { + if (messages.length) { + return Promise.resolve(messages.shift()); + } + return new Promise((resolve, reject) => { + waiters.push({ resolve, reject }); + }); + }, + }; +} + +async function receiveClose(socket: WebSocket): Promise<{ code: number; reason: string }> { + return await new Promise(resolve => { + socket.once('close', (code, reason) => { + resolve({ + code, + reason: reason.toString(), + }); + }); + }); +} + async function createClientWithFirstMessage(url: string) { const socket = new WebSocket(url); const firstMessage = receiveJSON(socket); @@ -123,4 +170,71 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { } }); + it('should handle websocket fetch controller in cluster', async () => { + const socket = new WebSocket(requestUrl(app, '/ws-fetch/cluster-fetch?name=bar&tag=cluster'), { + headers: { 'x-client-id': 'cluster-fetch-client' }, + }); + const receiver = createJSONReceiver(socket); + await new Promise((resolve, reject) => { + socket.once('open', resolve); + socket.once('error', reject); + }); + try { + const connection = await receiver.next(); + assert.equal(connection.type, 'connection'); + assert.equal(connection.header, 'cluster-fetch-client'); + assert.equal(connection.path, '/ws-fetch/cluster-fetch'); + assert.equal(typeof connection.pid, 'number'); + + const open = await receiver.next(); + assert.equal(open.type, 'open'); + assert.equal(open.path, '/ws-fetch/cluster-fetch'); + assert.equal(typeof open.pid, 'number'); + + socket.send(JSON.stringify({ + content: 'cluster-first', + })); + const first = await receiver.next(); + assert.equal(first.type, 'data'); + assert.equal(first.phase, 1); + assert.equal(first.id, 'cluster-fetch'); + assert.equal(first.name, 'bar'); + assert.deepEqual(first.tags, [ 'cluster' ]); + assert.equal(first.header, 'cluster-fetch-client'); + assert.equal(first.path, '/ws-fetch/cluster-fetch'); + assert.equal(first.content, 'cluster-first'); + assert.deepEqual(await receiver.next(), { + type: 'data', + phase: 2, + id: 'cluster-fetch', + content: 'cluster-first', + pid: first.pid, + }); + assert.equal(socket.readyState, WebSocket.OPEN); + + const closePromise = receiveClose(socket); + socket.send(JSON.stringify({ + content: 'cluster-done', + close: true, + })); + const closeData = await receiver.next(); + assert.equal(closeData.type, 'data'); + assert.equal(closeData.phase, 1); + assert.equal(closeData.content, 'cluster-done'); + assert.deepEqual(await receiver.next(), { + type: 'data', + phase: 2, + id: 'cluster-fetch', + content: 'cluster-done', + pid: closeData.pid, + }); + assert.deepEqual(await closePromise, { + code: 1000, + reason: 'server done', + }); + } finally { + await closeClient(socket); + } + }); + }); From ab2eb7153fbbc5f42f2970cf11079024ab116436 Mon Sep 17 00:00:00 2001 From: akitasummer Date: Thu, 9 Jul 2026 18:08:50 +0800 Subject: [PATCH 3/9] feat(standalone): support websocket controllers --- core/test-util/StandaloneTestUtil.ts | 39 +- core/types/standalone/index.ts | 1 + core/types/standalone/websocket.ts | 8 + standalone/service-worker/index.ts | 3 + standalone/service-worker/package.json | 4 +- .../service-worker/src/ServiceWorkerApp.ts | 5 +- .../src/hook/ControllerLoadUnitHook.ts | 10 + .../ServiceWorkerWebSocketContext.ts | 38 ++ .../websocket/WebSocketControllerRegister.ts | 549 ++++++++++++++++++ .../src/websocket/WebSocketEventHandler.ts | 45 ++ .../test/fixtures/http/GetController.ts | 21 + .../fixtures/websocket/WebSocketController.ts | 195 +++++++ .../test/fixtures/websocket/package.json | 6 + .../service-worker/test/http/response.test.ts | 20 + .../test/websocket/websocket.test.ts | 269 +++++++++ 15 files changed, 1206 insertions(+), 7 deletions(-) create mode 100644 core/types/standalone/websocket.ts create mode 100644 standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts create mode 100644 standalone/service-worker/src/websocket/WebSocketControllerRegister.ts create mode 100644 standalone/service-worker/src/websocket/WebSocketEventHandler.ts create mode 100644 standalone/service-worker/test/fixtures/websocket/WebSocketController.ts create mode 100644 standalone/service-worker/test/fixtures/websocket/package.json create mode 100644 standalone/service-worker/test/websocket/websocket.test.ts diff --git a/core/test-util/StandaloneTestUtil.ts b/core/test-util/StandaloneTestUtil.ts index 83e013ee9..c37f98f31 100644 --- a/core/test-util/StandaloneTestUtil.ts +++ b/core/test-util/StandaloneTestUtil.ts @@ -1,12 +1,15 @@ import { createServer, IncomingMessage, OutgoingHttpHeaders, Server, ServerOptions, ServerResponse } from 'node:http'; -import { pipeline } from 'node:stream'; +import { Duplex, pipeline } from 'node:stream'; import { Headers, BodyInit, Request, Response } from 'undici'; -import { FetchEvent } from '@eggjs/tegg-types/standalone'; +import { FetchEvent, WebSocketUpgradeEvent } from '@eggjs/tegg-types/standalone'; export type FetchEventListener = (event: FetchEvent) => Promise; +export type WebSocketUpgradeEventListener = (event: WebSocketUpgradeEvent) => Promise; +export type ServiceWorkerEvent = FetchEvent | WebSocketUpgradeEvent; +export type ServiceWorkerEventListener = FetchEventListener | WebSocketUpgradeEventListener | ((event: ServiceWorkerEvent) => Promise); export interface StartHTTPServerOptions extends ServerOptions { - listener: FetchEventListener; + listener: ServiceWorkerEventListener; } export class StandaloneTestUtil { @@ -42,13 +45,18 @@ export class StandaloneTestUtil { }); } - static #createHTTPServerListener(listener: FetchEventListener) { + static #createHTTPServerListener(listener: ServiceWorkerEventListener) { return async (req: IncomingMessage, res: ServerResponse) => { const request = StandaloneTestUtil.#buildRequest(req); // TODO currently fake FetchEvent const event: any = new Event('fetch'); event.request = request; - const response = await listener(event); + const response = await (listener as FetchEventListener)(event); + if (!response) { + res.writeHead(500); + res.end(); + return; + } const headers: OutgoingHttpHeaders = {}; for (const [ key, value ] of response.headers) { @@ -70,9 +78,30 @@ export class StandaloneTestUtil { }; } + static #createWebSocketUpgradeListener(listener: ServiceWorkerEventListener) { + return async (req: IncomingMessage, socket: Duplex, head: Buffer) => { + const event: any = new Event('websocket'); + event.request = req; + event.socket = socket; + event.head = head; + + try { + await listener(event); + } catch (error) { + if (!socket.destroyed) { + const message = 'Internal Server Error'; + socket.write(`HTTP/1.1 500 ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(message)}\r\n\r\n${message}`); + socket.destroy(); + } + console.error('websocket upgrade listener failed:', error); + } + }; + } + static startHTTPServer(host: string, port: number, { listener, ...options }: StartHTTPServerOptions) { const serverListener = StandaloneTestUtil.#createHTTPServerListener(listener); const server = createServer(options ?? {}, serverListener); + server.on('upgrade', StandaloneTestUtil.#createWebSocketUpgradeListener(listener)); return new Promise(resolve => { server.listen(port, host, () => resolve(server)); diff --git a/core/types/standalone/index.ts b/core/types/standalone/index.ts index c25290f76..d3fd286c8 100644 --- a/core/types/standalone/index.ts +++ b/core/types/standalone/index.ts @@ -1,2 +1,3 @@ export * from './fetch'; +export * from './websocket'; export * from './ServiceWorkerContext'; diff --git a/core/types/standalone/websocket.ts b/core/types/standalone/websocket.ts new file mode 100644 index 000000000..4829119df --- /dev/null +++ b/core/types/standalone/websocket.ts @@ -0,0 +1,8 @@ +import type { IncomingMessage } from 'node:http'; +import type { Duplex } from 'node:stream'; + +export interface WebSocketUpgradeEvent extends Event { + request: IncomingMessage; + socket: Duplex; + head: Buffer; +} diff --git a/standalone/service-worker/index.ts b/standalone/service-worker/index.ts index 58c3ae946..29d75facd 100644 --- a/standalone/service-worker/index.ts +++ b/standalone/service-worker/index.ts @@ -14,6 +14,9 @@ export * from './src/http/FetchRouter'; export * from './src/http/HTTPControllerRegister'; export * from './src/http/HTTPMethodRegister'; export * from './src/http/ServiceWorkerFetchContext'; +export * from './src/websocket/ServiceWorkerWebSocketContext'; +export * from './src/websocket/WebSocketControllerRegister'; +export * from './src/websocket/WebSocketEventHandler'; export * from './src/utils/RequestUtils'; export * from './src/utils/ResponseUtils'; export * from './src/mcp/AbstractControllerAdvice'; diff --git a/standalone/service-worker/package.json b/standalone/service-worker/package.json index 2667d277f..bb807e265 100644 --- a/standalone/service-worker/package.json +++ b/standalone/service-worker/package.json @@ -53,7 +53,8 @@ "egg-errors": "^2.3.0", "path-to-regexp": "^1.8.0", "type-is": "^1.6.18", - "urllib": "^4.0.0" + "urllib": "^4.0.0", + "ws": "^8.18.0" }, "peerDependencies": { "@eggjs/tegg": "^3" @@ -64,6 +65,7 @@ "@types/mocha": "^10.0.1", "@types/node": "^20.2.4", "@types/type-is": "^1.6.6", + "@types/ws": "^8.5.12", "cross-env": "^7.0.3", "mocha": "^10.2.0", "supertest": "^7.1.1", diff --git a/standalone/service-worker/src/ServiceWorkerApp.ts b/standalone/service-worker/src/ServiceWorkerApp.ts index ca2b09b3a..9d814908c 100644 --- a/standalone/service-worker/src/ServiceWorkerApp.ts +++ b/standalone/service-worker/src/ServiceWorkerApp.ts @@ -20,6 +20,8 @@ import { LoadUnitInnerClassHook } from './hook/LoadUnitInnerClassHook'; import { ServiceWorkerRunner } from './ServiceWorkerRunner'; import { StandaloneEggObjectFactory } from './StandaloneEggObjectFactory'; import { FetchEventHandler } from './http/FetchEventHandler'; +import { WebSocketEventHandler } from './websocket/WebSocketEventHandler'; +import { WebSocketControllerRegister } from './websocket/WebSocketControllerRegister'; export interface ServiceWorkerAppOptions { innerObjectHandlers?: RunnerOptions['innerObjectHandlers']; @@ -78,7 +80,7 @@ export class ServiceWorkerApp { innerObjectHandlers.httpclient = [{ obj: getDefaultHttpClient() }]; } - this.loadUnitInnerClassHook = new LoadUnitInnerClassHook([ StandaloneEggObjectFactory, ServiceWorkerRunner, FetchEventHandler ]); + this.loadUnitInnerClassHook = new LoadUnitInnerClassHook([ StandaloneEggObjectFactory, ServiceWorkerRunner, FetchEventHandler, WebSocketEventHandler ]); LoadUnitLifecycleUtil.registerLifecycle(this.loadUnitInnerClassHook); @@ -104,6 +106,7 @@ export class ServiceWorkerApp { // Clean up static singletons HTTPControllerRegister.clean(); MCPControllerRegister.clean(); + WebSocketControllerRegister.clean(); // Unregister lifecycle hooks LoadUnitLifecycleUtil.deleteLifecycle(this.contextProtoLoadUnitHook); diff --git a/standalone/service-worker/src/hook/ControllerLoadUnitHook.ts b/standalone/service-worker/src/hook/ControllerLoadUnitHook.ts index 8d7df7774..5d0b216b8 100644 --- a/standalone/service-worker/src/hook/ControllerLoadUnitHook.ts +++ b/standalone/service-worker/src/hook/ControllerLoadUnitHook.ts @@ -7,6 +7,7 @@ import { ControllerMetadataManager } from '../controller/ControllerMetadataManag import { FetchRouter } from '../http/FetchRouter'; import { HTTPControllerRegister } from '../http/HTTPControllerRegister'; import { MCPControllerRegister } from '../mcp/MCPControllerRegister'; +import { WebSocketControllerRegister } from '../websocket/WebSocketControllerRegister'; export class ControllerLoadUnitHook implements LifecycleHook { private readonly controllerRegisterFactory: ControllerRegisterFactory; @@ -34,6 +35,15 @@ export class ControllerLoadUnitHook implements LifecycleHook { return MCPControllerRegister.create(proto, controllerMeta, this.fetchRouter); }); + + // Register the WebSocket controller register creator + this.controllerRegisterFactory.registerControllerRegister(ControllerType.WEBSOCKET, (proto: EggPrototype, controllerMeta: ControllerMetadata) => { + return WebSocketControllerRegister.create(proto, controllerMeta); + }); + + this.controllerRegisterFactory.registerControllerRegister(ControllerType.WEBSOCKET_FETCH, (proto: EggPrototype, controllerMeta: ControllerMetadata) => { + return WebSocketControllerRegister.create(proto, controllerMeta); + }); } async postCreate(_: LoadUnitLifecycleContext, loadUnit: LoadUnit): Promise { diff --git a/standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts b/standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts new file mode 100644 index 000000000..79314d544 --- /dev/null +++ b/standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts @@ -0,0 +1,38 @@ +import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; +import type { WebSocket } from 'ws'; +import type { WebSocketContext } from '@eggjs/tegg'; + +export class ServiceWorkerWebSocketContext implements WebSocketContext { + readonly socket: WebSocket; + readonly request: IncomingMessage; + readonly url: URL; + readonly method = 'GET'; + readonly path: string; + readonly host: string; + readonly params: Record; + readonly query: Record; + readonly queries: Record; + readonly headers: IncomingHttpHeaders; + + constructor(options: { + socket: WebSocket; + request: IncomingMessage; + url: URL; + params: Record; + }) { + this.socket = options.socket; + this.request = options.request; + this.url = options.url; + this.path = options.url.pathname; + this.host = options.url.host; + this.params = options.params; + this.headers = options.request.headers; + this.query = {}; + this.queries = {}; + + for (const key of options.url.searchParams.keys()) { + this.query[key] = options.url.searchParams.get(key) ?? undefined; + this.queries[key] = options.url.searchParams.getAll(key); + } + } +} diff --git a/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts new file mode 100644 index 000000000..09e518719 --- /dev/null +++ b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts @@ -0,0 +1,549 @@ +import assert from 'node:assert'; +import type { IncomingMessage } from 'node:http'; +import { Duplex, pipeline } from 'node:stream'; +import pathToRegexp from 'path-to-regexp'; +import { Router as KoaRouter } from '@eggjs/router'; +import { FrameworkErrorFormater } from 'egg-errors'; +import { createWebSocketStream, RawData, WebSocket, WebSocketServer } from 'ws'; +import { + CONTROLLER_META_DATA, + ControllerMetadata, + ControllerType, + EggContext, + Next, + WebSocketControllerMeta, + WebSocketFetchClose, + WebSocketFetchControllerMeta, + WebSocketFetchMethodMeta, + WebSocketMethodMeta, + WebSocketParamType, + WebSocketPathParamMeta, + WebSocketQueriesParamMeta, + WebSocketQueryParamMeta, +} from '@eggjs/tegg'; +import { EggContainerFactory, EggPrototype } from '@eggjs/tegg/helper'; +import { WebSocketUpgradeEvent } from '@eggjs/tegg-types/standalone'; +import { ControllerRegister } from '../controller/ControllerRegister'; +import { RootProtoManager } from '../controller/RootProtoManager'; +import { ServiceWorkerWebSocketContext } from './ServiceWorkerWebSocketContext'; + +const noop = () => { + // noop +}; + +interface WebSocketRoute { + controllerProto: EggPrototype; + controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta; + methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta; + methodRealPath: string; + methodName: string; + host?: string; + keys: pathToRegexp.Key[]; + regexp: RegExp; +} + +interface WebSocketMethodPayload { + data?: RawData; + close?: WebSocketFetchClose; + error?: Error; + closeCode?: number; + closeReason?: Buffer; + getWebSocketStream?: () => Duplex; +} + +type Middleware = (ctx: ServiceWorkerWebSocketContext, next: Next) => Promise; + +export class WebSocketControllerRegister implements ControllerRegister { + static instance?: WebSocketControllerRegister; + + private readonly checkRouters: Map; + private readonly webSocketServer: WebSocketServer; + private controllerProtos: EggPrototype[] = []; + private routes: WebSocketRoute[] = []; + + static create(proto: EggPrototype, controllerMeta: ControllerMetadata) { + assert( + controllerMeta.type === ControllerType.WEBSOCKET || controllerMeta.type === ControllerType.WEBSOCKET_FETCH, + 'controller meta type is not WEBSOCKET or WEBSOCKET_FETCH', + ); + if (!WebSocketControllerRegister.instance) { + WebSocketControllerRegister.instance = new WebSocketControllerRegister(); + } + WebSocketControllerRegister.instance.controllerProtos.push(proto); + return WebSocketControllerRegister.instance; + } + + constructor() { + this.checkRouters = new Map(); + this.checkRouters.set('default', new KoaRouter({ sensitive: true })); + this.webSocketServer = new WebSocketServer({ noServer: true }); + } + + register(): Promise { + return Promise.resolve(); + } + + static clean() { + this.instance?.close(); + this.instance = undefined; + } + + doRegister(rootProtoManager: RootProtoManager) { + const methodMap = new Map(); + for (const proto of this.controllerProtos) { + const metadata = proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; + for (const method of metadata.methods) { + methodMap.set(method, proto); + } + } + const allMethods = Array.from(methodMap.keys()) + .sort((a, b) => b.priority - a.priority); + + for (const method of allMethods) { + const controllerProto = methodMap.get(method)!; + const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; + this.checkDuplicate(controllerMeta, method); + } + + this.routes = allMethods.flatMap(method => { + const controllerProto = methodMap.get(method)!; + const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; + const hosts = this.getMethodHosts(controllerMeta, method) || [ undefined ]; + return hosts.map(host => this.createRoute(controllerProto, controllerMeta, method, host)); + }); + + for (const route of this.routes) { + rootProtoManager.registerRootProto('GET', (ctx: EggContext) => { + if (route.regexp.test(ctx.path)) { + return route.controllerProto; + } + }, route.host || ''); + } + } + + close() { + this.webSocketServer.close(); + this.routes = []; + this.controllerProtos = []; + this.checkRouters.clear(); + } + + async handleUpgrade(event: WebSocketUpgradeEvent): Promise { + const url = this.createURL(event.request); + const matched = this.matchRoute(event.request, url); + if (!matched) { + return false; + } + + await new Promise(resolve => { + this.webSocketServer.handleUpgrade(event.request, event.socket as any, event.head, webSocket => { + this.handleConnection(matched.route, matched.params, url, event.request, webSocket) + .then(resolve) + .catch(error => { + console.error('[tegg/websocket] handle connection failed:', error); + if (webSocket.readyState === WebSocket.OPEN || webSocket.readyState === WebSocket.CONNECTING) { + webSocket.close(1011, 'Internal Server Error'); + } + resolve(); + }); + }); + }); + return true; + } + + static rejectSocket(socket: Duplex, status: number, message: string) { + if (socket.destroyed) { + return; + } + const body = message; + socket.write(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); + socket.destroy(); + } + + private createRoute( + controllerProto: EggPrototype, + controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, + methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, + host: string | undefined, + ): WebSocketRoute { + const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); + const keys: pathToRegexp.Key[] = []; + const regexp = pathToRegexp(methodRealPath, keys, { sensitive: true }); + return { + controllerProto, + controllerMeta, + methodMeta, + methodRealPath, + methodName: this.getMethodName(controllerMeta, methodMeta), + host, + keys, + regexp, + }; + } + + private checkDuplicate(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + let router = this.checkRouters.get('default')!; + const hosts = this.getMethodHosts(controllerMeta, methodMeta) || []; + if (!hosts.length) { + this.checkDuplicateInRouter(router, controllerMeta, methodMeta); + this.registerToRouter(router, controllerMeta, methodMeta); + return; + } + + hosts.forEach(host => { + router = this.checkRouters.get(host)!; + if (!router) { + router = new KoaRouter({ sensitive: true }); + this.checkRouters.set(host, router); + } + this.checkDuplicateInRouter(router, controllerMeta, methodMeta); + this.registerToRouter(router, controllerMeta, methodMeta); + }); + } + + private registerToRouter(router: KoaRouter, controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); + const methodName = this.getMethodName(controllerMeta, methodMeta); + Reflect.apply(router.get, router, [ methodName, methodRealPath, noop ]); + } + + private checkDuplicateInRouter(router: KoaRouter, controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); + const matched = router.match(methodRealPath, 'GET'); + const methodName = this.getMethodName(controllerMeta, methodMeta); + if (matched.route) { + const [ layer ] = matched.path; + const err = new Error(`register websocket controller ${methodName} failed, ${controllerMeta.type} ${methodRealPath} is conflict with exists rule ${layer.path}`); + throw FrameworkErrorFormater.format(err); + } + } + + private getMethodRealPath(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + return (controllerMeta as any).getMethodRealPath(methodMeta); + } + + private getMethodHosts(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): string[] | undefined { + return (controllerMeta as any).getMethodHosts(methodMeta); + } + + private getMethodName(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): string { + return (controllerMeta as any).getMethodName(methodMeta); + } + + private getMethodMiddlewares(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): Middleware[] { + return (controllerMeta as any).getMethodMiddlewares(methodMeta); + } + + private createURL(req: IncomingMessage): URL { + const host = req.headers.host || 'localhost'; + return new URL(req.url || '/', `http://${host}`); + } + + private matchRoute(req: IncomingMessage, url: URL): { route: WebSocketRoute; params: Record } | undefined { + for (const route of this.routes) { + if (route.host && !this.matchHost(route.host, req.headers.host)) { + continue; + } + const matched = route.regexp.exec(url.pathname); + if (!matched) { + continue; + } + const params: Record = {}; + route.keys.forEach((key, index) => { + params[String(key.name)] = decodeURIComponent(matched[index + 1]); + }); + return { route, params }; + } + } + + private matchHost(expectedHost: string, requestHost: string | undefined) { + if (!requestHost) { + return false; + } + return requestHost === expectedHost || requestHost.split(':')[0] === expectedHost; + } + + private async handleConnection( + route: WebSocketRoute, + params: Record, + url: URL, + req: IncomingMessage, + webSocket: WebSocket, + ) { + const webSocketCtx = new ServiceWorkerWebSocketContext({ + socket: webSocket, + request: req, + url, + params, + }); + const methodMiddlewares = this.getMethodMiddlewares(route.controllerMeta, route.methodMeta); + let invoked = false; + const handler = async (_ctx: ServiceWorkerWebSocketContext, next: Next) => { + invoked = true; + if (route.controllerMeta.type === ControllerType.WEBSOCKET_FETCH) { + await this.invokeFetchController(route, webSocketCtx); + } else { + await this.invokeController(route, webSocketCtx); + } + await next(); + }; + + await this.compose([ ...methodMiddlewares, handler ])(webSocketCtx, async () => { + // final middleware + }); + if (!invoked) { + webSocket.close(1008, 'WebSocket middleware did not call next'); + return; + } + await this.waitWebSocketClose(webSocket); + } + + private compose(middlewares: Middleware[]) { + return async (ctx: ServiceWorkerWebSocketContext, next: Next) => { + let index = -1; + const dispatch = async (i: number): Promise => { + if (i <= index) { + throw new Error('next() called multiple times'); + } + index = i; + const fn = middlewares[i] || next; + if (!fn) { + return; + } + await fn(ctx, () => dispatch(i + 1)); + }; + await dispatch(0); + }; + } + + private async invokeController(route: WebSocketRoute, webSocketCtx: ServiceWorkerWebSocketContext) { + const methodMeta = route.methodMeta as WebSocketMethodMeta; + const eggObj = await EggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); + const realObj = eggObj.obj; + const realMethod = realObj[methodMeta.name]; + let webSocketStream: Duplex | undefined; + const getWebSocketStream = () => { + if (!webSocketStream) { + webSocketStream = createWebSocketStream(webSocketCtx.socket); + } + return webSocketStream; + }; + const args = this.buildMethodArgs(methodMeta, webSocketCtx, { getWebSocketStream }); + const result = await Reflect.apply(realMethod, realObj, args); + this.pipeResponseStream(result, webSocketCtx.socket, getWebSocketStream); + } + + private buildMethodArgs( + methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, + webSocketCtx: ServiceWorkerWebSocketContext, + payload: WebSocketMethodPayload, + ) { + const argsLength = methodMeta.paramMap.size; + const hasContext = methodMeta.contextParamIndex !== undefined; + const contextIndex = methodMeta.contextParamIndex; + const methodArgsLength = argsLength + (hasContext ? 1 : 0); + const args: unknown[] = new Array(methodArgsLength); + if (hasContext) { + args[contextIndex!] = webSocketCtx; + } + for (const [ index, param ] of methodMeta.paramMap) { + switch (param.type) { + case WebSocketParamType.PARAM: { + const pathParam = param as WebSocketPathParamMeta; + args[index] = webSocketCtx.params[pathParam.name]; + break; + } + case WebSocketParamType.QUERY: { + const queryParam = param as WebSocketQueryParamMeta; + args[index] = webSocketCtx.query[queryParam.name]; + break; + } + case WebSocketParamType.QUERIES: { + const queryParam = param as WebSocketQueriesParamMeta; + args[index] = webSocketCtx.queries[queryParam.name] || []; + break; + } + case WebSocketParamType.HEADERS: { + args[index] = webSocketCtx.headers; + break; + } + case WebSocketParamType.REQUEST: { + args[index] = webSocketCtx.request; + break; + } + case WebSocketParamType.SOCKET: { + args[index] = webSocketCtx.socket; + break; + } + case WebSocketParamType.STREAM: { + assert(payload.getWebSocketStream, '@WebSocketStream can not be used here'); + args[index] = payload.getWebSocketStream(); + break; + } + case WebSocketParamType.DATA: { + args[index] = payload.data; + break; + } + case WebSocketParamType.CLOSE: { + args[index] = payload.close; + break; + } + case WebSocketParamType.ERROR: { + args[index] = payload.error; + break; + } + case WebSocketParamType.CLOSE_CODE: { + args[index] = payload.closeCode; + break; + } + case WebSocketParamType.CLOSE_REASON: { + args[index] = payload.closeReason; + break; + } + default: + assert.fail('never arrive'); + } + } + return args; + } + + private async invokeFetchController(route: WebSocketRoute, webSocketCtx: ServiceWorkerWebSocketContext) { + const controllerMeta = route.controllerMeta as WebSocketFetchControllerMeta; + const methodMeta = route.methodMeta as WebSocketFetchMethodMeta; + const eggObj = await EggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); + const realObj = eggObj.obj; + const close = this.createFetchClose(webSocketCtx.socket); + const invokeMethod = async ( + targetMethodMeta: WebSocketFetchMethodMeta | undefined, + payload: WebSocketMethodPayload = {}, + ) => { + if (!targetMethodMeta) { + return; + } + const realMethod = realObj[targetMethodMeta.name]; + const args = this.buildMethodArgs(targetMethodMeta, webSocketCtx, { + close, + ...payload, + }); + return await Reflect.apply(realMethod, realObj, args); + }; + const handleError = async (error: Error) => { + if (!controllerMeta.errorMethod) { + console.error('[tegg/websocket-fetch] handle error:', error); + return; + } + await invokeMethod(controllerMeta.errorMethod, { error }); + }; + const closeHandled = new Promise(resolve => { + webSocketCtx.socket.once('close', (code, reason) => { + invokeMethod(controllerMeta.closeMethod, { + closeCode: code, + closeReason: reason, + }).catch(error => { + console.error('[tegg/websocket-fetch] handle close failed:', error); + }).finally(resolve); + }); + }); + + webSocketCtx.socket.on('message', data => { + invokeMethod(methodMeta, { data }).then(result => { + this.sendFetchResponseStream(result, webSocketCtx.socket, handleError); + }).catch(error => { + handleError(error).catch(err => { + console.error('[tegg/websocket-fetch] handle message error failed:', err); + }); + }); + }); + webSocketCtx.socket.on('error', error => { + handleError(error).catch(err => { + console.error('[tegg/websocket-fetch] handle socket error failed:', err); + }); + }); + + await invokeMethod(controllerMeta.connectionMethod); + await invokeMethod(controllerMeta.openMethod); + await closeHandled; + } + + private createFetchClose(webSocket: WebSocket): WebSocketFetchClose { + return (code = 1000, reason?: string | Buffer) => { + if (webSocket.readyState !== WebSocket.OPEN && webSocket.readyState !== WebSocket.CONNECTING) { + return; + } + webSocket.close(code, reason); + }; + } + + private sendFetchResponseStream( + result: unknown, + webSocket: WebSocket, + handleError: (error: Error) => Promise, + ) { + if (result === undefined || result === null) { + return; + } + if (!this.isReadableStream(result)) { + handleError(new Error('WebSocketFetch method must return a readable stream or void')).catch(error => { + console.error('[tegg/websocket-fetch] handle invalid response failed:', error); + }); + return; + } + result.on('data', chunk => { + this.sendFetchChunk(webSocket, chunk, handleError); + }); + result.once('error', error => { + handleError(error).catch(err => { + console.error('[tegg/websocket-fetch] handle stream error failed:', err); + }); + }); + } + + private sendFetchChunk( + webSocket: WebSocket, + chunk: unknown, + handleError: (error: Error) => Promise, + ) { + if (webSocket.readyState !== WebSocket.OPEN) { + return; + } + webSocket.send(chunk as any, error => { + if (!error) { + return; + } + handleError(error).catch(err => { + console.error('[tegg/websocket-fetch] handle send error failed:', err); + }); + }); + } + + private pipeResponseStream(result: unknown, webSocket: WebSocket, getWebSocketStream: () => Duplex) { + if (!this.isReadableStream(result)) { + return; + } + + pipeline(result, getWebSocketStream(), error => { + if (!error) { + return; + } + if (webSocket.readyState === WebSocket.CLOSED || webSocket.readyState === WebSocket.CLOSING) { + return; + } + console.error('[tegg/websocket] pipe response stream failed:', error); + webSocket.close(1011, 'Internal Server Error'); + }); + } + + private isReadableStream(result: unknown): result is NodeJS.ReadableStream { + return !!result && typeof (result as NodeJS.ReadableStream).pipe === 'function'; + } + + private waitWebSocketClose(webSocket: WebSocket): Promise { + if (webSocket.readyState === WebSocket.CLOSED || webSocket.readyState === WebSocket.CLOSING) { + return Promise.resolve(); + } + return new Promise(resolve => { + webSocket.once('close', () => resolve()); + webSocket.once('error', () => resolve()); + }); + } +} diff --git a/standalone/service-worker/src/websocket/WebSocketEventHandler.ts b/standalone/service-worker/src/websocket/WebSocketEventHandler.ts new file mode 100644 index 000000000..92488c858 --- /dev/null +++ b/standalone/service-worker/src/websocket/WebSocketEventHandler.ts @@ -0,0 +1,45 @@ +import { AccessLevel, Inject } from '@eggjs/tegg'; +import { WebSocketUpgradeEvent } from '@eggjs/tegg-types/standalone'; +import { AbstractEventHandler, EventHandlerProto } from '@eggjs/tegg/standalone'; +import { RootProtoManager } from '../controller/RootProtoManager'; +import { WebSocketControllerRegister } from './WebSocketControllerRegister'; + +@EventHandlerProto('websocket', { accessLevel: AccessLevel.PUBLIC }) +export class WebSocketEventHandler extends AbstractEventHandler { + @Inject() + private readonly rootProtoManager: RootProtoManager; + + #initialized = false; + #initPromise?: Promise; + + private async initRoutes() { + if (this.#initialized) { + return; + } + if (!this.#initPromise) { + this.#initPromise = this.doInitRoutes().catch(err => { + this.#initPromise = undefined; + throw err; + }); + } + await this.#initPromise; + } + + private async doInitRoutes() { + WebSocketControllerRegister.instance?.doRegister(this.rootProtoManager); + this.#initialized = true; + } + + async handleEvent(event: WebSocketUpgradeEvent): Promise { + await this.initRoutes(); + try { + const handled = await WebSocketControllerRegister.instance?.handleUpgrade(event); + if (!handled) { + WebSocketControllerRegister.rejectSocket(event.socket, 404, 'Not Found'); + } + } catch (error) { + console.error('handle websocket event failed:', error); + WebSocketControllerRegister.rejectSocket(event.socket, 500, 'Internal Server Error'); + } + } +} diff --git a/standalone/service-worker/test/fixtures/http/GetController.ts b/standalone/service-worker/test/fixtures/http/GetController.ts index 0fd5a215b..0f70ea3a1 100644 --- a/standalone/service-worker/test/fixtures/http/GetController.ts +++ b/standalone/service-worker/test/fixtures/http/GetController.ts @@ -22,4 +22,25 @@ export class GetController { }, }); } + + @HTTPMethod({ method: HTTPMethodEnum.GET, path: '/api/sse' }) + async sse(@HTTPQuery() id: string) { + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(`event: ready\ndata: ${JSON.stringify({ id })}\n\n`)); + controller.enqueue(encoder.encode('event: done\ndata: ok\n\n')); + controller.close(); + }, + }); + + return new Response(stream, { + status: 200, + headers: { + 'content-type': 'text/event-stream', + 'cache-control': 'no-cache', + connection: 'keep-alive', + }, + }); + } } diff --git a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts new file mode 100644 index 000000000..f5d901213 --- /dev/null +++ b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts @@ -0,0 +1,195 @@ +import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; +import { PassThrough, type Readable } from 'node:stream'; +import type { RawData, WebSocket } from 'ws'; +import { + Context, + Headers, + Param, + Query, + Queries, + Request, + WebSocketClose, + WebSocketCloseCode, + WebSocketCloseReason, + WebSocketContext, + WebSocketController, + WebSocketData, + WebSocketError, + WebSocketFetchClose, + WebSocketFetchController, + WebSocketFetchMethod, + WebSocketFetchOnClose, + WebSocketFetchOnConnection, + WebSocketFetchOnError, + WebSocketFetchOnOpen, + WebSocketMethod, + WebSocketSocket, + WebSocketStream, +} from '@eggjs/tegg'; + +interface WebSocketFetchRequest { + content: string; + close?: boolean; + error?: boolean; + delayMs?: number; +} + +export const fetchCloseEvents: string[] = []; + +@WebSocketController({ + path: '/ws', +}) +export class StandaloneWebSocketController { + @WebSocketMethod({ + path: '/echo/:id', + }) + async echo( + @WebSocketSocket() socket: WebSocket, + @Param() id: string, + @Query() name: string, + @Queries({ name: 'tag' }) tags: string[], + @Headers() headers: IncomingHttpHeaders, + @Request() request: IncomingMessage, + @Context() ctx: WebSocketContext, + ) { + socket.send(JSON.stringify({ + type: 'ready', + id, + name, + tags, + header: headers['x-client-id'], + url: request.url, + path: ctx.path, + })); + + socket.on('message', data => { + socket.send(JSON.stringify({ + type: 'echo', + data: data.toString(), + })); + }); + } + + @WebSocketMethod({ + path: '/stream/:id', + }) + stream( + @WebSocketStream() input: Readable, + @Param() id: string, + @Query() name: string, + ) { + const output = new PassThrough(); + output.write(JSON.stringify({ + type: 'ready', + id, + name, + })); + + input.on('data', data => { + output.write(JSON.stringify({ + type: 'stream', + data: data.toString(), + })); + }); + input.on('end', () => output.end()); + input.on('error', error => output.destroy(error)); + + return output; + } +} + +@WebSocketFetchController({ + path: '/ws-fetch/:id', +}) +export class StandaloneWebSocketFetchController { + @WebSocketFetchOnConnection() + onConnection( + @Headers() headers: IncomingHttpHeaders, + @Context() ctx: WebSocketContext, + ) { + ctx.socket.send(JSON.stringify({ + type: 'connection', + header: headers['x-client-id'], + path: ctx.path, + })); + } + + @WebSocketFetchOnOpen() + onOpen( + @Context() ctx: WebSocketContext, + ) { + ctx.socket.send(JSON.stringify({ + type: 'open', + path: ctx.path, + })); + } + + @WebSocketFetchMethod() + onData( + @WebSocketData() data: RawData, + @WebSocketClose() close: WebSocketFetchClose, + @Param() id: string, + @Query() name: string, + @Queries({ name: 'tag' }) tags: string[], + @Headers() headers: IncomingHttpHeaders, + @Context() ctx: WebSocketContext, + ) { + const body = JSON.parse(data.toString()) as WebSocketFetchRequest; + const output = new PassThrough(); + output.write(JSON.stringify({ + type: 'data', + phase: 1, + id, + name, + tags, + header: headers['x-client-id'], + path: ctx.path, + content: body.content, + })); + + setTimeout(() => { + if (output.destroyed) { + return; + } + if (body.error) { + output.destroy(new Error(`fetch error: ${body.content}`)); + return; + } + output.write(JSON.stringify({ + type: 'data', + phase: 2, + id, + content: body.content, + })); + output.end(); + }, body.delayMs ?? 10); + + if (body.close) { + output.once('end', () => setImmediate(() => close(1000, 'server done'))); + } + return output; + } + + @WebSocketFetchOnError() + onError( + @WebSocketError() error: Error, + @Headers() headers: IncomingHttpHeaders, + @Context() ctx: WebSocketContext, + ) { + ctx.socket.send(JSON.stringify({ + type: 'error', + message: error.message, + header: headers['x-client-id'], + path: ctx.path, + })); + } + + @WebSocketFetchOnClose() + onClose( + @WebSocketCloseCode() code: number, + @WebSocketCloseReason() reason: Buffer, + @Param() id: string, + ) { + fetchCloseEvents.push(`${id}:${code}:${reason.toString()}`); + } +} diff --git a/standalone/service-worker/test/fixtures/websocket/package.json b/standalone/service-worker/test/fixtures/websocket/package.json new file mode 100644 index 000000000..962070300 --- /dev/null +++ b/standalone/service-worker/test/fixtures/websocket/package.json @@ -0,0 +1,6 @@ +{ + "name": "websocket-app", + "eggModule": { + "name": "websocketApp" + } +} diff --git a/standalone/service-worker/test/http/response.test.ts b/standalone/service-worker/test/http/response.test.ts index 218a88f97..360b26bc8 100644 --- a/standalone/service-worker/test/http/response.test.ts +++ b/standalone/service-worker/test/http/response.test.ts @@ -1,4 +1,5 @@ import { Server } from 'node:http'; +import assert from 'node:assert'; import httpRequest from 'supertest'; import { ServiceWorkerApp } from '../../src/ServiceWorkerApp'; import { StandaloneTestUtil } from '@eggjs/module-test-util/StandaloneTestUtil'; @@ -37,4 +38,23 @@ describe('standalone/service-worker/test/http/response.test.ts', () => { .get('/api/null-body?nil=1') .expect(204, ''); }); + + it('should return SSE stream response', async () => { + const res = await httpRequest(server) + .get('/api/sse?id=standalone') + .expect(200) + .expect('Cache-Control', 'no-cache') + .expect('Connection', 'keep-alive'); + + assert.equal(res.headers['content-type'], 'text/event-stream'); + assert.equal(res.text, [ + 'event: ready', + 'data: {"id":"standalone"}', + '', + 'event: done', + 'data: ok', + '', + '', + ].join('\n')); + }); }); diff --git a/standalone/service-worker/test/websocket/websocket.test.ts b/standalone/service-worker/test/websocket/websocket.test.ts new file mode 100644 index 000000000..01203151c --- /dev/null +++ b/standalone/service-worker/test/websocket/websocket.test.ts @@ -0,0 +1,269 @@ +import assert from 'node:assert'; +import { Server } from 'node:http'; +import { AddressInfo } from 'node:net'; +import { createWebSocketStream, WebSocket } from 'ws'; +import { ServiceWorkerApp } from '../../src/ServiceWorkerApp'; +import { StandaloneTestUtil } from '@eggjs/module-test-util/StandaloneTestUtil'; +import { TestUtils } from '../Utils'; +import { fetchCloseEvents } from '../fixtures/websocket/WebSocketController'; + +describe('standalone/service-worker/test/websocket/websocket.test.ts', () => { + let app: ServiceWorkerApp; + let server: Server; + let port: number; + + before(async function() { + if (StandaloneTestUtil.skipOnNode()) { + return this.skip(); + } + ({ app, server } = await TestUtils.createFetchApp('websocket')); + port = (server.address() as AddressInfo).port; + }); + + after(async () => { + server?.close(); + await app?.destroy(); + }); + + beforeEach(() => { + fetchCloseEvents.length = 0; + }); + + it('should handle websocket controller in standalone service worker', async () => { + const ws = createClient('/ws/echo/123?name=tegg&tag=a&tag=b', { + 'x-client-id': 'standalone-client', + }); + const ready = await waitJSONMessage(ws); + + assert.deepStrictEqual(ready, { + type: 'ready', + id: '123', + name: 'tegg', + tags: [ 'a', 'b' ], + header: 'standalone-client', + url: '/ws/echo/123?name=tegg&tag=a&tag=b', + path: '/ws/echo/123', + }); + + ws.send('hello'); + const echo = await waitJSONMessage(ws); + assert.deepStrictEqual(echo, { + type: 'echo', + data: 'hello', + }); + + await closeClient(ws); + }); + + it('should handle websocket stream controller in standalone service worker', async () => { + const ws = createClient('/ws/stream/stream-id?name=stream-name'); + const input = createWebSocketStream(ws); + const messages: unknown[] = []; + const done = new Promise((resolve, reject) => { + input.on('data', data => { + messages.push(JSON.parse(data.toString())); + if (messages.length === 2) { + resolve(); + } + }); + input.on('error', reject); + }); + + await waitOpen(ws); + input.write('first'); + await done; + + assert.deepStrictEqual(messages, [ + { + type: 'ready', + id: 'stream-id', + name: 'stream-name', + }, + { + type: 'stream', + data: 'first', + }, + ]); + + input.end(); + await closeClient(ws); + }); + + it('should isolate two websocket fetch connections with interleaved responses in standalone service worker', async () => { + const first = createClient('/ws-fetch/first?name=one&tag=a&tag=b', { + 'x-client-id': 'first-client', + }); + const second = createClient('/ws-fetch/second?name=two&tag=c', { + 'x-client-id': 'second-client', + }); + const firstMessages = createJSONMessageQueue(first); + const secondMessages = createJSONMessageQueue(second); + const firstInitialMessages = firstMessages.nextMany(2); + const secondInitialMessages = secondMessages.nextMany(2); + + assert.deepStrictEqual(await firstInitialMessages, [ + { + type: 'connection', + header: 'first-client', + path: '/ws-fetch/first', + }, + { + type: 'open', + path: '/ws-fetch/first', + }, + ]); + assert.deepStrictEqual(await secondInitialMessages, [ + { + type: 'connection', + header: 'second-client', + path: '/ws-fetch/second', + }, + { + type: 'open', + path: '/ws-fetch/second', + }, + ]); + + first.send(JSON.stringify({ content: 'slow', delayMs: 40 })); + const firstPhase1 = await firstMessages.next(); + second.send(JSON.stringify({ content: 'fast', delayMs: 5, close: true })); + const secondPhase1 = await secondMessages.next(); + const secondPhase2 = await secondMessages.next(); + const firstPhase2 = await firstMessages.next(); + + assert.deepStrictEqual([ + firstPhase1, + secondPhase1, + secondPhase2, + firstPhase2, + ], [ + { + type: 'data', + phase: 1, + id: 'first', + name: 'one', + tags: [ 'a', 'b' ], + header: 'first-client', + path: '/ws-fetch/first', + content: 'slow', + }, + { + type: 'data', + phase: 1, + id: 'second', + name: 'two', + tags: [ 'c' ], + header: 'second-client', + path: '/ws-fetch/second', + content: 'fast', + }, + { + type: 'data', + phase: 2, + id: 'second', + content: 'fast', + }, + { + type: 'data', + phase: 2, + id: 'first', + content: 'slow', + }, + ]); + + await waitClose(second); + await waitFor(() => fetchCloseEvents.includes('second:1000:server done')); + await closeClient(first); + }); + + function createClient(path: string, headers?: Record) { + return new WebSocket(`ws://127.0.0.1:${port}${path}`, { headers }); + } +}); + +function waitOpen(ws: WebSocket) { + if (ws.readyState === WebSocket.OPEN) { + return Promise.resolve(); + } + return new Promise((resolve, reject) => { + ws.once('open', () => resolve()); + ws.once('error', reject); + }); +} + +function waitJSONMessage(ws: WebSocket): Promise { + return new Promise((resolve, reject) => { + ws.once('message', data => resolve(JSON.parse(data.toString()))); + ws.once('error', reject); + }); +} + +async function waitJSONMessages(ws: WebSocket, count: number) { + const messages: unknown[] = []; + while (messages.length < count) { + messages.push(await waitJSONMessage(ws)); + } + return messages; +} + +function createJSONMessageQueue(ws: WebSocket) { + const messages: any[] = []; + const waiters: Array<(value: any) => void> = []; + + ws.on('message', data => { + const message = JSON.parse(data.toString()); + const waiter = waiters.shift(); + if (waiter) { + waiter(message); + return; + } + messages.push(message); + }); + + return { + next() { + const message = messages.shift(); + if (message) { + return Promise.resolve(message); + } + return new Promise(resolve => { + waiters.push(resolve); + }); + }, + async nextMany(count: number) { + const result: any[] = []; + while (result.length < count) { + result.push(await this.next()); + } + return result; + }, + }; +} + +function waitClose(ws: WebSocket) { + if (ws.readyState === WebSocket.CLOSED) { + return Promise.resolve(); + } + return new Promise(resolve => { + ws.once('close', () => resolve()); + }); +} + +async function closeClient(ws: WebSocket) { + if (ws.readyState === WebSocket.CLOSED || ws.readyState === WebSocket.CLOSING) { + return; + } + ws.close(); + await waitClose(ws); +} + +async function waitFor(predicate: () => boolean) { + const deadline = Date.now() + 1000; + while (Date.now() < deadline) { + if (predicate()) { + return; + } + await new Promise(resolve => setTimeout(resolve, 10)); + } + assert(predicate(), 'predicate should become true before timeout'); +} From 0f644cd510cfc07824364eb203742a29c93fb62e Mon Sep 17 00:00:00 2001 From: akitasummer Date: Fri, 10 Jul 2026 19:02:28 +0800 Subject: [PATCH 4/9] feat(controller): serialize websocket fetch streams --- .../websocket/WebSocketControllerRegister.ts | 112 ++++++++++++--- .../app/controller/AppController.ts | 12 ++ .../app/controller/WebSocketController.ts | 43 +++++- .../app/controller/WebSocketTestState.ts | 1 + .../test/websocket/websocket.test.ts | 118 +++++++++++++++ .../test/websocket/websocketCluster.test.ts | 136 ++++++++++++++++++ .../websocket/WebSocketControllerRegister.ts | 111 +++++++++++--- .../fixtures/websocket/WebSocketController.ts | 37 ++++- .../test/websocket/websocket.test.ts | 88 +++++++++++- 9 files changed, 608 insertions(+), 50 deletions(-) create mode 100644 plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketTestState.ts diff --git a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts index c3da3d308..c5cb3c461 100644 --- a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts +++ b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts @@ -415,6 +415,14 @@ export class WebSocketControllerRegister implements ControllerRegister { const eggObj = await this.eggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); const realObj = eggObj.obj; const close = this.createFetchClose(webSocketCtx.socket); + const responseStreams = new Set(); + let connectionClosed = false; + let controllerReady = false; + let resolveControllerReady!: () => void; + const controllerReadyHandled = new Promise(resolve => { + resolveControllerReady = resolve; + }); + let messageQueue = Promise.resolve(); const invokeMethod = async ( targetMethodMeta: WebSocketFetchMethodMeta | undefined, payload: WebSocketMethodPayload = {}, @@ -438,6 +446,8 @@ export class WebSocketControllerRegister implements ControllerRegister { }; const closeHandled = new Promise(resolve => { webSocketCtx.socket.once('close', (code, reason) => { + connectionClosed = true; + this.destroyFetchResponseStreams(responseStreams); this.app.ctxStorage.run(eggCtx, async () => { try { await invokeMethod(controllerMeta.closeMethod, { @@ -457,13 +467,22 @@ export class WebSocketControllerRegister implements ControllerRegister { }); webSocketCtx.socket.on('message', data => { - this.app.ctxStorage.run(eggCtx, async () => { - try { - const result = await invokeMethod(methodMeta, { data }); - this.sendFetchResponseStream(result, webSocketCtx.socket, handleError); - } catch (error) { - await handleError(error); + if (connectionClosed) { + return; + } + messageQueue = messageQueue.then(async () => { + await controllerReadyHandled; + if (!controllerReady || connectionClosed || webSocketCtx.socket.readyState !== WebSocket.OPEN) { + return; } + await this.app.ctxStorage.run(eggCtx, async () => { + try { + const result = await invokeMethod(methodMeta, { data }); + await this.sendFetchResponseStream(result, webSocketCtx.socket, responseStreams, handleError); + } catch (error) { + await handleError(error); + } + }); }).catch(error => { this.app.logger.error('[tegg/websocket-fetch] handle message failed: %s', error.stack || error.message); }); @@ -476,8 +495,13 @@ export class WebSocketControllerRegister implements ControllerRegister { }); }); - await invokeMethod(controllerMeta.connectionMethod); - await invokeMethod(controllerMeta.openMethod); + try { + await invokeMethod(controllerMeta.connectionMethod); + await invokeMethod(controllerMeta.openMethod); + controllerReady = true; + } finally { + resolveControllerReady(); + } await closeHandled; } @@ -490,30 +514,82 @@ export class WebSocketControllerRegister implements ControllerRegister { }; } - private sendFetchResponseStream( + private async sendFetchResponseStream( result: unknown, webSocket: WebSocket, + responseStreams: Set, handleError: (error: Error) => Promise, - ) { + ): Promise { if (result === undefined || result === null) { return; } if (!this.isReadableStream(result)) { - handleError(new Error('WebSocketFetch method must return a readable stream or void')).catch(error => { + await handleError(new Error('WebSocketFetch method must return a readable stream or void')).catch(error => { this.app.logger.error('[tegg/websocket-fetch] handle invalid response failed: %s', error.stack || error.message); }); return; } - result.on('data', chunk => { - this.sendFetchChunk(webSocket, chunk, handleError); - }); - result.once('error', error => { - handleError(error).catch(err => { - this.app.logger.error('[tegg/websocket-fetch] handle stream error failed: %s', err.stack || err.message); - }); + if (webSocket.readyState !== WebSocket.OPEN) { + this.destroyFetchResponseStream(result); + return; + } + responseStreams.add(result); + await new Promise(resolve => { + let settled = false; + const onData = (chunk: unknown) => { + this.sendFetchChunk(webSocket, chunk, handleError); + }; + const settle = (error?: Error) => { + if (settled) { + return; + } + settled = true; + result.removeListener('data', onData); + responseStreams.delete(result); + if (!error || webSocket.readyState !== WebSocket.OPEN) { + resolve(); + return; + } + handleError(error).catch(err => { + this.app.logger.error('[tegg/websocket-fetch] handle stream error failed: %s', err.stack || err.message); + }).finally(resolve); + }; + result.once('end', () => settle()); + result.once('close', () => settle()); + result.once('error', settle); + result.on('data', onData); + const state = result as NodeJS.ReadableStream & { + closed?: boolean; + errored?: Error | null; + readableEnded?: boolean; + }; + if (state.errored) { + settle(state.errored); + } else if (state.closed || state.readableEnded) { + settle(); + } }); } + private destroyFetchResponseStreams(responseStreams: Set) { + for (const responseStream of responseStreams) { + this.destroyFetchResponseStream(responseStream); + } + responseStreams.clear(); + } + + private destroyFetchResponseStream(responseStream: NodeJS.ReadableStream) { + const destroy = (responseStream as NodeJS.ReadableStream & { destroy?: () => void }).destroy; + if (typeof destroy !== 'function') { + return; + } + try { + Reflect.apply(destroy, responseStream, []); + } catch (error) { + this.app.logger.error('[tegg/websocket-fetch] destroy response stream failed: %s', error.stack || error.message); + } + } + private sendFetchChunk( webSocket: WebSocket, chunk: unknown, diff --git a/plugin/controller/test/fixtures/apps/controller-app/app/controller/AppController.ts b/plugin/controller/test/fixtures/apps/controller-app/app/controller/AppController.ts index e55736d35..4f543ff7d 100644 --- a/plugin/controller/test/fixtures/apps/controller-app/app/controller/AppController.ts +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/AppController.ts @@ -15,6 +15,7 @@ import { import AppService from '../../modules/multi-module-service/AppService'; import App from '../../modules/multi-module-common/model/App'; import { countMw } from '../middleware/count_mw'; +import { webSocketFetchStreamCloseEvents } from './WebSocketTestState'; @HTTPController({ path: '/apps', @@ -37,6 +38,17 @@ export class AppController { }; } + @HTTPMethod({ + method: HTTPMethodEnum.GET, + path: '/websocket-stream-events/:id', + }) + getWebSocketStreamEvent(@HTTPParam() id: string) { + return { + event: webSocketFetchStreamCloseEvents.get(id), + pid: process.pid, + }; + } + @HTTPMethod({ method: HTTPMethodEnum.GET, path: '', diff --git a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts index 9198921c6..379e90a82 100644 --- a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts @@ -1,5 +1,5 @@ import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; -import { PassThrough, type Readable } from 'node:stream'; +import { PassThrough, pipeline, type Readable } from 'node:stream'; import type { RawData, WebSocket } from 'ws'; import { Context, @@ -28,12 +28,16 @@ import { WebSocketStream, } from '@eggjs/tegg'; import AppService from '../../modules/multi-module-service/AppService'; +import { webSocketFetchStreamCloseEvents } from './WebSocketTestState'; interface WebSocketFetchRequest { content: string; close?: boolean; error?: boolean; delayMs?: number; + holdOpen?: boolean; + observeClose?: boolean; + pipeline?: boolean; } @WebSocketController({ @@ -154,8 +158,12 @@ export class AppWebSocketFetchController { @Context() ctx: WebSocketContext, ) { const body = JSON.parse(data.toString()) as WebSocketFetchRequest; - const output = new PassThrough(); - output.write(JSON.stringify({ + const source = new PassThrough(); + const output = body.pipeline ? new PassThrough() : source; + if (body.pipeline) { + pipeline(source, output, () => {}); + } + source.write(JSON.stringify({ type: 'data', phase: 1, id, @@ -167,22 +175,43 @@ export class AppWebSocketFetchController { pid: process.pid, })); + if (body.observeClose) { + output.once('close', () => { + webSocketFetchStreamCloseEvents.set( + `ws-fetch-stream-close-${id}-${body.content}`, + `${output.readableEnded}:${output.destroyed}`, + ); + }); + if (body.pipeline) { + source.once('close', () => { + webSocketFetchStreamCloseEvents.set( + `ws-fetch-source-close-${id}-${body.content}`, + `${source.readableEnded}:${source.destroyed}`, + ); + }); + } + } + + if (body.holdOpen) { + return output; + } + setTimeout(() => { - if (output.destroyed) { + if (source.destroyed) { return; } if (body.error) { - output.destroy(new Error(`fetch error: ${body.content}`)); + source.destroy(new Error(`fetch error: ${body.content}`)); return; } - output.write(JSON.stringify({ + source.write(JSON.stringify({ type: 'data', phase: 2, id, content: body.content, pid: process.pid, })); - output.end(); + source.end(); }, body.delayMs ?? 10); if (body.close) { diff --git a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketTestState.ts b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketTestState.ts new file mode 100644 index 000000000..3045506f7 --- /dev/null +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketTestState.ts @@ -0,0 +1 @@ +export const webSocketFetchStreamCloseEvents = new Map(); diff --git a/plugin/controller/test/websocket/websocket.test.ts b/plugin/controller/test/websocket/websocket.test.ts index 0145d7afa..66613c54d 100644 --- a/plugin/controller/test/websocket/websocket.test.ts +++ b/plugin/controller/test/websocket/websocket.test.ts @@ -137,6 +137,21 @@ async function closeClient(socket: WebSocket) { }); } +async function waitForStreamCloseEvent(app, name: string, timeout = 2000) { + const deadline = Date.now() + timeout; + while (Date.now() < deadline) { + const res = await app.httpRequest() + .get(`/apps/websocket-stream-events/${name}`) + .set('connection', 'close') + .expect(200); + if (res.body.event) { + return res.body.event; + } + await new Promise(resolve => setTimeout(resolve, 10)); + } + assert.fail(`stream close event ${name} should be stored before timeout`); +} + describe('plugin/controller/test/websocket/websocket.test.ts', () => { let app; @@ -280,6 +295,7 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { socket.send(JSON.stringify({ content: 'broken', error: true, + pipeline: true, })); const beforeError = await receiver.next(); assert.equal(beforeError.type, 'data'); @@ -326,6 +342,43 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { } }); + it('should serialize websocket fetch messages on the same connection', async () => { + const socket = new WebSocket(requestUrl(app, '/ws-fetch/serial')); + const receiver = createJSONReceiver(socket); + await new Promise((resolve, reject) => { + socket.once('open', resolve); + socket.once('error', reject); + }); + + try { + await receiver.next(); + await receiver.next(); + + socket.send(JSON.stringify({ + content: 'first', + delayMs: 60, + })); + socket.send(JSON.stringify({ + content: 'second', + delayMs: 5, + })); + + const messages = [ + await receiver.next(), + await receiver.next(), + await receiver.next(), + await receiver.next(), + ]; + assert.deepEqual( + messages.map(message => `${message.content}:${message.phase}`), + [ 'first:1', 'first:2', 'second:1', 'second:2' ], + ); + assert.equal(socket.readyState, WebSocket.OPEN); + } finally { + await closeClient(socket); + } + }); + it('should isolate two websocket fetch connections with interleaved data', async () => { const firstSocket = new WebSocket(requestUrl(app, '/ws-fetch/cross-1?name=first&tag=a'), { headers: { 'x-client-id': 'client-cross-1' }, @@ -437,6 +490,71 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { } }); + it('should destroy only the closed connection websocket fetch streams', async () => { + const firstSocket = new WebSocket(requestUrl(app, '/ws-fetch/cleanup-first')); + const secondSocket = new WebSocket(requestUrl(app, '/ws-fetch/cleanup-second')); + const firstMessages = createJSONReceiver(firstSocket); + const secondMessages = createJSONReceiver(secondSocket); + + await Promise.all([ + new Promise((resolve, reject) => { + firstSocket.once('open', resolve); + firstSocket.once('error', reject); + }), + new Promise((resolve, reject) => { + secondSocket.once('open', resolve); + secondSocket.once('error', reject); + }), + ]); + + try { + const [ firstInitial, secondInitial ] = await Promise.all([ + Promise.all([ firstMessages.next(), firstMessages.next() ]), + Promise.all([ secondMessages.next(), secondMessages.next() ]), + ]); + assert.deepEqual(firstInitial.map(message => message.type), [ 'connection', 'open' ]); + assert.deepEqual(secondInitial.map(message => message.type), [ 'connection', 'open' ]); + + firstSocket.send(JSON.stringify({ + content: 'hold', + holdOpen: true, + observeClose: true, + pipeline: true, + })); + secondSocket.send(JSON.stringify({ + content: 'complete', + delayMs: 20, + observeClose: true, + })); + + assert.equal((await firstMessages.next()).phase, 1); + assert.equal((await secondMessages.next()).phase, 1); + await closeClient(firstSocket); + + const secondPhase2 = await secondMessages.next(); + assert.equal(secondPhase2.phase, 2); + assert.equal(secondPhase2.content, 'complete'); + + assert.equal( + await waitForStreamCloseEvent(app, 'ws-fetch-stream-close-cleanup-first-hold'), + 'false:true', + ); + assert.equal( + await waitForStreamCloseEvent(app, 'ws-fetch-source-close-cleanup-first-hold'), + 'false:true', + ); + assert.equal( + await waitForStreamCloseEvent(app, 'ws-fetch-stream-close-cleanup-second-complete'), + 'true:true', + ); + } finally { + await Promise.all([ + closeClient(firstSocket), + closeClient(secondSocket), + ]); + } + }); + it('should reject unmatched websocket route', async () => { await assert.rejects( () => createClient(requestUrl(app, '/ws/not-found')), diff --git a/plugin/controller/test/websocket/websocketCluster.test.ts b/plugin/controller/test/websocket/websocketCluster.test.ts index 54b7c304c..139d753f6 100644 --- a/plugin/controller/test/websocket/websocketCluster.test.ts +++ b/plugin/controller/test/websocket/websocketCluster.test.ts @@ -87,6 +87,21 @@ async function closeClient(socket: WebSocket) { }); } +async function waitForStreamCloseEvent(app, name: string, timeout = 5000) { + const deadline = Date.now() + timeout; + while (Date.now() < deadline) { + const res = await app.httpRequest() + .get(`/apps/websocket-stream-events/${name}`) + .set('connection', 'close') + .expect(200); + if (res.body.event) { + return res.body.event; + } + await new Promise(resolve => setTimeout(resolve, 20)); + } + assert.fail(`stream close event ${name} should be stored before timeout`); +} + describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { let app; @@ -212,6 +227,25 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { }); assert.equal(socket.readyState, WebSocket.OPEN); + socket.send(JSON.stringify({ + content: 'cluster-broken', + error: true, + pipeline: true, + })); + const beforeError = await receiver.next(); + assert.equal(beforeError.type, 'data'); + assert.equal(beforeError.phase, 1); + assert.equal(beforeError.content, 'cluster-broken'); + const error = await receiver.next(); + assert.deepEqual(error, { + type: 'error', + message: 'fetch error: cluster-broken', + header: 'cluster-fetch-client', + path: '/ws-fetch/cluster-fetch', + pid: error.pid, + }); + assert.equal(socket.readyState, WebSocket.OPEN); + const closePromise = receiveClose(socket); socket.send(JSON.stringify({ content: 'cluster-done', @@ -237,4 +271,106 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { } }); + it('should serialize websocket fetch messages on the same connection in cluster', async () => { + const socket = new WebSocket(requestUrl(app, '/ws-fetch/cluster-serial')); + const receiver = createJSONReceiver(socket); + await new Promise((resolve, reject) => { + socket.once('open', resolve); + socket.once('error', reject); + }); + + try { + await receiver.next(); + await receiver.next(); + + socket.send(JSON.stringify({ + content: 'first', + delayMs: 60, + })); + socket.send(JSON.stringify({ + content: 'second', + delayMs: 5, + })); + + const messages = [ + await receiver.next(), + await receiver.next(), + await receiver.next(), + await receiver.next(), + ]; + assert.deepEqual( + messages.map(message => `${message.content}:${message.phase}`), + [ 'first:1', 'first:2', 'second:1', 'second:2' ], + ); + assert.equal(socket.readyState, WebSocket.OPEN); + } finally { + await closeClient(socket); + } + }); + + it('should destroy only the closed connection websocket fetch streams in cluster', async () => { + const firstSocket = new WebSocket(requestUrl(app, '/ws-fetch/cluster-cleanup-first')); + const secondSocket = new WebSocket(requestUrl(app, '/ws-fetch/cluster-cleanup-second')); + const firstMessages = createJSONReceiver(firstSocket); + const secondMessages = createJSONReceiver(secondSocket); + + await Promise.all([ + new Promise((resolve, reject) => { + firstSocket.once('open', resolve); + firstSocket.once('error', reject); + }), + new Promise((resolve, reject) => { + secondSocket.once('open', resolve); + secondSocket.once('error', reject); + }), + ]); + + try { + await Promise.all([ + firstMessages.next(), + firstMessages.next(), + secondMessages.next(), + secondMessages.next(), + ]); + + firstSocket.send(JSON.stringify({ + content: 'hold', + holdOpen: true, + observeClose: true, + pipeline: true, + })); + secondSocket.send(JSON.stringify({ + content: 'complete', + delayMs: 20, + observeClose: true, + })); + + assert.equal((await firstMessages.next()).phase, 1); + assert.equal((await secondMessages.next()).phase, 1); + await closeClient(firstSocket); + + const secondPhase2 = await secondMessages.next(); + assert.equal(secondPhase2.phase, 2); + assert.equal(secondPhase2.content, 'complete'); + + assert.equal( + await waitForStreamCloseEvent(app, 'ws-fetch-stream-close-cluster-cleanup-first-hold'), + 'false:true', + ); + assert.equal( + await waitForStreamCloseEvent(app, 'ws-fetch-source-close-cluster-cleanup-first-hold'), + 'false:true', + ); + assert.equal( + await waitForStreamCloseEvent(app, 'ws-fetch-stream-close-cluster-cleanup-second-complete'), + 'true:true', + ); + } finally { + await Promise.all([ + closeClient(firstSocket), + closeClient(secondSocket), + ]); + } + }); + }); diff --git a/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts index 09e518719..8f86ce15a 100644 --- a/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts +++ b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts @@ -413,6 +413,14 @@ export class WebSocketControllerRegister implements ControllerRegister { const eggObj = await EggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); const realObj = eggObj.obj; const close = this.createFetchClose(webSocketCtx.socket); + const responseStreams = new Set(); + let connectionClosed = false; + let controllerReady = false; + let resolveControllerReady!: () => void; + const controllerReadyHandled = new Promise(resolve => { + resolveControllerReady = resolve; + }); + let messageQueue = Promise.resolve(); const invokeMethod = async ( targetMethodMeta: WebSocketFetchMethodMeta | undefined, payload: WebSocketMethodPayload = {}, @@ -436,6 +444,8 @@ export class WebSocketControllerRegister implements ControllerRegister { }; const closeHandled = new Promise(resolve => { webSocketCtx.socket.once('close', (code, reason) => { + connectionClosed = true; + this.destroyFetchResponseStreams(responseStreams); invokeMethod(controllerMeta.closeMethod, { closeCode: code, closeReason: reason, @@ -446,12 +456,22 @@ export class WebSocketControllerRegister implements ControllerRegister { }); webSocketCtx.socket.on('message', data => { - invokeMethod(methodMeta, { data }).then(result => { - this.sendFetchResponseStream(result, webSocketCtx.socket, handleError); + if (connectionClosed) { + return; + } + messageQueue = messageQueue.then(async () => { + await controllerReadyHandled; + if (!controllerReady || connectionClosed || webSocketCtx.socket.readyState !== WebSocket.OPEN) { + return; + } + try { + const result = await invokeMethod(methodMeta, { data }); + await this.sendFetchResponseStream(result, webSocketCtx.socket, responseStreams, handleError); + } catch (error) { + await handleError(error); + } }).catch(error => { - handleError(error).catch(err => { - console.error('[tegg/websocket-fetch] handle message error failed:', err); - }); + console.error('[tegg/websocket-fetch] handle message error failed:', error); }); }); webSocketCtx.socket.on('error', error => { @@ -460,8 +480,13 @@ export class WebSocketControllerRegister implements ControllerRegister { }); }); - await invokeMethod(controllerMeta.connectionMethod); - await invokeMethod(controllerMeta.openMethod); + try { + await invokeMethod(controllerMeta.connectionMethod); + await invokeMethod(controllerMeta.openMethod); + controllerReady = true; + } finally { + resolveControllerReady(); + } await closeHandled; } @@ -474,30 +499,82 @@ export class WebSocketControllerRegister implements ControllerRegister { }; } - private sendFetchResponseStream( + private async sendFetchResponseStream( result: unknown, webSocket: WebSocket, + responseStreams: Set, handleError: (error: Error) => Promise, - ) { + ): Promise { if (result === undefined || result === null) { return; } if (!this.isReadableStream(result)) { - handleError(new Error('WebSocketFetch method must return a readable stream or void')).catch(error => { + await handleError(new Error('WebSocketFetch method must return a readable stream or void')).catch(error => { console.error('[tegg/websocket-fetch] handle invalid response failed:', error); }); return; } - result.on('data', chunk => { - this.sendFetchChunk(webSocket, chunk, handleError); - }); - result.once('error', error => { - handleError(error).catch(err => { - console.error('[tegg/websocket-fetch] handle stream error failed:', err); - }); + if (webSocket.readyState !== WebSocket.OPEN) { + this.destroyFetchResponseStream(result); + return; + } + responseStreams.add(result); + await new Promise(resolve => { + let settled = false; + const onData = (chunk: unknown) => { + this.sendFetchChunk(webSocket, chunk, handleError); + }; + const settle = (error?: Error) => { + if (settled) { + return; + } + settled = true; + result.removeListener('data', onData); + responseStreams.delete(result); + if (!error || webSocket.readyState !== WebSocket.OPEN) { + resolve(); + return; + } + handleError(error).catch(err => { + console.error('[tegg/websocket-fetch] handle stream error failed:', err); + }).finally(resolve); + }; + result.once('end', () => settle()); + result.once('close', () => settle()); + result.once('error', settle); + result.on('data', onData); + const state = result as NodeJS.ReadableStream & { + closed?: boolean; + errored?: Error | null; + readableEnded?: boolean; + }; + if (state.errored) { + settle(state.errored); + } else if (state.closed || state.readableEnded) { + settle(); + } }); } + private destroyFetchResponseStreams(responseStreams: Set) { + for (const responseStream of responseStreams) { + this.destroyFetchResponseStream(responseStream); + } + responseStreams.clear(); + } + + private destroyFetchResponseStream(responseStream: NodeJS.ReadableStream) { + const destroy = (responseStream as NodeJS.ReadableStream & { destroy?: () => void }).destroy; + if (typeof destroy !== 'function') { + return; + } + try { + Reflect.apply(destroy, responseStream, []); + } catch (error) { + console.error('[tegg/websocket-fetch] destroy response stream failed:', error); + } + } + private sendFetchChunk( webSocket: WebSocket, chunk: unknown, diff --git a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts index f5d901213..d656996ce 100644 --- a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts +++ b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts @@ -1,5 +1,5 @@ import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; -import { PassThrough, type Readable } from 'node:stream'; +import { PassThrough, pipeline, type Readable } from 'node:stream'; import type { RawData, WebSocket } from 'ws'; import { Context, @@ -32,9 +32,13 @@ interface WebSocketFetchRequest { close?: boolean; error?: boolean; delayMs?: number; + holdOpen?: boolean; + observeClose?: boolean; + pipeline?: boolean; } export const fetchCloseEvents: string[] = []; +export const fetchStreamCloseEvents: string[] = []; @WebSocketController({ path: '/ws', @@ -135,8 +139,12 @@ export class StandaloneWebSocketFetchController { @Context() ctx: WebSocketContext, ) { const body = JSON.parse(data.toString()) as WebSocketFetchRequest; - const output = new PassThrough(); - output.write(JSON.stringify({ + const source = new PassThrough(); + const output = body.pipeline ? new PassThrough() : source; + if (body.pipeline) { + pipeline(source, output, () => {}); + } + source.write(JSON.stringify({ type: 'data', phase: 1, id, @@ -147,21 +155,36 @@ export class StandaloneWebSocketFetchController { content: body.content, })); + if (body.observeClose) { + output.once('close', () => { + fetchStreamCloseEvents.push(`${id}:${body.content}:${output.readableEnded}:${output.destroyed}`); + }); + if (body.pipeline) { + source.once('close', () => { + fetchStreamCloseEvents.push(`source:${id}:${body.content}:${source.readableEnded}:${source.destroyed}`); + }); + } + } + + if (body.holdOpen) { + return output; + } + setTimeout(() => { - if (output.destroyed) { + if (source.destroyed) { return; } if (body.error) { - output.destroy(new Error(`fetch error: ${body.content}`)); + source.destroy(new Error(`fetch error: ${body.content}`)); return; } - output.write(JSON.stringify({ + source.write(JSON.stringify({ type: 'data', phase: 2, id, content: body.content, })); - output.end(); + source.end(); }, body.delayMs ?? 10); if (body.close) { diff --git a/standalone/service-worker/test/websocket/websocket.test.ts b/standalone/service-worker/test/websocket/websocket.test.ts index 01203151c..04cdc2b9a 100644 --- a/standalone/service-worker/test/websocket/websocket.test.ts +++ b/standalone/service-worker/test/websocket/websocket.test.ts @@ -5,7 +5,7 @@ import { createWebSocketStream, WebSocket } from 'ws'; import { ServiceWorkerApp } from '../../src/ServiceWorkerApp'; import { StandaloneTestUtil } from '@eggjs/module-test-util/StandaloneTestUtil'; import { TestUtils } from '../Utils'; -import { fetchCloseEvents } from '../fixtures/websocket/WebSocketController'; +import { fetchCloseEvents, fetchStreamCloseEvents } from '../fixtures/websocket/WebSocketController'; describe('standalone/service-worker/test/websocket/websocket.test.ts', () => { let app: ServiceWorkerApp; @@ -27,6 +27,7 @@ describe('standalone/service-worker/test/websocket/websocket.test.ts', () => { beforeEach(() => { fetchCloseEvents.length = 0; + fetchStreamCloseEvents.length = 0; }); it('should handle websocket controller in standalone service worker', async () => { @@ -173,7 +174,92 @@ describe('standalone/service-worker/test/websocket/websocket.test.ts', () => { await waitClose(second); await waitFor(() => fetchCloseEvents.includes('second:1000:server done')); + + first.send(JSON.stringify({ + content: 'broken', + error: true, + pipeline: true, + })); + assert.deepStrictEqual(await firstMessages.next(), { + type: 'data', + phase: 1, + id: 'first', + name: 'one', + tags: [ 'a', 'b' ], + header: 'first-client', + path: '/ws-fetch/first', + content: 'broken', + }); + assert.deepStrictEqual(await firstMessages.next(), { + type: 'error', + message: 'fetch error: broken', + header: 'first-client', + path: '/ws-fetch/first', + }); + await closeClient(first); + }); + + it('should serialize websocket fetch messages on the same connection in standalone service worker', async () => { + const socket = createClient('/ws-fetch/serial'); + const messages = createJSONMessageQueue(socket); + await messages.nextMany(2); + + socket.send(JSON.stringify({ + content: 'first', + delayMs: 60, + })); + socket.send(JSON.stringify({ + content: 'second', + delayMs: 5, + })); + + const responses = await messages.nextMany(4); + assert.deepStrictEqual( + responses.map(response => `${response.content}:${response.phase}`), + [ 'first:1', 'first:2', 'second:1', 'second:2' ], + ); + assert.equal(socket.readyState, WebSocket.OPEN); + await closeClient(socket); + }); + + it('should destroy only the closed connection websocket fetch streams in standalone service worker', async () => { + const first = createClient('/ws-fetch/cleanup-first'); + const second = createClient('/ws-fetch/cleanup-second'); + const firstMessages = createJSONMessageQueue(first); + const secondMessages = createJSONMessageQueue(second); + + await Promise.all([ + firstMessages.nextMany(2), + secondMessages.nextMany(2), + ]); + + first.send(JSON.stringify({ + content: 'hold', + holdOpen: true, + observeClose: true, + pipeline: true, + })); + second.send(JSON.stringify({ + content: 'complete', + delayMs: 20, + observeClose: true, + })); + + assert.equal((await firstMessages.next()).phase, 1); + assert.equal((await secondMessages.next()).phase, 1); await closeClient(first); + + assert.deepStrictEqual(await secondMessages.next(), { + type: 'data', + phase: 2, + id: 'cleanup-second', + content: 'complete', + }); + await waitFor(() => fetchStreamCloseEvents.includes('cleanup-first:hold:false:true')); + await waitFor(() => fetchStreamCloseEvents.includes('source:cleanup-first:hold:false:true')); + await waitFor(() => fetchStreamCloseEvents.includes('cleanup-second:complete:true:true')); + + await closeClient(second); }); function createClient(path: string, headers?: Record) { From d9182e51190a47ef0ae9d0b3624715c3abce615c Mon Sep 17 00:00:00 2001 From: akitasummer Date: Fri, 10 Jul 2026 19:09:40 +0800 Subject: [PATCH 5/9] fix(controller): resolve websocket lint errors --- .../controller-app/app/controller/WebSocketController.ts | 4 +++- .../test/fixtures/websocket/WebSocketController.ts | 4 +++- .../service-worker/test/websocket/websocket.test.ts | 8 -------- 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts index 379e90a82..5f1a42344 100644 --- a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts @@ -161,7 +161,9 @@ export class AppWebSocketFetchController { const source = new PassThrough(); const output = body.pipeline ? new PassThrough() : source; if (body.pipeline) { - pipeline(source, output, () => {}); + pipeline(source, output, () => { + // WebSocketFetch observes errors from the returned output stream. + }); } source.write(JSON.stringify({ type: 'data', diff --git a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts index d656996ce..1e314d36a 100644 --- a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts +++ b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts @@ -142,7 +142,9 @@ export class StandaloneWebSocketFetchController { const source = new PassThrough(); const output = body.pipeline ? new PassThrough() : source; if (body.pipeline) { - pipeline(source, output, () => {}); + pipeline(source, output, () => { + // WebSocketFetch observes errors from the returned output stream. + }); } source.write(JSON.stringify({ type: 'data', diff --git a/standalone/service-worker/test/websocket/websocket.test.ts b/standalone/service-worker/test/websocket/websocket.test.ts index 04cdc2b9a..cbe7a211e 100644 --- a/standalone/service-worker/test/websocket/websocket.test.ts +++ b/standalone/service-worker/test/websocket/websocket.test.ts @@ -284,14 +284,6 @@ function waitJSONMessage(ws: WebSocket): Promise { }); } -async function waitJSONMessages(ws: WebSocket, count: number) { - const messages: unknown[] = []; - while (messages.length < count) { - messages.push(await waitJSONMessage(ws)); - } - return messages; -} - function createJSONMessageQueue(ws: WebSocket) { const messages: any[] = []; const waiters: Array<(value: any) => void> = []; From 52be2fdea74c6472da1e85caee9b782343c61ca4 Mon Sep 17 00:00:00 2001 From: akitasummer Date: Fri, 10 Jul 2026 19:44:42 +0800 Subject: [PATCH 6/9] fix(controller): address websocket review feedback --- .../src/decorator/http/HTTPParam.ts | 49 --------------- .../src/decorator/websocket/WebSocketParam.ts | 62 +------------------ .../WebSocketFetchControllerMetaBuilder.ts | 3 +- ...bSocketFetchControllerMethodMetaBuilder.ts | 2 +- .../WebSocketControllerMetaBuilder.ts | 2 +- .../WebSocketControllerMethodMetaBuilder.ts | 2 +- .../test/websocket/WebSocketMeta.test.ts | 23 +++++++ .../controller-decorator/WebSocketParam.ts | 11 ---- core/types/controller-decorator/index.ts | 1 - .../websocket/WebSocketControllerRegister.ts | 55 +++++++++++++--- plugin/controller/package.json | 2 +- .../app/controller/WebSocketController.ts | 38 +++++++----- .../test/websocket/websocket.test.ts | 55 ++++++++++++++-- standalone/service-worker/package.json | 2 +- .../websocket/WebSocketControllerRegister.ts | 54 +++++++++++++--- .../fixtures/websocket/WebSocketController.ts | 38 +++++++----- .../test/websocket/websocket.test.ts | 17 +++++ 17 files changed, 232 insertions(+), 184 deletions(-) create mode 100644 core/controller-decorator/test/websocket/WebSocketMeta.test.ts delete mode 100644 core/types/controller-decorator/WebSocketParam.ts diff --git a/core/controller-decorator/src/decorator/http/HTTPParam.ts b/core/controller-decorator/src/decorator/http/HTTPParam.ts index e74dbf30e..38765253e 100644 --- a/core/controller-decorator/src/decorator/http/HTTPParam.ts +++ b/core/controller-decorator/src/decorator/http/HTTPParam.ts @@ -20,16 +20,6 @@ export function HTTPBody() { } export function HTTPHeaders() { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.HEADERS, parameterIndex, controllerClazz, methodName); - }; -} - -export function Headers() { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { assert(typeof propertyKey === 'string', `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); @@ -41,19 +31,6 @@ export function Headers() { } export function HTTPQuery(param?: HTTPQueryParams) { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); - const name = param?.name || argNames[parameterIndex]; - HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.QUERY, parameterIndex, controllerClazz, methodName); - HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); - }; -} - -export function Query(param?: HTTPQueryParams) { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { assert(typeof propertyKey === 'string', `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); @@ -69,19 +46,6 @@ export function Query(param?: HTTPQueryParams) { } export function HTTPQueries(param?: HTTPQueriesParams) { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); - const name = param?.name || argNames[parameterIndex]; - HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.QUERIES, parameterIndex, controllerClazz, methodName); - HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); - }; -} - -export function Queries(param?: HTTPQueriesParams) { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { assert(typeof propertyKey === 'string', `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); @@ -97,19 +61,6 @@ export function Queries(param?: HTTPQueriesParams) { } export function HTTPParam(param?: HTTPParamParams) { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); - const name = param?.name || argNames[parameterIndex]; - HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.PARAM, parameterIndex, controllerClazz, methodName); - HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); - }; -} - -export function Param(param?: HTTPParamParams) { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { assert(typeof propertyKey === 'string', `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); diff --git a/core/controller-decorator/src/decorator/websocket/WebSocketParam.ts b/core/controller-decorator/src/decorator/websocket/WebSocketParam.ts index 510785c1e..18182f112 100644 --- a/core/controller-decorator/src/decorator/websocket/WebSocketParam.ts +++ b/core/controller-decorator/src/decorator/websocket/WebSocketParam.ts @@ -1,68 +1,8 @@ import assert from 'node:assert'; import { WebSocketParamType } from '@eggjs/tegg-types'; -import type { EggProtoImplClass, WebSocketParamParams, WebSocketQueriesParams, WebSocketQueryParams } from '@eggjs/tegg-types'; -import { ObjectUtils } from '@eggjs/tegg-common-util'; +import type { EggProtoImplClass } from '@eggjs/tegg-types'; import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; -export function WebSocketHeaders() { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.HEADERS, parameterIndex, controllerClazz, methodName); - }; -} - -export function WebSocketQuery(param?: WebSocketQueryParams) { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); - const name = param?.name || argNames[parameterIndex]; - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.QUERY, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); - }; -} - -export function WebSocketQueries(param?: WebSocketQueriesParams) { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); - const name = param?.name || argNames[parameterIndex]; - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.QUERIES, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); - }; -} - -export function WebSocketParam(param?: WebSocketParamParams) { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - const argNames = ObjectUtils.getFunctionArgNameList(target[propertyKey]); - const name = param?.name || argNames[parameterIndex]; - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.PARAM, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); - }; -} - -export function WebSocketRequest() { - return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { - assert(typeof propertyKey === 'string', - `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); - const methodName = propertyKey as string; - const controllerClazz = target.constructor as EggProtoImplClass; - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.REQUEST, parameterIndex, controllerClazz, methodName); - }; -} - export function WebSocketSocket() { return function(target: any, propertyKey: PropertyKey, parameterIndex: number) { assert(typeof propertyKey === 'string', diff --git a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts index 7fa07eb99..0b48a1780 100644 --- a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts @@ -76,8 +76,9 @@ export class WebSocketFetchControllerMetaBuilder { build(): WebSocketFetchControllerMeta { ControllerValidator.validate(this.clazz); const controllerType = ControllerInfoUtil.getControllerType(this.clazz); - assert(controllerType === ControllerType.WEBSOCKET_FETCH, 'invalidate controller type'); + assert(controllerType === ControllerType.WEBSOCKET_FETCH, 'invalid controller type'); const webSocketPath = WebSocketInfoUtil.getWebSocketPath(this.clazz); + assert(webSocketPath, `build websocket fetch controller ${ClassUtil.classDescription(this.clazz)} failed: path is required`); const middlewares = ControllerInfoUtil.getControllerMiddlewares(this.clazz); const { methods, connectionMethod, openMethod, errorMethod, closeMethod } = this.buildMethods(); assert(methods.length === 1, 'websocket fetch controller must has exactly one @WebSocketFetchMethod'); diff --git a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts index 3197e223d..a3f7bdf9f 100644 --- a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts @@ -34,7 +34,7 @@ export class WebSocketFetchControllerMethodMetaBuilder { const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(i, this.clazz, this.methodName); if (!paramType) { const classDesc = ClassUtil.classDescription(this.clazz); - throw new Error(`${classDesc}:${this.methodName} param ${i} has no websocket fetch param type, Please add @WebSocketData, @WebSocketClose, @WebSocketError, @WebSocketCloseCode, @WebSocketCloseReason, @WebSocketParam, @WebSocketQuery, @WebSocketQueries, @WebSocketHeaders, @WebSocketRequest, @WebSocketSocket or @Context`); + throw new Error(`${classDesc}:${this.methodName} param ${i} has no websocket fetch param type, Please add @WebSocketData, @WebSocketClose, @WebSocketError, @WebSocketCloseCode, @WebSocketCloseReason, @HTTPParam, @HTTPQuery, @HTTPQueries, @HTTPHeaders, @Request, @WebSocketSocket or @Context`); } if (!this.isAllowedParamType(methodType, paramType)) { const classDesc = ClassUtil.classDescription(this.clazz); diff --git a/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts b/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts index 5aef7a594..eabfd3b3e 100644 --- a/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts @@ -35,7 +35,7 @@ export class WebSocketControllerMetaBuilder { build(): WebSocketControllerMeta { ControllerValidator.validate(this.clazz); const controllerType = ControllerInfoUtil.getControllerType(this.clazz); - assert(controllerType === ControllerType.WEBSOCKET, 'invalidate controller type'); + assert(controllerType === ControllerType.WEBSOCKET, 'invalid controller type'); const webSocketPath = WebSocketInfoUtil.getWebSocketPath(this.clazz); const middlewares = ControllerInfoUtil.getControllerMiddlewares(this.clazz); const methods = this.buildMethod(); diff --git a/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts index 213109efd..1d19907e3 100644 --- a/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts @@ -34,7 +34,7 @@ export class WebSocketControllerMethodMetaBuilder { const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(i, this.clazz, this.methodName); if (!paramType) { const classDesc = ClassUtil.classDescription(this.clazz); - throw new Error(`${classDesc}:${this.methodName} param ${i} has no websocket param type, Please add @WebSocketParam, @WebSocketQuery, @WebSocketQueries, @WebSocketHeaders, @WebSocketRequest, @WebSocketSocket, @WebSocketStream or @Context`); + throw new Error(`${classDesc}:${this.methodName} param ${i} has no websocket param type, Please add @HTTPParam, @HTTPQuery, @HTTPQueries, @HTTPHeaders, @Request, @WebSocketSocket, @WebSocketStream or @Context`); } } } diff --git a/core/controller-decorator/test/websocket/WebSocketMeta.test.ts b/core/controller-decorator/test/websocket/WebSocketMeta.test.ts new file mode 100644 index 000000000..f2c1792bb --- /dev/null +++ b/core/controller-decorator/test/websocket/WebSocketMeta.test.ts @@ -0,0 +1,23 @@ +import assert from 'node:assert'; +import { + ControllerMetaBuilderFactory, + WebSocketFetchController, + WebSocketFetchMethod, +} from '../..'; + +@WebSocketFetchController() +class MissingPathWebSocketFetchController { + @WebSocketFetchMethod() + onData() { + // test fixture + } +} + +describe('core/controller-decorator/test/websocket/WebSocketMeta.test.ts', () => { + it('should require websocket fetch controller path', () => { + assert.throws( + () => ControllerMetaBuilderFactory.build(MissingPathWebSocketFetchController), + /build websocket fetch controller .* failed: path is required/, + ); + }); +}); diff --git a/core/types/controller-decorator/WebSocketParam.ts b/core/types/controller-decorator/WebSocketParam.ts deleted file mode 100644 index 25111609a..000000000 --- a/core/types/controller-decorator/WebSocketParam.ts +++ /dev/null @@ -1,11 +0,0 @@ -export interface WebSocketQueryParams { - name?: string; -} - -export interface WebSocketQueriesParams { - name?: string; -} - -export interface WebSocketParamParams { - name?: string; -} diff --git a/core/types/controller-decorator/index.ts b/core/types/controller-decorator/index.ts index 98dda1a6b..6a688922a 100644 --- a/core/types/controller-decorator/index.ts +++ b/core/types/controller-decorator/index.ts @@ -7,7 +7,6 @@ export * from './HTTPMethod'; export * from './HTTPParam'; export * from './WebSocketController'; export * from './WebSocketMethod'; -export * from './WebSocketParam'; export * from './WebSocketContext'; export * from './WebSocketFetchController'; export * from './WebSocketFetchMethod'; diff --git a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts index c5cb3c461..8f2f1103e 100644 --- a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts +++ b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts @@ -230,7 +230,6 @@ export class WebSocketControllerRegister implements ControllerRegister { const url = this.createURL(req); const matched = this.matchRoute(req, url); if (!matched) { - this.rejectSocket(socket, 404, 'Not Found'); return; } @@ -495,14 +494,26 @@ export class WebSocketControllerRegister implements ControllerRegister { }); }); + let lifecycleError: Error | undefined; try { await invokeMethod(controllerMeta.connectionMethod); await invokeMethod(controllerMeta.openMethod); controllerReady = true; + } catch (error) { + lifecycleError = error; + await handleError(error).catch(err => { + this.app.logger.error('[tegg/websocket-fetch] handle lifecycle error failed: %s', err.stack || err.message); + }); + if (webSocketCtx.socket.readyState === WebSocket.OPEN || webSocketCtx.socket.readyState === WebSocket.CONNECTING) { + webSocketCtx.socket.close(1011, 'Internal Server Error'); + } } finally { resolveControllerReady(); } await closeHandled; + if (lifecycleError) { + throw lifecycleError; + } } private createFetchClose(webSocket: WebSocket): WebSocketFetchClose { @@ -539,12 +550,18 @@ export class WebSocketControllerRegister implements ControllerRegister { const onData = (chunk: unknown) => { this.sendFetchChunk(webSocket, chunk, handleError); }; + let onEnd!: () => void; + let onClose!: () => void; + let onError!: (error: Error) => void; const settle = (error?: Error) => { if (settled) { return; } settled = true; result.removeListener('data', onData); + result.removeListener('end', onEnd); + result.removeListener('close', onClose); + result.removeListener('error', onError); responseStreams.delete(result); if (!error || webSocket.readyState !== WebSocket.OPEN) { resolve(); @@ -554,9 +571,12 @@ export class WebSocketControllerRegister implements ControllerRegister { this.app.logger.error('[tegg/websocket-fetch] handle stream error failed: %s', err.stack || err.message); }).finally(resolve); }; - result.once('end', () => settle()); - result.once('close', () => settle()); - result.once('error', settle); + onEnd = () => settle(); + onClose = () => settle(); + onError = error => settle(error); + result.once('end', onEnd); + result.once('close', onClose); + result.once('error', onError); result.on('data', onData); const state = result as NodeJS.ReadableStream & { closed?: boolean; @@ -564,7 +584,8 @@ export class WebSocketControllerRegister implements ControllerRegister { readableEnded?: boolean; }; if (state.errored) { - settle(state.errored); + const streamError = state.errored; + setImmediate(() => settle(streamError)); } else if (state.closed || state.readableEnded) { settle(); } @@ -634,8 +655,25 @@ export class WebSocketControllerRegister implements ControllerRegister { return Promise.resolve(); } return new Promise(resolve => { - webSocket.once('close', () => resolve()); - webSocket.once('error', () => resolve()); + let settled = false; + let onClose!: () => void; + let onError!: (error: Error) => void; + const settle = () => { + if (settled) { + return; + } + settled = true; + webSocket.removeListener('close', onClose); + webSocket.removeListener('error', onError); + resolve(); + }; + onClose = () => settle(); + onError = error => { + this.app.logger.error('[tegg/websocket] socket error while waiting for close: %s', error.stack || error.message); + settle(); + }; + webSocket.once('close', onClose); + webSocket.once('error', onError); }); } @@ -644,7 +682,6 @@ export class WebSocketControllerRegister implements ControllerRegister { return; } const body = message; - socket.write(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); - socket.destroy(); + socket.end(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); } } diff --git a/plugin/controller/package.json b/plugin/controller/package.json index 2025b00f4..c0a9cb0f1 100644 --- a/plugin/controller/package.json +++ b/plugin/controller/package.json @@ -65,7 +65,7 @@ "path-to-regexp": "^1.8.0", "raw-body": "^2.5.2", "sdk-base": "^4.2.0", - "ws": "^8.18.0", + "ws": "^8.21.0", "zod": "^4.0.0" }, "devDependencies": { diff --git a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts index 5f1a42344..9ec24c27a 100644 --- a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts @@ -3,11 +3,11 @@ import { PassThrough, pipeline, type Readable } from 'node:stream'; import type { RawData, WebSocket } from 'ws'; import { Context, - Headers, + HTTPHeaders, + HTTPParam, + HTTPQueries, + HTTPQuery, Inject, - Param, - Queries, - Query, Request, WebSocketContext, WebSocketController, @@ -52,10 +52,10 @@ export class AppWebSocketController { }) async echo( @WebSocketSocket() socket: WebSocket, - @Param() id: string, - @Query() name: string, - @Queries({ name: 'tag' }) tags: string[], - @Headers() headers: IncomingHttpHeaders, + @HTTPParam() id: string, + @HTTPQuery() name: string, + @HTTPQueries({ name: 'tag' }) tags: string[], + @HTTPHeaders() headers: IncomingHttpHeaders, @Request() request: IncomingMessage, @Context() ctx: WebSocketContext, ) { @@ -91,8 +91,8 @@ export class AppWebSocketController { }) async stream( @WebSocketStream() input: Readable, - @Param() id: string, - @Query() name: string, + @HTTPParam() id: string, + @HTTPQuery() name: string, ) { const output = new PassThrough(); output.write(JSON.stringify({ @@ -125,9 +125,13 @@ export class AppWebSocketFetchController { @WebSocketFetchOnConnection() onConnection( - @Headers() headers: IncomingHttpHeaders, + @HTTPParam() id: string, + @HTTPHeaders() headers: IncomingHttpHeaders, @Context() ctx: WebSocketContext, ) { + if (id === 'connection-error') { + throw new Error('fetch connection error'); + } ctx.socket.send(JSON.stringify({ type: 'connection', header: headers['x-client-id'], @@ -151,10 +155,10 @@ export class AppWebSocketFetchController { onData( @WebSocketData() data: RawData, @WebSocketClose() close: WebSocketFetchClose, - @Param() id: string, - @Query() name: string, - @Queries({ name: 'tag' }) tags: string[], - @Headers() headers: IncomingHttpHeaders, + @HTTPParam() id: string, + @HTTPQuery() name: string, + @HTTPQueries({ name: 'tag' }) tags: string[], + @HTTPHeaders() headers: IncomingHttpHeaders, @Context() ctx: WebSocketContext, ) { const body = JSON.parse(data.toString()) as WebSocketFetchRequest; @@ -225,7 +229,7 @@ export class AppWebSocketFetchController { @WebSocketFetchOnError() onError( @WebSocketError() error: Error, - @Headers() headers: IncomingHttpHeaders, + @HTTPHeaders() headers: IncomingHttpHeaders, @Context() ctx: WebSocketContext, ) { ctx.socket.send(JSON.stringify({ @@ -241,7 +245,7 @@ export class AppWebSocketFetchController { async onClose( @WebSocketCloseCode() code: number, @WebSocketCloseReason() reason: Buffer, - @Param() id: string, + @HTTPParam() id: string, ) { await this.appService.save({ name: `ws-fetch-close-${id}`, diff --git a/plugin/controller/test/websocket/websocket.test.ts b/plugin/controller/test/websocket/websocket.test.ts index 66613c54d..e1e2380d4 100644 --- a/plugin/controller/test/websocket/websocket.test.ts +++ b/plugin/controller/test/websocket/websocket.test.ts @@ -1,5 +1,7 @@ import { strict as assert } from 'node:assert'; +import type { IncomingMessage } from 'node:http'; import path from 'node:path'; +import type { Duplex } from 'node:stream'; import mm from 'egg-mock'; import WebSocket from 'ws'; import type { RawData } from 'ws'; @@ -342,6 +344,37 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { } }); + it('should handle websocket fetch connection lifecycle errors before closing', async () => { + const socket = new WebSocket(requestUrl(app, '/ws-fetch/connection-error'), { + headers: { 'x-client-id': 'connection-error-client' }, + }); + const receiver = createJSONReceiver(socket); + const closePromise = receiveClose(socket); + await new Promise((resolve, reject) => { + socket.once('open', resolve); + socket.once('error', reject); + }); + + const error = await receiver.next(); + assert.deepEqual(error, { + type: 'error', + message: 'fetch connection error', + header: 'connection-error-client', + path: '/ws-fetch/connection-error', + pid: error.pid, + }); + assert.deepEqual(await closePromise, { + code: 1011, + reason: 'Internal Server Error', + }); + + const res = await app.httpRequest() + .get('/apps/ws-fetch-close-connection-error') + .expect(200); + assert.equal(res.body.app.name, 'ws-fetch-close-connection-error'); + assert.equal(res.body.app.desc, '1011:Internal Server Error'); + }); + it('should serialize websocket fetch messages on the same connection', async () => { const socket = new WebSocket(requestUrl(app, '/ws-fetch/serial')); const receiver = createJSONReceiver(socket); @@ -555,10 +588,22 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { } }); - it('should reject unmatched websocket route', async () => { - await assert.rejects( - () => createClient(requestUrl(app, '/ws/not-found')), - /Unexpected server response: 404/, - ); + it('should let another upgrade listener handle an unmatched websocket route', async () => { + const handleUpgrade = (request: IncomingMessage, socket: Duplex) => { + if (request.url !== '/ws/not-found') { + return; + } + const body = 'handled by another upgrade listener'; + socket.end(`HTTP/1.1 418 I'm a Teapot\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); + }; + app.server.on('upgrade', handleUpgrade); + try { + await assert.rejects( + () => createClient(requestUrl(app, '/ws/not-found')), + /Unexpected server response: 418/, + ); + } finally { + app.server.removeListener('upgrade', handleUpgrade); + } }); }); diff --git a/standalone/service-worker/package.json b/standalone/service-worker/package.json index bb807e265..9e559e1de 100644 --- a/standalone/service-worker/package.json +++ b/standalone/service-worker/package.json @@ -54,7 +54,7 @@ "path-to-regexp": "^1.8.0", "type-is": "^1.6.18", "urllib": "^4.0.0", - "ws": "^8.18.0" + "ws": "^8.21.0" }, "peerDependencies": { "@eggjs/tegg": "^3" diff --git a/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts index 8f86ce15a..16906e8cb 100644 --- a/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts +++ b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts @@ -156,8 +156,7 @@ export class WebSocketControllerRegister implements ControllerRegister { return; } const body = message; - socket.write(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); - socket.destroy(); + socket.end(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); } private createRoute( @@ -480,14 +479,26 @@ export class WebSocketControllerRegister implements ControllerRegister { }); }); + let lifecycleError: Error | undefined; try { await invokeMethod(controllerMeta.connectionMethod); await invokeMethod(controllerMeta.openMethod); controllerReady = true; + } catch (error) { + lifecycleError = error; + await handleError(error).catch(err => { + console.error('[tegg/websocket-fetch] handle lifecycle error failed:', err); + }); + if (webSocketCtx.socket.readyState === WebSocket.OPEN || webSocketCtx.socket.readyState === WebSocket.CONNECTING) { + webSocketCtx.socket.close(1011, 'Internal Server Error'); + } } finally { resolveControllerReady(); } await closeHandled; + if (lifecycleError) { + throw lifecycleError; + } } private createFetchClose(webSocket: WebSocket): WebSocketFetchClose { @@ -524,12 +535,18 @@ export class WebSocketControllerRegister implements ControllerRegister { const onData = (chunk: unknown) => { this.sendFetchChunk(webSocket, chunk, handleError); }; + let onEnd!: () => void; + let onClose!: () => void; + let onError!: (error: Error) => void; const settle = (error?: Error) => { if (settled) { return; } settled = true; result.removeListener('data', onData); + result.removeListener('end', onEnd); + result.removeListener('close', onClose); + result.removeListener('error', onError); responseStreams.delete(result); if (!error || webSocket.readyState !== WebSocket.OPEN) { resolve(); @@ -539,9 +556,12 @@ export class WebSocketControllerRegister implements ControllerRegister { console.error('[tegg/websocket-fetch] handle stream error failed:', err); }).finally(resolve); }; - result.once('end', () => settle()); - result.once('close', () => settle()); - result.once('error', settle); + onEnd = () => settle(); + onClose = () => settle(); + onError = error => settle(error); + result.once('end', onEnd); + result.once('close', onClose); + result.once('error', onError); result.on('data', onData); const state = result as NodeJS.ReadableStream & { closed?: boolean; @@ -549,7 +569,8 @@ export class WebSocketControllerRegister implements ControllerRegister { readableEnded?: boolean; }; if (state.errored) { - settle(state.errored); + const streamError = state.errored; + setImmediate(() => settle(streamError)); } else if (state.closed || state.readableEnded) { settle(); } @@ -619,8 +640,25 @@ export class WebSocketControllerRegister implements ControllerRegister { return Promise.resolve(); } return new Promise(resolve => { - webSocket.once('close', () => resolve()); - webSocket.once('error', () => resolve()); + let settled = false; + let onClose!: () => void; + let onError!: (error: Error) => void; + const settle = () => { + if (settled) { + return; + } + settled = true; + webSocket.removeListener('close', onClose); + webSocket.removeListener('error', onError); + resolve(); + }; + onClose = () => settle(); + onError = error => { + console.error('[tegg/websocket] socket error while waiting for close:', error); + settle(); + }; + webSocket.once('close', onClose); + webSocket.once('error', onError); }); } } diff --git a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts index 1e314d36a..1042b2612 100644 --- a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts +++ b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts @@ -3,10 +3,10 @@ import { PassThrough, pipeline, type Readable } from 'node:stream'; import type { RawData, WebSocket } from 'ws'; import { Context, - Headers, - Param, - Query, - Queries, + HTTPHeaders, + HTTPParam, + HTTPQueries, + HTTPQuery, Request, WebSocketClose, WebSocketCloseCode, @@ -49,10 +49,10 @@ export class StandaloneWebSocketController { }) async echo( @WebSocketSocket() socket: WebSocket, - @Param() id: string, - @Query() name: string, - @Queries({ name: 'tag' }) tags: string[], - @Headers() headers: IncomingHttpHeaders, + @HTTPParam() id: string, + @HTTPQuery() name: string, + @HTTPQueries({ name: 'tag' }) tags: string[], + @HTTPHeaders() headers: IncomingHttpHeaders, @Request() request: IncomingMessage, @Context() ctx: WebSocketContext, ) { @@ -79,8 +79,8 @@ export class StandaloneWebSocketController { }) stream( @WebSocketStream() input: Readable, - @Param() id: string, - @Query() name: string, + @HTTPParam() id: string, + @HTTPQuery() name: string, ) { const output = new PassThrough(); output.write(JSON.stringify({ @@ -108,9 +108,13 @@ export class StandaloneWebSocketController { export class StandaloneWebSocketFetchController { @WebSocketFetchOnConnection() onConnection( - @Headers() headers: IncomingHttpHeaders, + @HTTPParam() id: string, + @HTTPHeaders() headers: IncomingHttpHeaders, @Context() ctx: WebSocketContext, ) { + if (id === 'connection-error') { + throw new Error('fetch connection error'); + } ctx.socket.send(JSON.stringify({ type: 'connection', header: headers['x-client-id'], @@ -132,10 +136,10 @@ export class StandaloneWebSocketFetchController { onData( @WebSocketData() data: RawData, @WebSocketClose() close: WebSocketFetchClose, - @Param() id: string, - @Query() name: string, - @Queries({ name: 'tag' }) tags: string[], - @Headers() headers: IncomingHttpHeaders, + @HTTPParam() id: string, + @HTTPQuery() name: string, + @HTTPQueries({ name: 'tag' }) tags: string[], + @HTTPHeaders() headers: IncomingHttpHeaders, @Context() ctx: WebSocketContext, ) { const body = JSON.parse(data.toString()) as WebSocketFetchRequest; @@ -198,7 +202,7 @@ export class StandaloneWebSocketFetchController { @WebSocketFetchOnError() onError( @WebSocketError() error: Error, - @Headers() headers: IncomingHttpHeaders, + @HTTPHeaders() headers: IncomingHttpHeaders, @Context() ctx: WebSocketContext, ) { ctx.socket.send(JSON.stringify({ @@ -213,7 +217,7 @@ export class StandaloneWebSocketFetchController { onClose( @WebSocketCloseCode() code: number, @WebSocketCloseReason() reason: Buffer, - @Param() id: string, + @HTTPParam() id: string, ) { fetchCloseEvents.push(`${id}:${code}:${reason.toString()}`); } diff --git a/standalone/service-worker/test/websocket/websocket.test.ts b/standalone/service-worker/test/websocket/websocket.test.ts index cbe7a211e..131f02ae0 100644 --- a/standalone/service-worker/test/websocket/websocket.test.ts +++ b/standalone/service-worker/test/websocket/websocket.test.ts @@ -90,6 +90,23 @@ describe('standalone/service-worker/test/websocket/websocket.test.ts', () => { await closeClient(ws); }); + it('should handle websocket fetch connection lifecycle errors before closing', async () => { + const socket = createClient('/ws-fetch/connection-error', { + 'x-client-id': 'connection-error-client', + }); + const messages = createJSONMessageQueue(socket); + const closePromise = waitClose(socket); + + assert.deepStrictEqual(await messages.next(), { + type: 'error', + message: 'fetch connection error', + header: 'connection-error-client', + path: '/ws-fetch/connection-error', + }); + await closePromise; + await waitFor(() => fetchCloseEvents.includes('connection-error:1011:Internal Server Error')); + }); + it('should isolate two websocket fetch connections with interleaved responses in standalone service worker', async () => { const first = createClient('/ws-fetch/first?name=one&tag=a&tag=b', { 'x-client-id': 'first-client', From 3ae5b85f76876787006d47614727a8093e2f3953 Mon Sep 17 00:00:00 2001 From: akitasummer Date: Fri, 10 Jul 2026 22:48:13 +0800 Subject: [PATCH 7/9] test(controller): fix websocket ci compatibility --- plugin/controller/test/lib/EggControllerLoader.test.ts | 2 +- plugin/controller/test/websocket/websocketCluster.test.ts | 3 ++- standalone/service-worker/test/Utils.ts | 3 ++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/plugin/controller/test/lib/EggControllerLoader.test.ts b/plugin/controller/test/lib/EggControllerLoader.test.ts index b93bf16bc..2ebd556cc 100644 --- a/plugin/controller/test/lib/EggControllerLoader.test.ts +++ b/plugin/controller/test/lib/EggControllerLoader.test.ts @@ -17,7 +17,7 @@ describe('plugin/controller/test/lib/EggModuleLoader.test.ts', () => { const controllerDir = path.join(__dirname, '../fixtures/apps/controller-app/app/controller'); const loader = new EggControllerLoader(controllerDir); const clazzs = loader.load(); - assert.strictEqual(clazzs.length, 8); + assert.strictEqual(clazzs.length, 10); const AppController = clazzs[0]; const metadata = ControllerMetadataUtil.getControllerMetadata(AppController); assert(metadata); diff --git a/plugin/controller/test/websocket/websocketCluster.test.ts b/plugin/controller/test/websocket/websocketCluster.test.ts index 139d753f6..3deac95de 100644 --- a/plugin/controller/test/websocket/websocketCluster.test.ts +++ b/plugin/controller/test/websocket/websocketCluster.test.ts @@ -110,7 +110,8 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { mm(process, 'cwd', () => { return pluginRoot; }); - app = mm.cluster({ + const cluster = mm.cluster as any; + app = cluster({ baseDir: path.join(pluginRoot, 'test/fixtures/apps/controller-app'), framework: eggFramework, workers: 2, diff --git a/standalone/service-worker/test/Utils.ts b/standalone/service-worker/test/Utils.ts index ac82bcddb..86cb3b1e6 100644 --- a/standalone/service-worker/test/Utils.ts +++ b/standalone/service-worker/test/Utils.ts @@ -1,4 +1,5 @@ import path from 'node:path'; +import type { Response as UndiciResponse } from 'undici'; import { ServiceWorkerApp, ServiceWorkerAppOptions } from '../src/ServiceWorkerApp'; import { StandaloneTestUtil } from '@eggjs/module-test-util/StandaloneTestUtil'; @@ -22,7 +23,7 @@ export class TestUtils { const app = await TestUtils.createApp(name, init); // Use port 0 to let the OS assign a free port const server = await StandaloneTestUtil.startHTTPServer('127.0.0.1', 0, { - listener: e => app.handleEvent(e), + listener: e => app.handleEvent(e), }); return { app, server }; From 2782dd7c6f50c2fbda05d74e577a65dbe4659dba Mon Sep 17 00:00:00 2001 From: akitasummer Date: Fri, 10 Jul 2026 23:07:51 +0800 Subject: [PATCH 8/9] test(controller): guard mcp imports on node 16 --- plugin/controller/test/mcp/mcp.test.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/plugin/controller/test/mcp/mcp.test.ts b/plugin/controller/test/mcp/mcp.test.ts index a01f164ad..019317eb6 100644 --- a/plugin/controller/test/mcp/mcp.test.ts +++ b/plugin/controller/test/mcp/mcp.test.ts @@ -6,7 +6,6 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js'; import { CallToolResultSchema, ListToolsResultSchema, LoggingMessageNotificationSchema } from '@modelcontextprotocol/sdk/types.js'; import type { CallToolRequest, ListToolsRequest, JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js'; import assert from 'assert'; -import { MCPControllerRegister } from '../../lib/impl/mcp/MCPControllerRegister'; async function listTools(client: Client) { const toolsRequest: ListToolsRequest = { @@ -71,6 +70,8 @@ describe('plugin/controller/test/mcp/mcp.test.ts', () => { if (parseInt(process.version.slice(1, 3)) > 17) { // eslint-disable-next-line @typescript-eslint/no-var-requires const { StreamableHTTPClientTransport } = require('@modelcontextprotocol/sdk/client/streamableHttp.js'); + // eslint-disable-next-line @typescript-eslint/no-var-requires + const { MCPControllerRegister } = require('../../lib/impl/mcp/MCPControllerRegister'); let app; after(async () => { From 348592e0386711f1a09bf63b63c7986bdad04f8b Mon Sep 17 00:00:00 2001 From: akitasummer Date: Mon, 13 Jul 2026 00:35:48 +0800 Subject: [PATCH 9/9] fix(controller): address websocket review feedback --- .../src/decorator/http/HTTPParam.ts | 11 +- .../WebSocketFetchController.ts | 3 + .../websocket-fetch/WebSocketFetchMethod.ts | 23 +- .../websocket/WebSocketController.ts | 3 + .../decorator/websocket/WebSocketMethod.ts | 3 + .../WebSocketFetchControllerMetaBuilder.ts | 5 +- ...bSocketFetchControllerMethodMetaBuilder.ts | 13 +- .../WebSocketControllerMetaBuilder.ts | 3 +- .../WebSocketControllerMethodMetaBuilder.ts | 13 +- .../src/model/WebSocketControllerMeta.ts | 7 + .../src/model/WebSocketFetchControllerMeta.ts | 7 + .../src/model/WebSocketFetchMethodMeta.ts | 3 + .../src/model/WebSocketMethodMeta.ts | 3 + .../src/util/WebSocketInfoUtil.ts | 40 +- .../test/websocket/WebSocketMeta.test.ts | 117 ++++ .../controller-decorator/WebSocketContext.ts | 14 +- .../WebSocketController.ts | 1 + .../WebSocketFetchController.ts | 1 + .../WebSocketFetchMethod.ts | 5 + .../controller-decorator/WebSocketMethod.ts | 1 + core/websocket-runtime/index.ts | 3 + core/websocket-runtime/package.json | 54 ++ .../src/WebSocketControllerRuntime.ts | 229 +++++++ .../src/WebSocketFetchSession.ts | 347 ++++++++++ .../src/WebSocketRouteRegistry.ts | 140 ++++ .../test/WebSocketFetchSession.test.ts | 217 ++++++ .../test/WebSocketRouteRegistry.test.ts | 64 ++ core/websocket-runtime/tsconfig.json | 12 + core/websocket-runtime/tsconfig.pub.json | 3 + plugin/controller/app.ts | 1 - .../lib/impl/websocket/WebSocketContext.ts | 44 +- .../websocket/WebSocketControllerRegister.ts | 624 ++++-------------- plugin/controller/package.json | 1 + .../app/controller/WebSocketController.ts | 60 +- .../controller-app/config/config.default.js | 2 + .../test/websocket/websocket.test.ts | 72 ++ .../test/websocket/websocketCluster.test.ts | 17 +- standalone/service-worker/package.json | 1 + .../ServiceWorkerWebSocketContext.ts | 5 +- .../websocket/WebSocketControllerRegister.ts | 553 +++------------- .../fixtures/websocket/WebSocketController.ts | 58 +- .../test/websocket/websocket.test.ts | 49 ++ 42 files changed, 1773 insertions(+), 1059 deletions(-) create mode 100644 core/websocket-runtime/index.ts create mode 100644 core/websocket-runtime/package.json create mode 100644 core/websocket-runtime/src/WebSocketControllerRuntime.ts create mode 100644 core/websocket-runtime/src/WebSocketFetchSession.ts create mode 100644 core/websocket-runtime/src/WebSocketRouteRegistry.ts create mode 100644 core/websocket-runtime/test/WebSocketFetchSession.test.ts create mode 100644 core/websocket-runtime/test/WebSocketRouteRegistry.test.ts create mode 100644 core/websocket-runtime/tsconfig.json create mode 100644 core/websocket-runtime/tsconfig.pub.json diff --git a/core/controller-decorator/src/decorator/http/HTTPParam.ts b/core/controller-decorator/src/decorator/http/HTTPParam.ts index 38765253e..4161decf5 100644 --- a/core/controller-decorator/src/decorator/http/HTTPParam.ts +++ b/core/controller-decorator/src/decorator/http/HTTPParam.ts @@ -1,9 +1,8 @@ import assert from 'node:assert'; -import { HTTPParamType, WebSocketParamType } from '@eggjs/tegg-types'; +import { HTTPParamType } from '@eggjs/tegg-types'; import type { EggProtoImplClass, HTTPParamParams, HTTPQueriesParams, HTTPQueryParams } from '@eggjs/tegg-types'; import HTTPInfoUtil from '../../util/HTTPInfoUtil'; import { ObjectUtils } from '@eggjs/tegg-common-util'; -import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; // TODO url params // /foo/:id @@ -26,7 +25,6 @@ export function HTTPHeaders() { const methodName = propertyKey as string; const controllerClazz = target.constructor as EggProtoImplClass; HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.HEADERS, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.HEADERS, parameterIndex, controllerClazz, methodName); }; } @@ -40,8 +38,6 @@ export function HTTPQuery(param?: HTTPQueryParams) { const name = param?.name || argNames[parameterIndex]; HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.QUERY, parameterIndex, controllerClazz, methodName); HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.QUERY, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); }; } @@ -55,8 +51,6 @@ export function HTTPQueries(param?: HTTPQueriesParams) { const name = param?.name || argNames[parameterIndex]; HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.QUERIES, parameterIndex, controllerClazz, methodName); HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.QUERIES, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); }; } @@ -70,8 +64,6 @@ export function HTTPParam(param?: HTTPParamParams) { const name = param?.name || argNames[parameterIndex]; HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.PARAM, parameterIndex, controllerClazz, methodName); HTTPInfoUtil.setHTTPMethodParamName(name, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.PARAM, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamName(name, parameterIndex, controllerClazz, methodName); }; } @@ -85,7 +77,6 @@ export function Request() { const methodName = propertyKey as string; const controllerClazz = target.constructor as EggProtoImplClass; HTTPInfoUtil.setHTTPMethodParamType(HTTPParamType.REQUEST, parameterIndex, controllerClazz, methodName); - WebSocketInfoUtil.setWebSocketMethodParamType(WebSocketParamType.REQUEST, parameterIndex, controllerClazz, methodName); }; } diff --git a/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts index b0a8a6214..289666cdd 100644 --- a/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts +++ b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts @@ -14,6 +14,9 @@ export function WebSocketFetchController(param?: WebSocketFetchControllerParams) if (param?.path) { WebSocketInfoUtil.setWebSocketPath(param.path, constructor); } + if (param?.timeout !== undefined) { + ControllerInfoUtil.setControllerTimeout(param.timeout, constructor); + } const func = SingletonProto({ accessLevel: AccessLevel.PUBLIC, diff --git a/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts index 809c6a13a..910a7b7dd 100644 --- a/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts +++ b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts @@ -1,6 +1,6 @@ import assert from 'node:assert'; import { ControllerType, WebSocketFetchMethodType } from '@eggjs/tegg-types'; -import type { EggProtoImplClass, WebSocketFetchMethodParams } from '@eggjs/tegg-types'; +import type { EggProtoImplClass, WebSocketFetchLifecycleMethodParams, WebSocketFetchMethodParams } from '@eggjs/tegg-types'; import MethodInfoUtil from '../../util/MethodInfoUtil'; import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; @@ -8,6 +8,7 @@ function setWebSocketFetchMethodType( target: any, propertyKey: PropertyKey, methodType: WebSocketFetchMethodType, + timeout?: number, ) { assert(typeof propertyKey === 'string', `[controller/${target.name}] expect method name be typeof string, but now is ${String(propertyKey)}`); @@ -15,6 +16,9 @@ function setWebSocketFetchMethodType( const methodName = propertyKey as string; MethodInfoUtil.setMethodControllerType(controllerClazz, methodName, ControllerType.WEBSOCKET_FETCH); WebSocketInfoUtil.setWebSocketFetchMethodType(methodType, controllerClazz, methodName); + if (timeout !== undefined) { + MethodInfoUtil.setMethodTimeout(timeout, controllerClazz, methodName); + } return { controllerClazz, methodName }; } @@ -24,6 +28,7 @@ export function WebSocketFetchMethod(param?: WebSocketFetchMethodParams) { target, propertyKey, WebSocketFetchMethodType.DATA, + param?.timeout, ); WebSocketInfoUtil.setWebSocketMethodPath('', controllerClazz, methodName); if (param?.priority !== undefined) { @@ -32,26 +37,26 @@ export function WebSocketFetchMethod(param?: WebSocketFetchMethodParams) { }; } -export function WebSocketFetchOnConnection() { +export function WebSocketFetchOnConnection(param?: WebSocketFetchLifecycleMethodParams) { return function(target: any, propertyKey: PropertyKey) { - setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CONNECTION); + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CONNECTION, param?.timeout); }; } -export function WebSocketFetchOnOpen() { +export function WebSocketFetchOnOpen(param?: WebSocketFetchLifecycleMethodParams) { return function(target: any, propertyKey: PropertyKey) { - setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.OPEN); + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.OPEN, param?.timeout); }; } -export function WebSocketFetchOnError() { +export function WebSocketFetchOnError(param?: WebSocketFetchLifecycleMethodParams) { return function(target: any, propertyKey: PropertyKey) { - setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.ERROR); + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.ERROR, param?.timeout); }; } -export function WebSocketFetchOnClose() { +export function WebSocketFetchOnClose(param?: WebSocketFetchLifecycleMethodParams) { return function(target: any, propertyKey: PropertyKey) { - setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CLOSE); + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CLOSE, param?.timeout); }; } diff --git a/core/controller-decorator/src/decorator/websocket/WebSocketController.ts b/core/controller-decorator/src/decorator/websocket/WebSocketController.ts index 1672a66e7..b9772c424 100644 --- a/core/controller-decorator/src/decorator/websocket/WebSocketController.ts +++ b/core/controller-decorator/src/decorator/websocket/WebSocketController.ts @@ -14,6 +14,9 @@ export function WebSocketController(param?: WebSocketControllerParams) { if (param?.path) { WebSocketInfoUtil.setWebSocketPath(param.path, constructor); } + if (param?.timeout !== undefined) { + ControllerInfoUtil.setControllerTimeout(param.timeout, constructor); + } const func = SingletonProto({ accessLevel: AccessLevel.PUBLIC, diff --git a/core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts b/core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts index b2fbc6d99..ad7a71716 100644 --- a/core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts +++ b/core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts @@ -15,5 +15,8 @@ export function WebSocketMethod(param: WebSocketMethodParams) { if (param.priority !== undefined) { WebSocketInfoUtil.setWebSocketMethodPriority(param.priority, controllerClazz, methodName); } + if (param.timeout !== undefined) { + MethodInfoUtil.setMethodTimeout(param.timeout, controllerClazz, methodName); + } }; } diff --git a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts index 0b48a1780..ae995e05e 100644 --- a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts @@ -81,15 +81,16 @@ export class WebSocketFetchControllerMetaBuilder { assert(webSocketPath, `build websocket fetch controller ${ClassUtil.classDescription(this.clazz)} failed: path is required`); const middlewares = ControllerInfoUtil.getControllerMiddlewares(this.clazz); const { methods, connectionMethod, openMethod, errorMethod, closeMethod } = this.buildMethods(); - assert(methods.length === 1, 'websocket fetch controller must has exactly one @WebSocketFetchMethod'); + assert(methods.length === 1, 'websocket fetch controller must have exactly one @WebSocketFetchMethod'); const clazzName = this.clazz.name; const controllerName = ControllerInfoUtil.getControllerName(this.clazz) || clazzName; const property = PrototypeUtil.getProperty(this.clazz); const protoName = property!.name as string; const hosts = ControllerInfoUtil.getControllerHosts(this.clazz); + const timeout = ControllerInfoUtil.getControllerTimeout(this.clazz); const metadata = new WebSocketFetchControllerMeta( clazzName, protoName, controllerName, webSocketPath, middlewares, methods, hosts, - connectionMethod, openMethod, errorMethod, closeMethod); + connectionMethod, openMethod, errorMethod, closeMethod, timeout); ControllerMetadataUtil.setControllerMetadata(this.clazz, metadata); for (const method of metadata.methods) { const realPath = metadata.getMethodRealPath(method); diff --git a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts index a3f7bdf9f..5013cda0e 100644 --- a/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts @@ -20,7 +20,7 @@ export class WebSocketFetchControllerMethodMetaBuilder { private checkParamDecorators(methodType: WebSocketFetchMethodType) { const method = this.clazz.prototype[this.methodName]; const functionLength = method.length; - const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); + const paramIndexList = WebSocketInfoUtil.getCompatibleParamIndexList(this.clazz, this.methodName); const contextIndex = MethodInfoUtil.getMethodContextIndex(this.clazz, this.methodName); const hasAnnotationParamCount = typeof contextIndex === 'undefined' ? paramIndexList.length @@ -31,7 +31,7 @@ export class WebSocketFetchControllerMethodMetaBuilder { if (i === contextIndex) { continue; } - const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(i, this.clazz, this.methodName); + const paramType = WebSocketInfoUtil.getCompatibleMethodParamType(i, this.clazz, this.methodName); if (!paramType) { const classDesc = ClassUtil.classDescription(this.clazz); throw new Error(`${classDesc}:${this.methodName} param ${i} has no websocket fetch param type, Please add @WebSocketData, @WebSocketClose, @WebSocketError, @WebSocketCloseCode, @WebSocketCloseReason, @HTTPParam, @HTTPQuery, @HTTPQueries, @HTTPHeaders, @Request, @WebSocketSocket or @Context`); @@ -86,10 +86,10 @@ export class WebSocketFetchControllerMethodMetaBuilder { this.checkParamDecorators(methodType); const paramTypeMap = new Map(); - const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); + const paramIndexList = WebSocketInfoUtil.getCompatibleParamIndexList(this.clazz, this.methodName); for (const paramIndex of paramIndexList) { - const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(paramIndex, this.clazz, this.methodName)!; - const paramName = WebSocketInfoUtil.getWebSocketMethodParamName(paramIndex, this.clazz, this.methodName); + const paramType = WebSocketInfoUtil.getCompatibleMethodParamType(paramIndex, this.clazz, this.methodName)!; + const paramName = WebSocketInfoUtil.getCompatibleMethodParamName(paramIndex, this.clazz, this.methodName); const paramMeta = WebSocketParamMetaUtil.createParam(paramType, paramName); if (methodType === WebSocketFetchMethodType.DATA) { @@ -144,7 +144,8 @@ export class WebSocketFetchControllerMethodMetaBuilder { : webSocketPath; const paramTypeMap = this.buildParamType(realPath, methodType); const priority = methodType === WebSocketFetchMethodType.DATA ? this.getPriority() : HTTPPriorityUtil.DEFAULT_PRIORITY; + const timeout = MethodInfoUtil.getMethodTimeout(this.clazz, this.methodName); return new WebSocketFetchMethodMeta( - this.methodName, webSocketPath, middlewares, contextIndex, paramTypeMap, priority, hosts, methodType); + this.methodName, webSocketPath, middlewares, contextIndex, paramTypeMap, priority, hosts, methodType, timeout); } } diff --git a/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts b/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts index eabfd3b3e..1560eac45 100644 --- a/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts @@ -44,8 +44,9 @@ export class WebSocketControllerMetaBuilder { const property = PrototypeUtil.getProperty(this.clazz); const protoName = property!.name as string; const hosts = ControllerInfoUtil.getControllerHosts(this.clazz); + const timeout = ControllerInfoUtil.getControllerTimeout(this.clazz); const metadata = new WebSocketControllerMeta( - clazzName, protoName, controllerName, webSocketPath, middlewares, methods, hosts); + clazzName, protoName, controllerName, webSocketPath, middlewares, methods, hosts, timeout); ControllerMetadataUtil.setControllerMetadata(this.clazz, metadata); for (const method of metadata.methods) { const realPath = metadata.getMethodRealPath(method); diff --git a/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts index 1d19907e3..ac88f9de9 100644 --- a/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts @@ -20,7 +20,7 @@ export class WebSocketControllerMethodMetaBuilder { private checkParamDecorators() { const method = this.clazz.prototype[this.methodName]; const functionLength = method.length; - const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); + const paramIndexList = WebSocketInfoUtil.getCompatibleParamIndexList(this.clazz, this.methodName); const contextIndex = MethodInfoUtil.getMethodContextIndex(this.clazz, this.methodName); const hasAnnotationParamCount = typeof contextIndex === 'undefined' ? paramIndexList.length @@ -31,7 +31,7 @@ export class WebSocketControllerMethodMetaBuilder { if (i === contextIndex) { continue; } - const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(i, this.clazz, this.methodName); + const paramType = WebSocketInfoUtil.getCompatibleMethodParamType(i, this.clazz, this.methodName); if (!paramType) { const classDesc = ClassUtil.classDescription(this.clazz); throw new Error(`${classDesc}:${this.methodName} param ${i} has no websocket param type, Please add @HTTPParam, @HTTPQuery, @HTTPQueries, @HTTPHeaders, @Request, @WebSocketSocket, @WebSocketStream or @Context`); @@ -43,14 +43,14 @@ export class WebSocketControllerMethodMetaBuilder { this.checkParamDecorators(); const paramTypeMap = new Map(); - const paramIndexList = WebSocketInfoUtil.getParamIndexList(this.clazz, this.methodName); + const paramIndexList = WebSocketInfoUtil.getCompatibleParamIndexList(this.clazz, this.methodName); for (const paramIndex of paramIndexList) { - const paramType = WebSocketInfoUtil.getWebSocketMethodParamType(paramIndex, this.clazz, this.methodName)!; + const paramType = WebSocketInfoUtil.getCompatibleMethodParamType(paramIndex, this.clazz, this.methodName)!; if (this.isFetchOnlyParamType(paramType)) { const classDesc = ClassUtil.classDescription(this.clazz); throw new Error(`${classDesc}:${this.methodName} param ${paramIndex} is websocket fetch only`); } - const paramName = WebSocketInfoUtil.getWebSocketMethodParamName(paramIndex, this.clazz, this.methodName); + const paramName = WebSocketInfoUtil.getCompatibleMethodParamName(paramIndex, this.clazz, this.methodName); const paramMeta = WebSocketParamMetaUtil.createParam(paramType, paramName); try { @@ -105,7 +105,8 @@ export class WebSocketControllerMethodMetaBuilder { : webSocketPath; const paramTypeMap = this.buildParamType(realPath); const priority = this.getPriority(); + const timeout = MethodInfoUtil.getMethodTimeout(this.clazz, this.methodName); return new WebSocketMethodMeta( - this.methodName, webSocketPath, middlewares, contextIndex, paramTypeMap, priority, hosts); + this.methodName, webSocketPath, middlewares, contextIndex, paramTypeMap, priority, hosts, timeout); } } diff --git a/core/controller-decorator/src/model/WebSocketControllerMeta.ts b/core/controller-decorator/src/model/WebSocketControllerMeta.ts index c54535b6a..009ca6ab2 100644 --- a/core/controller-decorator/src/model/WebSocketControllerMeta.ts +++ b/core/controller-decorator/src/model/WebSocketControllerMeta.ts @@ -12,6 +12,7 @@ export class WebSocketControllerMeta implements ControllerMetadata { public readonly middlewares: readonly MiddlewareFunc[]; public readonly methods: readonly WebSocketMethodMeta[]; public readonly hosts?: string[]; + public readonly timeout?: number; constructor( className: string, @@ -21,6 +22,7 @@ export class WebSocketControllerMeta implements ControllerMetadata { middlewares: MiddlewareFunc[], methods: WebSocketMethodMeta[], hosts: string[] | undefined, + timeout?: number, ) { this.protoName = protoName; this.controllerName = controllerName; @@ -29,6 +31,7 @@ export class WebSocketControllerMeta implements ControllerMetadata { this.middlewares = middlewares; this.methods = methods; this.hosts = hosts; + this.timeout = timeout; } getMethodRealPath(method: WebSocketMethodMeta) { @@ -58,4 +61,8 @@ export class WebSocketControllerMeta implements ControllerMetadata { } return [ ...method.middlewares ]; } + + getMethodTimeout(method: WebSocketMethodMeta) { + return method.timeout ?? this.timeout; + } } diff --git a/core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts b/core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts index 823f90fe3..b5cc049d6 100644 --- a/core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts +++ b/core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts @@ -12,6 +12,7 @@ export class WebSocketFetchControllerMeta implements ControllerMetadata { public readonly middlewares: readonly MiddlewareFunc[]; public readonly methods: readonly WebSocketFetchMethodMeta[]; public readonly hosts?: string[]; + public readonly timeout?: number; public readonly connectionMethod?: WebSocketFetchMethodMeta; public readonly openMethod?: WebSocketFetchMethodMeta; public readonly errorMethod?: WebSocketFetchMethodMeta; @@ -29,6 +30,7 @@ export class WebSocketFetchControllerMeta implements ControllerMetadata { openMethod: WebSocketFetchMethodMeta | undefined, errorMethod: WebSocketFetchMethodMeta | undefined, closeMethod: WebSocketFetchMethodMeta | undefined, + timeout?: number, ) { this.protoName = protoName; this.controllerName = controllerName; @@ -41,6 +43,7 @@ export class WebSocketFetchControllerMeta implements ControllerMetadata { this.openMethod = openMethod; this.errorMethod = errorMethod; this.closeMethod = closeMethod; + this.timeout = timeout; } getMethodRealPath(method: WebSocketFetchMethodMeta) { @@ -70,4 +73,8 @@ export class WebSocketFetchControllerMeta implements ControllerMetadata { } return [ ...method.middlewares ]; } + + getMethodTimeout(method: WebSocketFetchMethodMeta) { + return method.timeout ?? this.timeout; + } } diff --git a/core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts b/core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts index ec0546a0b..44d72c91f 100644 --- a/core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts +++ b/core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts @@ -11,6 +11,7 @@ export class WebSocketFetchMethodMeta implements MethodMeta { public readonly priority: number; public readonly hosts: string[] | undefined; public readonly methodType: WebSocketFetchMethodType; + public readonly timeout: number | undefined; constructor( name: string, @@ -21,6 +22,7 @@ export class WebSocketFetchMethodMeta implements MethodMeta { priority: number, hosts: string[] | undefined, methodType: WebSocketFetchMethodType, + timeout?: number, ) { this.name = name; this.path = path; @@ -30,5 +32,6 @@ export class WebSocketFetchMethodMeta implements MethodMeta { this.priority = priority; this.hosts = hosts; this.methodType = methodType; + this.timeout = timeout; } } diff --git a/core/controller-decorator/src/model/WebSocketMethodMeta.ts b/core/controller-decorator/src/model/WebSocketMethodMeta.ts index 527783a9b..805bbb6a7 100644 --- a/core/controller-decorator/src/model/WebSocketMethodMeta.ts +++ b/core/controller-decorator/src/model/WebSocketMethodMeta.ts @@ -135,6 +135,7 @@ export class WebSocketMethodMeta implements MethodMeta { public readonly paramMap: Map; public readonly priority: number; public readonly hosts: string[] | undefined; + public readonly timeout: number | undefined; constructor( name: string, @@ -144,6 +145,7 @@ export class WebSocketMethodMeta implements MethodMeta { paramMap: Map, priority: number, hosts: string[] | undefined, + timeout?: number, ) { this.name = name; this.path = path; @@ -152,6 +154,7 @@ export class WebSocketMethodMeta implements MethodMeta { this.paramMap = paramMap; this.priority = priority; this.hosts = hosts; + this.timeout = timeout; } } diff --git a/core/controller-decorator/src/util/WebSocketInfoUtil.ts b/core/controller-decorator/src/util/WebSocketInfoUtil.ts index 0a9fed55b..482874ec6 100644 --- a/core/controller-decorator/src/util/WebSocketInfoUtil.ts +++ b/core/controller-decorator/src/util/WebSocketInfoUtil.ts @@ -6,9 +6,13 @@ import { CONTROLLER_WEBSOCKET_METHOD_PRIORITY, CONTROLLER_WEBSOCKET_FETCH_METHOD_TYPE_MAP, CONTROLLER_WEBSOCKET_PATH, + HTTPParamType, + WebSocketParamType, + type EggProtoImplClass, + type WebSocketFetchMethodType, } from '@eggjs/tegg-types'; -import type { EggProtoImplClass, WebSocketFetchMethodType, WebSocketParamType } from '@eggjs/tegg-types'; import { MapUtil } from '@eggjs/tegg-common-util'; +import HTTPInfoUtil from './HTTPInfoUtil'; type WebSocketMethodPathMap = Map; type WebSocketMethodParamTypeMap = Map>; @@ -50,12 +54,24 @@ export default class WebSocketInfoUtil { return Array.from(paramMap.keys()); } + static getCompatibleParamIndexList(clazz: EggProtoImplClass, methodName: string): number[] { + return Array.from(new Set([ + ...this.getParamIndexList(clazz, methodName), + ...HTTPInfoUtil.getParamIndexList(clazz, methodName), + ])); + } + static getWebSocketMethodParamType(parameterIndex: number, clazz: EggProtoImplClass, methodName: string): WebSocketParamType | undefined { const methodParamMap: WebSocketMethodParamTypeMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP, clazz); const paramMap = methodParamMap?.get(methodName); return paramMap?.get(parameterIndex); } + static getCompatibleMethodParamType(parameterIndex: number, clazz: EggProtoImplClass, methodName: string): WebSocketParamType | undefined { + return this.getWebSocketMethodParamType(parameterIndex, clazz, methodName) ?? + this.fromHTTPParamType(HTTPInfoUtil.getHTTPMethodParamType(parameterIndex, clazz, methodName)); + } + static setWebSocketMethodParamName(paramName: string, parameterIndex: number, clazz: EggProtoImplClass, methodName: string) { const methodParamNameMap: WebSocketMethodParamNameMap = MetadataUtil.initOwnMapMetaData(CONTROLLER_WEBSOCKET_METHOD_PARAM_NAME_MAP, clazz, new Map()); const paramMap = MapUtil.getOrStore(methodParamNameMap, methodName, new Map()); @@ -68,6 +84,11 @@ export default class WebSocketInfoUtil { return paramMap?.get(parameterIndex); } + static getCompatibleMethodParamName(parameterIndex: number, clazz: EggProtoImplClass, methodName: string): string | undefined { + return this.getWebSocketMethodParamName(parameterIndex, clazz, methodName) ?? + HTTPInfoUtil.getHTTPMethodParamName(parameterIndex, clazz, methodName); + } + static getWebSocketMethodPriority(clazz: EggProtoImplClass, methodName: string): number | undefined { const methodPriorityMap: WebSocketMethodPriorityMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_METHOD_PRIORITY, clazz); return methodPriorityMap?.get(methodName); @@ -87,4 +108,21 @@ export default class WebSocketInfoUtil { const methodTypeMap: WebSocketFetchMethodTypeMap | undefined = MetadataUtil.getMetaData(CONTROLLER_WEBSOCKET_FETCH_METHOD_TYPE_MAP, clazz); return methodTypeMap?.get(methodName); } + + private static fromHTTPParamType(paramType: HTTPParamType | undefined): WebSocketParamType | undefined { + switch (paramType) { + case HTTPParamType.PARAM: + return WebSocketParamType.PARAM; + case HTTPParamType.QUERY: + return WebSocketParamType.QUERY; + case HTTPParamType.QUERIES: + return WebSocketParamType.QUERIES; + case HTTPParamType.HEADERS: + return WebSocketParamType.HEADERS; + case HTTPParamType.REQUEST: + return WebSocketParamType.REQUEST; + default: + return undefined; + } + } } diff --git a/core/controller-decorator/test/websocket/WebSocketMeta.test.ts b/core/controller-decorator/test/websocket/WebSocketMeta.test.ts index f2c1792bb..9fac64a26 100644 --- a/core/controller-decorator/test/websocket/WebSocketMeta.test.ts +++ b/core/controller-decorator/test/websocket/WebSocketMeta.test.ts @@ -1,10 +1,82 @@ import assert from 'node:assert'; +import { MetadataUtil } from '@eggjs/core-decorator'; +import { + CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP, + ControllerType, + HTTPMethodEnum, +} from '@eggjs/tegg-types'; import { ControllerMetaBuilderFactory, + HTTPController, + HTTPMethod, + HTTPParam, + HTTPQuery, + WebSocketController, + WebSocketControllerMeta, WebSocketFetchController, + WebSocketFetchOnOpen, WebSocketFetchMethod, + WebSocketData, + WebSocketStream, + WebSocketMethod, + WebSocketPathParamMeta, + WebSocketQueryParamMeta, + WebSocketInfoUtil, } from '../..'; +@HTTPController({ path: '/http-only' }) +class HTTPOnlyController { + @HTTPMethod({ path: '/:id', method: HTTPMethodEnum.GET }) + get(@HTTPParam() id: string) { + return id; + } +} + +@WebSocketController({ path: '/websocket', timeout: 1000 }) +class CommonParamWebSocketController { + @WebSocketMethod({ path: '/:id', timeout: 50 }) + handle( + @HTTPParam() id: string, + @HTTPQuery({ name: 'name' }) name: string, + ) { + return { id, name }; + } +} + +@WebSocketFetchController({ path: '/duplicate-lifecycle' }) +class DuplicateLifecycleWebSocketFetchController { + @WebSocketFetchMethod() + onData() { + // test fixture + } + + @WebSocketFetchOnOpen() + firstOpen() { + // test fixture + } + + @WebSocketFetchOnOpen() + secondOpen() { + // test fixture + } +} + +@WebSocketFetchController({ path: '/invalid-stream' }) +class InvalidStreamWebSocketFetchController { + @WebSocketFetchMethod() + onData(@WebSocketStream() _stream: NodeJS.ReadableStream) { + void _stream; + } +} + +@WebSocketController({ path: '/invalid-data' }) +class InvalidDataWebSocketController { + @WebSocketMethod({ path: '/' }) + handle(@WebSocketData() _data: unknown) { + void _data; + } +} + @WebSocketFetchController() class MissingPathWebSocketFetchController { @WebSocketFetchMethod() @@ -20,4 +92,49 @@ describe('core/controller-decorator/test/websocket/WebSocketMeta.test.ts', () => /build websocket fetch controller .* failed: path is required/, ); }); + + it('should not add websocket metadata from HTTP parameter decorators', () => { + ControllerMetaBuilderFactory.build(HTTPOnlyController, ControllerType.HTTP); + assert.equal( + MetadataUtil.hasMetaData(CONTROLLER_WEBSOCKET_METHOD_PARAM_TYPE_MAP, HTTPOnlyController), + false, + ); + assert.deepEqual(WebSocketInfoUtil.getParamIndexList(HTTPOnlyController, 'get'), []); + }); + + it('should map common HTTP parameters while building websocket metadata', () => { + const metadata = ControllerMetaBuilderFactory.build( + CommonParamWebSocketController, + ControllerType.WEBSOCKET, + ) as WebSocketControllerMeta; + const method = metadata.methods[0]; + + assert.equal(metadata.timeout, 1000); + assert.equal(method.timeout, 50); + assert.deepEqual(method.paramMap, new Map([ + [ 0, new WebSocketPathParamMeta('id') ], + [ 1, new WebSocketQueryParamMeta('name') ], + ])); + }); + + it('should reject duplicate websocket fetch lifecycle methods', () => { + assert.throws( + () => ControllerMetaBuilderFactory.build(DuplicateLifecycleWebSocketFetchController), + /duplicate OPEN lifecycle method/, + ); + }); + + it('should reject websocket stream params in fetch methods', () => { + assert.throws( + () => ControllerMetaBuilderFactory.build(InvalidStreamWebSocketFetchController), + /type STREAM is not allowed in websocket fetch DATA method/, + ); + }); + + it('should reject websocket fetch params in ordinary websocket methods', () => { + assert.throws( + () => ControllerMetaBuilderFactory.build(InvalidDataWebSocketController), + /param 0 is websocket fetch only/, + ); + }); }); diff --git a/core/types/controller-decorator/WebSocketContext.ts b/core/types/controller-decorator/WebSocketContext.ts index eb82f5d7e..d9d426750 100644 --- a/core/types/controller-decorator/WebSocketContext.ts +++ b/core/types/controller-decorator/WebSocketContext.ts @@ -1,20 +1,14 @@ -import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; +import type { Context } from 'egg'; export interface WebSocketLike { on(event: string, listener: (...args: any[]) => void): this; once(event: string, listener: (...args: any[]) => void): this; + removeListener(event: string, listener: (...args: any[]) => void): this; send(data: any, cb?: (err?: Error) => void): void; close(code?: number, reason?: string | Buffer): void; readonly readyState?: number; } -export interface WebSocketContext { - readonly socket: Socket; - readonly request: IncomingMessage; - readonly path: string; - readonly host: string; - readonly params: Record; - readonly query: Record; - readonly queries: Record; - readonly headers: IncomingHttpHeaders; +export interface WebSocketContext extends Context { + readonly webSocket: Socket; } diff --git a/core/types/controller-decorator/WebSocketController.ts b/core/types/controller-decorator/WebSocketController.ts index bec77b2e4..a657cfbda 100644 --- a/core/types/controller-decorator/WebSocketController.ts +++ b/core/types/controller-decorator/WebSocketController.ts @@ -2,4 +2,5 @@ export interface WebSocketControllerParams { path?: string; controllerName?: string; protoName?: string; + timeout?: number; } diff --git a/core/types/controller-decorator/WebSocketFetchController.ts b/core/types/controller-decorator/WebSocketFetchController.ts index 5cb488dc9..51d42e6ca 100644 --- a/core/types/controller-decorator/WebSocketFetchController.ts +++ b/core/types/controller-decorator/WebSocketFetchController.ts @@ -2,4 +2,5 @@ export interface WebSocketFetchControllerParams { path?: string; controllerName?: string; protoName?: string; + timeout?: number; } diff --git a/core/types/controller-decorator/WebSocketFetchMethod.ts b/core/types/controller-decorator/WebSocketFetchMethod.ts index df43f5ee7..cdd9320a0 100644 --- a/core/types/controller-decorator/WebSocketFetchMethod.ts +++ b/core/types/controller-decorator/WebSocketFetchMethod.ts @@ -1,3 +1,8 @@ export interface WebSocketFetchMethodParams { priority?: number; + timeout?: number; +} + +export interface WebSocketFetchLifecycleMethodParams { + timeout?: number; } diff --git a/core/types/controller-decorator/WebSocketMethod.ts b/core/types/controller-decorator/WebSocketMethod.ts index 9019ce114..c31fb913e 100644 --- a/core/types/controller-decorator/WebSocketMethod.ts +++ b/core/types/controller-decorator/WebSocketMethod.ts @@ -1,4 +1,5 @@ export interface WebSocketMethodParams { path: string; priority?: number; + timeout?: number; } diff --git a/core/websocket-runtime/index.ts b/core/websocket-runtime/index.ts new file mode 100644 index 000000000..585df9dd5 --- /dev/null +++ b/core/websocket-runtime/index.ts @@ -0,0 +1,3 @@ +export * from './src/WebSocketControllerRuntime'; +export * from './src/WebSocketFetchSession'; +export * from './src/WebSocketRouteRegistry'; diff --git a/core/websocket-runtime/package.json b/core/websocket-runtime/package.json new file mode 100644 index 000000000..e017e8672 --- /dev/null +++ b/core/websocket-runtime/package.json @@ -0,0 +1,54 @@ +{ + "name": "@eggjs/tegg-websocket-runtime", + "description": "WebSocket controller runtime for tegg", + "version": "3.84.2", + "keywords": [ + "egg", + "typescript", + "tegg", + "websocket" + ], + "main": "dist/index.js", + "files": [ + "dist/**/*.js", + "dist/**/*.d.ts" + ], + "typings": "dist/index.d.ts", + "scripts": { + "test": "cross-env NODE_ENV=test NODE_OPTIONS='--no-deprecation' mocha", + "clean": "tsc -b --clean", + "tsc": "ut run clean && tsc -p ./tsconfig.json", + "tsc:pub": "ut run clean && tsc -p ./tsconfig.pub.json", + "prepublishOnly": "ut tsc:pub" + }, + "engines": { + "node": ">=14.0.0" + }, + "homepage": "https://github.com/eggjs/tegg", + "bugs": { + "url": "https://github.com/eggjs/tegg/issues" + }, + "author": "killagu ", + "license": "MIT", + "repository": { + "type": "git", + "url": "git@github.com:eggjs/tegg.git", + "directory": "core/websocket-runtime" + }, + "publishConfig": { + "access": "public" + }, + "dependencies": { + "@eggjs/controller-decorator": "^3.84.2", + "@eggjs/tegg-common-util": "^3.84.2", + "path-to-regexp": "^1.8.0" + }, + "devDependencies": { + "@types/mocha": "^10.0.1", + "@types/node": "^20.2.4", + "cross-env": "^7.0.3", + "mocha": "^10.2.0", + "ts-node": "^10.9.1", + "typescript": "^5.0.4" + } +} diff --git a/core/websocket-runtime/src/WebSocketControllerRuntime.ts b/core/websocket-runtime/src/WebSocketControllerRuntime.ts new file mode 100644 index 000000000..eac202178 --- /dev/null +++ b/core/websocket-runtime/src/WebSocketControllerRuntime.ts @@ -0,0 +1,229 @@ +import assert from 'node:assert'; +import { Duplex, pipeline } from 'node:stream'; +import { + WebSocketControllerMeta, + WebSocketFetchControllerMeta, + WebSocketFetchMethodMeta, + WebSocketMethodMeta, + WebSocketParamType, + WebSocketPathParamMeta, + WebSocketQueriesParamMeta, + WebSocketQueryParamMeta, +} from '@eggjs/controller-decorator'; +import type { WebSocketFetchClose } from '@eggjs/controller-decorator'; +import { TimerUtil } from '@eggjs/tegg-common-util'; +import { + WebSocketEventStream, + WebSocketFetchSession, + WebSocketSessionSocket, + WEBSOCKET_INTERNAL_ERROR_CODE, + WEBSOCKET_INTERNAL_ERROR_REASON, +} from './WebSocketFetchSession'; + +export interface WebSocketRuntimeContext { + context: unknown; + socket: Socket; + params: Record; + query: Record; + queries: Record; + headers: unknown; + request: unknown; +} + +export interface WebSocketRuntimeLogger { + debug(message: string): void; + error(message: string, error?: Error): void; +} + +export interface WebSocketControllerRuntimeOptions { + context: WebSocketRuntimeContext; + createWebSocketStream(socket: Socket): Duplex; + logger: WebSocketRuntimeLogger; + runInContext?(callback: () => Promise): Promise; +} + +interface WebSocketMethodPayload { + data?: Data; + close?: WebSocketFetchClose; + error?: Error; + closeCode?: number; + closeReason?: CloseReason; + getWebSocketStream?: () => Duplex; +} + +type WebSocketControllerObject = object; + +export class WebSocketControllerRuntime< + Socket extends WebSocketSessionSocket, + Data = unknown, + CloseReason = Buffer, +> { + constructor(private readonly options: WebSocketControllerRuntimeOptions) {} + + async invokeController( + realObj: WebSocketControllerObject, + controllerMeta: WebSocketControllerMeta, + methodMeta: WebSocketMethodMeta, + ) { + let webSocketStream: Duplex | undefined; + const getWebSocketStream = () => { + if (!webSocketStream) { + webSocketStream = this.options.createWebSocketStream(this.options.context.socket); + } + return webSocketStream; + }; + const result = await this.invokeMethod(realObj, controllerMeta, methodMeta, { getWebSocketStream }); + this.pipeResponseStream(result, getWebSocketStream); + } + + async invokeFetchController( + realObj: WebSocketControllerObject, + controllerMeta: WebSocketFetchControllerMeta, + methodMeta: WebSocketFetchMethodMeta, + events: WebSocketEventStream, + ) { + let session!: WebSocketFetchSession; + const invokeMethod = ( + targetMethodMeta: WebSocketFetchMethodMeta | undefined, + payload: WebSocketMethodPayload = {}, + ) => { + if (!targetMethodMeta) { + return undefined; + } + return this.invokeMethod(realObj, controllerMeta, targetMethodMeta, { + close: session.close, + ...payload, + }); + }; + + session = new WebSocketFetchSession({ + socket: this.options.context.socket, + events, + runInContext: this.options.runInContext, + onConnection: controllerMeta.connectionMethod + ? () => invokeMethod(controllerMeta.connectionMethod) + : undefined, + onOpen: controllerMeta.openMethod + ? () => invokeMethod(controllerMeta.openMethod) + : undefined, + onData: data => invokeMethod(methodMeta, { data }), + onError: controllerMeta.errorMethod + ? error => invokeMethod(controllerMeta.errorMethod, { error }) + : undefined, + onClose: controllerMeta.closeMethod + ? (closeCode, closeReason) => invokeMethod(controllerMeta.closeMethod, { closeCode, closeReason }) + : undefined, + logError: (message, error) => this.options.logger.error(message, error), + logDebug: message => this.options.logger.debug(message), + isFatalError: error => error instanceof TimerUtil.TimeoutError, + }); + await session.run(); + } + + private async invokeMethod( + realObj: WebSocketControllerObject, + controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, + methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, + payload: WebSocketMethodPayload, + ) { + const realMethod = (realObj as Record)[methodMeta.name]; + const args = this.buildMethodArgs(methodMeta, payload); + const timeout = methodMeta.timeout ?? controllerMeta.timeout; + try { + return await TimerUtil.timeout( + async () => await Reflect.apply(realMethod, realObj, args), + timeout, + ); + } catch (error) { + if (error instanceof TimerUtil.TimeoutError) { + this.options.logger.error( + `${controllerMeta.type} ${controllerMeta.controllerName}.${methodMeta.name} timed out after ${timeout}ms`, + error, + ); + } + throw error; + } + } + + private buildMethodArgs( + methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, + payload: WebSocketMethodPayload, + ) { + const indexes = Array.from(methodMeta.paramMap.keys()); + if (methodMeta.contextParamIndex !== undefined) { + indexes.push(methodMeta.contextParamIndex); + } + const args: unknown[] = new Array(indexes.length ? Math.max(...indexes) + 1 : 0); + if (methodMeta.contextParamIndex !== undefined) { + args[methodMeta.contextParamIndex] = this.options.context.context; + } + + for (const [ index, param ] of methodMeta.paramMap) { + switch (param.type) { + case WebSocketParamType.PARAM: { + const pathParam = param as WebSocketPathParamMeta; + args[index] = this.options.context.params[pathParam.name]; + break; + } + case WebSocketParamType.QUERY: { + const queryParam = param as WebSocketQueryParamMeta; + args[index] = this.options.context.query[queryParam.name]; + break; + } + case WebSocketParamType.QUERIES: { + const queryParam = param as WebSocketQueriesParamMeta; + args[index] = this.options.context.queries[queryParam.name] || []; + break; + } + case WebSocketParamType.HEADERS: + args[index] = this.options.context.headers; + break; + case WebSocketParamType.REQUEST: + args[index] = this.options.context.request; + break; + case WebSocketParamType.SOCKET: + args[index] = this.options.context.socket; + break; + case WebSocketParamType.STREAM: + assert(payload.getWebSocketStream, '@WebSocketStream can not be used here'); + args[index] = payload.getWebSocketStream(); + break; + case WebSocketParamType.DATA: + args[index] = payload.data; + break; + case WebSocketParamType.CLOSE: + args[index] = payload.close; + break; + case WebSocketParamType.ERROR: + args[index] = payload.error; + break; + case WebSocketParamType.CLOSE_CODE: + args[index] = payload.closeCode; + break; + case WebSocketParamType.CLOSE_REASON: + args[index] = payload.closeReason; + break; + default: + assert.fail('never arrive'); + } + } + return args; + } + + private pipeResponseStream(result: unknown, getWebSocketStream: () => Duplex) { + if (!this.isReadableStream(result)) { + return; + } + pipeline(result, getWebSocketStream(), error => { + if (!error || this.options.context.socket.readyState >= 2) { + return; + } + this.options.logger.error('WebSocket response stream pipeline failed', error); + this.options.context.socket.close(WEBSOCKET_INTERNAL_ERROR_CODE, WEBSOCKET_INTERNAL_ERROR_REASON); + }); + } + + private isReadableStream(result: unknown): result is NodeJS.ReadableStream { + return !!result && typeof (result as NodeJS.ReadableStream).pipe === 'function'; + } +} diff --git a/core/websocket-runtime/src/WebSocketFetchSession.ts b/core/websocket-runtime/src/WebSocketFetchSession.ts new file mode 100644 index 000000000..dbeb72fa0 --- /dev/null +++ b/core/websocket-runtime/src/WebSocketFetchSession.ts @@ -0,0 +1,347 @@ +export const WEBSOCKET_OPEN_STATE = 1; +export const WEBSOCKET_CONNECTING_STATE = 0; +export const WEBSOCKET_INTERNAL_ERROR_CODE = 1011; +export const WEBSOCKET_INTERNAL_ERROR_REASON = 'Internal Server Error'; + +export interface WebSocketSessionSocket { + on(event: string, listener: (...args: any[]) => void): this; + removeListener(event: string, listener: (...args: any[]) => void): this; + send(data: any, callback?: (error?: Error) => void): void; + close(code?: number, reason?: string | Buffer): void; + readonly readyState: number; +} + +export function waitForWebSocketClose( + socket: WebSocketSessionSocket, + onError: (error: Error) => void, +): Promise { + if (socket.readyState >= 2) { + return Promise.resolve(); + } + return new Promise(resolve => { + const handleError = (error: Error) => onError(error); + const handleClose = () => { + socket.removeListener('close', handleClose); + socket.removeListener('error', handleError); + resolve(); + }; + socket.on('close', handleClose); + socket.on('error', handleError); + }); +} + +export type WebSocketSessionEvent = + | { type: 'connection' } + | { type: 'open' } + | { type: 'message'; data: Data } + | { type: 'error'; error: Error } + | { type: 'close'; code: number; reason: CloseReason }; + +type EventIteratorResult = IteratorResult>; + +/** Captures socket events immediately and exposes them as a single ordered event stream. */ +export class WebSocketEventStream +implements AsyncIterable>, AsyncIterator> { + private events: Array> = []; + private waiters: Array<(result: EventIteratorResult) => void> = []; + private closeListeners = new Set<() => void>(); + private disposed = false; + private closeYielded = false; + private droppedMessages = 0; + private readonly onMessage = (data: Data) => { + this.enqueue({ type: 'message', data }); + }; + private readonly onError = (error: Error) => { + this.enqueue({ type: 'error', error }); + }; + private readonly onClose = (code: number, reason: CloseReason) => { + if (this.disposed) { + return; + } + for (const listener of this.closeListeners) { + listener(); + } + this.closeListeners.clear(); + this.droppedMessages += this.events.filter(event => event.type === 'message').length; + this.events = this.events.filter(event => event.type === 'connection' || event.type === 'open'); + this.socket.removeListener('message', this.onMessage); + this.socket.removeListener('error', this.onError); + this.socket.removeListener('close', this.onClose); + this.enqueue({ type: 'close', code, reason }); + }; + + constructor(private readonly socket: WebSocketSessionSocket) { + this.socket.on('message', this.onMessage); + this.socket.on('error', this.onError); + this.socket.on('close', this.onClose); + this.enqueue({ type: 'connection' }); + this.enqueue({ type: 'open' }); + } + + get closed(): boolean { + return this.closeYielded || this.events.some(event => event.type === 'close'); + } + + get droppedMessageCount(): number { + return this.droppedMessages; + } + + [Symbol.asyncIterator](): AsyncIterator> { + return this; + } + + next(): Promise> { + if (this.disposed || this.closeYielded) { + return Promise.resolve({ done: true, value: undefined }); + } + const event = this.events.shift(); + if (event) { + if (event.type === 'close') { + this.closeYielded = true; + } + return Promise.resolve({ done: false, value: event }); + } + return new Promise(resolve => this.waiters.push(resolve)); + } + + return(): Promise> { + this.dispose(); + return Promise.resolve({ done: true, value: undefined }); + } + + onSocketClose(listener: () => void): () => void { + if (this.closed) { + listener(); + return () => undefined; + } + this.closeListeners.add(listener); + return () => this.closeListeners.delete(listener); + } + + dispose() { + if (this.disposed) { + return; + } + this.disposed = true; + this.socket.removeListener('message', this.onMessage); + this.socket.removeListener('error', this.onError); + this.socket.removeListener('close', this.onClose); + this.events = []; + this.closeListeners.clear(); + for (const resolve of this.waiters.splice(0)) { + resolve({ done: true, value: undefined }); + } + } + + private enqueue(event: WebSocketSessionEvent) { + if (this.disposed || this.closeYielded) { + return; + } + const resolve = this.waiters.shift(); + if (resolve) { + if (event.type === 'close') { + this.closeYielded = true; + } + resolve({ done: false, value: event }); + return; + } + this.events.push(event); + } +} + +export interface WebSocketFetchSessionOptions { + socket: WebSocketSessionSocket; + events: WebSocketEventStream; + runInContext?(callback: () => Promise): Promise; + onConnection?(): unknown | Promise; + onOpen?(): unknown | Promise; + onData(data: Data): unknown | Promise; + onError?(error: Error): unknown | Promise; + onClose?(code: number, reason: CloseReason): unknown | Promise; + logError(message: string, error: Error): void; + logDebug?(message: string): void; + isFatalError?(error: Error): boolean; +} + +interface AsyncReadableStream extends AsyncIterable { + pipe(...args: any[]): unknown; + destroy(error?: Error): void; +} + +export class WebSocketFetchSession { + private acceptingData = true; + + readonly close = (code = 1000, reason?: string | Buffer) => { + const { socket } = this.options; + if (socket.readyState !== WEBSOCKET_OPEN_STATE && socket.readyState !== WEBSOCKET_CONNECTING_STATE) { + return; + } + socket.close(code, reason); + }; + + constructor(private readonly options: WebSocketFetchSessionOptions) {} + + async run(): Promise { + try { + for await (const event of this.options.events) { + await this.dispatch(event); + } + } finally { + this.options.events.dispose(); + } + } + + private async dispatch(event: WebSocketSessionEvent) { + switch (event.type) { + case 'connection': + await this.invokeLifecycle('connection', this.options.onConnection); + break; + case 'open': + await this.invokeLifecycle('open', this.options.onOpen); + break; + case 'message': + await this.invokeData(event.data); + break; + case 'error': + await this.handleError(event.error); + break; + case 'close': + if (this.options.events.droppedMessageCount > 0) { + this.options.logDebug?.( + `Discarded ${this.options.events.droppedMessageCount} queued WebSocket fetch message(s) after close`, + ); + } + await this.invokeClose(event.code, event.reason); + break; + default: + break; + } + } + + private async invokeLifecycle(name: 'connection' | 'open', callback: (() => unknown | Promise) | undefined) { + if (!callback || !this.acceptingData || this.options.events.closed) { + return; + } + try { + await this.runInContext(async () => await callback()); + } catch (error) { + const normalizedError = this.toError(error); + this.acceptingData = false; + this.options.logError(`WebSocket fetch ${name} handler failed`, normalizedError); + await this.handleError(normalizedError); + this.close(WEBSOCKET_INTERNAL_ERROR_CODE, WEBSOCKET_INTERNAL_ERROR_REASON); + } + } + + private async invokeData(data: Data) { + if (!this.acceptingData || this.options.events.closed || this.options.socket.readyState !== WEBSOCKET_OPEN_STATE) { + this.options.logDebug?.('Discarded a WebSocket fetch message because the connection is not open'); + return; + } + try { + const result = await this.runInContext(async () => await this.options.onData(data)); + await this.sendResponse(result); + } catch (error) { + if (!this.options.events.closed) { + const normalizedError = this.toError(error); + await this.handleError(normalizedError); + if (this.options.isFatalError?.(normalizedError)) { + this.close(WEBSOCKET_INTERNAL_ERROR_CODE, WEBSOCKET_INTERNAL_ERROR_REASON); + } + } + } + } + + private async invokeClose(code: number, reason: CloseReason) { + if (!this.options.onClose) { + return; + } + try { + await this.runInContext(async () => await this.options.onClose!(code, reason)); + } catch (error) { + this.options.logError('WebSocket fetch close handler failed', this.toError(error)); + } + } + + private async handleError(error: Error) { + if (!this.options.onError) { + this.options.logError('WebSocket fetch request failed', error); + return; + } + try { + await this.runInContext(async () => await this.options.onError!(error)); + } catch (handlerError) { + this.options.logError('WebSocket fetch error handler failed', this.toError(handlerError)); + } + } + + private async sendResponse(result: unknown) { + if (result === undefined || result === null) { + return; + } + if (!this.isReadableStream(result)) { + await this.handleError(new Error('WebSocketFetch method must return a readable stream or void')); + return; + } + if (this.options.events.closed || this.options.socket.readyState !== WEBSOCKET_OPEN_STATE) { + this.destroyResponseStream(result); + return; + } + + const removeCloseListener = this.options.events.onSocketClose(() => this.destroyResponseStream(result)); + try { + for await (const chunk of result) { + if (this.options.events.closed || this.options.socket.readyState !== WEBSOCKET_OPEN_STATE) { + this.destroyResponseStream(result); + return; + } + await this.sendChunk(chunk); + } + } finally { + removeCloseListener(); + } + } + + private sendChunk(chunk: unknown): Promise { + return new Promise((resolve, reject) => { + if (this.options.events.closed || this.options.socket.readyState !== WEBSOCKET_OPEN_STATE) { + resolve(); + return; + } + this.options.socket.send(chunk, error => { + if (error) { + reject(error); + return; + } + resolve(); + }); + }); + } + + private destroyResponseStream(stream: AsyncReadableStream) { + try { + stream.destroy(); + } catch (error) { + this.options.logError('WebSocket fetch response stream destroy failed', this.toError(error)); + } + } + + private runInContext(callback: () => Promise): Promise { + if (this.options.runInContext) { + return this.options.runInContext(callback); + } + return callback(); + } + + private isReadableStream(result: unknown): result is AsyncReadableStream { + const stream = result as Partial | undefined; + return !!stream && + typeof stream.pipe === 'function' && + typeof stream.destroy === 'function' && + typeof stream[Symbol.asyncIterator] === 'function'; + } + + private toError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)); + } +} diff --git a/core/websocket-runtime/src/WebSocketRouteRegistry.ts b/core/websocket-runtime/src/WebSocketRouteRegistry.ts new file mode 100644 index 000000000..3af6a815a --- /dev/null +++ b/core/websocket-runtime/src/WebSocketRouteRegistry.ts @@ -0,0 +1,140 @@ +import pathToRegexp from 'path-to-regexp'; +import { + ControllerType, + WebSocketControllerMeta, + WebSocketFetchControllerMeta, + WebSocketFetchMethodMeta, + WebSocketMethodMeta, +} from '@eggjs/controller-decorator'; + +export type WebSocketControllerMetadata = WebSocketControllerMeta | WebSocketFetchControllerMeta; +export type WebSocketRouteMethodMeta = WebSocketMethodMeta | WebSocketFetchMethodMeta; + +export interface WebSocketRoute { + controllerProto: ControllerProto; + controllerMeta: WebSocketControllerMetadata; + methodMeta: WebSocketRouteMethodMeta; + methodRealPath: string; + methodName: string; + host?: string; + keys: pathToRegexp.Key[]; + regexp: RegExp; +} + +export function createWebSocketRoutes( + controllerProtos: ControllerProto[], + getMetadata: (proto: ControllerProto) => WebSocketControllerMetadata, + validateMethod: (controllerMeta: WebSocketControllerMetadata, methodMeta: WebSocketRouteMethodMeta) => void, +): Array> { + const methodMap = new Map(); + for (const proto of controllerProtos) { + for (const method of getMetadata(proto).methods) { + methodMap.set(method, proto); + } + } + const methods = Array.from(methodMap.keys()).sort((a, b) => b.priority - a.priority); + + for (const method of methods) { + const controllerMeta = getMetadata(methodMap.get(method)!); + validateMethod(controllerMeta, method); + } + + return methods.flatMap(method => { + const controllerProto = methodMap.get(method)!; + const controllerMeta = getMetadata(controllerProto); + const methodHosts = getWebSocketMethodHosts(controllerMeta, method); + const hosts = methodHosts?.length ? methodHosts : [ undefined ]; + return hosts.map(host => createRoute(controllerProto, controllerMeta, method, host)); + }); +} + +export function matchWebSocketRoute( + routes: Array>, + pathname: string, + host: string | undefined, +): { route: WebSocketRoute; params: Record } | undefined { + for (const route of routes) { + if (route.host && !matchWebSocketHost(route.host, host)) { + continue; + } + const matched = route.regexp.exec(pathname); + if (!matched) { + continue; + } + const params: Record = {}; + route.keys.forEach((key, index) => { + const value = matched[index + 1]; + if (value !== undefined) { + params[String(key.name)] = decodeURIComponent(value); + } + }); + return { route, params }; + } +} + +export function getWebSocketMethodRealPath( + controllerMeta: WebSocketControllerMetadata, + methodMeta: WebSocketRouteMethodMeta, +) { + if (controllerMeta.type === ControllerType.WEBSOCKET) { + return controllerMeta.getMethodRealPath(methodMeta as WebSocketMethodMeta); + } + return controllerMeta.getMethodRealPath(methodMeta as WebSocketFetchMethodMeta); +} + +export function getWebSocketMethodHosts( + controllerMeta: WebSocketControllerMetadata, + methodMeta: WebSocketRouteMethodMeta, +): string[] | undefined { + if (controllerMeta.type === ControllerType.WEBSOCKET) { + return controllerMeta.getMethodHosts(methodMeta as WebSocketMethodMeta); + } + return controllerMeta.getMethodHosts(methodMeta as WebSocketFetchMethodMeta); +} + +export function getWebSocketMethodName( + controllerMeta: WebSocketControllerMetadata, + methodMeta: WebSocketRouteMethodMeta, +): string { + if (controllerMeta.type === ControllerType.WEBSOCKET) { + return controllerMeta.getMethodName(methodMeta as WebSocketMethodMeta); + } + return controllerMeta.getMethodName(methodMeta as WebSocketFetchMethodMeta); +} + +export function getWebSocketMethodMiddlewares( + controllerMeta: WebSocketControllerMetadata, + methodMeta: WebSocketRouteMethodMeta, +) { + if (controllerMeta.type === ControllerType.WEBSOCKET) { + return controllerMeta.getMethodMiddlewares(methodMeta as WebSocketMethodMeta); + } + return controllerMeta.getMethodMiddlewares(methodMeta as WebSocketFetchMethodMeta); +} + +function createRoute( + controllerProto: ControllerProto, + controllerMeta: WebSocketControllerMetadata, + methodMeta: WebSocketRouteMethodMeta, + host: string | undefined, +): WebSocketRoute { + const methodRealPath = getWebSocketMethodRealPath(controllerMeta, methodMeta); + const keys: pathToRegexp.Key[] = []; + return { + controllerProto, + controllerMeta, + methodMeta, + methodRealPath, + methodName: getWebSocketMethodName(controllerMeta, methodMeta), + host, + keys, + regexp: pathToRegexp(methodRealPath, keys, { sensitive: true }), + }; +} + +function matchWebSocketHost(expectedHost: string, requestHost: string | undefined) { + if (!requestHost) { + return false; + } + return requestHost === expectedHost || requestHost.split(':')[0] === expectedHost; +} diff --git a/core/websocket-runtime/test/WebSocketFetchSession.test.ts b/core/websocket-runtime/test/WebSocketFetchSession.test.ts new file mode 100644 index 000000000..f0557afb2 --- /dev/null +++ b/core/websocket-runtime/test/WebSocketFetchSession.test.ts @@ -0,0 +1,217 @@ +import { strict as assert } from 'node:assert'; +import { EventEmitter } from 'node:events'; +import { PassThrough, Readable } from 'node:stream'; +import { + WebSocketEventStream, + WebSocketFetchSession, + waitForWebSocketClose, +} from '..'; + +class FakeWebSocket extends EventEmitter { + readyState = 1; + readonly sent: unknown[] = []; + sendsInFlight = 0; + maxSendsInFlight = 0; + closeCode?: number; + closeReason?: string | Buffer; + onSend?: (data: unknown, callback: (error?: Error) => void) => void; + + send(data: unknown, callback: (error?: Error) => void = () => undefined) { + this.sent.push(data); + this.sendsInFlight++; + this.maxSendsInFlight = Math.max(this.maxSendsInFlight, this.sendsInFlight); + const done = (error?: Error) => { + this.sendsInFlight--; + callback(error); + }; + if (this.onSend) { + this.onSend(data, done); + return; + } + setImmediate(done); + } + + close(code = 1000, reason: string | Buffer = Buffer.alloc(0)) { + if (this.readyState > 1) { + return; + } + this.readyState = 3; + this.closeCode = code; + this.closeReason = reason; + this.emit('close', code, reason); + } +} + +describe('test/WebSocketFetchSession.test.ts', () => { + it('should keep waiting for close after a socket error', async () => { + const socket = new FakeWebSocket(); + const errors: Error[] = []; + let resolved = false; + const closed = waitForWebSocketClose(socket, error => errors.push(error)); + closed.then(() => { + resolved = true; + }); + + const error = new Error('socket failed'); + socket.emit('error', error); + await new Promise(resolve => setImmediate(resolve)); + assert.deepEqual(errors, [ error ]); + assert.equal(resolved, false); + + socket.close(); + await closed; + assert.equal(resolved, true); + }); + + it('should process response streams in request order with send backpressure', async () => { + const socket = new FakeWebSocket(); + const events = new WebSocketEventStream(socket); + const calls: string[] = []; + socket.onSend = (data, callback) => { + calls.push(`send:${String(data)}`); + setTimeout(() => { + callback(); + if (socket.sent.length === 4) { + socket.close(); + } + }, 5); + }; + const session = new WebSocketFetchSession({ + socket, + events, + onConnection: () => calls.push('connection'), + onOpen: () => calls.push('open'), + onData: data => { + calls.push(`data:${data}`); + return Readable.from([ `${data}:1`, `${data}:2` ]); + }, + onClose: () => calls.push('close'), + logError: (_message, error) => assert.fail(error), + }); + + socket.emit('message', 'first'); + socket.emit('message', 'second'); + await session.run(); + + assert.deepEqual(calls, [ + 'connection', + 'open', + 'data:first', + 'send:first:1', + 'send:first:2', + 'data:second', + 'send:second:1', + 'send:second:2', + 'close', + ]); + assert.equal(socket.maxSendsInFlight, 1); + }); + + it('should destroy the active response and discard queued messages on close', async () => { + const socket = new FakeWebSocket(); + const events = new WebSocketEventStream(socket); + const output = new PassThrough(); + const messages: string[] = []; + let closed = false; + const session = new WebSocketFetchSession({ + socket, + events, + onData: data => { + messages.push(data); + setImmediate(() => socket.close(1001, Buffer.from('gone'))); + return output; + }, + onClose: (code, reason) => { + assert.equal(code, 1001); + assert.equal(reason.toString(), 'gone'); + closed = true; + }, + logError: (_message, error) => assert.fail(error), + }); + + socket.emit('message', 'first'); + socket.emit('message', 'second'); + await session.run(); + + assert.deepEqual(messages, [ 'first' ]); + assert.equal(output.destroyed, true); + assert.equal(closed, true); + }); + + it('should invoke lifecycle and error handlers through the context runner', async () => { + const socket = new FakeWebSocket(); + const events = new WebSocketEventStream(socket); + let inContext = false; + const handlers: string[] = []; + const session = new WebSocketFetchSession({ + socket, + events, + runInContext: async callback => { + assert.equal(inContext, false); + inContext = true; + try { + return await callback(); + } finally { + inContext = false; + } + }, + onConnection: () => { + assert.equal(inContext, true); + handlers.push('connection'); + }, + onOpen: () => { + assert.equal(inContext, true); + handlers.push('open'); + }, + onData: () => { + assert.equal(inContext, true); + handlers.push('data'); + throw new Error('data failed'); + }, + onError: error => { + assert.equal(inContext, true); + assert.equal(error.message, 'data failed'); + handlers.push('error'); + socket.close(); + }, + onClose: () => { + assert.equal(inContext, true); + handlers.push('close'); + }, + logError: (_message, error) => assert.fail(error), + }); + + socket.emit('message', 'request'); + await session.run(); + + assert.deepEqual(handlers, [ 'connection', 'open', 'data', 'error', 'close' ]); + }); + + it('should close the connection after handling a fatal data error', async () => { + const socket = new FakeWebSocket(); + const events = new WebSocketEventStream(socket); + const fatalError = new Error('timeout'); + const handlers: string[] = []; + const session = new WebSocketFetchSession({ + socket, + events, + onData: () => { + throw fatalError; + }, + onError: error => { + assert.equal(error, fatalError); + handlers.push('error'); + }, + onClose: () => handlers.push('close'), + isFatalError: error => error === fatalError, + logError: (_message, error) => assert.fail(error), + }); + + socket.emit('message', 'request'); + await session.run(); + + assert.deepEqual(handlers, [ 'error', 'close' ]); + assert.equal(socket.closeCode, 1011); + assert.equal(socket.closeReason, 'Internal Server Error'); + }); +}); diff --git a/core/websocket-runtime/test/WebSocketRouteRegistry.test.ts b/core/websocket-runtime/test/WebSocketRouteRegistry.test.ts new file mode 100644 index 000000000..73f344096 --- /dev/null +++ b/core/websocket-runtime/test/WebSocketRouteRegistry.test.ts @@ -0,0 +1,64 @@ +import { strict as assert } from 'node:assert'; +import { + WebSocketControllerMeta, + WebSocketMethodMeta, +} from '@eggjs/controller-decorator'; +import { + createWebSocketRoutes, + matchWebSocketRoute, +} from '..'; + +describe('test/WebSocketRouteRegistry.test.ts', () => { + it('should sort routes, match hosts and preserve optional params', () => { + const optionalMethod = new WebSocketMethodMeta( + 'optional', + '/optional/:id?', + [], + undefined, + new Map(), + 10, + undefined, + ); + const lowerPriorityMethod = new WebSocketMethodMeta( + 'lowerPriority', + '/lower', + [], + undefined, + new Map(), + 1, + undefined, + ); + const metadata = new WebSocketControllerMeta( + 'TestController', + 'testController', + 'TestController', + '/ws', + [], + [ lowerPriorityMethod, optionalMethod ], + [ 'example.com' ], + ); + const proto = { metadata }; + const validated: string[] = []; + const routes = createWebSocketRoutes( + [ proto ], + item => item.metadata, + (_controllerMeta, methodMeta) => validated.push(methodMeta.name), + ); + + assert.deepEqual(validated, [ 'optional', 'lowerPriority' ]); + assert.equal(routes[0].methodMeta.name, 'optional'); + assert.equal(matchWebSocketRoute(routes, '/ws/optional', 'other.example.com'), undefined); + assert.deepEqual( + matchWebSocketRoute(routes, '/ws/optional', 'example.com')?.params, + {}, + ); + assert.deepEqual( + matchWebSocketRoute(routes, '/ws/optional/value', 'example.com')?.params, + { id: 'value' }, + ); + assert.throws( + () => matchWebSocketRoute(routes, '/ws/optional/%E0%A4%A', 'example.com'), + URIError, + ); + }); +}); diff --git a/core/websocket-runtime/tsconfig.json b/core/websocket-runtime/tsconfig.json new file mode 100644 index 000000000..64b224050 --- /dev/null +++ b/core/websocket-runtime/tsconfig.json @@ -0,0 +1,12 @@ +{ + "extends": "../../tsconfig.json", + "compilerOptions": { + "outDir": "dist", + "baseUrl": "./" + }, + "exclude": [ + "dist", + "node_modules", + "test" + ] +} diff --git a/core/websocket-runtime/tsconfig.pub.json b/core/websocket-runtime/tsconfig.pub.json new file mode 100644 index 000000000..fc8520e73 --- /dev/null +++ b/core/websocket-runtime/tsconfig.pub.json @@ -0,0 +1,3 @@ +{ + "extends": "./tsconfig.json" +} diff --git a/plugin/controller/app.ts b/plugin/controller/app.ts index 3bcca2760..b22835ee0 100644 --- a/plugin/controller/app.ts +++ b/plugin/controller/app.ts @@ -152,7 +152,6 @@ export default class ControllerAppBootHook { // and register methods after collect is done. HTTPControllerRegister.instance?.doRegister(this.app.rootProtoManager); WebSocketControllerRegister.instance?.doRegister(); - WebSocketControllerRegister.instance?.listen(); this.app.config.mcp.hooks = this.mcpControllerRegister?.hooks; } diff --git a/plugin/controller/lib/impl/websocket/WebSocketContext.ts b/plugin/controller/lib/impl/websocket/WebSocketContext.ts index 937e52eb6..e5590d12a 100644 --- a/plugin/controller/lib/impl/websocket/WebSocketContext.ts +++ b/plugin/controller/lib/impl/websocket/WebSocketContext.ts @@ -1,35 +1,17 @@ -import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; +import type { Context } from 'egg'; import type { WebSocket } from 'ws'; import type { WebSocketContext } from '@eggjs/tegg'; -export class WebSocketContextImpl implements WebSocketContext { - readonly socket: WebSocket; - readonly request: IncomingMessage; - readonly path: string; - readonly host: string; - readonly params: Record; - readonly query: Record; - readonly queries: Record; - readonly headers: IncomingHttpHeaders; - - constructor(options: { - socket: WebSocket; - request: IncomingMessage; - url: URL; - params: Record; - }) { - this.socket = options.socket; - this.request = options.request; - this.path = options.url.pathname; - this.host = options.url.host; - this.params = options.params; - this.headers = options.request.headers; - this.query = {}; - this.queries = {}; - - for (const key of options.url.searchParams.keys()) { - this.query[key] = options.url.searchParams.get(key) ?? undefined; - this.queries[key] = options.url.searchParams.getAll(key); - } - } +export function extendWebSocketContext( + ctx: Context, + webSocket: WebSocket, + params: Record, +): WebSocketContext { + Object.defineProperty(ctx, 'webSocket', { + configurable: true, + enumerable: false, + value: webSocket, + }); + ctx.params = params; + return ctx as WebSocketContext; } diff --git a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts index 8f2f1103e..69c4f7700 100644 --- a/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts +++ b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts @@ -1,8 +1,7 @@ import assert from 'node:assert'; import { IncomingMessage, ServerResponse } from 'node:http'; -import { Duplex, pipeline } from 'node:stream'; +import { Duplex } from 'node:stream'; import compose from 'koa-compose'; -import pathToRegexp from 'path-to-regexp'; import { Application, Context } from 'egg'; import { EggRouter } from '@eggjs/router'; import { FrameworkErrorFormater } from 'egg-errors'; @@ -13,46 +12,35 @@ import { ControllerType, Next, WebSocketControllerMeta, - WebSocketFetchClose, WebSocketFetchControllerMeta, WebSocketFetchMethodMeta, WebSocketMethodMeta, - WebSocketParamType, - WebSocketPathParamMeta, - WebSocketQueriesParamMeta, - WebSocketQueryParamMeta, } from '@eggjs/tegg'; import { EggContainerFactory } from '@eggjs/tegg-runtime'; import { EggPrototype } from '@eggjs/tegg-metadata'; +import { + WebSocketControllerRuntime, + WebSocketEventStream, + WebSocketRoute, + createWebSocketRoutes, + getWebSocketMethodHosts, + getWebSocketMethodMiddlewares, + getWebSocketMethodName, + getWebSocketMethodRealPath, + matchWebSocketRoute, + waitForWebSocketClose, + WEBSOCKET_INTERNAL_ERROR_CODE, + WEBSOCKET_INTERNAL_ERROR_REASON, +} from '@eggjs/tegg-websocket-runtime'; import { ROOT_PROTO } from '@eggjs/egg-module-common'; import { ControllerRegister } from '../../ControllerRegister'; import { RouterConflictError } from '../../errors'; -import { WebSocketContextImpl } from './WebSocketContext'; +import { extendWebSocketContext } from './WebSocketContext'; const noop = () => { // noop }; -interface WebSocketRoute { - controllerProto: EggPrototype; - controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta; - methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta; - methodRealPath: string; - methodName: string; - host?: string; - keys: pathToRegexp.Key[]; - regexp: RegExp; -} - -interface WebSocketMethodPayload { - data?: RawData; - close?: WebSocketFetchClose; - error?: Error; - closeCode?: number; - closeReason?: Buffer; - getWebSocketStream?: () => Duplex; -} - export class WebSocketControllerRegister implements ControllerRegister { static instance?: WebSocketControllerRegister; @@ -60,7 +48,7 @@ export class WebSocketControllerRegister implements ControllerRegister { private readonly eggContainerFactory: typeof EggContainerFactory; private readonly checkRouters: Map; private controllerProtos: EggPrototype[] = []; - private routes: WebSocketRoute[] = []; + private routes: Array> = []; private webSocketServer?: WebSocketServer; private upgradeHandler?: (req: IncomingMessage, socket: Duplex, head: Buffer) => void; private listening = false; @@ -94,28 +82,11 @@ export class WebSocketControllerRegister implements ControllerRegister { } doRegister() { - const methodMap = new Map(); - for (const proto of this.controllerProtos) { - const metadata = proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; - for (const method of metadata.methods) { - methodMap.set(method, proto); - } - } - const allMethods = Array.from(methodMap.keys()) - .sort((a, b) => b.priority - a.priority); - - for (const method of allMethods) { - const controllerProto = methodMap.get(method)!; - const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; - this.checkDuplicate(controllerMeta, method); - } - - this.routes = allMethods.flatMap(method => { - const controllerProto = methodMap.get(method)!; - const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; - const hosts = this.getMethodHosts(controllerMeta, method) || [ undefined ]; - return hosts.map(host => this.createRoute(controllerProto, controllerMeta, method, host)); - }); + this.routes = createWebSocketRoutes( + this.controllerProtos, + proto => proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta, + (controllerMeta, methodMeta) => this.checkDuplicate(controllerMeta, methodMeta), + ); } listen() { @@ -152,30 +123,9 @@ export class WebSocketControllerRegister implements ControllerRegister { this.checkRouters.clear(); } - private createRoute( - controllerProto: EggPrototype, - controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, - methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, - host: string | undefined, - ): WebSocketRoute { - const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); - const keys: pathToRegexp.Key[] = []; - const regexp = pathToRegexp(methodRealPath, keys, { sensitive: true }); - return { - controllerProto, - controllerMeta, - methodMeta, - methodRealPath, - methodName: this.getMethodName(controllerMeta, methodMeta), - host, - keys, - regexp, - }; - } - private checkDuplicate(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { let router = this.checkRouters.get('default')!; - const hosts = this.getMethodHosts(controllerMeta, methodMeta) || []; + const hosts = getWebSocketMethodHosts(controllerMeta, methodMeta) || []; if (!hosts.length) { this.checkDuplicateInRouter(router, controllerMeta, methodMeta); this.registerToRouter(router, controllerMeta, methodMeta); @@ -194,15 +144,15 @@ export class WebSocketControllerRegister implements ControllerRegister { } private registerToRouter(router: EggRouter, controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { - const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); - const methodName = this.getMethodName(controllerMeta, methodMeta); + const methodRealPath = getWebSocketMethodRealPath(controllerMeta, methodMeta); + const methodName = getWebSocketMethodName(controllerMeta, methodMeta); Reflect.apply(router.get, router, [ methodName, methodRealPath, noop ]); } private checkDuplicateInRouter(router: EggRouter, controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { - const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); + const methodRealPath = getWebSocketMethodRealPath(controllerMeta, methodMeta); const matched = router.match(methodRealPath, 'GET'); - const methodName = this.getMethodName(controllerMeta, methodMeta); + const methodName = getWebSocketMethodName(controllerMeta, methodMeta); if (matched.route) { const [ layer ] = matched.path; const err = new RouterConflictError(`register websocket controller ${methodName} failed, ${controllerMeta.type} ${methodRealPath} is conflict with exists rule ${layer.path}`); @@ -210,478 +160,158 @@ export class WebSocketControllerRegister implements ControllerRegister { } } - private getMethodRealPath(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { - return (controllerMeta as any).getMethodRealPath(methodMeta); - } - - private getMethodHosts(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): string[] | undefined { - return (controllerMeta as any).getMethodHosts(methodMeta); - } - - private getMethodName(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): string { - return (controllerMeta as any).getMethodName(methodMeta); - } - - private getMethodMiddlewares(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { - return (controllerMeta as any).getMethodMiddlewares(methodMeta); - } - private async handleUpgrade(req: IncomingMessage, socket: Duplex, head: Buffer) { - const url = this.createURL(req); - const matched = this.matchRoute(req, url); + const res = new ServerResponse(req); + const eggCtx = this.app.createContext(req, res) as unknown as Context; + const url = this.createURL(eggCtx); + let matched: { route: WebSocketRoute; params: Record } | undefined; + try { + matched = matchWebSocketRoute(this.routes, url.pathname, eggCtx.host); + } catch (error) { + if (error instanceof URIError) { + this.app.logger.warn('[tegg/websocket] reject malformed request path: %s', req.url); + this.rejectSocket(socket, 400, 'Bad Request'); + return; + } + throw error; + } if (!matched) { + const initialBytesWritten = this.getSocketBytesWritten(socket); + setImmediate(() => { + if ( + socket.destroyed || + socket.writableEnded || + this.getSocketBytesWritten(socket) > initialBytesWritten || + this.hasWebSocketOwner(socket) + ) { + return; + } + this.app.logger.warn('[tegg/websocket] no route matched upgrade request: %s', req.url); + this.rejectSocket(socket, 404, 'Not Found'); + }); return; } this.webSocketServer!.handleUpgrade(req, socket, head, webSocket => { - this.handleConnection(matched.route, matched.params, url, req, webSocket) + this.handleConnection(matched.route, matched.params, eggCtx, webSocket) .catch(error => { this.app.logger.error('[tegg/websocket] handle connection failed: %s', error.stack || error.message); if (webSocket.readyState === WebSocket.OPEN || webSocket.readyState === WebSocket.CONNECTING) { - webSocket.close(1011, 'Internal Server Error'); + webSocket.close(WEBSOCKET_INTERNAL_ERROR_CODE, WEBSOCKET_INTERNAL_ERROR_REASON); } }); }); } - private createURL(req: IncomingMessage): URL { - const host = req.headers.host || 'localhost'; - return new URL(req.url || '/', `http://${host}`); - } - - private matchRoute(req: IncomingMessage, url: URL): { route: WebSocketRoute; params: Record } | undefined { - for (const route of this.routes) { - if (route.host && !this.matchHost(route.host, req.headers.host)) { - continue; - } - const matched = route.regexp.exec(url.pathname); - if (!matched) { - continue; - } - const params: Record = {}; - route.keys.forEach((key, index) => { - params[String(key.name)] = decodeURIComponent(matched[index + 1]); - }); - return { route, params }; - } - } - - private matchHost(expectedHost: string, requestHost: string | undefined) { - if (!requestHost) { - return false; - } - return requestHost === expectedHost || requestHost.split(':')[0] === expectedHost; + private createURL(ctx: Context): URL { + const host = ctx.host || 'localhost'; + return new URL(ctx.url || '/', `${ctx.protocol}://${host}`); } private async handleConnection( - route: WebSocketRoute, + route: WebSocketRoute, params: Record, - url: URL, - req: IncomingMessage, + eggCtx: Context, webSocket: WebSocket, ) { - const res = new ServerResponse(req); - const eggCtx = this.app.createContext(req, res) as unknown as Context; - (eggCtx as any)[ROOT_PROTO] = route.controllerProto; - const webSocketCtx = new WebSocketContextImpl({ - socket: webSocket, - request: req, - url, - params, + const closeHandled = waitForWebSocketClose(webSocket, error => { + this.app.logger.error('[tegg/websocket] socket error while waiting for close: %s', error.stack || error.message); + }); + Reflect.set(eggCtx, ROOT_PROTO, route.controllerProto); + const webSocketCtx = extendWebSocketContext(eggCtx, webSocket, params); + const fetchEvents = route.controllerMeta.type === ControllerType.WEBSOCKET_FETCH + ? new WebSocketEventStream(webSocket) + : undefined; + const runtime = new WebSocketControllerRuntime({ + context: { + context: webSocketCtx, + socket: webSocket, + params: webSocketCtx.params, + query: webSocketCtx.query, + queries: webSocketCtx.queries, + headers: webSocketCtx.headers, + request: webSocketCtx.req, + }, + createWebSocketStream, + runInContext: callback => this.app.ctxStorage.run(eggCtx, callback), + logger: { + debug: message => this.app.logger.debug('[tegg/websocket] %s', message), + error: (message, error) => { + if (error) { + this.app.logger.error('[tegg/websocket] %s: %s', message, error.stack || error.message); + return; + } + this.app.logger.error('[tegg/websocket] %s', message); + }, + }, }); const lifecycleMiddleware = this.app.middleware.teggCtxLifecycleMiddleware(); - const methodMiddlewares = this.getMethodMiddlewares(route.controllerMeta, route.methodMeta); + const methodMiddlewares = getWebSocketMethodMiddlewares(route.controllerMeta, route.methodMeta); let invoked = false; const handler = async (_ctx: Context, next: Next) => { invoked = true; if (route.controllerMeta.type === ControllerType.WEBSOCKET_FETCH) { - await this.invokeFetchController(route, webSocketCtx, eggCtx); + await this.invokeFetchController(route, runtime, fetchEvents!); } else { - await this.invokeController(route, webSocketCtx); + await this.invokeController(route, runtime); } await next(); }; const composed = compose([ ...methodMiddlewares, handler ]); - await this.app.ctxStorage.run(eggCtx, async () => { - await lifecycleMiddleware(eggCtx, async () => { - await composed(eggCtx, async () => { - // final middleware + try { + await this.app.ctxStorage.run(eggCtx, async () => { + await lifecycleMiddleware(eggCtx, async () => { + await composed(eggCtx, async () => { + // final middleware + }); + if (!invoked) { + this.app.logger.debug('[tegg/websocket] %s was short-circuited by middleware', route.methodName); + } + await closeHandled; }); - if (!invoked) { - webSocket.close(1008, 'WebSocket middleware did not call next'); - return; - } - await this.waitWebSocketClose(webSocket); }); - }); + } finally { + fetchEvents?.dispose(); + } } - private async invokeController(route: WebSocketRoute, webSocketCtx: WebSocketContextImpl) { + private async invokeController( + route: WebSocketRoute, + runtime: WebSocketControllerRuntime, + ) { const methodMeta = route.methodMeta as WebSocketMethodMeta; + const controllerMeta = route.controllerMeta as WebSocketControllerMeta; const eggObj = await this.eggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); - const realObj = eggObj.obj; - const realMethod = realObj[methodMeta.name]; - let webSocketStream: Duplex | undefined; - const getWebSocketStream = () => { - if (!webSocketStream) { - webSocketStream = createWebSocketStream(webSocketCtx.socket); - } - return webSocketStream; - }; - const args = this.buildMethodArgs(methodMeta, webSocketCtx, { getWebSocketStream }); - const result = await Reflect.apply(realMethod, realObj, args); - this.pipeResponseStream(result, webSocketCtx.socket, getWebSocketStream); + await runtime.invokeController(eggObj.obj, controllerMeta, methodMeta); } - private buildMethodArgs( - methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, - webSocketCtx: WebSocketContextImpl, - payload: WebSocketMethodPayload, + private async invokeFetchController( + route: WebSocketRoute, + runtime: WebSocketControllerRuntime, + events: WebSocketEventStream, ) { - const argsLength = methodMeta.paramMap.size; - const hasContext = methodMeta.contextParamIndex !== undefined; - const contextIndex = methodMeta.contextParamIndex; - const methodArgsLength = argsLength + (hasContext ? 1 : 0); - const args: unknown[] = new Array(methodArgsLength); - if (hasContext) { - args[contextIndex!] = webSocketCtx; - } - for (const [ index, param ] of methodMeta.paramMap) { - switch (param.type) { - case WebSocketParamType.PARAM: { - const pathParam = param as WebSocketPathParamMeta; - args[index] = webSocketCtx.params[pathParam.name]; - break; - } - case WebSocketParamType.QUERY: { - const queryParam = param as WebSocketQueryParamMeta; - args[index] = webSocketCtx.query[queryParam.name]; - break; - } - case WebSocketParamType.QUERIES: { - const queryParam = param as WebSocketQueriesParamMeta; - args[index] = webSocketCtx.queries[queryParam.name] || []; - break; - } - case WebSocketParamType.HEADERS: { - args[index] = webSocketCtx.headers; - break; - } - case WebSocketParamType.REQUEST: { - args[index] = webSocketCtx.request; - break; - } - case WebSocketParamType.SOCKET: { - args[index] = webSocketCtx.socket; - break; - } - case WebSocketParamType.STREAM: { - assert(payload.getWebSocketStream, '@WebSocketStream can not be used here'); - args[index] = payload.getWebSocketStream(); - break; - } - case WebSocketParamType.DATA: { - args[index] = payload.data; - break; - } - case WebSocketParamType.CLOSE: { - args[index] = payload.close; - break; - } - case WebSocketParamType.ERROR: { - args[index] = payload.error; - break; - } - case WebSocketParamType.CLOSE_CODE: { - args[index] = payload.closeCode; - break; - } - case WebSocketParamType.CLOSE_REASON: { - args[index] = payload.closeReason; - break; - } - default: - assert.fail('never arrive'); - } - } - return args; - } - - private async invokeFetchController(route: WebSocketRoute, webSocketCtx: WebSocketContextImpl, eggCtx: Context) { const controllerMeta = route.controllerMeta as WebSocketFetchControllerMeta; const methodMeta = route.methodMeta as WebSocketFetchMethodMeta; const eggObj = await this.eggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); - const realObj = eggObj.obj; - const close = this.createFetchClose(webSocketCtx.socket); - const responseStreams = new Set(); - let connectionClosed = false; - let controllerReady = false; - let resolveControllerReady!: () => void; - const controllerReadyHandled = new Promise(resolve => { - resolveControllerReady = resolve; - }); - let messageQueue = Promise.resolve(); - const invokeMethod = async ( - targetMethodMeta: WebSocketFetchMethodMeta | undefined, - payload: WebSocketMethodPayload = {}, - ) => { - if (!targetMethodMeta) { - return; - } - const realMethod = realObj[targetMethodMeta.name]; - const args = this.buildMethodArgs(targetMethodMeta, webSocketCtx, { - close, - ...payload, - }); - return await Reflect.apply(realMethod, realObj, args); - }; - const handleError = async (error: Error) => { - if (!controllerMeta.errorMethod) { - this.app.logger.error('[tegg/websocket-fetch] handle error: %s', error.stack || error.message); - return; - } - await invokeMethod(controllerMeta.errorMethod, { error }); - }; - const closeHandled = new Promise(resolve => { - webSocketCtx.socket.once('close', (code, reason) => { - connectionClosed = true; - this.destroyFetchResponseStreams(responseStreams); - this.app.ctxStorage.run(eggCtx, async () => { - try { - await invokeMethod(controllerMeta.closeMethod, { - closeCode: code, - closeReason: reason, - }); - } catch (error) { - this.app.logger.error('[tegg/websocket-fetch] handle close failed: %s', error.stack || error.message); - } finally { - resolve(); - } - }).catch(error => { - this.app.logger.error('[tegg/websocket-fetch] handle close failed: %s', error.stack || error.message); - resolve(); - }); - }); - }); - - webSocketCtx.socket.on('message', data => { - if (connectionClosed) { - return; - } - messageQueue = messageQueue.then(async () => { - await controllerReadyHandled; - if (!controllerReady || connectionClosed || webSocketCtx.socket.readyState !== WebSocket.OPEN) { - return; - } - await this.app.ctxStorage.run(eggCtx, async () => { - try { - const result = await invokeMethod(methodMeta, { data }); - await this.sendFetchResponseStream(result, webSocketCtx.socket, responseStreams, handleError); - } catch (error) { - await handleError(error); - } - }); - }).catch(error => { - this.app.logger.error('[tegg/websocket-fetch] handle message failed: %s', error.stack || error.message); - }); - }); - webSocketCtx.socket.on('error', error => { - this.app.ctxStorage.run(eggCtx, async () => { - await handleError(error); - }).catch(error => { - this.app.logger.error('[tegg/websocket-fetch] handle socket error failed: %s', error.stack || error.message); - }); - }); - - let lifecycleError: Error | undefined; - try { - await invokeMethod(controllerMeta.connectionMethod); - await invokeMethod(controllerMeta.openMethod); - controllerReady = true; - } catch (error) { - lifecycleError = error; - await handleError(error).catch(err => { - this.app.logger.error('[tegg/websocket-fetch] handle lifecycle error failed: %s', err.stack || err.message); - }); - if (webSocketCtx.socket.readyState === WebSocket.OPEN || webSocketCtx.socket.readyState === WebSocket.CONNECTING) { - webSocketCtx.socket.close(1011, 'Internal Server Error'); - } - } finally { - resolveControllerReady(); - } - await closeHandled; - if (lifecycleError) { - throw lifecycleError; - } - } - - private createFetchClose(webSocket: WebSocket): WebSocketFetchClose { - return (code = 1000, reason?: string | Buffer) => { - if (webSocket.readyState !== WebSocket.OPEN && webSocket.readyState !== WebSocket.CONNECTING) { - return; - } - webSocket.close(code, reason); - }; - } - - private async sendFetchResponseStream( - result: unknown, - webSocket: WebSocket, - responseStreams: Set, - handleError: (error: Error) => Promise, - ): Promise { - if (result === undefined || result === null) { - return; - } - if (!this.isReadableStream(result)) { - await handleError(new Error('WebSocketFetch method must return a readable stream or void')).catch(error => { - this.app.logger.error('[tegg/websocket-fetch] handle invalid response failed: %s', error.stack || error.message); - }); - return; - } - if (webSocket.readyState !== WebSocket.OPEN) { - this.destroyFetchResponseStream(result); - return; - } - responseStreams.add(result); - await new Promise(resolve => { - let settled = false; - const onData = (chunk: unknown) => { - this.sendFetchChunk(webSocket, chunk, handleError); - }; - let onEnd!: () => void; - let onClose!: () => void; - let onError!: (error: Error) => void; - const settle = (error?: Error) => { - if (settled) { - return; - } - settled = true; - result.removeListener('data', onData); - result.removeListener('end', onEnd); - result.removeListener('close', onClose); - result.removeListener('error', onError); - responseStreams.delete(result); - if (!error || webSocket.readyState !== WebSocket.OPEN) { - resolve(); - return; - } - handleError(error).catch(err => { - this.app.logger.error('[tegg/websocket-fetch] handle stream error failed: %s', err.stack || err.message); - }).finally(resolve); - }; - onEnd = () => settle(); - onClose = () => settle(); - onError = error => settle(error); - result.once('end', onEnd); - result.once('close', onClose); - result.once('error', onError); - result.on('data', onData); - const state = result as NodeJS.ReadableStream & { - closed?: boolean; - errored?: Error | null; - readableEnded?: boolean; - }; - if (state.errored) { - const streamError = state.errored; - setImmediate(() => settle(streamError)); - } else if (state.closed || state.readableEnded) { - settle(); - } - }); + await runtime.invokeFetchController(eggObj.obj, controllerMeta, methodMeta, events); } - private destroyFetchResponseStreams(responseStreams: Set) { - for (const responseStream of responseStreams) { - this.destroyFetchResponseStream(responseStream); - } - responseStreams.clear(); - } - - private destroyFetchResponseStream(responseStream: NodeJS.ReadableStream) { - const destroy = (responseStream as NodeJS.ReadableStream & { destroy?: () => void }).destroy; - if (typeof destroy !== 'function') { - return; - } - try { - Reflect.apply(destroy, responseStream, []); - } catch (error) { - this.app.logger.error('[tegg/websocket-fetch] destroy response stream failed: %s', error.stack || error.message); - } - } - - private sendFetchChunk( - webSocket: WebSocket, - chunk: unknown, - handleError: (error: Error) => Promise, - ) { - if (webSocket.readyState !== WebSocket.OPEN) { - return; - } - webSocket.send(chunk as any, error => { - if (!error) { - return; - } - handleError(error).catch(err => { - this.app.logger.error('[tegg/websocket-fetch] handle send error failed: %s', err.stack || err.message); - }); - }); - } - - private pipeResponseStream(result: unknown, webSocket: WebSocket, getWebSocketStream: () => Duplex) { - if (!this.isReadableStream(result)) { + private rejectSocket(socket: Duplex, status: number, message: string) { + if (socket.destroyed) { return; } - - pipeline(result, getWebSocketStream(), error => { - if (!error) { - return; - } - if (webSocket.readyState === WebSocket.CLOSED || webSocket.readyState === WebSocket.CLOSING) { - return; - } - this.app.logger.error('[tegg/websocket] pipe response stream failed: %s', error.stack || error.message); - webSocket.close(1011, 'Internal Server Error'); - }); + const body = message; + socket.end(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); } - private isReadableStream(result: unknown): result is NodeJS.ReadableStream { - return !!result && typeof (result as NodeJS.ReadableStream).pipe === 'function'; + private getSocketBytesWritten(socket: Duplex): number { + return (socket as Duplex & { bytesWritten?: number }).bytesWritten || 0; } - private waitWebSocketClose(webSocket: WebSocket): Promise { - if (webSocket.readyState === WebSocket.CLOSED || webSocket.readyState === WebSocket.CLOSING) { - return Promise.resolve(); - } - return new Promise(resolve => { - let settled = false; - let onClose!: () => void; - let onError!: (error: Error) => void; - const settle = () => { - if (settled) { - return; - } - settled = true; - webSocket.removeListener('close', onClose); - webSocket.removeListener('error', onError); - resolve(); - }; - onClose = () => settle(); - onError = error => { - this.app.logger.error('[tegg/websocket] socket error while waiting for close: %s', error.stack || error.message); - settle(); - }; - webSocket.once('close', onClose); - webSocket.once('error', onError); + private hasWebSocketOwner(socket: Duplex): boolean { + return Object.getOwnPropertySymbols(socket).some(symbol => { + return symbol.description === 'websocket' && Boolean(Reflect.get(socket, symbol)); }); } - - private rejectSocket(socket: Duplex, status: number, message: string) { - if (socket.destroyed) { - return; - } - const body = message; - socket.end(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); - } } diff --git a/plugin/controller/package.json b/plugin/controller/package.json index c0a9cb0f1..08a36af2a 100644 --- a/plugin/controller/package.json +++ b/plugin/controller/package.json @@ -56,6 +56,7 @@ "@eggjs/tegg-loader": "^3.84.2", "@eggjs/tegg-metadata": "^3.84.2", "@eggjs/tegg-runtime": "^3.84.2", + "@eggjs/tegg-websocket-runtime": "^3.84.2", "@modelcontextprotocol/sdk": "^1.23.0", "await-event": "^2.1.0", "content-type": "^1.0.5", diff --git a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts index 9ec24c27a..6217a7578 100644 --- a/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts @@ -1,5 +1,6 @@ import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; import { PassThrough, pipeline, type Readable } from 'node:stream'; +import type { Context as EggContext } from 'egg'; import type { RawData, WebSocket } from 'ws'; import { Context, @@ -7,7 +8,9 @@ import { HTTPParam, HTTPQueries, HTTPQuery, + Host, Inject, + Middleware, Request, WebSocketContext, WebSocketController, @@ -40,6 +43,14 @@ interface WebSocketFetchRequest { pipeline?: boolean; } +async function webSocketShortCircuit(ctx: EggContext) { + const webSocketCtx = ctx as WebSocketContext; + webSocketCtx.webSocket.send(JSON.stringify({ + type: 'middleware', + path: ctx.path, + })); +} + @WebSocketController({ path: '/ws', }) @@ -74,6 +85,8 @@ export class AppWebSocketController { header: headers['x-client-id'], url: request.url, path: ctx.path, + sameSocket: ctx.webSocket === socket, + sameRequest: ctx.req === request, pid: process.pid, })); @@ -114,6 +127,47 @@ export class AppWebSocketController { return output; } + + @Host('proxy.example.com') + @WebSocketMethod({ + path: '/proxy', + }) + proxy(@Context() ctx: WebSocketContext) { + ctx.webSocket.send(JSON.stringify({ + type: 'proxy', + host: ctx.host, + protocol: ctx.protocol, + })); + } + + @WebSocketMethod({ + path: '/optional/:id?', + }) + optional( + @HTTPParam() id: string | undefined, + @WebSocketSocket() socket: WebSocket, + ) { + socket.send(JSON.stringify({ + type: 'optional', + id: id ?? null, + })); + } + + @Middleware(webSocketShortCircuit) + @WebSocketMethod({ + path: '/middleware-short-circuit', + }) + middlewareShortCircuit() { + throw new Error('middleware should not invoke this method'); + } + + @WebSocketMethod({ + path: '/timeout', + timeout: 20, + }) + async timeout() { + await new Promise(resolve => setTimeout(resolve, 100)); + } } @WebSocketFetchController({ @@ -132,7 +186,7 @@ export class AppWebSocketFetchController { if (id === 'connection-error') { throw new Error('fetch connection error'); } - ctx.socket.send(JSON.stringify({ + ctx.webSocket.send(JSON.stringify({ type: 'connection', header: headers['x-client-id'], path: ctx.path, @@ -144,7 +198,7 @@ export class AppWebSocketFetchController { onOpen( @Context() ctx: WebSocketContext, ) { - ctx.socket.send(JSON.stringify({ + ctx.webSocket.send(JSON.stringify({ type: 'open', path: ctx.path, pid: process.pid, @@ -232,7 +286,7 @@ export class AppWebSocketFetchController { @HTTPHeaders() headers: IncomingHttpHeaders, @Context() ctx: WebSocketContext, ) { - ctx.socket.send(JSON.stringify({ + ctx.webSocket.send(JSON.stringify({ type: 'error', message: error.message, header: headers['x-client-id'], diff --git a/plugin/controller/test/fixtures/apps/controller-app/config/config.default.js b/plugin/controller/test/fixtures/apps/controller-app/config/config.default.js index 31ec8ab4f..581037f79 100644 --- a/plugin/controller/test/fixtures/apps/controller-app/config/config.default.js +++ b/plugin/controller/test/fixtures/apps/controller-app/config/config.default.js @@ -3,6 +3,8 @@ module.exports = function() { const config = { keys: 'test key', + proxy: true, + hostHeaders: 'x-forwarded-host', security: { csrf: { ignoreJSON: false, diff --git a/plugin/controller/test/websocket/websocket.test.ts b/plugin/controller/test/websocket/websocket.test.ts index e1e2380d4..ec8ae66af 100644 --- a/plugin/controller/test/websocket/websocket.test.ts +++ b/plugin/controller/test/websocket/websocket.test.ts @@ -200,6 +200,8 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { header: 'client-1', url: '/ws/echo/foo?name=bar&tag=a&tag=b', path: '/ws/echo/foo', + sameSocket: true, + sameRequest: true, pid: ready.pid, }); assert.equal(typeof ready.pid, 'number'); @@ -242,6 +244,69 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { } }); + it('should use Egg proxy host and protocol for websocket routing', async () => { + const { socket, firstMessage } = await createClientWithFirstMessage( + requestUrl(app, '/ws/proxy'), + { + 'x-forwarded-host': 'proxy.example.com', + 'x-forwarded-proto': 'https', + }, + ); + try { + assert.deepEqual(await firstMessage, { + type: 'proxy', + host: 'proxy.example.com', + protocol: 'https', + }); + } finally { + await closeClient(socket); + } + }); + + it('should keep an omitted optional websocket path parameter undefined', async () => { + const { socket, firstMessage } = await createClientWithFirstMessage( + requestUrl(app, '/ws/optional'), + ); + try { + assert.deepEqual(await firstMessage, { + type: 'optional', + id: null, + }); + } finally { + await closeClient(socket); + } + }); + + it('should reject malformed websocket path encoding with 400', async () => { + await assert.rejects( + () => createClient(requestUrl(app, '/ws/echo/%E0%A4%A')), + /Unexpected server response: 400/, + ); + }); + + it('should allow websocket middleware to short circuit without closing', async () => { + const { socket, firstMessage } = await createClientWithFirstMessage( + requestUrl(app, '/ws/middleware-short-circuit'), + ); + try { + assert.deepEqual(await firstMessage, { + type: 'middleware', + path: '/ws/middleware-short-circuit', + }); + assert.equal(socket.readyState, WebSocket.OPEN); + } finally { + await closeClient(socket); + } + }); + + it('should close a timed out websocket controller method', async () => { + const socket = await createClient(requestUrl(app, '/ws/timeout')); + assert.deepEqual(await receiveClose(socket), { + code: 1011, + reason: 'Internal Server Error', + }); + }); + it('should handle websocket fetch controller', async () => { const socket = new WebSocket(requestUrl(app, '/ws-fetch/fetch-1?name=bar&tag=a&tag=b'), { headers: { 'x-client-id': 'client-fetch' }, @@ -606,4 +671,11 @@ describe('plugin/controller/test/websocket/websocket.test.ts', () => { app.server.removeListener('upgrade', handleUpgrade); } }); + + it('should reject an unmatched websocket route when no other listener handles it', async () => { + await assert.rejects( + () => createClient(requestUrl(app, '/ws/not-found-without-listener')), + /Unexpected server response: 404/, + ); + }); }); diff --git a/plugin/controller/test/websocket/websocketCluster.test.ts b/plugin/controller/test/websocket/websocketCluster.test.ts index 3deac95de..98629ff71 100644 --- a/plugin/controller/test/websocket/websocketCluster.test.ts +++ b/plugin/controller/test/websocket/websocketCluster.test.ts @@ -152,7 +152,7 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { for (const [ index, echo ] of echoMessages.entries()) { assert.equal(echo.type, 'echo'); assert.equal(echo.data, `hello-${index}`); - assert.equal(typeof echo.pid, 'number'); + assert.equal(echo.pid, readyMessages[index].pid); } } finally { await Promise.all(clients.map(socket => closeClient(socket))); @@ -179,7 +179,7 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { for (const [ index, echo ] of echoMessages.entries()) { assert.equal(echo.type, 'stream'); assert.equal(echo.data, `stream-${index}`); - assert.equal(typeof echo.pid, 'number'); + assert.equal(echo.pid, readyMessages[index].pid); } } finally { await Promise.all(clients.map(socket => closeClient(socket))); @@ -205,7 +205,7 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { const open = await receiver.next(); assert.equal(open.type, 'open'); assert.equal(open.path, '/ws-fetch/cluster-fetch'); - assert.equal(typeof open.pid, 'number'); + assert.equal(open.pid, connection.pid); socket.send(JSON.stringify({ content: 'cluster-first', @@ -219,6 +219,7 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { assert.equal(first.header, 'cluster-fetch-client'); assert.equal(first.path, '/ws-fetch/cluster-fetch'); assert.equal(first.content, 'cluster-first'); + assert.equal(first.pid, connection.pid); assert.deepEqual(await receiver.next(), { type: 'data', phase: 2, @@ -237,13 +238,14 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { assert.equal(beforeError.type, 'data'); assert.equal(beforeError.phase, 1); assert.equal(beforeError.content, 'cluster-broken'); + assert.equal(beforeError.pid, connection.pid); const error = await receiver.next(); assert.deepEqual(error, { type: 'error', message: 'fetch error: cluster-broken', header: 'cluster-fetch-client', path: '/ws-fetch/cluster-fetch', - pid: error.pid, + pid: connection.pid, }); assert.equal(socket.readyState, WebSocket.OPEN); @@ -256,6 +258,7 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { assert.equal(closeData.type, 'data'); assert.equal(closeData.phase, 1); assert.equal(closeData.content, 'cluster-done'); + assert.equal(closeData.pid, connection.pid); assert.deepEqual(await receiver.next(), { type: 'data', phase: 2, @@ -281,8 +284,9 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { }); try { - await receiver.next(); - await receiver.next(); + const connection = await receiver.next(); + const open = await receiver.next(); + assert.equal(open.pid, connection.pid); socket.send(JSON.stringify({ content: 'first', @@ -303,6 +307,7 @@ describe('plugin/controller/test/websocket/websocketCluster.test.ts', () => { messages.map(message => `${message.content}:${message.phase}`), [ 'first:1', 'first:2', 'second:1', 'second:2' ], ); + assert.deepEqual(new Set(messages.map(message => message.pid)), new Set([ connection.pid ])); assert.equal(socket.readyState, WebSocket.OPEN); } finally { await closeClient(socket); diff --git a/standalone/service-worker/package.json b/standalone/service-worker/package.json index 9e559e1de..e5ab1e335 100644 --- a/standalone/service-worker/package.json +++ b/standalone/service-worker/package.json @@ -49,6 +49,7 @@ "@eggjs/tegg-metadata": "^3.84.2", "@eggjs/tegg-standalone": "^3.84.2", "@eggjs/tegg-types": "^3.84.2", + "@eggjs/tegg-websocket-runtime": "^3.84.2", "@modelcontextprotocol/sdk": "1.24.3", "egg-errors": "^2.3.0", "path-to-regexp": "^1.8.0", diff --git a/standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts b/standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts index 79314d544..29cc8883c 100644 --- a/standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts +++ b/standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts @@ -1,9 +1,9 @@ import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; import type { WebSocket } from 'ws'; -import type { WebSocketContext } from '@eggjs/tegg'; -export class ServiceWorkerWebSocketContext implements WebSocketContext { +export class ServiceWorkerWebSocketContext { readonly socket: WebSocket; + readonly webSocket: WebSocket; readonly request: IncomingMessage; readonly url: URL; readonly method = 'GET'; @@ -21,6 +21,7 @@ export class ServiceWorkerWebSocketContext implements WebSocketContext; }) { this.socket = options.socket; + this.webSocket = options.socket; this.request = options.request; this.url = options.url; this.path = options.url.pathname; diff --git a/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts index 16906e8cb..445b70615 100644 --- a/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts +++ b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts @@ -1,7 +1,6 @@ import assert from 'node:assert'; import type { IncomingMessage } from 'node:http'; -import { Duplex, pipeline } from 'node:stream'; -import pathToRegexp from 'path-to-regexp'; +import { Duplex } from 'node:stream'; import { Router as KoaRouter } from '@eggjs/router'; import { FrameworkErrorFormater } from 'egg-errors'; import { createWebSocketStream, RawData, WebSocket, WebSocketServer } from 'ws'; @@ -12,16 +11,25 @@ import { EggContext, Next, WebSocketControllerMeta, - WebSocketFetchClose, WebSocketFetchControllerMeta, WebSocketFetchMethodMeta, WebSocketMethodMeta, - WebSocketParamType, - WebSocketPathParamMeta, - WebSocketQueriesParamMeta, - WebSocketQueryParamMeta, } from '@eggjs/tegg'; import { EggContainerFactory, EggPrototype } from '@eggjs/tegg/helper'; +import { + WebSocketControllerRuntime, + WebSocketEventStream, + WebSocketRoute, + createWebSocketRoutes, + getWebSocketMethodHosts, + getWebSocketMethodMiddlewares, + getWebSocketMethodName, + getWebSocketMethodRealPath, + matchWebSocketRoute, + waitForWebSocketClose, + WEBSOCKET_INTERNAL_ERROR_CODE, + WEBSOCKET_INTERNAL_ERROR_REASON, +} from '@eggjs/tegg-websocket-runtime'; import { WebSocketUpgradeEvent } from '@eggjs/tegg-types/standalone'; import { ControllerRegister } from '../controller/ControllerRegister'; import { RootProtoManager } from '../controller/RootProtoManager'; @@ -31,26 +39,6 @@ const noop = () => { // noop }; -interface WebSocketRoute { - controllerProto: EggPrototype; - controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta; - methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta; - methodRealPath: string; - methodName: string; - host?: string; - keys: pathToRegexp.Key[]; - regexp: RegExp; -} - -interface WebSocketMethodPayload { - data?: RawData; - close?: WebSocketFetchClose; - error?: Error; - closeCode?: number; - closeReason?: Buffer; - getWebSocketStream?: () => Duplex; -} - type Middleware = (ctx: ServiceWorkerWebSocketContext, next: Next) => Promise; export class WebSocketControllerRegister implements ControllerRegister { @@ -59,7 +47,7 @@ export class WebSocketControllerRegister implements ControllerRegister { private readonly checkRouters: Map; private readonly webSocketServer: WebSocketServer; private controllerProtos: EggPrototype[] = []; - private routes: WebSocketRoute[] = []; + private routes: Array> = []; static create(proto: EggPrototype, controllerMeta: ControllerMetadata) { assert( @@ -89,28 +77,11 @@ export class WebSocketControllerRegister implements ControllerRegister { } doRegister(rootProtoManager: RootProtoManager) { - const methodMap = new Map(); - for (const proto of this.controllerProtos) { - const metadata = proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; - for (const method of metadata.methods) { - methodMap.set(method, proto); - } - } - const allMethods = Array.from(methodMap.keys()) - .sort((a, b) => b.priority - a.priority); - - for (const method of allMethods) { - const controllerProto = methodMap.get(method)!; - const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; - this.checkDuplicate(controllerMeta, method); - } - - this.routes = allMethods.flatMap(method => { - const controllerProto = methodMap.get(method)!; - const controllerMeta = controllerProto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta; - const hosts = this.getMethodHosts(controllerMeta, method) || [ undefined ]; - return hosts.map(host => this.createRoute(controllerProto, controllerMeta, method, host)); - }); + this.routes = createWebSocketRoutes( + this.controllerProtos, + proto => proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta, + (controllerMeta, methodMeta) => this.checkDuplicate(controllerMeta, methodMeta), + ); for (const route of this.routes) { rootProtoManager.registerRootProto('GET', (ctx: EggContext) => { @@ -130,7 +101,17 @@ export class WebSocketControllerRegister implements ControllerRegister { async handleUpgrade(event: WebSocketUpgradeEvent): Promise { const url = this.createURL(event.request); - const matched = this.matchRoute(event.request, url); + let matched: { route: WebSocketRoute; params: Record } | undefined; + try { + matched = matchWebSocketRoute(this.routes, url.pathname, event.request.headers.host); + } catch (error) { + if (error instanceof URIError) { + console.warn('[tegg/websocket] reject malformed request path:', event.request.url); + WebSocketControllerRegister.rejectSocket(event.socket, 400, 'Bad Request'); + return true; + } + throw error; + } if (!matched) { return false; } @@ -142,7 +123,7 @@ export class WebSocketControllerRegister implements ControllerRegister { .catch(error => { console.error('[tegg/websocket] handle connection failed:', error); if (webSocket.readyState === WebSocket.OPEN || webSocket.readyState === WebSocket.CONNECTING) { - webSocket.close(1011, 'Internal Server Error'); + webSocket.close(WEBSOCKET_INTERNAL_ERROR_CODE, WEBSOCKET_INTERNAL_ERROR_REASON); } resolve(); }); @@ -159,30 +140,9 @@ export class WebSocketControllerRegister implements ControllerRegister { socket.end(`HTTP/1.1 ${status} ${message}\r\nConnection: close\r\nContent-Length: ${Buffer.byteLength(body)}\r\n\r\n${body}`); } - private createRoute( - controllerProto: EggPrototype, - controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, - methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, - host: string | undefined, - ): WebSocketRoute { - const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); - const keys: pathToRegexp.Key[] = []; - const regexp = pathToRegexp(methodRealPath, keys, { sensitive: true }); - return { - controllerProto, - controllerMeta, - methodMeta, - methodRealPath, - methodName: this.getMethodName(controllerMeta, methodMeta), - host, - keys, - regexp, - }; - } - private checkDuplicate(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { let router = this.checkRouters.get('default')!; - const hosts = this.getMethodHosts(controllerMeta, methodMeta) || []; + const hosts = getWebSocketMethodHosts(controllerMeta, methodMeta) || []; if (!hosts.length) { this.checkDuplicateInRouter(router, controllerMeta, methodMeta); this.registerToRouter(router, controllerMeta, methodMeta); @@ -201,15 +161,15 @@ export class WebSocketControllerRegister implements ControllerRegister { } private registerToRouter(router: KoaRouter, controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { - const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); - const methodName = this.getMethodName(controllerMeta, methodMeta); + const methodRealPath = getWebSocketMethodRealPath(controllerMeta, methodMeta); + const methodName = getWebSocketMethodName(controllerMeta, methodMeta); Reflect.apply(router.get, router, [ methodName, methodRealPath, noop ]); } private checkDuplicateInRouter(router: KoaRouter, controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { - const methodRealPath = this.getMethodRealPath(controllerMeta, methodMeta); + const methodRealPath = getWebSocketMethodRealPath(controllerMeta, methodMeta); const matched = router.match(methodRealPath, 'GET'); - const methodName = this.getMethodName(controllerMeta, methodMeta); + const methodName = getWebSocketMethodName(controllerMeta, methodMeta); if (matched.route) { const [ layer ] = matched.path; const err = new Error(`register websocket controller ${methodName} failed, ${controllerMeta.type} ${methodRealPath} is conflict with exists rule ${layer.path}`); @@ -217,84 +177,72 @@ export class WebSocketControllerRegister implements ControllerRegister { } } - private getMethodRealPath(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { - return (controllerMeta as any).getMethodRealPath(methodMeta); - } - - private getMethodHosts(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): string[] | undefined { - return (controllerMeta as any).getMethodHosts(methodMeta); - } - - private getMethodName(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): string { - return (controllerMeta as any).getMethodName(methodMeta); - } - - private getMethodMiddlewares(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta): Middleware[] { - return (controllerMeta as any).getMethodMiddlewares(methodMeta); - } - private createURL(req: IncomingMessage): URL { const host = req.headers.host || 'localhost'; return new URL(req.url || '/', `http://${host}`); } - private matchRoute(req: IncomingMessage, url: URL): { route: WebSocketRoute; params: Record } | undefined { - for (const route of this.routes) { - if (route.host && !this.matchHost(route.host, req.headers.host)) { - continue; - } - const matched = route.regexp.exec(url.pathname); - if (!matched) { - continue; - } - const params: Record = {}; - route.keys.forEach((key, index) => { - params[String(key.name)] = decodeURIComponent(matched[index + 1]); - }); - return { route, params }; - } - } - - private matchHost(expectedHost: string, requestHost: string | undefined) { - if (!requestHost) { - return false; - } - return requestHost === expectedHost || requestHost.split(':')[0] === expectedHost; - } - private async handleConnection( - route: WebSocketRoute, + route: WebSocketRoute, params: Record, url: URL, req: IncomingMessage, webSocket: WebSocket, ) { + const closeHandled = waitForWebSocketClose(webSocket, error => { + console.error('[tegg/websocket] socket error while waiting for close:', error); + }); const webSocketCtx = new ServiceWorkerWebSocketContext({ socket: webSocket, request: req, url, params, }); - const methodMiddlewares = this.getMethodMiddlewares(route.controllerMeta, route.methodMeta); + const fetchEvents = route.controllerMeta.type === ControllerType.WEBSOCKET_FETCH + ? new WebSocketEventStream(webSocket) + : undefined; + const runtime = new WebSocketControllerRuntime({ + context: { + context: webSocketCtx, + socket: webSocket, + params: webSocketCtx.params, + query: webSocketCtx.query, + queries: webSocketCtx.queries, + headers: webSocketCtx.headers, + request: webSocketCtx.request, + }, + createWebSocketStream, + logger: { + debug: message => console.debug(`[tegg/websocket] ${message}`), + error: (message, error) => console.error(`[tegg/websocket] ${message}:`, error || ''), + }, + }); + const methodMiddlewares = getWebSocketMethodMiddlewares( + route.controllerMeta, + route.methodMeta, + ) as unknown as Middleware[]; let invoked = false; const handler = async (_ctx: ServiceWorkerWebSocketContext, next: Next) => { invoked = true; if (route.controllerMeta.type === ControllerType.WEBSOCKET_FETCH) { - await this.invokeFetchController(route, webSocketCtx); + await this.invokeFetchController(route, runtime, fetchEvents!); } else { - await this.invokeController(route, webSocketCtx); + await this.invokeController(route, runtime); } await next(); }; - await this.compose([ ...methodMiddlewares, handler ])(webSocketCtx, async () => { - // final middleware - }); - if (!invoked) { - webSocket.close(1008, 'WebSocket middleware did not call next'); - return; + try { + await this.compose([ ...methodMiddlewares, handler ])(webSocketCtx, async () => { + // final middleware + }); + if (!invoked) { + console.debug(`[tegg/websocket] ${route.methodName} was short-circuited by middleware`); + } + await closeHandled; + } finally { + fetchEvents?.dispose(); } - await this.waitWebSocketClose(webSocket); } private compose(middlewares: Middleware[]) { @@ -315,350 +263,25 @@ export class WebSocketControllerRegister implements ControllerRegister { }; } - private async invokeController(route: WebSocketRoute, webSocketCtx: ServiceWorkerWebSocketContext) { + private async invokeController( + route: WebSocketRoute, + runtime: WebSocketControllerRuntime, + ) { + const controllerMeta = route.controllerMeta as WebSocketControllerMeta; const methodMeta = route.methodMeta as WebSocketMethodMeta; const eggObj = await EggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); - const realObj = eggObj.obj; - const realMethod = realObj[methodMeta.name]; - let webSocketStream: Duplex | undefined; - const getWebSocketStream = () => { - if (!webSocketStream) { - webSocketStream = createWebSocketStream(webSocketCtx.socket); - } - return webSocketStream; - }; - const args = this.buildMethodArgs(methodMeta, webSocketCtx, { getWebSocketStream }); - const result = await Reflect.apply(realMethod, realObj, args); - this.pipeResponseStream(result, webSocketCtx.socket, getWebSocketStream); + await runtime.invokeController(eggObj.obj, controllerMeta, methodMeta); } - private buildMethodArgs( - methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta, - webSocketCtx: ServiceWorkerWebSocketContext, - payload: WebSocketMethodPayload, + private async invokeFetchController( + route: WebSocketRoute, + runtime: WebSocketControllerRuntime, + events: WebSocketEventStream, ) { - const argsLength = methodMeta.paramMap.size; - const hasContext = methodMeta.contextParamIndex !== undefined; - const contextIndex = methodMeta.contextParamIndex; - const methodArgsLength = argsLength + (hasContext ? 1 : 0); - const args: unknown[] = new Array(methodArgsLength); - if (hasContext) { - args[contextIndex!] = webSocketCtx; - } - for (const [ index, param ] of methodMeta.paramMap) { - switch (param.type) { - case WebSocketParamType.PARAM: { - const pathParam = param as WebSocketPathParamMeta; - args[index] = webSocketCtx.params[pathParam.name]; - break; - } - case WebSocketParamType.QUERY: { - const queryParam = param as WebSocketQueryParamMeta; - args[index] = webSocketCtx.query[queryParam.name]; - break; - } - case WebSocketParamType.QUERIES: { - const queryParam = param as WebSocketQueriesParamMeta; - args[index] = webSocketCtx.queries[queryParam.name] || []; - break; - } - case WebSocketParamType.HEADERS: { - args[index] = webSocketCtx.headers; - break; - } - case WebSocketParamType.REQUEST: { - args[index] = webSocketCtx.request; - break; - } - case WebSocketParamType.SOCKET: { - args[index] = webSocketCtx.socket; - break; - } - case WebSocketParamType.STREAM: { - assert(payload.getWebSocketStream, '@WebSocketStream can not be used here'); - args[index] = payload.getWebSocketStream(); - break; - } - case WebSocketParamType.DATA: { - args[index] = payload.data; - break; - } - case WebSocketParamType.CLOSE: { - args[index] = payload.close; - break; - } - case WebSocketParamType.ERROR: { - args[index] = payload.error; - break; - } - case WebSocketParamType.CLOSE_CODE: { - args[index] = payload.closeCode; - break; - } - case WebSocketParamType.CLOSE_REASON: { - args[index] = payload.closeReason; - break; - } - default: - assert.fail('never arrive'); - } - } - return args; - } - - private async invokeFetchController(route: WebSocketRoute, webSocketCtx: ServiceWorkerWebSocketContext) { const controllerMeta = route.controllerMeta as WebSocketFetchControllerMeta; const methodMeta = route.methodMeta as WebSocketFetchMethodMeta; const eggObj = await EggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); - const realObj = eggObj.obj; - const close = this.createFetchClose(webSocketCtx.socket); - const responseStreams = new Set(); - let connectionClosed = false; - let controllerReady = false; - let resolveControllerReady!: () => void; - const controllerReadyHandled = new Promise(resolve => { - resolveControllerReady = resolve; - }); - let messageQueue = Promise.resolve(); - const invokeMethod = async ( - targetMethodMeta: WebSocketFetchMethodMeta | undefined, - payload: WebSocketMethodPayload = {}, - ) => { - if (!targetMethodMeta) { - return; - } - const realMethod = realObj[targetMethodMeta.name]; - const args = this.buildMethodArgs(targetMethodMeta, webSocketCtx, { - close, - ...payload, - }); - return await Reflect.apply(realMethod, realObj, args); - }; - const handleError = async (error: Error) => { - if (!controllerMeta.errorMethod) { - console.error('[tegg/websocket-fetch] handle error:', error); - return; - } - await invokeMethod(controllerMeta.errorMethod, { error }); - }; - const closeHandled = new Promise(resolve => { - webSocketCtx.socket.once('close', (code, reason) => { - connectionClosed = true; - this.destroyFetchResponseStreams(responseStreams); - invokeMethod(controllerMeta.closeMethod, { - closeCode: code, - closeReason: reason, - }).catch(error => { - console.error('[tegg/websocket-fetch] handle close failed:', error); - }).finally(resolve); - }); - }); - - webSocketCtx.socket.on('message', data => { - if (connectionClosed) { - return; - } - messageQueue = messageQueue.then(async () => { - await controllerReadyHandled; - if (!controllerReady || connectionClosed || webSocketCtx.socket.readyState !== WebSocket.OPEN) { - return; - } - try { - const result = await invokeMethod(methodMeta, { data }); - await this.sendFetchResponseStream(result, webSocketCtx.socket, responseStreams, handleError); - } catch (error) { - await handleError(error); - } - }).catch(error => { - console.error('[tegg/websocket-fetch] handle message error failed:', error); - }); - }); - webSocketCtx.socket.on('error', error => { - handleError(error).catch(err => { - console.error('[tegg/websocket-fetch] handle socket error failed:', err); - }); - }); - - let lifecycleError: Error | undefined; - try { - await invokeMethod(controllerMeta.connectionMethod); - await invokeMethod(controllerMeta.openMethod); - controllerReady = true; - } catch (error) { - lifecycleError = error; - await handleError(error).catch(err => { - console.error('[tegg/websocket-fetch] handle lifecycle error failed:', err); - }); - if (webSocketCtx.socket.readyState === WebSocket.OPEN || webSocketCtx.socket.readyState === WebSocket.CONNECTING) { - webSocketCtx.socket.close(1011, 'Internal Server Error'); - } - } finally { - resolveControllerReady(); - } - await closeHandled; - if (lifecycleError) { - throw lifecycleError; - } - } - - private createFetchClose(webSocket: WebSocket): WebSocketFetchClose { - return (code = 1000, reason?: string | Buffer) => { - if (webSocket.readyState !== WebSocket.OPEN && webSocket.readyState !== WebSocket.CONNECTING) { - return; - } - webSocket.close(code, reason); - }; + await runtime.invokeFetchController(eggObj.obj, controllerMeta, methodMeta, events); } - private async sendFetchResponseStream( - result: unknown, - webSocket: WebSocket, - responseStreams: Set, - handleError: (error: Error) => Promise, - ): Promise { - if (result === undefined || result === null) { - return; - } - if (!this.isReadableStream(result)) { - await handleError(new Error('WebSocketFetch method must return a readable stream or void')).catch(error => { - console.error('[tegg/websocket-fetch] handle invalid response failed:', error); - }); - return; - } - if (webSocket.readyState !== WebSocket.OPEN) { - this.destroyFetchResponseStream(result); - return; - } - responseStreams.add(result); - await new Promise(resolve => { - let settled = false; - const onData = (chunk: unknown) => { - this.sendFetchChunk(webSocket, chunk, handleError); - }; - let onEnd!: () => void; - let onClose!: () => void; - let onError!: (error: Error) => void; - const settle = (error?: Error) => { - if (settled) { - return; - } - settled = true; - result.removeListener('data', onData); - result.removeListener('end', onEnd); - result.removeListener('close', onClose); - result.removeListener('error', onError); - responseStreams.delete(result); - if (!error || webSocket.readyState !== WebSocket.OPEN) { - resolve(); - return; - } - handleError(error).catch(err => { - console.error('[tegg/websocket-fetch] handle stream error failed:', err); - }).finally(resolve); - }; - onEnd = () => settle(); - onClose = () => settle(); - onError = error => settle(error); - result.once('end', onEnd); - result.once('close', onClose); - result.once('error', onError); - result.on('data', onData); - const state = result as NodeJS.ReadableStream & { - closed?: boolean; - errored?: Error | null; - readableEnded?: boolean; - }; - if (state.errored) { - const streamError = state.errored; - setImmediate(() => settle(streamError)); - } else if (state.closed || state.readableEnded) { - settle(); - } - }); - } - - private destroyFetchResponseStreams(responseStreams: Set) { - for (const responseStream of responseStreams) { - this.destroyFetchResponseStream(responseStream); - } - responseStreams.clear(); - } - - private destroyFetchResponseStream(responseStream: NodeJS.ReadableStream) { - const destroy = (responseStream as NodeJS.ReadableStream & { destroy?: () => void }).destroy; - if (typeof destroy !== 'function') { - return; - } - try { - Reflect.apply(destroy, responseStream, []); - } catch (error) { - console.error('[tegg/websocket-fetch] destroy response stream failed:', error); - } - } - - private sendFetchChunk( - webSocket: WebSocket, - chunk: unknown, - handleError: (error: Error) => Promise, - ) { - if (webSocket.readyState !== WebSocket.OPEN) { - return; - } - webSocket.send(chunk as any, error => { - if (!error) { - return; - } - handleError(error).catch(err => { - console.error('[tegg/websocket-fetch] handle send error failed:', err); - }); - }); - } - - private pipeResponseStream(result: unknown, webSocket: WebSocket, getWebSocketStream: () => Duplex) { - if (!this.isReadableStream(result)) { - return; - } - - pipeline(result, getWebSocketStream(), error => { - if (!error) { - return; - } - if (webSocket.readyState === WebSocket.CLOSED || webSocket.readyState === WebSocket.CLOSING) { - return; - } - console.error('[tegg/websocket] pipe response stream failed:', error); - webSocket.close(1011, 'Internal Server Error'); - }); - } - - private isReadableStream(result: unknown): result is NodeJS.ReadableStream { - return !!result && typeof (result as NodeJS.ReadableStream).pipe === 'function'; - } - - private waitWebSocketClose(webSocket: WebSocket): Promise { - if (webSocket.readyState === WebSocket.CLOSED || webSocket.readyState === WebSocket.CLOSING) { - return Promise.resolve(); - } - return new Promise(resolve => { - let settled = false; - let onClose!: () => void; - let onError!: (error: Error) => void; - const settle = () => { - if (settled) { - return; - } - settled = true; - webSocket.removeListener('close', onClose); - webSocket.removeListener('error', onError); - resolve(); - }; - onClose = () => settle(); - onError = error => { - console.error('[tegg/websocket] socket error while waiting for close:', error); - settle(); - }; - webSocket.once('close', onClose); - webSocket.once('error', onError); - }); - } } diff --git a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts index 1042b2612..628b63103 100644 --- a/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts +++ b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts @@ -1,5 +1,6 @@ import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; import { PassThrough, pipeline, type Readable } from 'node:stream'; +import type { Context as EggContext } from 'egg'; import type { RawData, WebSocket } from 'ws'; import { Context, @@ -7,11 +8,11 @@ import { HTTPParam, HTTPQueries, HTTPQuery, + Middleware, Request, WebSocketClose, WebSocketCloseCode, WebSocketCloseReason, - WebSocketContext, WebSocketController, WebSocketData, WebSocketError, @@ -26,6 +27,7 @@ import { WebSocketSocket, WebSocketStream, } from '@eggjs/tegg'; +import type { ServiceWorkerWebSocketContext } from '../../../src/websocket/ServiceWorkerWebSocketContext'; interface WebSocketFetchRequest { content: string; @@ -40,6 +42,14 @@ interface WebSocketFetchRequest { export const fetchCloseEvents: string[] = []; export const fetchStreamCloseEvents: string[] = []; +async function webSocketShortCircuit(ctx: EggContext) { + const webSocketCtx = ctx as unknown as ServiceWorkerWebSocketContext; + webSocketCtx.webSocket.send(JSON.stringify({ + type: 'middleware', + path: ctx.path, + })); +} + @WebSocketController({ path: '/ws', }) @@ -54,7 +64,7 @@ export class StandaloneWebSocketController { @HTTPQueries({ name: 'tag' }) tags: string[], @HTTPHeaders() headers: IncomingHttpHeaders, @Request() request: IncomingMessage, - @Context() ctx: WebSocketContext, + @Context() ctx: ServiceWorkerWebSocketContext, ) { socket.send(JSON.stringify({ type: 'ready', @@ -64,6 +74,7 @@ export class StandaloneWebSocketController { header: headers['x-client-id'], url: request.url, path: ctx.path, + sameSocket: ctx.webSocket === socket, })); socket.on('message', data => { @@ -100,6 +111,35 @@ export class StandaloneWebSocketController { return output; } + + @WebSocketMethod({ + path: '/optional/:id?', + }) + optional( + @HTTPParam() id: string | undefined, + @WebSocketSocket() socket: WebSocket, + ) { + socket.send(JSON.stringify({ + type: 'optional', + id: id ?? null, + })); + } + + @Middleware(webSocketShortCircuit) + @WebSocketMethod({ + path: '/middleware-short-circuit', + }) + middlewareShortCircuit() { + throw new Error('middleware should not invoke this method'); + } + + @WebSocketMethod({ + path: '/timeout', + timeout: 20, + }) + async timeout() { + await new Promise(resolve => setTimeout(resolve, 100)); + } } @WebSocketFetchController({ @@ -110,12 +150,12 @@ export class StandaloneWebSocketFetchController { onConnection( @HTTPParam() id: string, @HTTPHeaders() headers: IncomingHttpHeaders, - @Context() ctx: WebSocketContext, + @Context() ctx: ServiceWorkerWebSocketContext, ) { if (id === 'connection-error') { throw new Error('fetch connection error'); } - ctx.socket.send(JSON.stringify({ + ctx.webSocket.send(JSON.stringify({ type: 'connection', header: headers['x-client-id'], path: ctx.path, @@ -124,9 +164,9 @@ export class StandaloneWebSocketFetchController { @WebSocketFetchOnOpen() onOpen( - @Context() ctx: WebSocketContext, + @Context() ctx: ServiceWorkerWebSocketContext, ) { - ctx.socket.send(JSON.stringify({ + ctx.webSocket.send(JSON.stringify({ type: 'open', path: ctx.path, })); @@ -140,7 +180,7 @@ export class StandaloneWebSocketFetchController { @HTTPQuery() name: string, @HTTPQueries({ name: 'tag' }) tags: string[], @HTTPHeaders() headers: IncomingHttpHeaders, - @Context() ctx: WebSocketContext, + @Context() ctx: ServiceWorkerWebSocketContext, ) { const body = JSON.parse(data.toString()) as WebSocketFetchRequest; const source = new PassThrough(); @@ -203,9 +243,9 @@ export class StandaloneWebSocketFetchController { onError( @WebSocketError() error: Error, @HTTPHeaders() headers: IncomingHttpHeaders, - @Context() ctx: WebSocketContext, + @Context() ctx: ServiceWorkerWebSocketContext, ) { - ctx.socket.send(JSON.stringify({ + ctx.webSocket.send(JSON.stringify({ type: 'error', message: error.message, header: headers['x-client-id'], diff --git a/standalone/service-worker/test/websocket/websocket.test.ts b/standalone/service-worker/test/websocket/websocket.test.ts index 131f02ae0..5e963a330 100644 --- a/standalone/service-worker/test/websocket/websocket.test.ts +++ b/standalone/service-worker/test/websocket/websocket.test.ts @@ -44,6 +44,7 @@ describe('standalone/service-worker/test/websocket/websocket.test.ts', () => { header: 'standalone-client', url: '/ws/echo/123?name=tegg&tag=a&tag=b', path: '/ws/echo/123', + sameSocket: true, }); ws.send('hello'); @@ -90,6 +91,45 @@ describe('standalone/service-worker/test/websocket/websocket.test.ts', () => { await closeClient(ws); }); + it('should handle optional websocket path parameters in standalone service worker', async () => { + const ws = createClient('/ws/optional'); + assert.deepStrictEqual(await waitJSONMessage(ws), { + type: 'optional', + id: null, + }); + await closeClient(ws); + }); + + it('should allow standalone websocket middleware to short circuit', async () => { + const ws = createClient('/ws/middleware-short-circuit'); + assert.deepStrictEqual(await waitJSONMessage(ws), { + type: 'middleware', + path: '/ws/middleware-short-circuit', + }); + assert.equal(ws.readyState, WebSocket.OPEN); + await closeClient(ws); + }); + + it('should close timed out standalone websocket controller methods', async () => { + const ws = createClient('/ws/timeout'); + await waitOpen(ws); + assert.deepStrictEqual(await waitCloseInfo(ws), { + code: 1011, + reason: 'Internal Server Error', + }); + }); + + it('should reject malformed and unmatched standalone websocket routes', async () => { + await assert.rejects( + () => waitOpen(createClient('/ws/echo/%E0%A4%A')), + /Unexpected server response: 400/, + ); + await assert.rejects( + () => waitOpen(createClient('/ws/not-found')), + /Unexpected server response: 404/, + ); + }); + it('should handle websocket fetch connection lifecycle errors before closing', async () => { const socket = createClient('/ws-fetch/connection-error', { 'x-client-id': 'connection-error-client', @@ -344,6 +384,15 @@ function waitClose(ws: WebSocket) { }); } +function waitCloseInfo(ws: WebSocket): Promise<{ code: number; reason: string }> { + return new Promise(resolve => { + ws.once('close', (code, reason) => resolve({ + code, + reason: reason.toString(), + })); + }); +} + async function closeClient(ws: WebSocket) { if (ws.readyState === WebSocket.CLOSED || ws.readyState === WebSocket.CLOSING) { return;