framework: add event stream support

framework/framework/base.ts

@@ -27,8 +27,8 @@ export default (logger, xff, xhost) => async (ctx: KoaContext, next: Next) => {
         template: null,
         redirect: null,
         attachment: (name, streamOrBuffer) => {
-            ctx.attachment(name);
+            if (name) ctx.attachment(name);
-            if (streamOrBuffer instanceof Buffer) {
+            if (streamOrBuffer instanceof Buffer || streamOrBuffer instanceof PassThrough) {
                 response.body = null;
                 ctx.body = streamOrBuffer;
             } else {

framework/framework/package.json

@@ -10,7 +10,7 @@
         "@hydrooj/utils": "workspace:^",
         "emoji-regex": "^10.4.0",
         "fs-extra": "^11.3.0",
-        "koa": "^2.15.3",
+        "koa": "^2.15.4",
         "koa-body": "^6.0.1",
         "koa-compress": "^5.1.1",
         "koa-router": "^13.0.1",

framework/framework/server.ts

@@ -1,6 +1,7 @@
 import http from 'http';
 import { tmpdir } from 'os';
 import { join } from 'path';
+import { PassThrough } from 'stream';
 import { Context, Service } from 'cordis';
 import type { Files } from 'formidable';
 import fs from 'fs-extra';
@@ -345,7 +346,7 @@ export class WebService extends Service {
     ConnectionHandler = ConnectionHandler;
 
     constructor(ctx: Context, public config: WebServiceConfig) {
-        super(ctx, 'server', true);
+        super(ctx, 'server');
         ctx.mixin('server', ['Route', 'Connection', 'withHandlerClass']);
         this.server.keys = this.config.keys;
         this.server.proxy = this.config.proxy;
@@ -353,13 +354,17 @@ export class WebService extends Service {
         this.server.use(Compress());
         this.server.use(async (c, next) => {
             if (c.request.headers.origin && this.config.cors) {
-                const host = new URL(c.request.headers.origin).host;
+                try {
-                if (host !== c.request.headers.host && `,${this.config.cors},`.includes(`,${host},`)) {
+                    const host = new URL(c.request.headers.origin).host;
-                    c.set('Access-Control-Allow-Credentials', 'true');
+                    if (host !== c.request.headers.host && `,${this.config.cors},`.includes(`,${host},`)) {
-                    c.set('Access-Control-Allow-Origin', c.request.headers.origin);
+                        c.set('Access-Control-Allow-Credentials', 'true');
-                    c.set('Access-Control-Allow-Headers', corsAllowHeaders);
+                        c.set('Access-Control-Allow-Origin', c.request.headers.origin);
-                    c.set('Vary', 'Origin');
+                        c.set('Access-Control-Allow-Headers', corsAllowHeaders);
-                    c.cors = true;
+                        c.set('Vary', 'Origin');
+                        c.cors = true;
+                    }
+                } catch (e) {
+                    // invalid origin header, ignore
                 }
             }
             if (c.request.method.toLowerCase() === 'options') {
@@ -378,7 +383,7 @@ export class WebService extends Service {
                     await next();
                 } finally {
                     const endTime = Date.now();
-                    if (!(c.nolog || c.response.headers.nolog)) {
+                    if (!c.nolog && !c.response.headers.nolog) {
                         logger.debug(`${c.request.method} /${c.domainId || 'system'}${c.request.path} \
 ${c.response.status} ${endTime - startTime}ms ${c.response.length}`);
                     }
@@ -547,10 +552,33 @@ ${c.response.status} ${endTime - startTime}ms ${c.response.length}`);
         }
     }
 
-    private async handleWS(ctx: KoaContext, HandlerClass, checker, conn, layer) {
+    private async handleWS(ctx: KoaContext, HandlerClass, checker, conn, layer?) {
         const { args } = ctx.HydroContext;
         const h = new HandlerClass(ctx, this.ctx);
         await this.ctx.parallel('connection/create', h);
+        const stream = new PassThrough();
+        if (!conn) {
+            // By HTTP
+            ctx.request.socket.setTimeout(0);
+            ctx.req.socket.setNoDelay(true);
+            ctx.req.socket.setKeepAlive(true);
+            ctx.set({
+                'X-Accel-Buffering': 'no',
+                'Content-Type': 'text/event-stream',
+                'Cache-Control': 'no-cache',
+                Connection: 'keep-alive',
+            });
+            ctx.HydroContext.request.websocket = true;
+            ctx.compress = false;
+            conn = {
+                close() {
+                    stream.end();
+                },
+                send(data: any) {
+                    stream.write(`${args.sse ? 'data: ' : ''}${data}\n${args.sse ? '\n' : ''}`);
+                },
+            };
+        }
         ctx.handler = h;
         h.conn = conn;
         const disposables = [];
@@ -562,54 +590,59 @@ ${c.response.status} ${endTime - startTime}ms ${c.response.length}`);
             if (h.prepare) await h.prepare(args);
             // eslint-disable-next-line @typescript-eslint/no-shadow
             for (const { name, target } of h.__subscribe || []) disposables.push(this.ctx.on(name, target.bind(h)));
-            let lastHeartbeat = Date.now();
             let closed = false;
             let interval: NodeJS.Timeout;
             const clean = () => {
                 if (closed) return;
                 closed = true;
                 this.ctx.emit('connection/close', h);
-                layer.clients.delete(conn);
+                if (layer) layer.clients.delete(conn);
                 if (interval) clearInterval(interval);
                 for (const d of disposables) d();
                 h.cleanup?.(args);
             };
-            interval = setInterval(() => {
+            if (layer) {
-                if (Date.now() - lastHeartbeat > 80000) {
+                let lastHeartbeat = Date.now();
-                    clean();
+                interval = setInterval(() => {
-                    conn.terminate();
+                    if (Date.now() - lastHeartbeat > 80000) {
-                }
+                        clean();
-                if (Date.now() - lastHeartbeat > 30000) conn.send('ping');
+                        conn.terminate();
-            }, 40000);
+                    }
-            conn.on('pong', () => {
+                    if (Date.now() - lastHeartbeat > 30000) conn.send('ping');
-                lastHeartbeat = Date.now();
+                }, 40000);
-            });
+                conn.on('pong', () => {
-            conn.onmessage = (e) => {
+                    lastHeartbeat = Date.now();
-                lastHeartbeat = Date.now();
+                });
-                if (e.data === 'pong') return;
+                conn.onmessage = (e) => {
-                if (e.data === 'ping') {
+                    lastHeartbeat = Date.now();
-                    conn.send('pong');
+                    if (e.data === 'pong') return;
-                    return;
+                    if (e.data === 'ping') {
-                }
+                        conn.send('pong');
-                let payload;
+                        return;
-                try {
+                    }
-                    payload = JSON.parse(e.data.toString());
+                    let payload;
-                } catch {
+                    try {
-                    conn.close();
+                        payload = JSON.parse(e.data.toString());
-                }
+                    } catch {
-                try {
+                        conn.close();
-                    h.message?.(payload);
+                    }
-                } catch (err) {
+                    try {
-                    logger.error(e);
+                        h.message?.(payload);
-                }
+                    } catch (err) {
-            };
+                        logger.error(e);
+                    }
+                };
+            } else ctx.body = stream;
             await this.ctx.parallel('connection/active', h);
-            if (conn.readyState === conn.OPEN) {
+            if (layer) {
-                conn.on('close', clean);
+                if (conn.readyState === conn.OPEN) {
-                conn.resume();
+                    conn.on('close', clean);
-            } else clean();
+                    conn.resume();
+                } else clean();
+            } else stream.on('close', clean);
         } catch (e) {
             await h.onerror(e);
+            if (!layer) ctx.status = 500;
         }
     }
 
@@ -628,6 +661,7 @@ ${c.response.status} ${endTime - startTime}ms ${c.response.length}`);
             const layer = router.ws(path, async (conn, _req, ctx) => {
                 await this.handleWS(ctx as any, HandlerClass, checker, conn, layer);
             });
+            router.get(path, (ctx) => this.handleWS(ctx as any, HandlerClass, checker, null, null));
         }
         const dispose = router.disposeLastOp;
         // @ts-ignore