From 1e62a9b027d095815b14f9998246101f829902b7 Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Sun, 6 Nov 2022 06:34:39 +0800
Subject: [PATCH] core: handle Handler registration on multi paths

---
 packages/hydrooj/src/plugin-api.ts     |  1 +
 packages/hydrooj/src/service/server.ts | 17 +++++++++++------
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/packages/hydrooj/src/plugin-api.ts b/packages/hydrooj/src/plugin-api.ts
index f8611033a..720249a1a 100644
--- a/packages/hydrooj/src/plugin-api.ts
+++ b/packages/hydrooj/src/plugin-api.ts
@@ -8,6 +8,7 @@ import superagent from 'superagent';
 import { Context } from './context';
 import db from './service/db';
 export { ObjectID, ObjectId, FilterQuery } from 'mongodb';
+export { WebSocket, WebSocketServer } from 'ws';
 export * from './utils';
 export * from './interface';
 export * from './pipelineUtils';
diff --git a/packages/hydrooj/src/service/server.ts b/packages/hydrooj/src/service/server.ts
index 5143710f6..825292ffc 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 { isClass, parseMemoryMB } from '@hydrooj/utils/lib/utils';
+import { Counter, isClass, parseMemoryMB } from '@hydrooj/utils/lib/utils';
 import { Context, Service } from '../context';
 import {
     HydroError, InvalidOperationError, MethodNotAllowedError,
@@ -389,6 +389,7 @@ class NotFoundHandler extends Handler {
 export class RouteService extends Service {
     static readonly methods = ['Route', 'Connection', 'withHandlerClass'];
     private registry = {};
+    private registrationCount = Counter();
 
     constructor(ctx) {
         super(ctx, 'server', true);
@@ -396,15 +397,19 @@ export class RouteService extends Service {
 
     private register(func: typeof Route | typeof Connection, ...args: Parameters<typeof Route>) {
         const HandlerClass = args[2];
+        const name = HandlerClass.name;
         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];
+        if (this.registrationCount[name] && this.registry[name] !== HandlerClass) {
+            logger.warn('Route with name %s already exists.', name);
+        }
+        this.registry[name] = HandlerClass;
+        this.registrationCount[name]++;
         const res = func(...args);
-        this.ctx.parallel(`handler/register/${HandlerClass.name}`, HandlerClass);
+        this.ctx.parallel(`handler/register/${name}`, HandlerClass);
         this.caller?.on('dispose', () => {
+            this.registrationCount[name]--;
+            if (!this.registrationCount[name]) delete this.registry[name];
             res();
-            dispose();
         });
     }