web-terminal/agent/src/transport/runTunnel.ts

/**
 * Long-running tunnel supervisor — PLAN_RELAY_PHASE1 C2. Ports the PROVEN cafeDemo assembly
 * (dialRelay → holdTunnel → createStreamRouter → dialLoopback, plus heartbeat) into a supervised
 * loop that survives disconnects: exponential backoff reconnect (T10 policy), heartbeat liveness
 * (T9), and GOAWAY/revocation-aware teardown (T14, INV12 — a revoked host NEVER reconnects).
 *
 * All IO is injectable via `RunTunnelDeps` so the loop is unit-testable with fakes (see cafeDemo);
 * the two-arg `runTunnel(cfg, ks)` default path wires the real `ws` sockets. INV2 is preserved: the
 * router splices OPAQUE bytes (identityTransform) — no terminal parsing happens here.
 */
import { WebSocket } from 'ws'
import type { AgentConfig } from '../config/agentConfig.js'
import type { Keystore } from '../keys/keystore.js'
import { createLogger, type Logger } from '../log/logger.js'
import { createRevocationState, applyGoAway } from '../lifecycle/revocation.js'
import { dialRelay, type TlsWsConstructor } from './dial.js'
import { dialLoopback, type DialLoopback, type WsConstructor } from './loopback.js'
import { holdTunnel, type Tunnel } from './tunnel.js'
import { createStreamRouter, identityTransform } from './streamRouter.js'
import { createHeartbeat } from './heartbeat.js'
import { createBackoff, reconnectLoop, type BackoffPolicy, type Sleep } from './backoff.js'
import type { TimerLike, WsLike } from './seams.js'

/** Handle returned by `runTunnel`: stop the supervisor, or await its terminal exit code. */
export interface TunnelHandle {
  /** Request graceful shutdown; resolves once the supervisor loop has fully stopped. */
  stop(): Promise<void>
  /** Resolves with a process exit code when the loop ends (stopped or host revoked ⇒ 0). */
  readonly done: Promise<number>
}

/** Injectable seams for the supervisor. All optional; unset fields default to real `ws` IO. */
export interface RunTunnelDeps {
  connectRelay(): Promise<WsLike>
  dialLoopback: DialLoopback
  logger: Logger
  timer: TimerLike
  sleep: Sleep
  backoff: BackoffPolicy
}

const realTimer: TimerLike = {
  setTimeout: (cb, ms) => setTimeout(cb, ms),
  clearTimeout: (h) => clearTimeout(h as ReturnType<typeof setTimeout>),
  setInterval: (cb, ms) => setInterval(cb, ms),
  clearInterval: (h) => clearInterval(h as ReturnType<typeof setInterval>),
}
const realSleep: Sleep = (ms) => new Promise<void>((r) => setTimeout(r, ms))

function resolveDeps(cfg: AgentConfig, ks: Keystore, o?: Partial<RunTunnelDeps>): RunTunnelDeps {
  return {
    connectRelay:
      o?.connectRelay ?? (() => dialRelay(cfg, ks, { Ctor: WebSocket as unknown as TlsWsConstructor })),
    dialLoopback: o?.dialLoopback ?? dialLoopback(cfg.localTargetUrl, WebSocket as unknown as WsConstructor),
    logger: o?.logger ?? createLogger('info'),
    timer: o?.timer ?? realTimer,
    sleep: o?.sleep ?? realSleep,
    backoff: o?.backoff ?? createBackoff({ jitter: true }),
  }
}

/**
 * Start the supervised tunnel. Returns immediately with a handle; the reconnect loop runs in the
 * background. Never awaits the first connection (an unreachable relay would otherwise hang the
 * caller with no way to `stop()`).
 */
export async function runTunnel(
  cfg: AgentConfig,
  ks: Keystore,
  overrides?: Partial<RunTunnelDeps>,
): Promise<TunnelHandle> {
  const { connectRelay, dialLoopback: dialLb, logger, timer, sleep, backoff } = resolveDeps(cfg, ks, overrides)

  let stopped = false
  let currentTunnel: Tunnel | null = null
  let currentSocket: WsLike | null = null

  // INV12: revocation tears the live tunnel down immediately and suppresses all reconnects.
  const revocation = createRevocationState(() => currentTunnel?.close())
  const shouldStop = (): boolean => stopped || revocation.isRevoked()

  async function dialTunnel(): Promise<Tunnel> {
    const socket = await connectRelay()
    currentSocket = socket
    return holdTunnel(socket)
  }

  /** Run ONE tunnel session; resolves when this session ends (dead heartbeat / close / GOAWAY). */
  function runSession(tunnel: Tunnel, socket: WsLike): Promise<void> {
    return new Promise<void>((resolve) => {
      const router = createStreamRouter(cfg, tunnel, dialLb, identityTransform, logger)
      const heartbeat = createHeartbeat(tunnel, { timer })
      tunnel.dispatchTo(router, heartbeat)

      let settled = false
      const endSession = (): void => {
        if (settled) return
        settled = true
        heartbeat.stop()
        resolve()
      }

      heartbeat.onDead(() => {
        logger.log('warn', 'heartbeat missed — tunnel presumed down, will reconnect')
        tunnel.close()
        endSession()
      })
      socket.on('close', () => endSession())
      socket.on('error', () => {
        tunnel.close()
        endSession()
      })
      tunnel.onGoAway((reason) => {
        const action = applyGoAway(reason, revocation) // 'revoked' ⇒ no reconnect (INV12)
        logger.log('info', 'received GOAWAY', { action })
        tunnel.close()
        endSession()
      })

      heartbeat.start()
    })
  }

  async function supervise(): Promise<number> {
    while (!shouldStop()) {
      const tunnel = await reconnectLoop(dialTunnel, backoff, shouldStop, sleep)
      if (tunnel === null || currentSocket === null) break
      if (shouldStop()) {
        // stop()/revoke raced with the in-flight dial — discard the fresh tunnel.
        tunnel.close()
        break
      }
      currentTunnel = tunnel
      logger.log('info', 'relay tunnel established')
      await runSession(tunnel, currentSocket)
      currentTunnel = null
      currentSocket = null
    }
    return 0
  }

  const done = supervise()

  return {
    async stop(): Promise<void> {
      stopped = true
      currentTunnel?.close()
      await done
    },
    done,
  }
}