framework: add event stream support

This commit is contained in:
undefined committed 2025-02-19 03:50:58 +08:00
1 parent fc0fd0fe29
commit 6f744e9f97
3 files changed
+82 -48

No files matched your search

+2 -2
View File
@@ -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 {
+1 -1
View File
@@ -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",
+79 -45
View File
@@ -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