clawtap/server/transport/websocket-transport.ts

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 in noServer mode. Call handleUpgrade() from your server upgrade handler. */
  setup(_server: HttpServer | HttpsServer): void {
    this.wss = new WebSocketServer({
      noServer: true,
      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();
  }

  /** Route an HTTP upgrade request into this WSS (called from server's upgrade handler). */
  handleUpgrade(req: import('http').IncomingMessage, socket: import('stream').Duplex, head: Buffer): void {
    this.wss!.handleUpgrade(req, socket, head, (ws) => {
      this.wss!.emit('connection', ws, req);
    });
  }

  /** 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;
    }
  }
}