web-terminal/term-relay/test/relay-node.test.ts

import { describe, test, expect, vi } from 'vitest'
import { createRelayNode } from '../data-plane/relay-node.js'
import type { AgentListener, AgentTunnel } from '../data-plane/agent-listener.js'
import type { MuxSession, MuxStreamHandle } from '../mux/mux-session.js'
import type { WebSocketLike } from '../data-plane/ws-like.js'
import type { UpgradeDecision, UpgradeRequest } from '../data-plane/upgrade.js'
import { loadDataPlaneConfig } from '../data-plane/config.js'

const config = loadDataPlaneConfig({
  BASE_DOMAIN: 'term.example.com',
  TLS_CERT_PATH: '/c',
  TLS_KEY_PATH: '/k',
  AGENT_CA_CERT_PATH: '/ca',
  RELAY_NODE_ID: 'node-1',
})

interface FakeStream {
  handle: MuxStreamHandle
  emitData(b: Uint8Array): void
  emitClose(): void
  writes: Uint8Array[]
  rstClosed: boolean
}

function fakeTunnel(hostId: string) {
  let nextId = 0
  const fakeStreams: FakeStream[] = []
  const session = {
    openStream: () => {
      let dataCb: ((b: Uint8Array) => void) | null = null
      let closeCb: ((rst: boolean) => void) | null = null
      const fs: FakeStream = { writes: [], rstClosed: false, emitData: () => {}, emitClose: () => {}, handle: undefined as unknown as MuxStreamHandle }
      const handle: MuxStreamHandle = {
        streamId: ++nextId,
        writeData: (b) => {
          fs.writes.push(b)
          return true
        },
        close: (rst = false) => {
          fs.rstClosed = rst
        },
        onData: (cb) => (dataCb = cb),
        onClose: (cb) => (closeCb = cb),
      }
      fs.handle = handle
      fs.emitData = (b) => dataCb?.(b)
      fs.emitClose = () => closeCb?.(false)
      fakeStreams.push(fs)
      return handle
    },
    onWire: () => {},
    drain: vi.fn(),
    close: vi.fn(),
  } satisfies MuxSession
  const tunnel: AgentTunnel = { hostId, accountId: 'acct-1', session, closeTunnel: vi.fn() }
  return { tunnel, fakeStreams }
}

function fakeWs() {
  let onMsg: ((b: Uint8Array) => void) | null = null
  let onClose: (() => void) | null = null
  return {
    sent: [] as Uint8Array[],
    closedWith: undefined as number | undefined,
    send(b: Uint8Array) {
      this.sent.push(b)
    },
    close(code?: number) {
      this.closedWith = code
    },
    onMessage(h: (b: Uint8Array) => void) {
      onMsg = h
    },
    onClose(h: () => void) {
      onClose = h
    },
    emitMessage: (b: Uint8Array) => onMsg?.(b),
    emitClose: () => onClose?.(),
  }
}

function listenerWith(tunnels: Map<string, AgentTunnel>): AgentListener {
  return { attach: vi.fn(), tunnels: () => tunnels, close: vi.fn() }
}

function decisionFor(hostId: string, jti: string): UpgradeDecision {
  return {
    ok: true,
    hostId,
    acceptedSubprotocol: 'term.relay.v1',
    open: { streamId: 0, subdomain: 'alice', requestPath: '/term', originHeader: 'o', remoteAddrHash: 'h', capabilityTokenRef: jti },
  }
}

const req = {} as UpgradeRequest

describe('relay node opaque splice (T10, INV2/INV5/INV7/INV11/INV12)', () => {
  test('authorized upgrade → bytes splice byte-identical in both directions', async () => {
    const { tunnel, fakeStreams } = fakeTunnel('host-1')
    const node = createRelayNode({ config, listener: listenerWith(new Map([['host-1', tunnel]])), authorize: async () => decisionFor('host-1', 'jti-a') })
    const ws = fakeWs()
    await node.handleBrowserUpgrade(req, ws)
    const fs = fakeStreams[0]!
    const up = new Uint8Array([1, 2, 3])
    ws.emitMessage(up) // browser → agent
    expect([...fs.writes[0]!]).toEqual([1, 2, 3])
    const down = new Uint8Array([9, 8, 7])
    fs.emitData(down) // agent → browser
    expect([...ws.sent[0]!]).toEqual([9, 8, 7])
  })

  test('unauthorized upgrade → ws closed with the status; no stream, no tunnel touched', async () => {
    const { tunnel } = fakeTunnel('host-1')
    const node = createRelayNode({ config, listener: listenerWith(new Map([['host-1', tunnel]])), authorize: async () => ({ ok: false, status: 403 }) })
    const ws = fakeWs()
    await node.handleBrowserUpgrade(req, ws)
    expect(ws.closedWith).toBe(403)
  })

  test('agent offline (no tunnel) → ws closed 1013, no crash', async () => {
    const node = createRelayNode({ config, listener: listenerWith(new Map()), authorize: async () => decisionFor('host-x', 'jti') })
    const ws = fakeWs()
    await node.handleBrowserUpgrade(req, ws)
    expect(ws.closedWith).toBe(1013)
  })

  test('single-device revocation (Finding-3): closeStream by jti RSTs only that jti', async () => {
    const { tunnel, fakeStreams } = fakeTunnel('host-1')
    const node = createRelayNode({ config, listener: listenerWith(new Map([['host-1', tunnel]])), authorize: async () => decisionFor('host-1', jtiRef()) })
    let currentJti = 'jtiA'
    function jtiRef() {
      return currentJti
    }
    const wsA = fakeWs()
    await node.handleBrowserUpgrade(req, wsA)
    currentJti = 'jtiB'
    const wsB = fakeWs()
    await node.handleBrowserUpgrade(req, wsB)

    const count = node.closeStream('host-1', { jti: 'jtiA' })
    expect(count).toBe(1)
    expect(fakeStreams[0]!.rstClosed).toBe(true) // A RST
    expect(wsA.closedWith).toBe(4403)
    expect(fakeStreams[1]!.rstClosed).toBe(false) // B untouched
    // B still flows
    fakeStreams[1]!.emitData(new Uint8Array([5]))
    expect([...wsB.sent[0]!]).toEqual([5])
  })

  test('closeStream with two splices under the SAME jti returns 2; unknown jti returns 0', async () => {
    const { tunnel } = fakeTunnel('host-1')
    const node = createRelayNode({ config, listener: listenerWith(new Map([['host-1', tunnel]])), authorize: async () => decisionFor('host-1', 'jtiSame') })
    await node.handleBrowserUpgrade(req, fakeWs())
    await node.handleBrowserUpgrade(req, fakeWs())
    expect(node.closeStream('host-1', { jti: 'jtiSame' })).toBe(2)
    expect(node.closeStream('host-1', { jti: 'unknown' })).toBe(0)
  })

  test('after a browser ws closes normally, its splice entry is gone (no stale handle)', async () => {
    const { tunnel } = fakeTunnel('host-1')
    const node = createRelayNode({ config, listener: listenerWith(new Map([['host-1', tunnel]])), authorize: async () => decisionFor('host-1', 'jtiZ') })
    const ws = fakeWs()
    await node.handleBrowserUpgrade(req, ws)
    ws.emitClose()
    expect(node.closeStream('host-1', { jti: 'jtiZ' })).toBe(0)
  })
})