diff --git a/framework/framework/base.ts b/framework/framework/base.ts index 2dded4986..b7413117b 100644 --- a/framework/framework/base.ts +++ b/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 (streamOrBuffer instanceof Buffer) { + if (name) ctx.attachment(name); + if (streamOrBuffer instanceof Buffer || streamOrBuffer instanceof PassThrough) { response.body = null; ctx.body = streamOrBuffer; } else { diff --git a/framework/framework/package.json b/framework/framework/package.json index 955bfe371..af4bf1d25 100644 --- a/framework/framework/package.json +++ b/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", diff --git a/framework/framework/server.ts b/framework/framework/server.ts index 350a1a314..4b278664c 100644 --- a/framework/framework/server.ts +++ b/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; - if (host !== c.request.headers.host && `,${this.config.cors},`.includes(`,${host},`)) { - c.set('Access-Control-Allow-Credentials', 'true'); - c.set('Access-Control-Allow-Origin', c.request.headers.origin); - c.set('Access-Control-Allow-Headers', corsAllowHeaders); - c.set('Vary', 'Origin'); - c.cors = true; + try { + const host = new URL(c.request.headers.origin).host; + if (host !== c.request.headers.host && `,${this.config.cors},`.includes(`,${host},`)) { + c.set('Access-Control-Allow-Credentials', 'true'); + c.set('Access-Control-Allow-Origin', c.request.headers.origin); + c.set('Access-Control-Allow-Headers', corsAllowHeaders); + 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 (Date.now() - lastHeartbeat > 80000) { - clean(); - conn.terminate(); - } - if (Date.now() - lastHeartbeat > 30000) conn.send('ping'); - }, 40000); - conn.on('pong', () => { - lastHeartbeat = Date.now(); - }); - conn.onmessage = (e) => { - lastHeartbeat = Date.now(); - if (e.data === 'pong') return; - if (e.data === 'ping') { - conn.send('pong'); - return; - } - let payload; - try { - payload = JSON.parse(e.data.toString()); - } catch { - conn.close(); - } - try { - h.message?.(payload); - } catch (err) { - logger.error(e); - } - }; + if (layer) { + let lastHeartbeat = Date.now(); + interval = setInterval(() => { + if (Date.now() - lastHeartbeat > 80000) { + clean(); + conn.terminate(); + } + if (Date.now() - lastHeartbeat > 30000) conn.send('ping'); + }, 40000); + conn.on('pong', () => { + lastHeartbeat = Date.now(); + }); + conn.onmessage = (e) => { + lastHeartbeat = Date.now(); + if (e.data === 'pong') return; + if (e.data === 'ping') { + conn.send('pong'); + return; + } + let payload; + try { + payload = JSON.parse(e.data.toString()); + } catch { + conn.close(); + } + try { + h.message?.(payload); + } catch (err) { + logger.error(e); + } + }; + } else ctx.body = stream; await this.ctx.parallel('connection/active', h); - if (conn.readyState === conn.OPEN) { - conn.on('close', clean); - conn.resume(); - } else clean(); + if (layer) { + if (conn.readyState === conn.OPEN) { + conn.on('close', 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