85 lines
2.7 KiB
JavaScript
85 lines
2.7 KiB
JavaScript
const { WebSocketServer } = require('ws');
|
|
const sessionMiddleware = require('./session');
|
|
const queue = require('./queue');
|
|
const log = require('./logger');
|
|
|
|
let wss = null;
|
|
// Map jobId → Set of connected clients
|
|
const subscribers = new Map();
|
|
|
|
// Minimal response stub so express-session can run during the HTTP upgrade handshake.
|
|
const resStub = { getHeader() {}, setHeader() {}, writeHead() {}, on() {}, once() {}, removeListener() {}, emit() {}, end() {} };
|
|
|
|
function attach(server) {
|
|
// noServer + manual upgrade so we can authenticate before accepting the socket.
|
|
wss = new WebSocketServer({ noServer: true, maxPayload: 8 * 1024 });
|
|
|
|
server.on('upgrade', (req, socket, head) => {
|
|
let pathname;
|
|
try { pathname = new URL(req.url, 'http://localhost').pathname; } catch { pathname = req.url; }
|
|
if (pathname !== '/ws') return; // not ours — leave for others / default handling
|
|
|
|
sessionMiddleware(req, resStub, () => {
|
|
if (!req.session || !req.session.user) {
|
|
socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
|
|
socket.destroy();
|
|
return;
|
|
}
|
|
wss.handleUpgrade(req, socket, head, (ws) => {
|
|
ws._user = req.session.user;
|
|
wss.emit('connection', ws, req);
|
|
});
|
|
});
|
|
});
|
|
|
|
wss.on('connection', (ws) => {
|
|
ws._jobId = null;
|
|
|
|
ws.on('message', (data) => {
|
|
try {
|
|
const msg = JSON.parse(data);
|
|
if (msg.type === 'subscribe' && msg.jobId) {
|
|
// Only allow subscribing to one's OWN job (prevents cross-user IDOR)
|
|
const job = queue.getJob(String(msg.jobId));
|
|
if (!job || String(job.user_id) !== String(ws._user && ws._user.id)) return;
|
|
// Unsubscribe from previous
|
|
if (ws._jobId) {
|
|
const prev = subscribers.get(ws._jobId);
|
|
if (prev) { prev.delete(ws); if (prev.size === 0) subscribers.delete(ws._jobId); }
|
|
}
|
|
// Subscribe to new job
|
|
ws._jobId = String(msg.jobId);
|
|
if (!subscribers.has(ws._jobId)) subscribers.set(ws._jobId, new Set());
|
|
subscribers.get(ws._jobId).add(ws);
|
|
}
|
|
} catch {}
|
|
});
|
|
|
|
ws.on('close', () => {
|
|
if (ws._jobId) {
|
|
const subs = subscribers.get(ws._jobId);
|
|
if (subs) { subs.delete(ws); if (subs.size === 0) subscribers.delete(ws._jobId); }
|
|
}
|
|
});
|
|
});
|
|
|
|
log.info('WebSocket server attached at /ws (authenticated)');
|
|
}
|
|
|
|
function notify(jobId, data) {
|
|
const subs = subscribers.get(String(jobId));
|
|
if (!subs || subs.size === 0) return;
|
|
const msg = JSON.stringify({ jobId, ...data });
|
|
for (const ws of subs) {
|
|
if (ws.readyState === 1) { // OPEN
|
|
ws.send(msg);
|
|
}
|
|
}
|
|
}
|
|
|
|
function getConnectionCount() {
|
|
return wss ? wss.clients.size : 0;
|
|
}
|
|
|
|
module.exports = { attach, notify, getConnectionCount };
|