core&ui: add events channel

packages/hydrooj/src/handler/connection.ts

@@ -0,0 +1,89 @@
+import { Context } from '../context';
+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';
+
+class WebsocketEventsConnectionManagerHandler extends ConnectionHandler {
+    category = '#gateway';
+    private privileged = false;
+    private channels: Set<string> = new Set();
+
+    async prepare() {
+        const secret = SystemModel.get('websocket.secret');
+        if (this.request.headers['x-hydro-websocket-gateway'] === secret) {
+            this.privileged = true;
+            this.ctx.on('user/message', async (uid, mdoc) => {
+                await this.notifyMessage(uid, mdoc);
+            });
+        }
+    }
+
+    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 },
+        });
+    }
+
+    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);
+            });
+        }
+    }
+
+    async message(payload: any) {
+        try {
+            payload = JSON.parse(payload);
+        } catch (e) {
+            console.error('Invalid payload', e);
+            return;
+        }
+        const success = [];
+        const failed = [];
+        const session = payload?.credential
+            ? await TokenModel.get(payload.credential, TokenModel.TYPE_SESSION)
+            : null;
+        if (!session) return;
+        const user = await UserModel.getById('system', session.uid);
+        for (const channel of payload.channels) {
+            if (channel === 'message' && user.hasPriv(PRIV.PRIV_USER_PROFILE)) {
+                success.push(`message:${user._id}`);
+                // eslint-disable-next-line no-await-in-loop
+                await this.accept(`message:${user._id}`);
+                continue;
+            }
+            failed.push(channel);
+        }
+        if (success.length) {
+            this.send({
+                operation: 'accept',
+                channels: success,
+                connection_id: payload.connection_id,
+            });
+        }
+        if (failed.length) {
+            this.send({
+                operation: 'reject',
+                channels: failed,
+                connection_id: payload.connection_id,
+            });
+        }
+    }
+}
+
+export function apply(ctx: Context) {
+    ctx.Connection('websocket_gateway', '/websocket', WebsocketEventsConnectionManagerHandler);
+}

packages/hydrooj/src/model/message.ts

@@ -18,17 +18,18 @@ class MessageModel {
 
     @ArgMethod
     static async send(
-        from: number, to: number,
+        from: number, to: number | number[],
         content: string, flag: number = MessageModel.FLAG_UNREAD,
     ) {
         const _id = new ObjectId();
-        const mdoc: MessageDoc = {
-            _id, from, to, content, flag,
+        if (!Array.isArray(to)) to = [to];
+        const base = {
+            _id, from, content, flag,
         };
-        await MessageModel.coll.insertOne(mdoc);
-        if (from !== to) bus.broadcast('user/message', to, mdoc);
+        await MessageModel.coll.insertMany(to.map((t) => ({ ...base, to: t })));
+        bus.broadcast('user/message', to, base);
         if (flag & MessageModel.FLAG_UNREAD) await user.inc(to, 'unreadMsg', 1);
-        return mdoc;
+        return base;
     }
 
     static async sendInfo(to: number, content: string) {
@@ -36,7 +37,7 @@ class MessageModel {
         const mdoc: MessageDoc = {
             _id, from: 1, to, content, flag: MessageModel.FLAG_INFO | MessageModel.FLAG_I18N,
         };
-        bus.broadcast('user/message', to, mdoc);
+        bus.broadcast('user/message', [to], mdoc);
     }
 
     static async get(_id: ObjectId) {

packages/hydrooj/src/model/user.ts

@@ -295,13 +295,14 @@ class UserModel {
     }
 
     @ArgMethod
-    static async inc(_id: number, field: string, n: number = 1) {
-        if (_id < -999) return null;
-        const udoc = await coll.findOne({ _id });
-        if (!udoc) throw new UserNotFoundError(_id);
-        await coll.updateOne({ _id }, { $inc: { [field]: n } });
-        deleteUserCache(udoc);
-        return udoc;
+    static async inc(_id: number | number[], field: string, n: number = 1) {
+        const ids = (Array.isArray(_id) ? _id : [_id]).filter((i) => i >= -999);
+        if (!ids.length) return null;
+        const udocs = await coll.find({ _id: { $in: ids } }).toArray();
+        if (udocs.length !== ids.length) throw new UserNotFoundError(_id);
+        await coll.updateMany({ _id: { $in: ids } }, { $inc: { [field]: n } });
+        for (const udoc of udocs) deleteUserCache(udoc);
+        return udocs;
     }
 
     @ArgMethod

packages/hydrooj/src/service/bus.ts

@@ -36,7 +36,7 @@ export interface EventMap {
     'task/daily': () => void;
     'task/daily/finish': (pref: Record<string, number>) => void;
 
-    'user/message': (uid: number, mdoc: MessageDoc) => void
+    'user/message': (uid: number[], mdoc: Omit<MessageDoc, 'to'>) => void
     'user/get': (udoc: User) => void
     'user/delcache': (content: string | true) => void
 

packages/ui-default/components/message/worker.ts

@@ -28,6 +28,31 @@ const ack = {};
 function broadcastMsg(message: any) {
   for (const p of ports) p.postMessage(message);
 }
+
+function onMessage(payload: any) {
+  broadcastMsg({ type: 'message', payload });
+  let acked = false;
+  ack[payload.mdoc._id] = () => { acked = true; };
+  setTimeout(() => {
+    delete ack[payload.mdoc._id];
+    if (acked) return;
+    if (payload.mdoc.flag & FLAG_INFO) return;
+    if (Notification?.permission !== 'granted') {
+      console.log('Notification permission denied');
+      return;
+    }
+    // eslint-disable-next-line no-new
+    new Notification(
+      payload.udoc.uname || 'Hydro Notification',
+      {
+        tag: `message-${payload.mdoc._id}`,
+        icon: payload.udoc.avatarUrl || '/android-chrome-192x192.png',
+        body: payload.mdoc.content,
+      },
+    );
+  }, 5000);
+}
+
 function initConn(path: string, port: MessagePort, cookie: any) {
   ports.add(port);
   if (cookie !== lcookie) conn?.close();
@@ -55,27 +80,7 @@ function initConn(path: string, port: MessagePort, cookie: any) {
       broadcastMsg({ type: 'close', error: payload.error });
       conn.close();
     } else {
-      broadcastMsg({ type: 'message', payload });
-      let acked = false;
-      ack[payload.mdoc._id] = () => { acked = true; };
-      setTimeout(() => {
-        delete ack[payload.mdoc._id];
-        if (acked) return;
-        if (payload.mdoc.flag & FLAG_INFO) return;
-        if (Notification?.permission !== 'granted') {
-          console.log('Notification permission denied');
-          return;
-        }
-        // eslint-disable-next-line no-new
-        new Notification(
-          payload.udoc.uname || 'Hydro Notification',
-          {
-            tag: `message-${payload.mdoc._id}`,
-            icon: payload.udoc.avatarUrl || '/android-chrome-192x192.png',
-            body: payload.mdoc.content,
-          },
-        );
-      }, 5000);
+      onMessage(payload);
     }
   };
 }