adapt ws

.eslintignore

+/dist
\ No newline at end of file

README.md

@@ -92,8 +92,21 @@ Koishi-Nest 的配置项和 Koishi 配置项一致，参照 [Koishi 文档](http
 
 插件的使用可以参考 [Koishi 文档](https://koishi.js.org/v4/guide/plugin/plugin.html)。 `moduleSelection` 的使用见本文 **复用性** 部分。
 
+* `useWs`: `boolean` 默认 `false` 。是否启用 WebSocket 。**异步配置该项应写入异步配置项中。** 详见本文的 **WebSocket** 部分。
+
 * `isGlobal`: `boolean` 默认 `true` 。指示 Koishi-Nest 模块是否应被注册为全局模块，建议开启。当安装了其他模块的情况下，需要将 Koishi-Nest 注册为全局模块使得其他模块可以正常注入 Koishi-Nest 作为依赖项。 **异步配置该项应写入异步配置项中。** 关于全局模块请参考 [Nest.js 文档](https://docs.nestjs.cn/8/modules?id=%e5%85%a8%e5%b1%80%e6%a8%a1%e5%9d%97) 。
 
+## WebSocket
+
+Koishi-Nest 针对 Koishi 的 WebSocket 功能进行了针对 Nest.js 的适配。
+
+若要使用 Koishi 的与 WebSocket 服务器相关的插件或功能，需要在 Koishi-Nest 配置项中，把 `useWs` 设置为 `true` ，并在 `main.ts` 修改下列部分。
+
+```ts
+  const app = await NestFactory.create(AppModule);
+  app.useWebSocketAdapter(new KoishiWsAdapter(app));
+```
+
 ## 注入 Koishi 实例
 
 可以直接注入 Koishi 实例或上下文进行注册操作。这种情况下，建议让 Nest 提供者实现 `OnModuleInit` 接口，并在该事件方法中进行 Koishi 指令注册操作。

index.ts

@@ -3,3 +3,4 @@ export * from './src/koishi.interfaces';
 export * from './src/koishi.service';
 export * from './src/koishi.module';
 export * from './src/koishi-context.factory';
+export * from './src/koishi.ws-adapter';

package.json

@@ -32,15 +32,14 @@
   "peerDependencies": {
     "@nestjs/common": "^8.0.0",
     "@nestjs/core": "^8.0.0",
-    "koishi": "^4.0.0-alpha.11",
-    "proxy-agent": "^5.0.0",
-    "reflect-metadata": "^0.1.13"
+    "koishi": "^4.0.0-alpha.12",
+    "reflect-metadata": "^0.1.13",
+    "rxjs": "^7.4.0"
   },
   "devDependencies": {
     "@nestjs/common": "^8.0.9",
     "@nestjs/core": "^8.0.9",
     "@types/koa": "^2.13.4",
-    "@types/koa__router": "^8.0.8",
     "@types/koa-bodyparser": "^4.3.3",
     "@types/lodash": "^4.14.175",
     "@types/node": "^16.10.2",
@@ -50,16 +49,19 @@
     "eslint-config-prettier": "^8.3.0",
     "eslint-plugin-prettier": "^3.4.1",
     "express": "^4.17.1",
-    "koishi": "^4.0.0-alpha.11",
+    "koishi": "^4.0.0-alpha.12",
     "prettier": "^2.4.1",
-    "proxy-agent": "^5.0.0",
     "reflect-metadata": "^0.1.13",
+    "rxjs": "^7.4.0",
     "typescript": "^4.4.3"
   },
   "dependencies": {
-    "@koa/router": "^10.1.1",
+    "@nestjs/platform-ws": "^8.1.2",
+    "@nestjs/websockets": "^8.1.2",
+    "@types/ws": "^8.2.0",
     "koa": "^2.13.3",
     "koa-bodyparser": "^4.3.0",
-    "lodash": "^4.17.21"
+    "lodash": "^4.17.21",
+    "ws": "^8.2.3"
   }
 }

src/koishi-http-discovery/koishi-http-discovery.service.ts

+import { Injectable } from '@nestjs/common';
+import { AbstractHttpAdapter, HttpAdapterHost, ModuleRef } from '@nestjs/core';
+import { AbstractWsAdapter } from '@nestjs/websockets';
+
+@Injectable()
+export class KoishiHttpDiscoveryService {
+  constructor(private readonly moduleRef: ModuleRef) {}
+
+  getHttpAdapter(): AbstractHttpAdapter {
+    const apdaterHost = this.moduleRef.get(HttpAdapterHost, { strict: false });
+    if (apdaterHost) {
+      return apdaterHost.httpAdapter;
+    } else {
+      return null;
+    }
+  }
+}

src/koishi.interfaces.ts

@@ -52,11 +52,14 @@ export interface KoishiModuleSelection extends ContextSelector {
   module: Type<any>;
 }
 
-export interface WhetherGlobalOption {
+export interface KoishiModuleTopOptions {
   isGlobal?: boolean;
+  useWs?: boolean;
 }
 
-export interface KoishiModuleOptions extends App.Config, WhetherGlobalOption {
+export interface KoishiModuleOptions
+  extends App.Config,
+    KoishiModuleTopOptions {
   usePlugins?: KoishiModulePlugin<Plugin>[];
   loggerPrefix?: string;
   loggerColor?: number;
@@ -69,7 +72,7 @@ export interface KoishiModuleOptionsFactory {
 
 export interface KoishiModuleAsyncOptions
   extends Pick<ModuleMetadata, 'imports'>,
-    WhetherGlobalOption {
+    KoishiModuleTopOptions {
   useExisting?: Type<KoishiModuleOptionsFactory>;
   useClass?: Type<KoishiModuleOptionsFactory>;
   useFactory?: (

src/koishi.module.ts

@@ -26,6 +26,8 @@ import { Context } from 'koishi';
 import { defaultContextContainer } from './koishi-context.factory';
 import { KoishiInjectionService } from './providers/koishi-injection.service';
 import { KoishiContextService } from './providers/koishi-context.service';
+import { KoishiHttpDiscoveryService } from './koishi-http-discovery/koishi-http-discovery.service';
+import { KoishiWebsocketGateway } from './providers/koishi-websocket.gateway';
 
 const koishiContextProvider: Provider<Context> = {
   provide: KOISHI_CONTEXT,
@@ -44,6 +46,7 @@ const koishiContextProvider: Provider<Context> = {
     koishiContextProvider,
     KoishiContextService,
     KoishiInjectionService,
+    KoishiHttpDiscoveryService,
   ],
   exports: [KoishiService, koishiContextProvider],
 })
@@ -58,6 +61,7 @@ export class KoishiModule implements NestModule {
       providers: [
         { provide: KOISHI_MODULE_OPTIONS, useValue: options },
         ...defaultContextContainer.contextsToProvide,
+        ...(options.useWs ? [KoishiWebsocketGateway] : []),
       ],
       exports: defaultContextContainer.contextsToProvide,
       global: options.isGlobal != null ? options.isGlobal : true,
@@ -72,6 +76,7 @@ export class KoishiModule implements NestModule {
         ...this.createAsyncProviders(options),
         ...defaultContextContainer.contextsToProvide,
         ...(options.extraProviders || []),
+        ...(options.useWs ? [KoishiWebsocketGateway] : []),
       ],
       exports: defaultContextContainer.contextsToProvide,
       global: options.isGlobal != null ? options.isGlobal : true,

src/koishi.service.ts

-import { App } from 'koishi';
+import { App, Router } from 'koishi';
 import {
   Inject,
   Injectable,
@@ -9,38 +9,40 @@ import {
 import { KoishiModuleOptions } from './koishi.interfaces';
 import { Server } from 'http';
 import Koa from 'koa';
-import KoaRouter from '@koa/router';
 import KoaBodyParser from 'koa-bodyparser';
 import { KoishiMetascanService } from './providers/koishi-metascan.service';
 import { applySelector } from './utility/koishi.utility';
 import { KOISHI_MODULE_OPTIONS } from './utility/koishi.constants';
 import { KoishiLoggerService } from './providers/koishi-logger.service';
+import { KoishiHttpDiscoveryService } from './koishi-http-discovery/koishi-http-discovery.service';
 
 @Injectable()
 export class KoishiService
   extends App
-  implements OnModuleInit, OnApplicationBootstrap, OnModuleDestroy {
+  implements OnModuleInit, OnApplicationBootstrap, OnModuleDestroy
+{
   constructor(
     @Inject(KOISHI_MODULE_OPTIONS)
     private readonly koishiModuleOptions: KoishiModuleOptions,
     private readonly metascan: KoishiMetascanService,
+    private readonly httpDiscovery: KoishiHttpDiscoveryService,
     private readonly koishiLogger: KoishiLoggerService,
   ) {
     super({
       ...koishiModuleOptions,
       port: 0,
     });
-    this.router = new KoaRouter();
+    this.router = new Router();
+    this._nestKoaTmpInstance.use(KoaBodyParser());
     this._nestKoaTmpInstance.use(this.router.routes());
     this._nestKoaTmpInstance.use(this.router.allowedMethods());
-    this._nestKoaTmpInstance.use(KoaBodyParser());
     this.options.port = 1;
   }
 
   readonly _nestKoaTmpInstance = new Koa();
 
   private async setHttpServer() {
-    const httpAdapter = this.metascan.getHttpAdapter();
+    const httpAdapter = this.httpDiscovery.getHttpAdapter();
     if (!httpAdapter) {
       return;
     }

src/koishi.ws-adapter.ts

+import { WsAdapter } from '@nestjs/platform-ws';
+import { INestApplicationContext } from '@nestjs/common';
+import * as http from 'http';
+
+export class KoishiWsAdapter extends WsAdapter {
+  constructor(appOrHttpServer?: INestApplicationContext | any) {
+    super(appOrHttpServer);
+  }
+
+  override ensureHttpServerExists(
+    port: number,
+    httpServer = http.createServer(),
+  ) {
+    if (this.httpServersRegistry.has(port)) {
+      return;
+    }
+    this.httpServersRegistry.set(port, httpServer);
+
+    httpServer.on('upgrade', (request, socket, head) => {
+      const baseUrl = 'ws://' + request.headers.host + '/';
+      const pathname = new URL(request.url, baseUrl).pathname;
+      const wsServersCollection = this.wsServersRegistry.get(port);
+
+      let isRequestDelegated = false;
+      let fallbackWsServer: any;
+      for (const wsServer of wsServersCollection) {
+        if (pathname === wsServer.path) {
+          wsServer.handleUpgrade(request, socket, head, (ws: unknown) => {
+            wsServer.emit('connection', ws, request);
+          });
+          isRequestDelegated = true;
+          break;
+        }
+        if (wsServer.path === '__koishi_fallback') {
+          fallbackWsServer = wsServer;
+        }
+      }
+      if (!isRequestDelegated && fallbackWsServer) {
+        const wsServer = fallbackWsServer;
+        wsServer.handleUpgrade(request, socket, head, (ws: unknown) => {
+          wsServer.emit('connection', ws, request);
+        });
+        isRequestDelegated = true;
+      }
+      if (!isRequestDelegated) {
+        socket.destroy();
+      }
+    });
+    return httpServer;
+  }
+}

src/providers/koishi-metascan.service.ts

@@ -40,15 +40,6 @@ export class KoishiMetascanService {
     private readonly ctxService: KoishiContextService,
   ) {}
 
-  getHttpAdapter(): AbstractHttpAdapter {
-    const apdaterHost = this.moduleRef.get(HttpAdapterHost, { strict: false });
-    if (apdaterHost) {
-      return apdaterHost.httpAdapter;
-    } else {
-      return null;
-    }
-  }
-
   private preRegisterCommandActionArg(config: CommandPutConfig, cmd: Command) {
     if (!config) {
       return;

src/providers/koishi-websocket.gateway.ts

+import {
+  OnGatewayConnection,
+  OnGatewayInit,
+  WebSocketGateway,
+  WebSocketServer,
+} from '@nestjs/websockets';
+import { KoishiService } from '../koishi.service';
+import type WebSocket from 'ws';
+import type { Server } from 'ws';
+import { IncomingMessage } from 'http';
+
+@WebSocketGateway()
+export class KoishiWebsocketGateway
+  implements OnGatewayInit, OnGatewayConnection {
+  constructor(private readonly koishi: KoishiService) {}
+
+  @WebSocketServer()
+  wsServer: Server;
+
+  afterInit(server: any): any {
+    // console.log('Init ws server', server, server === this.wsServer);
+    this.wsServer.path = '__koishi_fallback';
+    this.koishi._wsServer = this.wsServer;
+  }
+
+  handleConnection(socket: WebSocket, request: IncomingMessage) {
+    // console.log(socket);
+    for (const manager of this.koishi.router.wsStack) {
+      if (manager.accept(socket, request)) return;
+    }
+    socket.close();
+  }
+}