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

/**
 * T11 · Graceful drain + GOAWAY (INV7/INV12). GOAWAY on streamId 0 to every (or one) tunnel,
 * then teardown. The grace window is REASON-DIFFERENTIATED (Finding-2):
 *   effectiveGraceMs = (reason === revoked) ? 0 : inFlightGraceMs
 * `revoked` is a SECURITY action against a compromised host/device — it forces immediate teardown
 * (RST/close now, deregister now) regardless of any caller-supplied grace, so INV12's "its live
 * tunnel drops within seconds" holds even when operator-drain grace is tuned to tens of seconds.
 * `operatorDrain`/`shutdown` keep a grace window purely for clean migration.
 */
import type { AgentTunnel, RouteRegistrar } from './agent-listener.js'

/** Drain reason ⇄ §4.1 GOAWAY wire code (frozen in relay-contracts as GOAWAY_REASON_TO_CODE). */
export const DRAIN_REASON = { operatorDrain: 1, revoked: 2, shutdown: 3 } as const
export type DrainReason = (typeof DRAIN_REASON)[keyof typeof DRAIN_REASON]

export interface DrainDeps {
  tunnels(): ReadonlyMap<string, AgentTunnel>
  registrar: RouteRegistrar
  reason: DrainReason
  inFlightGraceMs: number
  // Injected timer for the grace window (default global setTimeout); testability seam.
  setTimeoutFn?: (fn: () => void, ms: number) => unknown
}

export interface DrainHostDeps extends DrainDeps {
  // single-host target resolved by drainHost's first arg
}

function effectiveGrace(reason: DrainReason, inFlightGraceMs: number): number {
  return reason === DRAIN_REASON.revoked ? 0 : Math.max(0, inFlightGraceMs)
}

async function tearDownTunnel(t: AgentTunnel, registrar: RouteRegistrar): Promise<void> {
  t.closeTunnel() // socket close (agent side sees the streams die)
  await registrar.deregister(t.hostId) // ingress stops routing to this host on this node
}

/** GOAWAY then teardown one tunnel, honoring the reason-differentiated grace. */
function drainOne(
  t: AgentTunnel,
  reason: DrainReason,
  inFlightGraceMs: number,
  registrar: RouteRegistrar,
  setTimeoutFn: (fn: () => void, ms: number) => unknown,
): Promise<void> {
  t.session.drain(0, reason) // §4.1 GOAWAY on streamId 0; refuses new streams thereafter
  const grace = effectiveGrace(reason, inFlightGraceMs)
  if (grace === 0) {
    return tearDownTunnel(t, registrar) // immediate (revoked, or zero-grace operator drain)
  }
  setTimeoutFn(() => {
    void tearDownTunnel(t, registrar)
  }, grace)
  return Promise.resolve() // grace scheduled; drain initiated
}

/** Drain EVERY tunnel on the node (operator drain / shutdown / node-wide revocation). */
export async function drainNode(deps: DrainDeps): Promise<void> {
  const setTimeoutFn = deps.setTimeoutFn ?? ((fn, ms) => setTimeout(fn, ms))
  await Promise.all(
    [...deps.tunnels().values()].map((t) =>
      drainOne(t, deps.reason, deps.inFlightGraceMs, deps.registrar, setTimeoutFn),
    ),
  )
}

/** Drain exactly ONE host, leaving siblings running (whole-host revocation / operator drain, INV12). */
export async function drainHost(hostId: string, deps: DrainHostDeps): Promise<void> {
  const setTimeoutFn = deps.setTimeoutFn ?? ((fn, ms) => setTimeout(fn, ms))
  const t = deps.tunnels().get(hostId)
  if (t === undefined) return // no-op: already gone (deny-by-default, never a broader kill)
  await drainOne(t, deps.reason, deps.inFlightGraceMs, deps.registrar, setTimeoutFn)
}