From 5cadd87712b2e48dc14bf56ae1984a15f76666ac Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Tue, 17 Jun 2025 04:47:59 +0800
Subject: [PATCH] core: connection: add unsubscribe callback

---
 packages/hydrooj/src/handler/connection.ts | 29 ++++++++++++++++++++--
 packages/hydrooj/src/service/bus.ts        |  4 ++-
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/packages/hydrooj/src/handler/connection.ts b/packages/hydrooj/src/handler/connection.ts
index 9b17143f8..7094a2310 100644
--- a/packages/hydrooj/src/handler/connection.ts
+++ b/packages/hydrooj/src/handler/connection.ts
@@ -14,6 +14,7 @@ class WebsocketEventsConnectionManagerHandler extends ConnectionHandler {
     category = '#gateway';
     private privileged = false;
     private channels: Set<string> = new Set();
+    private disposables: Map<string, () => void> = new Map();
 
     async prepare() {
         const secret = SystemModel.get('websocket.secret');
@@ -29,11 +30,35 @@ class WebsocketEventsConnectionManagerHandler extends ConnectionHandler {
     async accept(channel: string) {
         if (this.channels.has(channel)) return;
         this.channels.add(channel);
-        await this.ctx.parallel('subscription/enable', channel, this, this.privileged);
+        const d: (() => void)[] = [];
+        const disposeAll = () => {
+            for (const dispose of d) dispose();
+        };
+        this.disposables.set(channel, disposeAll);
+        await this.ctx.parallel('subscription/enable', channel, this, this.privileged, (disposable) => {
+            d.push(disposable);
+        });
     }
 
     async message(payload: any) {
-        if (!['resume', 'subscribe'].includes(payload.operation)) return;
+        if (['resume', 'subscribe'].includes(payload.operation)) await this.subscribe(payload);
+        if (payload.operation === 'unsubscribe') await this.unsubscribe(payload);
+    }
+
+    async unsubscribe(payload: any) {
+        for (const channel of payload.channels || []) {
+            try {
+                if (!this.channels.has(channel)) continue;
+                this.channels.delete(channel);
+                this.disposables.get(channel)?.();
+                this.disposables.delete(channel);
+            } catch (e) {
+                logger.error('Error unsubscribing from channel %s: %s', channel, e);
+            }
+        }
+    }
+
+    async subscribe(payload: any) {
         const accept = [];
         const reject = [];
         const session = payload.credential
diff --git a/packages/hydrooj/src/service/bus.ts b/packages/hydrooj/src/service/bus.ts
index f465bfdc8..1ecde1fb7 100644
--- a/packages/hydrooj/src/service/bus.ts
+++ b/packages/hydrooj/src/service/bus.ts
@@ -24,7 +24,9 @@ export interface EventMap {
 
     'subscription/init': (h: ConnectionHandler<Context>, privileged: boolean) => VoidReturn
     'subscription/subscribe': (channel: string, user: User, metadata: Record<string, string>) => VoidReturn
-    'subscription/enable': (channel: string, h: ConnectionHandler<Context>, privileged: boolean) => VoidReturn
+    'subscription/enable': (
+        channel: string, h: ConnectionHandler<Context>, privileged: boolean, onDispose: (disposable: () => void) => void
+    ) => VoidReturn
 
     'app/watch/change': (path: string) => VoidReturn
     'app/watch/unlink': (path: string) => VoidReturn