web-terminal/agent/test/streamRouter.test.ts

import { describe, expect, it, vi } from 'vitest'
import { decodeMuxFrame, encodeMuxFrame, type MuxOpen } from 'relay-contracts'
import type { AgentConfig } from '../src/config/agentConfig.js'
import { createLogger } from '../src/log/logger.js'
import { holdTunnel, dataHeader } from '../src/transport/tunnel.js'
import { createStreamRouter, identityTransform } from '../src/transport/streamRouter.js'
import { FakeWs } from './fixtures/fakes.js'

const CFG: AgentConfig = {
  relayUrl: 'wss://relay/agent',
  enrollUrl: 'https://x/enroll',
  stateDir: '/tmp/x',
  localTargetUrl: 'ws://127.0.0.1:3000',
  subdomain: 'host-42',
  hostId: 'h-1',
}

const OPEN: MuxOpen = {
  streamId: 5,
  subdomain: 'host-42',
  requestPath: '/term?join=abc',
  originHeader: 'https://host-42.term.example.com',
  remoteAddrHash: 'x',
  capabilityTokenRef: 'jti',
}

const silentLogger = createLogger('error', () => {})
const flush = () => new Promise((r) => setImmediate(r))

function setup(dialImpl?: (path: string, origin: string) => Promise<FakeWs>) {
  const upstream = new FakeWs()
  const tunnel = holdTunnel(upstream)
  const loopbacks: FakeWs[] = []
  const dial = vi.fn(
    dialImpl ??
      (async () => {
        const ws = new FakeWs()
        loopbacks.push(ws)
        return ws
      }),
  )
  const router = createStreamRouter(CFG, tunnel, dial as never, identityTransform, silentLogger)
  return { upstream, router, dial, loopbacks }
}

function lastFrame(ws: FakeWs) {
  return decodeMuxFrame(ws.sent.at(-1)!)
}

describe('StreamRouter (T8)', () => {
  it('dials loopback with the request path and replayed Origin', async () => {
    const { router, dial } = setup()
    router.handleOpen(OPEN)
    await flush()
    expect(dial).toHaveBeenCalledWith('/term?join=abc', 'https://host-42.term.example.com')
  })

  it('INV1 defense-in-depth: foreign subdomain is RST and NEVER dialed', async () => {
    const { router, dial, upstream } = setup()
    router.handleOpen({ ...OPEN, subdomain: 'host-999' })
    await flush()
    expect(dial).not.toHaveBeenCalled()
    const f = lastFrame(upstream)
    expect(f.header.type).toBe('close')
    expect(f.header.rst).toBe(true)
    expect(router.activeStreamCount()).toBe(0)
  })

  it('splices loopback output → tunnel DATA', async () => {
    const { router, upstream, loopbacks } = setup()
    router.handleOpen(OPEN)
    await flush()
    const bytes = new Uint8Array([9, 8, 7])
    loopbacks[0]!.emit('message', bytes)
    const f = lastFrame(upstream)
    expect(f.header.type).toBe('data')
    expect(Buffer.from(f.payload).equals(Buffer.from(bytes))).toBe(true)
  })

  it('forwards tunnel DATA → loopback socket (buffering pre-open)', async () => {
    const { router, loopbacks } = setup()
    router.handleOpen(OPEN)
    router.handleData(5, new Uint8Array([1, 2])) // arrives before dial resolves → buffered
    await flush()
    router.handleData(5, new Uint8Array([3, 4]))
    expect(loopbacks[0]!.sent.map((b) => [...b])).toEqual([[1, 2], [3, 4]])
  })

  it('DATA on an unknown stream is RST (illegal transition)', () => {
    const { router, upstream } = setup()
    router.handleData(999, new Uint8Array([1]))
    const f = lastFrame(upstream)
    expect(f.header.type).toBe('close')
    expect(f.header.rst).toBe(true)
  })

  it('two concurrent streams get separate sockets (no cross-stream bleed)', async () => {
    const { router, loopbacks } = setup()
    router.handleOpen(OPEN)
    router.handleOpen({ ...OPEN, streamId: 6 })
    await flush()
    router.handleData(5, new Uint8Array([0xaa]))
    router.handleData(6, new Uint8Array([0xbb]))
    expect(loopbacks[0]!.sent).toHaveLength(1)
    expect(loopbacks[1]!.sent).toHaveLength(1)
    expect([...loopbacks[0]!.sent[0]!]).toEqual([0xaa])
    expect([...loopbacks[1]!.sent[0]!]).toEqual([0xbb])
    expect(router.activeStreamCount()).toBe(2)
  })

  it('CLOSE frees the loopback socket', async () => {
    const { router, loopbacks } = setup()
    router.handleOpen(OPEN)
    await flush()
    router.handleClose(5, false)
    expect(loopbacks[0]!.closed).toBe(true)
    expect(router.activeStreamCount()).toBe(0)
  })

  it('loopback connect-refused → RST upstream (typed, no swallow)', async () => {
    const { router, upstream } = setup(async () => {
      throw new Error('ECONNREFUSED')
    })
    router.handleOpen(OPEN)
    await flush()
    const f = lastFrame(upstream)
    expect(f.header.type).toBe('close')
    expect(f.header.rst).toBe(true)
  })
})

describe('opaque splice sanity (INV11)', () => {
  it('does not import any terminal parser (identity transform)', () => {
    const bytes = new Uint8Array([0x1b, 0x5b, 0x41])
    expect(identityTransform.outbound(1, bytes)).toBe(bytes)
    // encode/decode round-trip stays byte-exact
    const framed = encodeMuxFrame(dataHeader(1, bytes.length), bytes)
    expect(Buffer.from(decodeMuxFrame(framed).payload).equals(Buffer.from(bytes))).toBe(true)
  })
})