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

import { describe, expect, it, vi } from 'vitest'
import { decodeMuxFrame, encodeGoaway, encodeMuxFrame, encodeOpen, type MuxOpen } from 'relay-contracts'
import type { AgentConfig } from '../src/config/agentConfig.js'
import type { Keystore } from '../src/keys/keystore.js'
import { createBackoff } from '../src/transport/backoff.js'
import { runTunnel, type RunTunnelDeps } from '../src/transport/runTunnel.js'
import { FakeTimer, FakeWs } from './fixtures/fakes.js'

/**
 * C2 supervisor: proves runTunnel ports the cafeDemo assembly into a supervised loop —
 * bytes splice both ways, a dead session reconnects, and a `revoked` GOAWAY stops for good (INV12).
 */
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 KS = {} as unknown as Keystore // unused when connectRelay/dialLoopback are injected
const flush = (): Promise<void> => new Promise((r) => setImmediate(r))

function emitOpen(upstream: FakeWs, open: MuxOpen): void {
  const payload = encodeOpen(open)
  upstream.emitMessage(
    encodeMuxFrame(
      { version: 1, type: 'open', fin: false, rst: false, streamId: open.streamId, payloadLen: payload.length },
      payload,
    ),
  )
}
function emitGoAwayRevoked(upstream: FakeWs): void {
  const payload = encodeGoaway(0, 'revoked')
  upstream.emitMessage(
    encodeMuxFrame(
      { version: 1, type: 'goaway', fin: false, rst: false, streamId: 0, payloadLen: payload.length },
      payload,
    ),
  )
}

function baseDeps(over: Partial<RunTunnelDeps>): Partial<RunTunnelDeps> {
  return {
    timer: new FakeTimer(), // never auto-fires ⇒ heartbeat is inert during the test
    sleep: async () => {},
    backoff: createBackoff(),
    ...over,
  }
}

describe('runTunnel supervisor (C2)', () => {
  it('splices bytes both ways through the loopback', async () => {
    const upstream = new FakeWs()
    const loopback = new FakeWs()
    const connectRelay = vi.fn(async () => upstream)
    const dialLoopback = vi.fn(async () => loopback)
    const handle = await runTunnel(CFG, KS, baseDeps({ connectRelay, dialLoopback: dialLoopback as never }))
    await flush()

    emitOpen(upstream, OPEN)
    await flush()
    expect(dialLoopback).toHaveBeenCalledWith('/term?join=abc', 'https://host-42.term.example.com')

    upstream.emitMessage(encodeMuxFrame({ version: 1, type: 'data', fin: false, rst: false, streamId: 5, payloadLen: 3 }, new Uint8Array([104, 105, 10])))
    expect(loopback.sent.at(-1)).toEqual(new Uint8Array([104, 105, 10]))

    loopback.emit('message', new Uint8Array([79, 75])) // "OK" echoes back upstream as DATA
    const last = decodeMuxFrame(upstream.sent.at(-1)!)
    expect(last.header.type).toBe('data')
    expect([...last.payload]).toEqual([79, 75])

    await handle.stop()
    expect(await handle.done).toBe(0)
  })

  it('reconnects after the tunnel dies', async () => {
    const sockets = [new FakeWs(), new FakeWs()]
    let i = 0
    const connectRelay = vi.fn(async () => sockets[i++]!)
    const handle = await runTunnel(CFG, KS, baseDeps({ connectRelay, dialLoopback: async () => new FakeWs() }))
    await flush()
    expect(connectRelay).toHaveBeenCalledTimes(1)

    sockets[0]!.emit('close') // first session dies ⇒ supervisor redials
    await flush()
    expect(connectRelay).toHaveBeenCalledTimes(2)

    await handle.stop()
  })

  it('a revoked GOAWAY tears down and NEVER reconnects (INV12)', async () => {
    const connectRelay = vi.fn(async () => new FakeWs())
    let socket: FakeWs | undefined
    const wrapped = vi.fn(async () => {
      socket = new FakeWs()
      return socket
    })
    const handle = await runTunnel(CFG, KS, baseDeps({ connectRelay: wrapped, dialLoopback: async () => new FakeWs() }))
    await flush()
    expect(wrapped).toHaveBeenCalledTimes(1)

    emitGoAwayRevoked(socket!)
    await flush()

    expect(await handle.done).toBe(0)
    expect(wrapped).toHaveBeenCalledTimes(1) // no reconnect after revocation
    expect(connectRelay).not.toHaveBeenCalled()
  })

  it('stop() ends the loop with exit code 0 and does not reconnect', async () => {
    const connectRelay = vi.fn(async () => new FakeWs())
    const handle = await runTunnel(CFG, KS, baseDeps({ connectRelay, dialLoopback: async () => new FakeWs() }))
    await flush()

    await handle.stop()
    expect(await handle.done).toBe(0)
    expect(connectRelay).toHaveBeenCalledTimes(1)
  })
})