From 84309e86fd8dfa84c0c5c4d4f41d1b1ae337df6a Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Sun, 6 Nov 2022 05:43:28 +0800
Subject: [PATCH] core: add withHandlerClass api

---
 packages/hydrooj/src/context.ts        |  1 +
 packages/hydrooj/src/service/bus.ts    |  1 +
 packages/hydrooj/src/service/server.ts | 51 ++++++++++++++++++++------
 3 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/packages/hydrooj/src/context.ts b/packages/hydrooj/src/context.ts
index 45420557d..f6995989a 100644
--- a/packages/hydrooj/src/context.ts
+++ b/packages/hydrooj/src/context.ts
@@ -25,6 +25,7 @@ export interface Context {
     loader: Loader;
     Route: typeof import('./service/server').Route;
     Connection: typeof import('./service/server').Connection;
+    withHandlerClass: import('./service/server').RouteService['withHandlerClass'];
     setTimeout: typeof setTimeout;
     setInterval: typeof setInterval;
     setImmediate: typeof setImmediate;
diff --git a/packages/hydrooj/src/service/bus.ts b/packages/hydrooj/src/service/bus.ts
index d4e1224c3..3b0343915 100644
--- a/packages/hydrooj/src/service/bus.ts
+++ b/packages/hydrooj/src/service/bus.ts
@@ -26,6 +26,7 @@ type HandlerEvents =
     KnownHandlerEvents
     & Record<`handler/${HookType}/${string}`, (thisArg: Handler & Record<string, any>) => VoidReturn>
     & Record<`handler/${HookType}`, (thisArg: Handler) => VoidReturn>
+    & Record<`handler/register/${string}`, (HandlerClass: typeof Handler) => VoidReturn>
     & Record<`connection/${'create' | 'active' | 'close'}`, (thisArg: ConnectionHandler) => VoidReturn>;
 
 /* eslint-disable @typescript-eslint/naming-convention */
diff --git a/packages/hydrooj/src/service/server.ts b/packages/hydrooj/src/service/server.ts
index 46705a3df..5143710f6 100644
--- a/packages/hydrooj/src/service/server.ts
+++ b/packages/hydrooj/src/service/server.ts
@@ -8,7 +8,7 @@ import Compress from 'koa-compress';
 import proxy from 'koa-proxies';
 import cache from 'koa-static-cache';
 import WebSocket from 'ws';
-import { parseMemoryMB } from '@hydrooj/utils/lib/utils';
+import { isClass, parseMemoryMB } from '@hydrooj/utils/lib/utils';
 import { Context, Service } from '../context';
 import {
     HydroError, InvalidOperationError, MethodNotAllowedError,
@@ -386,26 +386,54 @@ class NotFoundHandler extends Handler {
     all() { }
 }
 
-class RouteService extends Service {
-    static readonly methods = ['Route', 'Connection'];
+export class RouteService extends Service {
+    static readonly methods = ['Route', 'Connection', 'withHandlerClass'];
+    private registry = {};
+
     constructor(ctx) {
         super(ctx, 'server', true);
     }
 
-    // eslint-disable-next-line @typescript-eslint/naming-convention
-    Route(...args: Parameters<typeof Route>) {
-        const res = Route(...args);
-        this.caller?.on('dispose', () => res());
+    private register(func: typeof Route | typeof Connection, ...args: Parameters<typeof Route>) {
+        const HandlerClass = args[2];
+        if (!isClass(HandlerClass)) throw new Error('Invalid registration.');
+        if (this.registry[HandlerClass.name]) logger.warn('Route name with %s already exists.', HandlerClass.name);
+        this.registry[HandlerClass.name] = HandlerClass;
+        const dispose = () => delete this.registry[HandlerClass.name];
+        const res = func(...args);
+        this.ctx.parallel(`handler/register/${HandlerClass.name}`, HandlerClass);
+        this.caller?.on('dispose', () => {
+            res();
+            dispose();
+        });
+    }
+
+    public withHandlerClass(name: string, callback: (HandlerClass: typeof HandlerCommon) => any) {
+        if (this.registry[name]) callback(this.registry[name]);
+        this.ctx.on(`handler/register/${name}`, callback);
+        this.caller?.on('dispose', () => {
+            this.ctx.off(`handler/register/${name}`, callback);
+        });
     }
 
     // eslint-disable-next-line @typescript-eslint/naming-convention
-    Connection(...args: Parameters<typeof Connection>) {
-        const res = Connection(...args);
-        this.caller?.on('dispose', () => res());
+    public Route(...args: Parameters<typeof Route>) {
+        this.register(Route, ...args);
+    }
+
+    // eslint-disable-next-line @typescript-eslint/naming-convention
+    public Connection(...args: Parameters<typeof Connection>) {
+        this.register(Connection, ...args);
     }
 }
 
-export async function apply(pluginContext) {
+declare module '../context' {
+    interface Context {
+        server: RouteService;
+    }
+}
+
+export async function apply(pluginContext: Context) {
     Context.service('server', RouteService);
     pluginContext.server = new RouteService(pluginContext);
     app.keys = system.get('session.keys') as unknown as string[];
@@ -490,6 +518,7 @@ global.Hydro.service.server = {
     wsServer,
     router,
     captureAllRoutes,
+    RouteService,
     HandlerCommon,
     Handler,
     ConnectionHandler,