88 lines
2.9 KiB
TypeScript
88 lines
2.9 KiB
TypeScript
import { EventEmitter } from 'events';
|
|
import { WebSocketServer } from 'ws';
|
|
import type WebSocket from 'ws';
|
|
import type { Server as HttpServer } from 'http';
|
|
import type { Server as HttpsServer } from 'https';
|
|
import { verifyWebSocketToken } from '../auth.js';
|
|
import { WebSocketConnection } from './websocket-connection.js';
|
|
import type { ClientConnection } from './client-connection.js';
|
|
import type { ClientMessage } from '../types/messages.js';
|
|
|
|
/**
|
|
* WebSocketTransport — manages the WebSocket server lifecycle.
|
|
*
|
|
* Emits:
|
|
* 'connection' (conn: ClientConnection) — new authenticated client connected
|
|
* 'message' (conn: ClientConnection, msg: ClientMessage) — parsed message from client
|
|
*/
|
|
export class WebSocketTransport extends EventEmitter {
|
|
private wss: WebSocketServer | null = null;
|
|
private pingInterval: ReturnType<typeof setInterval> | null = null;
|
|
// Tracks whether each WS responded to the last ping. Initialized true so a
|
|
// connection is never terminated before it has a chance to respond.
|
|
private alive = new WeakMap<WebSocket, boolean>();
|
|
|
|
/** Create WebSocketServer on /ws path with JWT verification and ping/pong keepalive. */
|
|
setup(server: HttpServer | HttpsServer): void {
|
|
this.wss = new WebSocketServer({
|
|
server,
|
|
path: '/ws',
|
|
verifyClient: ({ req }, cb) => {
|
|
const url = new URL(req.url!, `http://${req.headers.host}`);
|
|
const token = url.searchParams.get('token');
|
|
if (!token || !verifyWebSocketToken(token)) {
|
|
cb(false, 401, 'Unauthorized');
|
|
return;
|
|
}
|
|
cb(true);
|
|
},
|
|
});
|
|
|
|
this.wss.on('connection', (ws: WebSocket) => {
|
|
this.alive.set(ws, true);
|
|
ws.on('pong', () => this.alive.set(ws, true));
|
|
|
|
const conn = new WebSocketConnection(ws);
|
|
|
|
this.emit('connection', conn);
|
|
|
|
ws.on('message', (raw: Buffer | ArrayBuffer | Buffer[]) => {
|
|
let msg: ClientMessage;
|
|
try {
|
|
msg = JSON.parse(raw.toString()) as ClientMessage;
|
|
} catch {
|
|
conn.send({ type: 'error', error: 'Invalid JSON' });
|
|
return;
|
|
}
|
|
this.emit('message', conn, msg);
|
|
});
|
|
});
|
|
|
|
// Ping/pong keepalive every 30s — terminate connections that miss a pong.
|
|
this.pingInterval = setInterval(() => {
|
|
if (!this.wss) return;
|
|
for (const ws of this.wss.clients) {
|
|
if (!this.alive.get(ws)) {
|
|
ws.terminate(); // no pong since last ping — dead connection
|
|
continue;
|
|
}
|
|
this.alive.set(ws, false);
|
|
if (ws.readyState === 1) ws.ping();
|
|
}
|
|
}, 30_000);
|
|
this.pingInterval.unref();
|
|
}
|
|
|
|
/** Shut down the WebSocket server and stop keepalive. */
|
|
destroy(): void {
|
|
if (this.pingInterval) {
|
|
clearInterval(this.pingInterval);
|
|
this.pingInterval = null;
|
|
}
|
|
if (this.wss) {
|
|
this.wss.close();
|
|
this.wss = null;
|
|
}
|
|
}
|
|
}
|