diff --git a/packages/hydrooj/src/handler/connection.ts b/packages/hydrooj/src/handler/connection.ts index 8605b5ac5..b33c0d4e8 100644 --- a/packages/hydrooj/src/handler/connection.ts +++ b/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}`, + }; + } + }); } diff --git a/packages/hydrooj/src/service/bus.ts b/packages/hydrooj/src/service/bus.ts index d499e6d11..8d8444a59 100644 --- a/packages/hydrooj/src/service/bus.ts +++ b/packages/hydrooj/src/service/bus.ts @@ -22,6 +22,10 @@ export interface EventMap { 'app/before-reload': (entries: Set) => VoidReturn 'app/reload': (entries: Set) => VoidReturn + 'subscription/init': (h: ConnectionHandler, privileged: boolean) => VoidReturn + 'subscription/subscribe': (channel: string, user: User) => VoidReturn + 'subscription/enable': (channel: string, h: ConnectionHandler, privileged: boolean) => VoidReturn + 'app/watch/change': (path: string) => VoidReturn 'app/watch/unlink': (path: string) => VoidReturn