dashcaddy/dashcaddy-api/src/websocket/dashboard-ws.js

/**
 * DC-076: WebSocket server for real-time dashboard updates
 *
 * Runs alongside the existing SSE endpoint (/api/v1/events/stream).
 * Shares the same event broadcasts but over a bidirectional WebSocket
 * connection, enabling client→server commands (e.g. "subscribe to
 * container X", "set alert threshold").
 *
 * Protocol: JSON messages with {type, data} envelope.
 * Server→client: {type: 'event', event: '<name>', data: {...}}
 * Client→server: {type: 'subscribe', events: ['resource-alert', ...]}
 *                {type: 'ping'} → {type: 'pong'}
 */
const { WebSocketServer } = require('ws');

/**
 * Parse the `Cookie` header into a plain `{name: value}` map.
 * WS upgrade requests don't go through Express's cookie-parser, so we
 * do it by hand here. We deliberately do NOT decode the values — the
 * session-cookie HMAC verifier reads the raw cookie string verbatim
 * (`payloadB64.sig` shape), so any decoding (e.g. url-decode) would
 * corrupt the signature. Single cookie-pair per call, no nesting.
 *
 * @param {string|undefined} header - Raw Cookie header value
 * @returns {Object<string, string>} name → value map (empty string for blanks)
 */
function parseCookieHeader(header) {
  const out = {};
  if (!header) return out;
  for (const part of header.split(';')) {
    const idx = part.indexOf('=');
    if (idx === -1) continue;
    const name = part.slice(0, idx).trim();
    if (!name) continue;
    const value = part.slice(idx + 1).trim();
    out[name] = value;
  }
  return out;
}

function createDashboardWS(server, deps = {}) {
  const wss = new WebSocketServer({ noServer: true });

  // Event broadcasters that the events.js SSE route already wires up.
  // We listen to the same EventEmitters and forward to WS clients.
  const {
    resourceMonitor,
    healthChecker,
    updateManager,
    dependencyManager,
    autoRestartManager,
    driftDetector,
    sslMonitor,
    dnsPropagationChecker,
    log,
  } = deps;

  // ── Auth verifier (injected by server.js from app.locals.ctx.session) ──
  // The WS upgrade path bypasses Express middleware, so the global
  // `totpAuthMiddleware` (which calls `isSessionValid(req)`) never runs.
  // We accept the SAME verifier here so a valid browser session cookie
  // grants access and nothing else does.
  //
  // The injected verifier receives the raw HTTP upgrade request (an
  // IncomingMessage with `.headers.cookie`). Production wires
  // `app.locals.ctx.session.isValid` directly — it accepts the same
  // shape, parses the Cookie header internally, and runs the HMAC
  // check. Tests inject a stub.
  //
  // DC-061 hardening: prior code only checked that the `dashcaddy_session`
  // SUBSTRING appeared in the Cookie header. That let an attacker set any
  // cookie named `dashcaddy_session=garbage` (or include the literal text
  // in another cookie's value) and bypass auth. The injected verifier
  // runs HMAC validation, so a present-but-invalid cookie now 401s.
  const authVerifier = typeof deps.authVerifier === 'function'
    ? deps.authVerifier
    : (req) => {
        // Last-resort fallback: presence-only check on a non-empty
        // session-cookie value. Used only when the caller didn't inject
        // a real verifier (e.g. tests, unusual boot paths). Production
        // wires the real one from app.locals.ctx.session.isValid.
        const parsed = parseCookieHeader(req && req.headers && req.headers.cookie);
        const raw = parsed.dashcaddy_session || parsed.sid;
        return typeof raw === 'string' && raw.length > 0;
      };

  // Track connected clients and their subscriptions
  const wsClients = new Set();

  // Track the listener functions we attach to shared EventEmitters so we
  // can detach exactly OUR listeners on close() — without disturbing the
  // SSE route's listeners on the same emitters. DC-061 critical fix:
  // the previous code called `resourceMonitor.removeAllListeners()` which
  // silently killed the SSE route's `alert`/`status-check`/etc subscribers
  // whenever close() ran (hot reload, graceful restart).
  const emitterListeners = [];

  function attachListener(emitter, event, handler) {
    if (!emitter || typeof emitter.on !== 'function') return;
    emitter.on(event, handler);
    emitterListeners.push({ emitter, event, handler });
  }

  function broadcast(event, data) {
    const msg = JSON.stringify({ type: 'event', event, data });
    for (const client of wsClients) {
      if (client.readyState !== 1) continue; // OPEN only
      // Check subscription filter
      if (client.subscribedEvents && !client.subscribedEvents.has(event)) continue;
      try {
        client.send(msg);
      } catch {
        wsClients.delete(client);
      }
    }
  }

  // ── Wire up EventEmitter listeners (same events as SSE) ──

  attachListener(resourceMonitor, 'alert', (data) => broadcast('resource-alert', data));
  attachListener(resourceMonitor, 'auto-restart', (data) => broadcast('auto-restart', data));

  attachListener(healthChecker, 'status-check', (data) => {
    broadcast('status-change', {
      serviceId: data.serviceId,
      name: data.name,
      status: data.status,
      responseTime: data.responseTime,
      timestamp: data.timestamp,
    });
  });
  attachListener(healthChecker, 'incident-created', (data) => broadcast('incident', { type: 'created', ...data }));
  attachListener(healthChecker, 'incident-resolved', (data) => broadcast('incident', { type: 'resolved', ...data }));

  attachListener(updateManager, 'update-available', (data) => broadcast('update-available', data));
  attachListener(updateManager, 'update-start', (data) => broadcast('update-start', data));
  attachListener(updateManager, 'update-complete', (data) => broadcast('update-complete', data));
  attachListener(updateManager, 'update-failed', (data) => broadcast('update-failed', data));
  attachListener(updateManager, 'auto-update-start', (data) => broadcast('auto-update-start', data));
  attachListener(updateManager, 'auto-update-complete', (data) => broadcast('auto-update-complete', data));

  attachListener(dependencyManager, 'dependency-restart-start', (data) => broadcast('dependency-restart-start', data));
  attachListener(dependencyManager, 'dependency-restart-progress', (data) => broadcast('dependency-restart-progress', data));
  attachListener(dependencyManager, 'dependency-restart-complete', (data) => broadcast('dependency-restart-complete', data));
  attachListener(dependencyManager, 'dependency-restart-failed', (data) => broadcast('dependency-restart-failed', data));

  attachListener(autoRestartManager, 'auto-restart-attempt', (data) => broadcast('auto-restart-attempt', data));
  attachListener(autoRestartManager, 'auto-restart-success', (data) => broadcast('auto-restart-success', data));
  attachListener(autoRestartManager, 'auto-restart-failed', (data) => broadcast('auto-restart-failed', data));
  attachListener(autoRestartManager, 'auto-restart-max-reached', (data) => broadcast('auto-restart-max-reached', data));

  attachListener(driftDetector, 'drift-detected', (data) => broadcast('drift-detected', data));

  attachListener(sslMonitor, 'cert-expiring', (data) => broadcast('cert-expiring', data));
  attachListener(sslMonitor, 'cert-critical', (data) => broadcast('cert-critical', data));

  attachListener(dnsPropagationChecker, 'propagation-check', (data) => broadcast('dns-propagation-check', data));
  attachListener(dnsPropagationChecker, 'propagation-complete', (data) => broadcast('dns-propagation-complete', data));
  attachListener(dnsPropagationChecker, 'propagation-timeout', (data) => broadcast('dns-propagation-timeout', data));

  // ── Handle upgrade requests at /api/v1/ws ──

  server.on('upgrade', (request, socket, head) => {
    const url = new URL(request.url, 'http://localhost');

    // Only handle exact /api/v1/ws path — the exec WS handler manages its own path
    if (url.pathname !== '/api/v1/ws' && url.pathname !== '/ws/dashboard') {
      return; // Let other upgrade handlers deal with it
    }

    // DC-061 auth gate: WS upgrade bypasses Express middleware, so we
    // must validate the session here. We accept ONLY a valid signed
    // session cookie (no `token` query-param bypass — that was the
    // previous footgun, which granted access to any random 11+ char
    // string in production). The verifier is injected from
    // app.locals.ctx.session.isValid in production.
    const ok = authVerifier(request);

    if (!ok) {
      const ip = (request.socket && request.socket.remoteAddress) || 'unknown';
      if (log && log.warn) {
        log.warn('websocket', 'WS upgrade rejected — no valid session', { ip, path: url.pathname });
      }
      // 401 + Connection: close so the client doesn't retry.
      socket.write('HTTP/1.1 401 Unauthorized\r\nConnection: close\r\n\r\n');
      socket.destroy();
      return;
    }

    wss.handleUpgrade(request, socket, head, (ws) => {
      wss.emit('connection', ws, request);
    });
  });

  // ── Connection handler ──

  wss.on('connection', (ws, req) => {
    ws.subscribedEvents = null; // null = receive all events
    wsClients.add(ws);

    if (log) {
      log.info('websocket', 'Client connected', { total: wsClients.size });
    }

    // Send welcome message
    ws.send(JSON.stringify({
      type: 'connected',
      data: { clients: wsClients.size },
    }));

    // Heartbeat every 30s
    ws.isAlive = true;
    const heartbeat = setInterval(() => {
      if (ws.readyState !== 1) {
        clearInterval(heartbeat);
        return;
      }
      ws.isAlive = false;
      try {
        ws.ping();
      } catch {
        clearInterval(heartbeat);
        wsClients.delete(ws);
      }
    }, 30000);

    ws.on('pong', () => { ws.isAlive = true; });

    ws.on('message', (raw) => {
      // Cap message size at 16 KB — defense-in-depth against a malicious
      // peer that exploits ws's message framing to flood our parser.
      // The `ws` library already enforces this via its constructor option,
      // but a second guard at the handler level catches any future
      // regressions (e.g. someone passing `maxPayload` differently).
      if (raw.length > 16 * 1024) {
        ws.send(JSON.stringify({ type: 'error', error: 'Message too large' }));
        try { ws.close(1009, 'Message too large'); } catch { /* already closed */ }
        return;
      }

      let msg;
      try {
        msg = JSON.parse(raw.toString());
      } catch {
        ws.send(JSON.stringify({ type: 'error', error: 'Invalid JSON' }));
        return;
      }

      switch (msg.type) {
        case 'subscribe':
          if (Array.isArray(msg.events)) {
            ws.subscribedEvents = new Set(msg.events);
            ws.send(JSON.stringify({ type: 'subscribed', events: msg.events }));
          }
          break;

        case 'unsubscribe':
          // Actually unsubscribe — set to empty set so no events are received
          ws.subscribedEvents = new Set();
          ws.send(JSON.stringify({ type: 'unsubscribed' }));
          break;

        case 'subscribe-all':
          // Reset to receive ALL events
          ws.subscribedEvents = null;
          ws.send(JSON.stringify({ type: 'subscribed-all' }));
          break;

        case 'ping':
          ws.send(JSON.stringify({ type: 'pong' }));
          break;

        case 'client-count':
          ws.send(JSON.stringify({ type: 'client-count', count: wsClients.size }));
          break;

        default:
          // Unknown message — ignore silently
          break;
      }
    });

    ws.on('close', () => {
      clearInterval(heartbeat);
      wsClients.delete(ws);
      if (log) {
        log.info('websocket', 'Client disconnected', { total: wsClients.size });
      }
    });

    ws.on('error', () => {
      clearInterval(heartbeat);
      wsClients.delete(ws);
    });
  });

  // Periodic sweep for dead connections
  const sweepInterval = setInterval(() => {
    for (const ws of wss.clients) {
      if (!ws.isAlive) {
        ws.terminate();
        wsClients.delete(ws);
      }
    }
  }, 60000);
  sweepInterval.unref();

  return {
    wss,
    getClientCount: () => wsClients.size,
    broadcast,
    close: () => {
      clearInterval(sweepInterval);
      for (const ws of wss.clients) {
        ws.terminate();
      }
      wsClients.clear();
      wss.close();
      // DC-061: detach ONLY the listeners we attached. Previously the
      // module called `resourceMonitor.removeAllListeners()` (and same
      // for healthChecker / updateManager), which silently wiped the
      // SSE route's listeners on the same shared emitters — the SSE
      // stream went dead the moment close() ran (hot reload, restart).
      for (const { emitter, event, handler } of emitterListeners) {
        if (emitter && typeof emitter.removeListener === 'function') {
          emitter.removeListener(event, handler);
        }
      }
      emitterListeners.length = 0;
    },
  };
}

module.exports = createDashboardWS;
module.exports.createDashboardWS = createDashboardWS;
module.exports.parseCookieHeader = parseCookieHeader;