core: connection: resume & support hooks

packages/hydrooj/src/handler/connection.ts

@@ -1,10 +1,14 @@
 import { Context } from '../context';
+import { ForbiddenError } from '../error';
 import avatar from '../lib/avatar';
 import { PRIV } from '../model/builtin';
 import * as SystemModel from '../model/system';
 import TokenModel from '../model/token';
 import UserModel from '../model/user';
 import { ConnectionHandler } from '../service/server';
+import { Logger } from '../utils';
+
+const logger = new Logger('connection');
 
 class WebsocketEventsConnectionManagerHandler extends ConnectionHandler {
     category = '#gateway';
@@ -13,36 +17,19 @@ class WebsocketEventsConnectionManagerHandler extends ConnectionHandler {
 
     async prepare() {
         const secret = SystemModel.get('websocket.secret');
-        if (secret && this.request.headers['x-hydro-websocket-gateway'] === secret) {
-            this.ctx.logger('connection').info('Gateway websocket connection from %s', this.request.ip);
+        const token = this.request.headers['x-hydro-websocket-gateway'];
+        if (token) {
+            if (!secret || token !== secret) throw new ForbiddenError('Invalid token');
+            logger.info('Gateway websocket connection from %s', this.request.ip);
             this.privileged = true;
-            this.ctx.on('user/message', async (uid, mdoc) => {
-                await this.notifyMessage(uid, mdoc);
-            });
         }
+        await this.ctx.parallel('subscription/init', this, this.privileged);
     }
 
-    async notifyMessage(uid: number[], mdoc: any) {
-        const udoc = (await UserModel.getById('system', mdoc.from))!;
-        udoc.avatarUrl = avatar(udoc.avatar, 64);
-        this.send({
-            operation: 'event',
-            channels: uid.map((u) => `message:${u}`),
-            payload: { udoc, mdoc },
-        });
-    }
-
-    accept(channel: string) {
+    async accept(channel: string) {
         if (this.channels.has(channel)) return;
         this.channels.add(channel);
-        if (this.privileged) return;
-        if (channel.startsWith('message:')) {
-            const [, uid] = channel.split(':');
-            this.ctx.on('user/message', async (uids, mdoc) => {
-                if (!uids.includes(+uid)) return;
-                await this.notifyMessage([+uid], mdoc);
-            });
-        }
+        await this.ctx.parallel('subscription/enable', channel, this, this.privileged);
     }
 
     async message(payload: any) {
@@ -51,15 +38,21 @@ class WebsocketEventsConnectionManagerHandler extends ConnectionHandler {
         const session = payload?.credential
             ? await TokenModel.get(payload.credential, TokenModel.TYPE_SESSION)
             : null;
-        if (!session) return;
-        const user = await UserModel.getById('system', session.uid);
+        if (!session && payload.operation !== 'resume') return;
+        if (payload.operation === 'resume' && !this.privileged) return;
+        const user = payload.operation === 'resume' ? null : await UserModel.getById('system', session.uid);
         for (const channel of payload.channels) {
-            if (channel === 'message' && user.hasPriv(PRIV.PRIV_USER_PROFILE)) {
-                accept.push(`message:${user._id}`);
-                this.accept(`message:${user._id}`);
-                continue;
+            try {
+                const result = payload.operation === 'resume'
+                    ? { ok: true, channel }
+                    // eslint-disable-next-line no-await-in-loop
+                    : await this.ctx.bail('subscription/subscribe', channel, user);
+                if (result?.ok) accept.push(result.channel);
+                else reject.push(channel);
+            } catch (e) {
+                logger.error('Error subscribing to channel %s for user %s: %s', channel, user._id, e);
+                reject.push(channel);
             }
-            reject.push(channel);
         }
         this.send({
             operation: 'verify',
@@ -72,4 +65,38 @@ class WebsocketEventsConnectionManagerHandler extends ConnectionHandler {
 
 export function apply(ctx: Context) {
     ctx.Connection('websocket_gateway', '/websocket', WebsocketEventsConnectionManagerHandler);
+
+    async function notifyMessage(uid: number[], mdoc: any, h) {
+        const udoc = (await UserModel.getById('system', mdoc.from))!;
+        return {
+            operation: 'event',
+            channels: uid.map((u) => `message:${u}`),
+            payload: { udoc: { ...udoc.serialize(h) as any, avatarUrl: avatar(udoc.avatar, 128) }, mdoc },
+        };
+    }
+
+    ctx.on('subscription/init', (h, privileged) => {
+        if (!privileged) return;
+        h.ctx.on('user/message', async (uid, mdoc) => {
+            h.send(await notifyMessage(uid, mdoc, h));
+        });
+    });
+
+    ctx.on('subscription/enable', (channel, h, privileged) => {
+        if (!channel.startsWith('message:') || privileged) return;
+        const uid = +channel.split(':')[1];
+        h.ctx.on('user/message', async (uids, mdoc) => {
+            if (!uids.includes(uid)) return;
+            h.send(await notifyMessage([uid], mdoc, h));
+        });
+    });
+
+    ctx.on('subscription/subscribe', (channel, user) => { // eslint-disable-line consistent-return
+        if (channel === 'message' && user.hasPriv(PRIV.PRIV_USER_PROFILE)) {
+            return {
+                ok: true,
+                channel: `message:${user._id}`,
+            };
+        }
+    });
 }

packages/hydrooj/src/service/bus.ts

@@ -22,6 +22,10 @@ export interface EventMap {
     'app/before-reload': (entries: Set<string>) => VoidReturn
     'app/reload': (entries: Set<string>) => VoidReturn
 
+    'subscription/init': (h: ConnectionHandler<Context>, privileged: boolean) => VoidReturn
+    'subscription/subscribe': (channel: string, user: User) => VoidReturn
+    'subscription/enable': (channel: string, h: ConnectionHandler<Context>, privileged: boolean) => VoidReturn
+
     'app/watch/change': (path: string) => VoidReturn
     'app/watch/unlink': (path: string) => VoidReturn