hyreal-gen-pic/napos/ws.js

'use strict';
/**
 * 零依赖的 WebSocket 客户端（RFC6455 子集，够 CDP 用）。
 *
 * 为什么不用 Node 内置的全局 WebSocket：
 *   Chromium 的 DevTools WebSocket 端点会拒绝带 Origin 头的握手。
 *   浏览器环境自动带 Origin，undici 的实现行为不确定。
 *   这里用 http.request 手动升级，完全掌控请求头 —— 不发 Origin，握手必定通过。
 *
 * 支持：文本/二进制帧、分片续帧、ping/pong、64 位长度、关闭帧。
 * 客户端发帧一律 mask（协议要求）。
 */

const http = require('http');
const crypto = require('crypto');
const { EventEmitter } = require('events');

const OP = { CONT: 0x0, TEXT: 0x1, BIN: 0x2, CLOSE: 0x8, PING: 0x9, PONG: 0xa };

function encodeFrame(payload, opcode = OP.TEXT) {
  const len = payload.length;
  let head;
  if (len < 126) {
    head = Buffer.alloc(2);
    head[1] = len;
  } else if (len < 65536) {
    head = Buffer.alloc(4);
    head[1] = 126;
    head.writeUInt16BE(len, 2);
  } else {
    head = Buffer.alloc(10);
    head[1] = 127;
    head.writeUInt32BE(Math.floor(len / 4294967296), 2);
    head.writeUInt32BE(len >>> 0, 6);
  }
  head[0] = 0x80 | opcode; // FIN + opcode
  head[1] |= 0x80;         // MASK

  const mask = crypto.randomBytes(4);
  const out = Buffer.alloc(head.length + 4 + len);
  head.copy(out, 0);
  mask.copy(out, head.length);
  const off = head.length + 4;
  for (let i = 0; i < len; i++) out[off + i] = payload[i] ^ mask[i & 3];
  return out;
}

class RawWebSocket extends EventEmitter {
  constructor(url) {
    super();
    this.url = new URL(url);
    this.sock = null;
    this.buf = Buffer.alloc(0);
    this.frag = [];        // 分片累积
    this.fragOp = null;
    this.readyState = 0;   // 0 connecting 1 open 2 closing 3 closed
  }

  connect(timeoutMs = 10000) {
    return new Promise((resolve, reject) => {
      const key = crypto.randomBytes(16).toString('base64');
      const req = http.request({
        host: this.url.hostname,
        port: this.url.port || 80,
        path: this.url.pathname + this.url.search,
        method: 'GET',
        headers: {
          Connection: 'Upgrade',
          Upgrade: 'websocket',
          'Sec-WebSocket-Key': key,
          'Sec-WebSocket-Version': '13',
        },
        timeout: timeoutMs,
      });

      let settled = false;
      const fail = (e) => { if (!settled) { settled = true; reject(e); } };

      req.on('upgrade', (res, socket, head) => {
        settled = true;
        this.sock = socket;
        this.readyState = 1;
        socket.setNoDelay(true);
        socket.on('data', (d) => this._feed(d));
        socket.on('close', () => {
          const wasOpen = this.readyState === 1;
          this.readyState = 3;
          this.emit('close');
          // 从未 emit 过 error 的连接断掉时，给上层一个明确的失败信号而不是静默
          if (wasOpen) this.emit('error', new Error('WebSocket 连接已断开'));
        });
        socket.on('error', (e) => {
          this.readyState = 3;
          this.emit('error', e);
        });
        if (head && head.length) this._feed(head);
        resolve(res.headers);
      });
      req.on('response', (res) => {
        // 握手被拒：必须销毁请求释放底层 socket，否则每次探测失败都泄漏一条连接
        res.resume();
        req.destroy();
        fail(new Error(`握手被拒: HTTP ${res.statusCode}`));
      });
      req.on('error', fail);
      req.on('timeout', () => req.destroy(new Error('WebSocket 握手超时')));
      req.end();
    });
  }

  _feed(chunk) {
    this.buf = Buffer.concat([this.buf, chunk]);
    for (;;) {
      if (this.buf.length < 2) return;
      const b0 = this.buf[0];
      const b1 = this.buf[1];
      const fin = (b0 & 0x80) !== 0;
      const opcode = b0 & 0x0f;
      const masked = (b1 & 0x80) !== 0;
      let len = b1 & 0x7f;
      let pos = 2;
      if (len === 126) {
        if (this.buf.length < pos + 2) return;
        len = this.buf.readUInt16BE(pos); pos += 2;
      } else if (len === 127) {
        if (this.buf.length < pos + 8) return;
        const hi = this.buf.readUInt32BE(pos);
        const lo = this.buf.readUInt32BE(pos + 4);
        len = hi * 4294967296 + lo; pos += 8;
      }
      let mask = null;
      if (masked) {
        if (this.buf.length < pos + 4) return;
        mask = this.buf.subarray(pos, pos + 4); pos += 4;
      }
      if (this.buf.length < pos + len) return;
      let payload = this.buf.subarray(pos, pos + len);
      this.buf = this.buf.subarray(pos + len);

      if (mask) {
        payload = Buffer.from(payload); // 复制，避免改动原缓冲
        for (let i = 0; i < payload.length; i++) payload[i] ^= mask[i & 3];
      }

      if (opcode === OP.PING) { this._sendFrame(payload, OP.PONG); continue; }
      if (opcode === OP.PONG) continue;
      if (opcode === OP.CLOSE) { this.close(); continue; }

      if (opcode === OP.CONT) {
        this.frag.push(payload);
      } else {
        this.frag = [payload];
        this.fragOp = opcode;
      }
      if (fin) {
        const full = this.frag.length === 1 ? this.frag[0] : Buffer.concat(this.frag);
        const startOp = this.fragOp;
        this.frag = []; this.fragOp = null;
        if (startOp === OP.TEXT) this.emit('message', full.toString('utf8'));
        else this.emit('binary', full);
      }
    }
  }

  _sendFrame(payload, opcode) {
    if (!this.sock || this.readyState !== 1) return;
    try { this.sock.write(encodeFrame(payload, opcode)); } catch (e) { this.emit('error', e); }
  }

  send(text) { this._sendFrame(Buffer.from(text, 'utf8'), OP.TEXT); }

  close() {
    if (this.readyState === 3) return;
    this.readyState = 2;
    try { this._sendFrame(Buffer.from([0x03, 0xe8]), OP.CLOSE); } catch { /* ignore */ }
    try { this.sock.end(); } catch { /* ignore */ }
    this.readyState = 3;
    this.emit('close');
  }
}

module.exports = { RawWebSocket, encodeFrame };