diff --git a/core/controller-decorator/index.ts b/core/controller-decorator/index.ts index 03cd745ac..910e0217e 100644 --- a/core/controller-decorator/index.ts +++ b/core/controller-decorator/index.ts @@ -1,4 +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'; @@ -10,6 +12,12 @@ 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/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'; @@ -19,6 +27,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/websocket-fetch/WebSocketFetchController.ts b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts new file mode 100644 index 000000000..289666cdd --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchController.ts @@ -0,0 +1,29 @@ +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); + } + if (param?.timeout !== undefined) { + ControllerInfoUtil.setControllerTimeout(param.timeout, 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..910a7b7dd --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket-fetch/WebSocketFetchMethod.ts @@ -0,0 +1,62 @@ +import assert from 'node:assert'; +import { ControllerType, WebSocketFetchMethodType } from '@eggjs/tegg-types'; +import type { EggProtoImplClass, WebSocketFetchLifecycleMethodParams, WebSocketFetchMethodParams } from '@eggjs/tegg-types'; +import MethodInfoUtil from '../../util/MethodInfoUtil'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +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)}`); + const controllerClazz = target.constructor as EggProtoImplClass; + 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 }; +} + +export function WebSocketFetchMethod(param?: WebSocketFetchMethodParams) { + return function(target: any, propertyKey: PropertyKey) { + const { controllerClazz, methodName } = setWebSocketFetchMethodType( + target, + propertyKey, + WebSocketFetchMethodType.DATA, + param?.timeout, + ); + WebSocketInfoUtil.setWebSocketMethodPath('', controllerClazz, methodName); + if (param?.priority !== undefined) { + WebSocketInfoUtil.setWebSocketMethodPriority(param.priority, controllerClazz, methodName); + } + }; +} + +export function WebSocketFetchOnConnection(param?: WebSocketFetchLifecycleMethodParams) { + return function(target: any, propertyKey: PropertyKey) { + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CONNECTION, param?.timeout); + }; +} + +export function WebSocketFetchOnOpen(param?: WebSocketFetchLifecycleMethodParams) { + return function(target: any, propertyKey: PropertyKey) { + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.OPEN, param?.timeout); + }; +} + +export function WebSocketFetchOnError(param?: WebSocketFetchLifecycleMethodParams) { + return function(target: any, propertyKey: PropertyKey) { + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.ERROR, param?.timeout); + }; +} + +export function WebSocketFetchOnClose(param?: WebSocketFetchLifecycleMethodParams) { + return function(target: any, propertyKey: PropertyKey) { + setWebSocketFetchMethodType(target, propertyKey, WebSocketFetchMethodType.CLOSE, param?.timeout); + }; +} 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/decorator/websocket/WebSocketController.ts b/core/controller-decorator/src/decorator/websocket/WebSocketController.ts new file mode 100644 index 000000000..b9772c424 --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket/WebSocketController.ts @@ -0,0 +1,29 @@ +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); + } + if (param?.timeout !== undefined) { + ControllerInfoUtil.setControllerTimeout(param.timeout, 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..ad7a71716 --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket/WebSocketMethod.ts @@ -0,0 +1,22 @@ +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); + } + if (param.timeout !== undefined) { + MethodInfoUtil.setMethodTimeout(param.timeout, 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..18182f112 --- /dev/null +++ b/core/controller-decorator/src/decorator/websocket/WebSocketParam.ts @@ -0,0 +1,24 @@ +import assert from 'node:assert'; +import { WebSocketParamType } from '@eggjs/tegg-types'; +import type { EggProtoImplClass } from '@eggjs/tegg-types'; +import WebSocketInfoUtil from '../../util/WebSocketInfoUtil'; + +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-fetch/WebSocketFetchControllerMetaBuilder.ts b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts new file mode 100644 index 000000000..ae995e05e --- /dev/null +++ b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMetaBuilder.ts @@ -0,0 +1,110 @@ +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, '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 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, timeout); + 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..5013cda0e --- /dev/null +++ b/core/controller-decorator/src/impl/websocket-fetch/WebSocketFetchControllerMethodMetaBuilder.ts @@ -0,0 +1,151 @@ +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.getCompatibleParamIndexList(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.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`); + } + 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.getCompatibleParamIndexList(this.clazz, this.methodName); + for (const paramIndex of paramIndexList) { + 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) { + 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; + const timeout = MethodInfoUtil.getMethodTimeout(this.clazz, this.methodName); + return new WebSocketFetchMethodMeta( + 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 new file mode 100644 index 000000000..1560eac45 --- /dev/null +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMetaBuilder.ts @@ -0,0 +1,66 @@ +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, 'invalid 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 timeout = ControllerInfoUtil.getControllerTimeout(this.clazz); + const metadata = new WebSocketControllerMeta( + clazzName, protoName, controllerName, webSocketPath, middlewares, methods, hosts, timeout); + 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..ac88f9de9 --- /dev/null +++ b/core/controller-decorator/src/impl/websocket/WebSocketControllerMethodMetaBuilder.ts @@ -0,0 +1,112 @@ +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'; +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.getCompatibleParamIndexList(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.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`); + } + } + } + + private buildParamType(webSocketPath: string): Map { + this.checkParamDecorators(); + + const paramTypeMap = new Map(); + const paramIndexList = WebSocketInfoUtil.getCompatibleParamIndexList(this.clazz, this.methodName); + for (const paramIndex of paramIndexList) { + 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.getCompatibleMethodParamName(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; + } + + 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) { + 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(); + const timeout = MethodInfoUtil.getMethodTimeout(this.clazz, this.methodName); + return new WebSocketMethodMeta( + 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 new file mode 100644 index 000000000..009ca6ab2 --- /dev/null +++ b/core/controller-decorator/src/model/WebSocketControllerMeta.ts @@ -0,0 +1,68 @@ +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[]; + public readonly timeout?: number; + + constructor( + className: string, + protoName: EggPrototypeName, + controllerName: string, + path: string | undefined, + middlewares: MiddlewareFunc[], + methods: WebSocketMethodMeta[], + hosts: string[] | undefined, + timeout?: number, + ) { + this.protoName = protoName; + this.controllerName = controllerName; + this.className = className; + this.path = path; + this.middlewares = middlewares; + this.methods = methods; + this.hosts = hosts; + this.timeout = timeout; + } + + 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 ]; + } + + 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 new file mode 100644 index 000000000..b5cc049d6 --- /dev/null +++ b/core/controller-decorator/src/model/WebSocketFetchControllerMeta.ts @@ -0,0 +1,80 @@ +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 timeout?: number; + 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, + timeout?: number, + ) { + 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; + this.timeout = timeout; + } + + 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 ]; + } + + 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 new file mode 100644 index 000000000..44d72c91f --- /dev/null +++ b/core/controller-decorator/src/model/WebSocketFetchMethodMeta.ts @@ -0,0 +1,37 @@ +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; + public readonly timeout: number | undefined; + + constructor( + name: string, + path: string, + middlewares: MiddlewareFunc[], + contextParamIndex: number | undefined, + paramMap: Map, + priority: number, + hosts: string[] | undefined, + methodType: WebSocketFetchMethodType, + timeout?: number, + ) { + this.name = name; + this.path = path; + this.middlewares = middlewares; + this.contextParamIndex = contextParamIndex; + this.paramMap = paramMap; + 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 new file mode 100644 index 000000000..805bbb6a7 --- /dev/null +++ b/core/controller-decorator/src/model/WebSocketMethodMeta.ts @@ -0,0 +1,207 @@ +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 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; + 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 timeout: number | undefined; + + constructor( + name: string, + path: string, + middlewares: MiddlewareFunc[], + contextParamIndex: number | undefined, + paramMap: Map, + priority: number, + hosts: string[] | undefined, + timeout?: number, + ) { + this.name = name; + this.path = path; + this.middlewares = middlewares; + this.contextParamIndex = contextParamIndex; + this.paramMap = paramMap; + this.priority = priority; + this.hosts = hosts; + this.timeout = timeout; + } +} + +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(); + } + 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 4d355ee7f..80ed915ae 100644 --- a/core/controller-decorator/src/model/index.ts +++ b/core/controller-decorator/src/model/index.ts @@ -3,6 +3,10 @@ export * from './HTTPControllerMeta'; export * from './HTTPRequest'; 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 new file mode 100644 index 000000000..482874ec6 --- /dev/null +++ b/core/controller-decorator/src/util/WebSocketInfoUtil.ts @@ -0,0 +1,128 @@ +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_FETCH_METHOD_TYPE_MAP, + CONTROLLER_WEBSOCKET_PATH, + HTTPParamType, + WebSocketParamType, + type EggProtoImplClass, + type WebSocketFetchMethodType, +} from '@eggjs/tegg-types'; +import { MapUtil } from '@eggjs/tegg-common-util'; +import HTTPInfoUtil from './HTTPInfoUtil'; + +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) { + 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 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()); + 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 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); + } + + static setWebSocketMethodPriority(priority: number, clazz: EggProtoImplClass, methodName: string) { + 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); + } + + 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 new file mode 100644 index 000000000..9fac64a26 --- /dev/null +++ b/core/controller-decorator/test/websocket/WebSocketMeta.test.ts @@ -0,0 +1,140 @@ +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() + 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/, + ); + }); + + 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/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/controller-decorator/MetadataKey.ts b/core/types/controller-decorator/MetadataKey.ts index 5eb77550e..901822b49 100644 --- a/core/types/controller-decorator/MetadataKey.ts +++ b/core/types/controller-decorator/MetadataKey.ts @@ -15,6 +15,13 @@ 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 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'); 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..d9d426750 --- /dev/null +++ b/core/types/controller-decorator/WebSocketContext.ts @@ -0,0 +1,14 @@ +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 extends Context { + readonly webSocket: Socket; +} diff --git a/core/types/controller-decorator/WebSocketController.ts b/core/types/controller-decorator/WebSocketController.ts new file mode 100644 index 000000000..a657cfbda --- /dev/null +++ b/core/types/controller-decorator/WebSocketController.ts @@ -0,0 +1,6 @@ +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 new file mode 100644 index 000000000..51d42e6ca --- /dev/null +++ b/core/types/controller-decorator/WebSocketFetchController.ts @@ -0,0 +1,6 @@ +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 new file mode 100644 index 000000000..cdd9320a0 --- /dev/null +++ b/core/types/controller-decorator/WebSocketFetchMethod.ts @@ -0,0 +1,8 @@ +export interface WebSocketFetchMethodParams { + priority?: number; + timeout?: number; +} + +export interface WebSocketFetchLifecycleMethodParams { + timeout?: 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/WebSocketMethod.ts b/core/types/controller-decorator/WebSocketMethod.ts new file mode 100644 index 000000000..c31fb913e --- /dev/null +++ b/core/types/controller-decorator/WebSocketMethod.ts @@ -0,0 +1,5 @@ +export interface WebSocketMethodParams { + path: string; + priority?: number; + timeout?: number; +} diff --git a/core/types/controller-decorator/index.ts b/core/types/controller-decorator/index.ts index 828f96679..6a688922a 100644 --- a/core/types/controller-decorator/index.ts +++ b/core/types/controller-decorator/index.ts @@ -5,6 +5,12 @@ export * from './builder'; export * from './HTTPController'; export * from './HTTPMethod'; export * from './HTTPParam'; +export * from './WebSocketController'; +export * from './WebSocketMethod'; +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 4ace7fe56..bfdb0805f 100644 --- a/core/types/controller-decorator/model/types.ts +++ b/core/types/controller-decorator/model/types.ts @@ -7,6 +7,8 @@ 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', @@ -23,6 +25,8 @@ 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', @@ -53,6 +57,29 @@ export enum HTTPParamType { COOKIES = 'COOKIES', } +export enum WebSocketParamType { + PARAM = 'PARAM', + QUERY = 'QUERY', + QUERIES = 'QUERIES', + HEADERS = 'HEADERS', + 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 { STDIO = 'STDIO', SSE = 'SSE', 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/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 5093705da..b22835ee0 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,8 @@ export default class ControllerAppBootHook { return new EggControllerLoader(unitPath); }); 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 => { @@ -148,10 +151,15 @@ 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(); 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..e5590d12a --- /dev/null +++ b/plugin/controller/lib/impl/websocket/WebSocketContext.ts @@ -0,0 +1,17 @@ +import type { Context } from 'egg'; +import type { WebSocket } from 'ws'; +import type { WebSocketContext } from '@eggjs/tegg'; + +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 new file mode 100644 index 000000000..69c4f7700 --- /dev/null +++ b/plugin/controller/lib/impl/websocket/WebSocketControllerRegister.ts @@ -0,0 +1,317 @@ +import assert from 'node:assert'; +import { IncomingMessage, ServerResponse } from 'node:http'; +import { Duplex } from 'node:stream'; +import compose from 'koa-compose'; +import { Application, Context } from 'egg'; +import { EggRouter } from '@eggjs/router'; +import { FrameworkErrorFormater } from 'egg-errors'; +import { createWebSocketStream, RawData, WebSocket, WebSocketServer } from 'ws'; +import { + CONTROLLER_META_DATA, + ControllerMetadata, + ControllerType, + Next, + WebSocketControllerMeta, + WebSocketFetchControllerMeta, + WebSocketFetchMethodMeta, + WebSocketMethodMeta, +} 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 { extendWebSocketContext } from './WebSocketContext'; + +const noop = () => { + // noop +}; + +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: Array> = []; + 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 || controllerMeta.type === ControllerType.WEBSOCKET_FETCH, + 'controller meta type is not WEBSOCKET or WEBSOCKET_FETCH', + ); + 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() { + this.routes = createWebSocketRoutes( + this.controllerProtos, + proto => proto.getMetaData(CONTROLLER_META_DATA) as WebSocketControllerMeta | WebSocketFetchControllerMeta, + (controllerMeta, methodMeta) => this.checkDuplicate(controllerMeta, methodMeta), + ); + } + + 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 checkDuplicate(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + let router = this.checkRouters.get('default')!; + const hosts = getWebSocketMethodHosts(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 | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + 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 = getWebSocketMethodRealPath(controllerMeta, methodMeta); + const matched = router.match(methodRealPath, 'GET'); + 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}`); + throw FrameworkErrorFormater.format(err); + } + } + + private async handleUpgrade(req: IncomingMessage, socket: Duplex, head: Buffer) { + 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, 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(WEBSOCKET_INTERNAL_ERROR_CODE, WEBSOCKET_INTERNAL_ERROR_REASON); + } + }); + }); + } + + private createURL(ctx: Context): URL { + const host = ctx.host || 'localhost'; + return new URL(ctx.url || '/', `${ctx.protocol}://${host}`); + } + + private async handleConnection( + route: WebSocketRoute, + params: Record, + eggCtx: Context, + webSocket: WebSocket, + ) { + 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 = 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, runtime, fetchEvents!); + } else { + await this.invokeController(route, runtime); + } + await next(); + }; + const composed = compose([ ...methodMiddlewares, handler ]); + + 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; + }); + }); + } finally { + fetchEvents?.dispose(); + } + } + + 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); + await runtime.invokeController(eggObj.obj, controllerMeta, methodMeta); + } + + private async invokeFetchController( + route: WebSocketRoute, + runtime: WebSocketControllerRuntime, + events: WebSocketEventStream, + ) { + const controllerMeta = route.controllerMeta as WebSocketFetchControllerMeta; + const methodMeta = route.methodMeta as WebSocketFetchMethodMeta; + const eggObj = await this.eggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); + await runtime.invokeFetchController(eggObj.obj, controllerMeta, methodMeta, events); + } + + 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}`); + } + + private getSocketBytesWritten(socket: Duplex): number { + return (socket as Duplex & { bytesWritten?: number }).bytesWritten || 0; + } + + private hasWebSocketOwner(socket: Duplex): boolean { + return Object.getOwnPropertySymbols(socket).some(symbol => { + return symbol.description === 'websocket' && Boolean(Reflect.get(socket, symbol)); + }); + } +} diff --git a/plugin/controller/package.json b/plugin/controller/package.json index 9019818a0..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", @@ -65,6 +66,7 @@ "path-to-regexp": "^1.8.0", "raw-body": "^2.5.2", "sdk-base": "^4.2.0", + "ws": "^8.21.0", "zod": "^4.0.0" }, "devDependencies": { @@ -73,6 +75,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/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 new file mode 100644 index 000000000..6217a7578 --- /dev/null +++ b/plugin/controller/test/fixtures/apps/controller-app/app/controller/WebSocketController.ts @@ -0,0 +1,309 @@ +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, + HTTPHeaders, + HTTPParam, + HTTPQueries, + HTTPQuery, + Host, + Inject, + Middleware, + 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'; +import { webSocketFetchStreamCloseEvents } from './WebSocketTestState'; + +interface WebSocketFetchRequest { + content: string; + close?: boolean; + error?: boolean; + delayMs?: number; + holdOpen?: boolean; + observeClose?: boolean; + 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', +}) +export class AppWebSocketController { + @Inject() + appService: AppService; + + @WebSocketMethod({ + path: '/echo/:id', + }) + async echo( + @WebSocketSocket() socket: WebSocket, + @HTTPParam() id: string, + @HTTPQuery() name: string, + @HTTPQueries({ name: 'tag' }) tags: string[], + @HTTPHeaders() 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, + sameSocket: ctx.webSocket === socket, + sameRequest: ctx.req === request, + 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, + @HTTPParam() id: string, + @HTTPQuery() 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; + } + + @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({ + path: '/ws-fetch/:id', +}) +export class AppWebSocketFetchController { + @Inject() + appService: AppService; + + @WebSocketFetchOnConnection() + onConnection( + @HTTPParam() id: string, + @HTTPHeaders() headers: IncomingHttpHeaders, + @Context() ctx: WebSocketContext, + ) { + if (id === 'connection-error') { + throw new Error('fetch connection error'); + } + ctx.webSocket.send(JSON.stringify({ + type: 'connection', + header: headers['x-client-id'], + path: ctx.path, + pid: process.pid, + })); + } + + @WebSocketFetchOnOpen() + onOpen( + @Context() ctx: WebSocketContext, + ) { + ctx.webSocket.send(JSON.stringify({ + type: 'open', + path: ctx.path, + pid: process.pid, + })); + } + + @WebSocketFetchMethod() + onData( + @WebSocketData() data: RawData, + @WebSocketClose() close: WebSocketFetchClose, + @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; + const source = new PassThrough(); + const output = body.pipeline ? new PassThrough() : source; + if (body.pipeline) { + pipeline(source, output, () => { + // WebSocketFetch observes errors from the returned output stream. + }); + } + source.write(JSON.stringify({ + type: 'data', + phase: 1, + id, + name, + tags, + header: headers['x-client-id'], + path: ctx.path, + content: body.content, + 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 (source.destroyed) { + return; + } + if (body.error) { + source.destroy(new Error(`fetch error: ${body.content}`)); + return; + } + source.write(JSON.stringify({ + type: 'data', + phase: 2, + id, + content: body.content, + pid: process.pid, + })); + source.end(); + }, body.delayMs ?? 10); + + if (body.close) { + output.once('end', () => setImmediate(() => close(1000, 'server done'))); + } + return output; + } + + @WebSocketFetchOnError() + onError( + @WebSocketError() error: Error, + @HTTPHeaders() headers: IncomingHttpHeaders, + @Context() ctx: WebSocketContext, + ) { + ctx.webSocket.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, + @HTTPParam() id: string, + ) { + await this.appService.save({ + name: `ws-fetch-close-${id}`, + desc: `${code}:${reason.toString()}`, + }); + } +} 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/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/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/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 () => { diff --git a/plugin/controller/test/websocket/websocket.test.ts b/plugin/controller/test/websocket/websocket.test.ts new file mode 100644 index 000000000..ec8ae66af --- /dev/null +++ b/plugin/controller/test/websocket/websocket.test.ts @@ -0,0 +1,681 @@ +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'; + +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()); +} + +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); + 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(); + }); +} + +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; + + 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', + sameSocket: true, + sameRequest: true, + 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 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' }, + }); + 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, + pipeline: 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 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); + 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' }, + }); + 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 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 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); + } + }); + + 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 new file mode 100644 index 000000000..98629ff71 --- /dev/null +++ b/plugin/controller/test/websocket/websocketCluster.test.ts @@ -0,0 +1,382 @@ +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()); +} + +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); + 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(); + }); +} + +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; + + before(async () => { + mm(process.env, 'EGG_TYPESCRIPT', true); + mm(process, 'cwd', () => { + return pluginRoot; + }); + const cluster = mm.cluster as any; + app = 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(echo.pid, readyMessages[index].pid); + } + } 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(echo.pid, readyMessages[index].pid); + } + } finally { + await Promise.all(clients.map(socket => closeClient(socket))); + } + }); + + 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(open.pid, connection.pid); + + 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.equal(first.pid, connection.pid); + assert.deepEqual(await receiver.next(), { + type: 'data', + phase: 2, + id: 'cluster-fetch', + content: 'cluster-first', + pid: first.pid, + }); + 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'); + 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: connection.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.equal(closeData.pid, connection.pid); + 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); + } + }); + + 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 { + const connection = await receiver.next(); + const open = await receiver.next(); + assert.equal(open.pid, connection.pid); + + 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.deepEqual(new Set(messages.map(message => message.pid)), new Set([ connection.pid ])); + 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/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..e5ab1e335 100644 --- a/standalone/service-worker/package.json +++ b/standalone/service-worker/package.json @@ -49,11 +49,13 @@ "@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", "type-is": "^1.6.18", - "urllib": "^4.0.0" + "urllib": "^4.0.0", + "ws": "^8.21.0" }, "peerDependencies": { "@eggjs/tegg": "^3" @@ -64,6 +66,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..29cc8883c --- /dev/null +++ b/standalone/service-worker/src/websocket/ServiceWorkerWebSocketContext.ts @@ -0,0 +1,39 @@ +import type { IncomingHttpHeaders, IncomingMessage } from 'node:http'; +import type { WebSocket } from 'ws'; + +export class ServiceWorkerWebSocketContext { + readonly socket: WebSocket; + readonly webSocket: 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.webSocket = 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..445b70615 --- /dev/null +++ b/standalone/service-worker/src/websocket/WebSocketControllerRegister.ts @@ -0,0 +1,287 @@ +import assert from 'node:assert'; +import type { IncomingMessage } from 'node:http'; +import { Duplex } from 'node:stream'; +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, + WebSocketFetchControllerMeta, + WebSocketFetchMethodMeta, + WebSocketMethodMeta, +} 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'; +import { ServiceWorkerWebSocketContext } from './ServiceWorkerWebSocketContext'; + +const noop = () => { + // noop +}; + +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: Array> = []; + + 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) { + 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) => { + 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); + 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; + } + + 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(WEBSOCKET_INTERNAL_ERROR_CODE, WEBSOCKET_INTERNAL_ERROR_REASON); + } + resolve(); + }); + }); + }); + return true; + } + + static 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}`); + } + + private checkDuplicate(controllerMeta: WebSocketControllerMeta | WebSocketFetchControllerMeta, methodMeta: WebSocketMethodMeta | WebSocketFetchMethodMeta) { + let router = this.checkRouters.get('default')!; + const hosts = getWebSocketMethodHosts(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 = 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 = getWebSocketMethodRealPath(controllerMeta, methodMeta); + const matched = router.match(methodRealPath, 'GET'); + 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}`); + throw FrameworkErrorFormater.format(err); + } + } + + private createURL(req: IncomingMessage): URL { + const host = req.headers.host || 'localhost'; + return new URL(req.url || '/', `http://${host}`); + } + + private async handleConnection( + 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 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, runtime, fetchEvents!); + } else { + await this.invokeController(route, runtime); + } + await next(); + }; + + 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(); + } + } + + 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, + runtime: WebSocketControllerRuntime, + ) { + const controllerMeta = route.controllerMeta as WebSocketControllerMeta; + const methodMeta = route.methodMeta as WebSocketMethodMeta; + const eggObj = await EggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); + await runtime.invokeController(eggObj.obj, controllerMeta, methodMeta); + } + + private async invokeFetchController( + route: WebSocketRoute, + runtime: WebSocketControllerRuntime, + events: WebSocketEventStream, + ) { + const controllerMeta = route.controllerMeta as WebSocketFetchControllerMeta; + const methodMeta = route.methodMeta as WebSocketFetchMethodMeta; + const eggObj = await EggContainerFactory.getOrCreateEggObject(route.controllerProto, route.controllerProto.name); + await runtime.invokeFetchController(eggObj.obj, controllerMeta, methodMeta, events); + } + +} 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/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 }; 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..628b63103 --- /dev/null +++ b/standalone/service-worker/test/fixtures/websocket/WebSocketController.ts @@ -0,0 +1,264 @@ +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, + HTTPHeaders, + HTTPParam, + HTTPQueries, + HTTPQuery, + Middleware, + Request, + WebSocketClose, + WebSocketCloseCode, + WebSocketCloseReason, + WebSocketController, + WebSocketData, + WebSocketError, + WebSocketFetchClose, + WebSocketFetchController, + WebSocketFetchMethod, + WebSocketFetchOnClose, + WebSocketFetchOnConnection, + WebSocketFetchOnError, + WebSocketFetchOnOpen, + WebSocketMethod, + WebSocketSocket, + WebSocketStream, +} from '@eggjs/tegg'; +import type { ServiceWorkerWebSocketContext } from '../../../src/websocket/ServiceWorkerWebSocketContext'; + +interface WebSocketFetchRequest { + content: string; + close?: boolean; + error?: boolean; + delayMs?: number; + holdOpen?: boolean; + observeClose?: boolean; + pipeline?: boolean; +} + +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', +}) +export class StandaloneWebSocketController { + @WebSocketMethod({ + path: '/echo/:id', + }) + async echo( + @WebSocketSocket() socket: WebSocket, + @HTTPParam() id: string, + @HTTPQuery() name: string, + @HTTPQueries({ name: 'tag' }) tags: string[], + @HTTPHeaders() headers: IncomingHttpHeaders, + @Request() request: IncomingMessage, + @Context() ctx: ServiceWorkerWebSocketContext, + ) { + socket.send(JSON.stringify({ + type: 'ready', + id, + name, + tags, + header: headers['x-client-id'], + url: request.url, + path: ctx.path, + sameSocket: ctx.webSocket === socket, + })); + + socket.on('message', data => { + socket.send(JSON.stringify({ + type: 'echo', + data: data.toString(), + })); + }); + } + + @WebSocketMethod({ + path: '/stream/:id', + }) + stream( + @WebSocketStream() input: Readable, + @HTTPParam() id: string, + @HTTPQuery() 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; + } + + @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({ + path: '/ws-fetch/:id', +}) +export class StandaloneWebSocketFetchController { + @WebSocketFetchOnConnection() + onConnection( + @HTTPParam() id: string, + @HTTPHeaders() headers: IncomingHttpHeaders, + @Context() ctx: ServiceWorkerWebSocketContext, + ) { + if (id === 'connection-error') { + throw new Error('fetch connection error'); + } + ctx.webSocket.send(JSON.stringify({ + type: 'connection', + header: headers['x-client-id'], + path: ctx.path, + })); + } + + @WebSocketFetchOnOpen() + onOpen( + @Context() ctx: ServiceWorkerWebSocketContext, + ) { + ctx.webSocket.send(JSON.stringify({ + type: 'open', + path: ctx.path, + })); + } + + @WebSocketFetchMethod() + onData( + @WebSocketData() data: RawData, + @WebSocketClose() close: WebSocketFetchClose, + @HTTPParam() id: string, + @HTTPQuery() name: string, + @HTTPQueries({ name: 'tag' }) tags: string[], + @HTTPHeaders() headers: IncomingHttpHeaders, + @Context() ctx: ServiceWorkerWebSocketContext, + ) { + const body = JSON.parse(data.toString()) as WebSocketFetchRequest; + const source = new PassThrough(); + const output = body.pipeline ? new PassThrough() : source; + if (body.pipeline) { + pipeline(source, output, () => { + // WebSocketFetch observes errors from the returned output stream. + }); + } + source.write(JSON.stringify({ + type: 'data', + phase: 1, + id, + name, + tags, + header: headers['x-client-id'], + path: ctx.path, + 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 (source.destroyed) { + return; + } + if (body.error) { + source.destroy(new Error(`fetch error: ${body.content}`)); + return; + } + source.write(JSON.stringify({ + type: 'data', + phase: 2, + id, + content: body.content, + })); + source.end(); + }, body.delayMs ?? 10); + + if (body.close) { + output.once('end', () => setImmediate(() => close(1000, 'server done'))); + } + return output; + } + + @WebSocketFetchOnError() + onError( + @WebSocketError() error: Error, + @HTTPHeaders() headers: IncomingHttpHeaders, + @Context() ctx: ServiceWorkerWebSocketContext, + ) { + ctx.webSocket.send(JSON.stringify({ + type: 'error', + message: error.message, + header: headers['x-client-id'], + path: ctx.path, + })); + } + + @WebSocketFetchOnClose() + onClose( + @WebSocketCloseCode() code: number, + @WebSocketCloseReason() reason: Buffer, + @HTTPParam() 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..5e963a330 --- /dev/null +++ b/standalone/service-worker/test/websocket/websocket.test.ts @@ -0,0 +1,413 @@ +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, fetchStreamCloseEvents } 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; + fetchStreamCloseEvents.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', + sameSocket: true, + }); + + 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 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', + }); + 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', + }); + 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')); + + 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) { + 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); + }); +} + +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()); + }); +} + +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; + } + 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'); +}