dsh-remote/lib/ws.mjs

// Minimal dependency-free WebSocket (RFC 6455) for the dsh-remote link.
//
// Why not `ws`: DSH plugins loaded through `link:` resolve imports from their own
// folder, and slaves install the plugin from a tarball without node_modules. The link
// only needs binary + text frames, ping/pong and close, so a small implementation keeps
// the plugin installable anywhere with zero dependencies.
import crypto from 'node:crypto'
import { EventEmitter } from 'node:events'
import http from 'node:http'
import https from 'node:https'

const GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
const MAX_FRAME = 64 * 1024 * 1024

export const acceptKey = (key) => crypto.createHash('sha1').update(key + GUID).digest('base64')

/** One WebSocket endpoint over a raw socket (server or client role). */
export class WsConn extends EventEmitter {
  constructor(socket, { client = false, head } = {}) {
    super()
    this.socket = socket
    this.client = client
    this.open = true
    this.buf = head && head.length ? Buffer.from(head) : Buffer.alloc(0)
    this.frag = null
    socket.setNoDelay(true)
    socket.on('data', (d) => this._onData(d))
    socket.on('close', () => this._closed(1006, 'socket closed'))
    // Never throw from an unhandled 'error' event: a reset peer must not crash the host.
    socket.on('error', (e) => { if (this.listenerCount('error')) this.emit('error', e); this._closed(1006, String(e?.message ?? e)) })
    if (this.buf.length) queueMicrotask(() => this._onData(Buffer.alloc(0)))
  }

  get bufferedAmount() { return this.socket.writableLength }

  _closed(code, reason) {
    if (!this.open) return
    this.open = false
    this.emit('close', code, reason)
  }

  _onData(chunk) {
    if (chunk.length) this.buf = this.buf.length ? Buffer.concat([this.buf, chunk]) : chunk
    for (;;) {
      const b = this.buf
      if (b.length < 2) return
      const fin = (b[0] & 0x80) !== 0
      const op = b[0] & 0x0f
      const masked = (b[1] & 0x80) !== 0
      let len = b[1] & 0x7f
      let off = 2
      if (len === 126) {
        if (b.length < 4) return
        len = b.readUInt16BE(2); off = 4
      } else if (len === 127) {
        if (b.length < 10) return
        const big = b.readBigUInt64BE(2)
        if (big > BigInt(MAX_FRAME)) return this.destroy(1009, 'frame too large')
        len = Number(big); off = 10
      }
      if (len > MAX_FRAME) return this.destroy(1009, 'frame too large')
      const maskOff = off
      if (masked) off += 4
      if (b.length < off + len) return
      let payload = b.subarray(off, off + len)
      if (masked) {
        const mask = b.subarray(maskOff, maskOff + 4)
        payload = Buffer.from(payload)
        for (let i = 0; i < payload.length; i++) payload[i] ^= mask[i & 3]
      }
      this.buf = b.subarray(off + len)
      this._frame(fin, op, payload)
      if (!this.open) return
    }
  }

  _frame(fin, op, payload) {
    if (op === 0x8) {
      const code = payload.length >= 2 ? payload.readUInt16BE(0) : 1005
      try { this._send(0x8, payload.subarray(0, 2)) } catch {}
      this.socket.end()
      return this._closed(code, payload.subarray(2).toString('utf8'))
    }
    if (op === 0x9) return this._send(0xa, payload)
    if (op === 0xa) return this.emit('pong', payload)
    if (op === 0x0) {
      if (!this.frag) return this.destroy(1002, 'unexpected continuation')
      this.frag.parts.push(payload)
      if (!fin) return
      const data = Buffer.concat(this.frag.parts)
      const binary = this.frag.op === 0x2
      this.frag = null
      return this.emit('message', binary ? data : data.toString('utf8'), binary)
    }
    if (op === 0x1 || op === 0x2) {
      if (!fin) { this.frag = { op, parts: [payload] }; return }
      return this.emit('message', op === 0x2 ? Buffer.from(payload) : payload.toString('utf8'), op === 0x2)
    }
    this.destroy(1002, 'bad opcode')
  }

  _send(op, data) {
    if (!this.socket.writable) return false
    const payload = Buffer.isBuffer(data) ? data : Buffer.from(data)
    const len = payload.length
    let header
    if (len < 126) { header = Buffer.alloc(2); header[1] = len }
    else if (len < 65536) { header = Buffer.alloc(4); header[1] = 126; header.writeUInt16BE(len, 2) }
    else { header = Buffer.alloc(10); header[1] = 127; header.writeBigUInt64BE(BigInt(len), 2) }
    header[0] = 0x80 | op
    if (this.client) {
      header[1] |= 0x80
      const mask = crypto.randomBytes(4)
      const body = Buffer.from(payload)
      for (let i = 0; i < body.length; i++) body[i] ^= mask[i & 3]
      return this.socket.write(Buffer.concat([header, mask, body]))
    }
    return this.socket.write(Buffer.concat([header, payload]))
  }

  send(data) {
    if (!this.open) return false
    return typeof data === 'string' ? this._send(0x1, Buffer.from(data, 'utf8')) : this._send(0x2, data)
  }

  ping(data = Buffer.alloc(0)) { if (this.open) this._send(0x9, data) }

  close(code = 1000, reason = '') {
    if (!this.open) return
    const p = Buffer.alloc(2 + Buffer.byteLength(reason))
    p.writeUInt16BE(code, 0); p.write(reason, 2)
    try { this._send(0x8, p) } catch {}
    this.socket.end()
    setTimeout(() => this.socket.destroy(), 2000).unref()
    this._closed(code, reason)
  }

  destroy(code = 1006, reason = '') {
    this._closed(code, reason)
    this.socket.destroy()
  }
}

/** Complete a server-side upgrade on (req, socket, head). Returns a WsConn or null. */
export function acceptUpgrade(req, socket, head, extraHeaders = {}) {
  const key = req.headers['sec-websocket-key']
  if (!key || String(req.headers.upgrade).toLowerCase() !== 'websocket') {
    socket.end('HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n')
    return null
  }
  const lines = [
    'HTTP/1.1 101 Switching Protocols',
    'Upgrade: websocket',
    'Connection: Upgrade',
    `Sec-WebSocket-Accept: ${acceptKey(key)}`,
    ...Object.entries(extraHeaders).map(([k, v]) => `${k}: ${v}`),
  ]
  socket.write(lines.join('\r\n') + '\r\n\r\n')
  return new WsConn(socket, { head })
}

export function rejectUpgrade(socket, status, message = '') {
  const text = { 400: 'Bad Request', 401: 'Unauthorized', 403: 'Forbidden', 404: 'Not Found', 409: 'Conflict', 502: 'Bad Gateway' }[status] || 'Error'
  try {
    socket.end(`HTTP/1.1 ${status} ${text}\r\nConnection: close\r\nContent-Type: text/plain\r\nContent-Length: ${Buffer.byteLength(message)}\r\n\r\n${message}`)
  } catch { socket.destroy() }
}

/** Client connect: resolves a WsConn, rejects with {status, body} on refusal. */
export function connect(url, { headers = {}, timeoutMs = 15000, rejectUnauthorized = true, onHeaders } = {}) {
  return new Promise((resolve, reject) => {
    const u = new URL(url)
    const secure = u.protocol === 'wss:' || u.protocol === 'https:'
    const key = crypto.randomBytes(16).toString('base64')
    const req = (secure ? https : http).request({
      host: u.hostname,
      port: u.port || (secure ? 443 : 80),
      path: u.pathname + u.search,
      method: 'GET',
      servername: u.hostname,
      rejectUnauthorized,
      headers: {
        ...headers,
        Host: u.host,
        Connection: 'Upgrade',
        Upgrade: 'websocket',
        'Sec-WebSocket-Version': '13',
        'Sec-WebSocket-Key': key,
      },
      timeout: timeoutMs,
    })
    req.on('timeout', () => { req.destroy(new Error(`connect timeout after ${timeoutMs} ms`)) })
    req.on('error', reject)
    req.on('response', (res) => {
      onHeaders?.(res.headers)
      let body = ''
      res.setEncoding('utf8')
      res.on('data', (d) => { if (body.length < 4096) body += d })
      res.on('end', () => {
        const loc = res.headers.location ? ` → ${String(res.headers.location).slice(0, 100)}` : ''
        const err = new Error(`upgrade refused: HTTP ${res.statusCode}${loc} ${res.headers.location ? '' : body.trim().slice(0, 200)}`.trim())
        err.status = res.statusCode
        err.location = res.headers.location
        err.contentType = String(res.headers['content-type'] || '')
        err.body = body
        reject(err)
      })
    })
    req.on('upgrade', (res, socket, head) => {
      if (res.headers['sec-websocket-accept'] !== acceptKey(key)) {
        socket.destroy()
        return reject(new Error('bad Sec-WebSocket-Accept'))
      }
      socket.setTimeout(0)
      resolve(new WsConn(socket, { client: true, head }))
    })
    req.end()
  })
}