feuerwerk/tools/mcp/websocket.mjs

import { createHash } from 'node:crypto';

/**
 * The server half of RFC 6455, as much as the relay needs (#429): text messages, fragmented or not,
 * ping and pong, close. Node brings a WebSocket client of its own but no server, and the relay keeps
 * without dependencies, like the rest of tools/.
 *
 * Binary messages are refused (1003): the tab only ever speaks JSON.
 */

const GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11';
const utf8 = new TextDecoder('utf-8', { fatal: true });

export function acceptKey(key) {
  return createHash('sha1')
    .update(key + GUID)
    .digest('base64');
}

/**
 * Answers an upgrade request with 101 and hands back the connection, or refuses it with 400 and
 * answers null.
 */
export function acceptWebSocket(req, socket, head, { maxMessageBytes }) {
  const key = req.headers['sec-websocket-key'];
  if (
    String(req.headers.upgrade).toLowerCase() !== 'websocket' ||
    typeof key !== 'string' ||
    !/^[A-Za-z0-9+/]{22}==$/.test(key) ||
    req.headers['sec-websocket-version'] !== '13'
  ) {
    socket.end('HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n');
    return null;
  }
  socket.write(
    [
      'HTTP/1.1 101 Switching Protocols',
      'Upgrade: websocket',
      'Connection: Upgrade',
      `Sec-WebSocket-Accept: ${acceptKey(key)}`,
      '',
      '',
    ].join('\r\n'),
  );
  return new WebSocketConnection(socket, head, maxMessageBytes);
}

export class WebSocketConnection {
  /** @type {(text: string) => void} */
  onmessage = () => {};
  /** @type {(code: number, reason: string) => void} */
  onclose = () => {};

  #socket;
  #max;
  #buffer = Buffer.alloc(0);
  #fragments = [];
  #fragmentBytes = 0;
  #closed = false;

  constructor(socket, head, maxMessageBytes) {
    this.#socket = socket;
    this.#max = maxMessageBytes;
    socket.setNoDelay(true);
    socket.on('data', (data) => this.#receive(data));
    socket.on('close', () => this.#gone(1006, ''));
    socket.on('error', () => socket.destroy());
    if (head.length > 0) this.#receive(head);
  }

  get closed() {
    return this.#closed;
  }

  send(text) {
    if (!this.#closed) this.#frame(0x1, Buffer.from(text, 'utf8'));
  }

  ping() {
    if (!this.#closed) this.#frame(0x9, Buffer.alloc(0));
  }

  close(code = 1000, reason = '') {
    if (this.#closed) return;
    const words = Buffer.from(reason.slice(0, 100), 'utf8');
    const payload = Buffer.alloc(2 + words.length);
    payload.writeUInt16BE(code, 0);
    words.copy(payload, 2);
    this.#frame(0x8, payload);
    this.#socket.end();
    this.#gone(code, reason);
  }

  #gone(code, reason) {
    if (this.#closed) return;
    this.#closed = true;
    this.onclose(code, reason);
  }

  #frame(opcode, payload) {
    const length = payload.length;
    let header;
    if (length < 126) {
      header = Buffer.from([0x80 | opcode, length]);
    } else if (length < 65536) {
      header = Buffer.alloc(4);
      header[0] = 0x80 | opcode;
      header[1] = 126;
      header.writeUInt16BE(length, 2);
    } else {
      header = Buffer.alloc(10);
      header[0] = 0x80 | opcode;
      header[1] = 127;
      header.writeBigUInt64BE(BigInt(length), 2);
    }
    this.#socket.write(Buffer.concat([header, payload]));
  }

  #fail(code, reason) {
    this.close(code, reason);
    this.#socket.destroy();
  }

  #receive(data) {
    if (this.#closed) return;
    this.#buffer = this.#buffer.length === 0 ? data : Buffer.concat([this.#buffer, data]);
    for (;;) {
      const buffer = this.#buffer;
      if (buffer.length < 2) return;
      const fin = (buffer[0] & 0x80) !== 0;
      const opcode = buffer[0] & 0x0f;
      if ((buffer[0] & 0x70) !== 0) return this.#fail(1002, 'reserved bits');
      // A client masks every frame it sends (RFC 6455, 5.1).
      if ((buffer[1] & 0x80) === 0) return this.#fail(1002, 'unmasked frame');
      let length = buffer[1] & 0x7f;
      let offset = 2;
      if (length === 126) {
        if (buffer.length < 4) return;
        length = buffer.readUInt16BE(2);
        offset = 4;
      } else if (length === 127) {
        if (buffer.length < 10) return;
        const long = buffer.readBigUInt64BE(2);
        if (long > BigInt(this.#max)) return this.#fail(1009, 'message too big');
        length = Number(long);
        offset = 10;
      }
      const control = opcode >= 0x8;
      if (control && (!fin || length > 125)) return this.#fail(1002, 'bad control frame');
      if (!control && this.#fragmentBytes + length > this.#max) return this.#fail(1009, 'message too big');
      if (buffer.length < offset + 4 + length) return;
      const mask = buffer.subarray(offset, offset + 4);
      const payload = Buffer.from(buffer.subarray(offset + 4, offset + 4 + length));
      for (let i = 0; i < payload.length; i++) payload[i] ^= mask[i & 3];
      this.#buffer = buffer.subarray(offset + 4 + length);
      this.#handle(opcode, fin, payload);
      if (this.#closed) return;
    }
  }

  #handle(opcode, fin, payload) {
    switch (opcode) {
      case 0x0:
      case 0x1:
        if ((opcode === 0x1) === this.#fragments.length > 0) return this.#fail(1002, 'bad continuation');
        this.#fragments.push(payload);
        this.#fragmentBytes += payload.length;
        if (!fin) return;
        {
          const whole = Buffer.concat(this.#fragments);
          this.#fragments = [];
          this.#fragmentBytes = 0;
          let text;
          try {
            text = utf8.decode(whole);
          } catch {
            return this.#fail(1007, 'not UTF-8');
          }
          this.onmessage(text);
        }
        return;
      case 0x2:
        return this.#fail(1003, 'text only');
      case 0x8: {
        // Echoed, unless it is one no endpoint may send.
        const code = payload.length >= 2 ? payload.readUInt16BE(0) : 1000;
        const sendable = code >= 1000 && code <= 4999 && ![1004, 1005, 1006, 1015].includes(code);
        this.close(sendable ? code : 1000);
        return;
      }
      case 0x9:
        this.#frame(0xa, payload);
        return;
      case 0xa:
        return;
      default:
        return this.#fail(1002, 'unknown opcode');
    }
  }
}