From 0a944b22dbdbfd8962716d81694db12de7a1f37a Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Tue, 20 Feb 2024 22:22:59 +0800
Subject: [PATCH] core: fix memory leak

---
 packages/hydrooj/src/service/router.ts |  9 ++++++---
 packages/hydrooj/src/service/server.ts | 11 +++++++----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/packages/hydrooj/src/service/router.ts b/packages/hydrooj/src/service/router.ts
index 1a64c612d..070d705c2 100644
--- a/packages/hydrooj/src/service/router.ts
+++ b/packages/hydrooj/src/service/router.ts
@@ -14,7 +14,7 @@ export class WebSocketLayer {
     clients = new Set<WebSocket>();
     regexp: RegExp;
 
-    constructor(private router: Router, path: Parameters<typeof pathToRegexp>[0], public callback?: WebSocketCallback) {
+    constructor(path: Parameters<typeof pathToRegexp>[0], public callback?: WebSocketCallback) {
         this.regexp = pathToRegexp(path);
     }
 
@@ -49,9 +49,12 @@ export class Router extends KoaRouter {
     }
 
     ws(path: Parameters<typeof pathToRegexp>[0], callback?: WebSocketCallback) {
-        const layer = new WebSocketLayer(this, path, callback);
+        const layer = new WebSocketLayer(path, callback);
         this.wsStack.push(layer);
-        this.disposeLastOp = () => remove(this.wsStack, layer);
+        this.disposeLastOp = () => {
+            layer.close();
+            remove(this.wsStack, layer);
+        };
         return layer;
     }
 }
diff --git a/packages/hydrooj/src/service/server.ts b/packages/hydrooj/src/service/server.ts
index c821771dd..a83b5081d 100644
--- a/packages/hydrooj/src/service/server.ts
+++ b/packages/hydrooj/src/service/server.ts
@@ -414,7 +414,7 @@ export function Connection(
     ...permPrivChecker: Array<number | bigint | Function>
 ) {
     const checker = Checker(permPrivChecker);
-    router.ws(prefix, async (conn, _, ctx) => {
+    const layer = router.ws(prefix, async (conn, _, ctx) => {
         const {
             args, request, response, user, domain, UiContext,
         } = ctx.HydroContext;
@@ -437,6 +437,7 @@ export function Connection(
                 if (closed) return;
                 closed = true;
                 bus.emit('connection/close', h);
+                layer.clients.delete(conn);
                 if (interval) clearInterval(interval);
                 for (const d of disposables) d();
                 h.cleanup?.(args);
@@ -446,7 +447,7 @@ export function Connection(
                     clean();
                     conn.terminate();
                 }
-                if (Date.now() - lastHeartbeat > 30000) conn.ping();
+                if (Date.now() - lastHeartbeat > 30000) conn.send('ping');
             }, 40000);
             conn.on('pong', () => {
                 lastHeartbeat = Date.now();
@@ -460,9 +461,11 @@ export function Connection(
                 }
                 h.message?.(JSON.parse(e.data.toString()));
             };
-            conn.onclose = clean;
-            conn.resume();
             await bus.parallel('connection/active', h);
+            if (conn.readyState === conn.OPEN) {
+                conn.on('close', clean);
+                conn.resume();
+            } else clean();
         } catch (e) {
             await h.onerror(e);
         }