dsh-remote/lib/mux.mjs

// Stream multiplexer over one WsConn.
//
// Frames:
//   text   = JSON control message  { t: <type>, ...fields }
//   binary = data frame            [u32 stream id][payload]
//
// Streams are opened with a control message `{t:'open', id, kind, meta}` and carry raw
// bytes in both directions. They implement backpressure-lite: a stream pauses its local
// source while the underlying socket buffer is above HIGH_WATER.
//
// Request/response control calls: `{t:'call', cid, method, params}` →
// `{t:'reply', cid, ok, result | error}`. Unsolicited `{t:'event', name, data}`.
import { EventEmitter } from 'node:events'

const HIGH_WATER = 8 * 1024 * 1024

export class Mux extends EventEmitter {
  /**
   * @param ws WsConn
   * @param {{ odd: boolean, handlers?: Record<string, Function>, onStream?: Function }} opts
   *   odd: this side allocates odd stream ids (master = even, slave = odd) so ids never collide.
   */
  constructor(ws, { odd, handlers = {}, onStream } = {}) {
    super()
    this.ws = ws
    this.nextId = odd ? 1 : 2
    this.nextCid = 1
    this.streams = new Map()
    this.calls = new Map()
    this.handlers = handlers
    this.onStream = onStream
    this.closed = false
    this.lastSeen = Date.now()
    ws.on('message', (data, binary) => {
      this.lastSeen = Date.now()
      if (binary) this._data(data)
      else this._control(data)
    })
    ws.on('pong', () => { this.lastSeen = Date.now() })
    ws.on('close', (code, reason) => this._teardown(code, reason))
    this.pinger = setInterval(() => {
      if (Date.now() - this.lastSeen > 45000) return ws.destroy(1001, 'ping timeout')
      ws.ping()
    }, 15000)
    this.pinger.unref?.()
  }

  _teardown(code, reason) {
    if (this.closed) return
    this.closed = true
    clearInterval(this.pinger)
    for (const s of this.streams.values()) s._remoteEnd(new Error('link closed'))
    this.streams.clear()
    for (const { reject, timer } of this.calls.values()) { clearTimeout(timer); reject(new Error('link closed')) }
    this.calls.clear()
    this.emit('close', code, reason)
  }

  close(code, reason) { this.ws.close(code, reason) }

  sendControl(msg) {
    if (this.closed) return false
    return this.ws.send(JSON.stringify(msg))
  }

  _control(text) {
    let msg
    try { msg = JSON.parse(text) } catch { return }
    switch (msg.t) {
      case 'open': {
        const s = new Stream(this, msg.id, msg.kind, msg.meta || {})
        this.streams.set(msg.id, s)
        Promise.resolve()
          .then(() => (this.onStream ? this.onStream(s) : s.fail('no stream handler')))
          .catch((e) => s.fail(String(e?.message ?? e)))
        return
      }
      case 'head': return this.streams.get(msg.id)?._head(msg.meta || {})
      case 'end': return this.streams.get(msg.id)?._remoteEnd()
      case 'fail': return this.streams.get(msg.id)?._remoteEnd(new Error(msg.error || 'stream failed'))
      case 'pause': return this.streams.get(msg.id)?.emit('remote-pause')
      case 'resume': return this.streams.get(msg.id)?.emit('remote-resume')
      case 'call': return this._call(msg)
      case 'reply': {
        const c = this.calls.get(msg.cid)
        if (!c) return
        this.calls.delete(msg.cid)
        clearTimeout(c.timer)
        if (msg.ok) c.resolve(msg.result)
        else c.reject(Object.assign(new Error(msg.error || 'remote error'), { status: msg.status }))
        return
      }
      case 'event': return this.emit('event', msg.name, msg.data)
      default: return this.emit('control', msg)
    }
  }

  async _call(msg) {
    const h = this.handlers[msg.method]
    try {
      if (!h) throw Object.assign(new Error(`unknown method ${msg.method}`), { status: 404 })
      const result = await h(msg.params ?? {}, this)
      this.sendControl({ t: 'reply', cid: msg.cid, ok: true, result: result ?? null })
    } catch (e) {
      this.sendControl({ t: 'reply', cid: msg.cid, ok: false, error: String(e?.message ?? e), status: e?.status })
    }
  }

  call(method, params = {}, timeoutMs = 30000) {
    if (this.closed) return Promise.reject(new Error('link closed'))
    const cid = this.nextCid++
    return new Promise((resolve, reject) => {
      const timer = setTimeout(() => {
        this.calls.delete(cid)
        reject(new Error(`call ${method} timed out`))
      }, timeoutMs)
      timer.unref?.()
      this.calls.set(cid, { resolve, reject, timer })
      this.sendControl({ t: 'call', cid, method, params })
    })
  }

  event(name, data) { this.sendControl({ t: 'event', name, data }) }

  _data(buf) {
    if (buf.length < 4) return
    const id = buf.readUInt32BE(0)
    this.streams.get(id)?._remoteData(buf.subarray(4))
  }

  /** Open an outgoing stream. */
  open(kind, meta = {}) {
    const id = this.nextId
    this.nextId += 2
    const s = new Stream(this, id, kind, meta)
    this.streams.set(id, s)
    this.sendControl({ t: 'open', id, kind, meta })
    // The opener attaches listeners synchronously; deliver from the next tick on.
    queueMicrotask(() => s.start())
    return s
  }

  _writeData(id, chunk) {
    const header = Buffer.alloc(4)
    header.writeUInt32BE(id, 0)
    // Split large chunks so one stream cannot monopolise the link.
    const MAX = 256 * 1024
    let ok = true
    for (let off = 0; off < chunk.length; off += MAX) {
      ok = this.ws.send(Buffer.concat([header, chunk.subarray(off, off + MAX)])) && ok
    }
    return ok && this.ws.bufferedAmount < HIGH_WATER
  }
}

/**
 * One bidirectional byte stream. Events: 'head'(meta), 'data'(buf), 'end', 'error'(err),
 * 'drain'. Writers should stop on write()===false and wait for 'drain'.
 */
export class Stream extends EventEmitter {
  constructor(mux, id, kind, meta) {
    super()
    this.mux = mux
    this.id = id
    this.kind = kind
    this.meta = meta
    this.localEnded = false
    this.remoteEnded = false
    this.done = false
    // Incoming events are queued until the consumer attaches its listeners and calls
    // start(): handlers often await something (cookie, upstream connect) first, and
    // data/end arriving in the meantime must not be lost.
    this.paused = true
    this.queue = []
  }

  /** Begin delivering queued + future incoming events. Idempotent. */
  start() {
    if (!this.paused) return this
    this.paused = false
    const q = this.queue
    this.queue = []
    for (const [ev, arg] of q) this._deliver(ev, arg)
    return this
  }

  _deliver(ev, arg) {
    if (this.paused) { this.queue.push([ev, arg]); return }
    if (ev === 'head') this.emit('head', arg)
    else if (ev === 'data') this.emit('data', arg)
    else if (ev === 'end') { this.emit('end'); this._maybeDone() }
    else if (ev === 'error') { if (this.listenerCount('error')) this.emit('error', arg); this._finish() }
  }

  head(meta) { this.mux.sendControl({ t: 'head', id: this.id, meta }) }

  write(chunk) {
    if (this.localEnded || this.mux.closed) return false
    const ok = this.mux._writeData(this.id, Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk))
    if (!ok && !this._drainWait) {
      this._drainWait = true
      const sock = this.mux.ws.socket
      sock.once('drain', () => { this._drainWait = false; this.emit('drain') })
    }
    return ok
  }

  end() {
    if (this.localEnded) return
    this.localEnded = true
    this.mux.sendControl({ t: 'end', id: this.id })
    this._maybeDone()
  }

  fail(error) {
    if (this.done) return
    this.mux.sendControl({ t: 'fail', id: this.id, error: String(error) })
    this.localEnded = true
    this.remoteEnded = true
    this._finish()
  }

  _head(meta) { this._deliver('head', meta) }
  _remoteData(buf) { if (!this.remoteEnded) this._deliver('data', buf) }

  _remoteEnd(err) {
    if (this.remoteEnded) return
    this.remoteEnded = true
    if (err) {
      this.localEnded = true
      this._deliver('error', err)
      return
    }
    this._deliver('end')
  }

  _maybeDone() { if (this.localEnded && this.remoteEnded) this._finish() }

  _finish() {
    if (this.done) return
    this.done = true
    this.mux.streams.delete(this.id)
    this.emit('close')
  }
}