web-terminal/term-relay/data-plane/relay-node.ts

/**
 * T10 · Stateless relay node + opaque splice (INV2/INV5/INV7/INV11 core). Browser upgrade →
 * authorize (P5 via T8) → openStream on the host's agent tunnel → OPAQUE byte splice. The node
 * NEVER decodes DATA (INV2/INV11): a byte is a byte. Buffers are per-connection (each ws↔stream
 * pair) — NO global mutable buffer pool (INV5, kills cross-tenant buffer bleed). All maps are
 * in-RAM only (INV7): a crash loses nothing; agents reconnect + re-register.
 *
 * closeStream (Finding-3/INV12): a per-`jti` splice index lets P5 RST exactly the stream(s) under
 * ONE revoked capability token, leaving the owner's other devices (different `jti`s) attached —
 * removing the "nuke everyone or nothing" false choice.
 */
import type { WebSocketLike } from './ws-like.js'
import type { AgentListener } from './agent-listener.js'
import type { DataPlaneConfig } from './config.js'
import type { UpgradeRequest, UpgradeDecision } from './upgrade.js'
import type { MuxStreamHandle } from '../mux/mux-session.js'

const WS_TRY_LATER = 1013 // agent offline for this host
const WS_REVOKED = 4403 // single-device revocation close code

interface Splice {
  readonly hostId: string
  readonly streamId: number
  readonly jti: string
  readonly ws: WebSocketLike
  readonly stream: MuxStreamHandle
}

export type StreamSelector = { readonly jti: string } | { readonly streamId: number }

export interface RelayNode {
  handleAgentTunnel(ws: WebSocketLike, peerCert: Uint8Array): void
  handleBrowserUpgrade(req: UpgradeRequest, ws: WebSocketLike): Promise<void>
  drain(reason: number): Promise<void>
  closeTunnel(hostId: string): void
  closeStream(hostId: string, selector: StreamSelector): number
}

export interface RelayNodeDeps {
  config: DataPlaneConfig
  listener: AgentListener
  authorize(req: UpgradeRequest): Promise<UpgradeDecision>
}

function selectorMatches(splice: Splice, selector: StreamSelector): boolean {
  return 'jti' in selector ? splice.jti === selector.jti : splice.streamId === selector.streamId
}

export function createRelayNode(deps: RelayNodeDeps): RelayNode {
  // In-RAM splice index (INV7). No durable buffers, no global pool (INV5).
  const splices = new Set<Splice>()

  function removeSplice(splice: Splice): void {
    splices.delete(splice)
  }

  return {
    handleAgentTunnel(ws, peerCert) {
      deps.listener.attach(ws, peerCert) // TLS+mTLS already enforced upstream (T9)
    },

    async handleBrowserUpgrade(req, ws) {
      const decision = await deps.authorize(req)
      if (!decision.ok) {
        ws.close(decision.status) // INV6/INV15: only P5's {ok:true} opens a stream
        return
      }
      const tunnel = deps.listener.tunnels().get(decision.hostId) // strictly by token-derived hostId (INV1/INV3)
      if (tunnel === undefined) {
        ws.close(WS_TRY_LATER) // agent offline; no crash
        return
      }
      const stream = tunnel.session.openStream(decision.open) // relay allocates streamId
      const splice: Splice = {
        hostId: decision.hostId,
        streamId: stream.streamId,
        jti: decision.open.capabilityTokenRef,
        ws,
        stream,
      }
      splices.add(splice)

      // OPAQUE splice — verbatim Uint8Array both ways (INV2), per-connection only (INV5).
      stream.onData((bytes) => ws.send(bytes)) // agent → browser
      stream.onClose(() => {
        ws.close()
        removeSplice(splice)
      })
      ws.onMessage((bytes) => stream.writeData(bytes)) // browser → agent
      ws.onClose(() => {
        stream.close()
        removeSplice(splice)
      })
    },

    async drain(reason) {
      // Whole-node drain: GOAWAY every tunnel, then close (deregisters via the listener's teardown).
      for (const t of [...deps.listener.tunnels().values()]) {
        t.session.drain(0, reason)
        t.closeTunnel()
      }
      await Promise.resolve()
    },

    closeTunnel(hostId) {
      // Whole-host revocation lever — kicks every device on the host.
      for (const splice of [...splices]) {
        if (splice.hostId === hostId) {
          splice.ws.close(WS_REVOKED)
          removeSplice(splice)
        }
      }
      deps.listener.tunnels().get(hostId)?.closeTunnel()
    },

    closeStream(hostId, selector) {
      // Single-device revocation (Finding-3): RST exactly the matching stream(s); unknown → 0 (no-op).
      let count = 0
      for (const splice of [...splices]) {
        if (splice.hostId !== hostId) continue
        if (!selectorMatches(splice, selector)) continue
        splice.stream.close(true) // RST just this stream via the mux (tunnel stays up)
        splice.ws.close(WS_REVOKED)
        removeSplice(splice)
        count += 1
      }
      return count
    },
  }
}