web-terminal/term-relay/test/revocation-subscriber.test.ts

import { describe, test, expect, vi } from 'vitest'
import { RELAY_REVOCATIONS_CHANNEL, REVOCATION_PUSH_BUDGET_MS, type KillSignal } from 'relay-contracts'
import { startRevocationSubscriber, hostsForScope, type Subscription } from '../data-plane/revocation-subscriber.js'
import { DRAIN_REASON } from '../data-plane/drain.js'
import type { AgentTunnel } from '../data-plane/agent-listener.js'

const UUID_H = '11111111-1111-4111-8111-111111111111'
const UUID_H2 = '22222222-2222-4222-8222-222222222222'
const UUID_A = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa'

function tunnel(hostId: string, accountId: string): AgentTunnel {
  return { hostId, accountId, session: {} as AgentTunnel['session'], closeTunnel: () => {} }
}

function fakeBus() {
  let handler: ((raw: string) => void) | null = null
  const sub: Subscription = { close: vi.fn() }
  return {
    subscribe: (channel: string, onMessage: (raw: string) => void) => {
      expect(channel).toBe(RELAY_REVOCATIONS_CHANNEL)
      handler = onMessage
      return sub
    },
    publish: (signal: KillSignal | string) => handler?.(typeof signal === 'string' ? signal : JSON.stringify(signal)),
    sub,
  }
}

describe('hostsForScope (T14, pure blast-radius selector)', () => {
  const live = new Map([
    [UUID_H, tunnel(UUID_H, UUID_A)],
    [UUID_H2, tunnel(UUID_H2, 'other-acct')],
  ])
  const accountOf = (h: string) => live.get(h)?.accountId

  test('host scope: only if this node serves it', () => {
    expect(hostsForScope({ kind: 'host', hostId: UUID_H }, live, accountOf)).toEqual([UUID_H])
    expect(hostsForScope({ kind: 'host', hostId: 'nope' }, live, accountOf)).toEqual([])
  })
  test('account scope: only the account’s hosts on this node', () => {
    expect(hostsForScope({ kind: 'account', accountId: UUID_A }, live, accountOf)).toEqual([UUID_H])
  })
  test('global scope: every live host', () => {
    expect(hostsForScope({ kind: 'global' }, live, accountOf)).toEqual([UUID_H, UUID_H2])
  })
})

describe('revocation subscriber (T14, INV12/FIX 4)', () => {
  function setup(tunnels: Map<string, AgentTunnel>) {
    const bus = fakeBus()
    const drainHost = vi.fn(async () => {})
    const onApplied = vi.fn()
    const onDropped = vi.fn()
    const clock = { t: 0 }
    startRevocationSubscriber({
      subscribe: bus.subscribe,
      drainHost,
      tunnels: () => tunnels,
      now: () => clock.t,
      onApplied,
      onDropped,
    })
    return { bus, drainHost, onApplied, onDropped, clock }
  }

  test('host signal → drainHost(H, revoked) once, within the push budget', async () => {
    const { bus, drainHost } = setup(new Map([[UUID_H, tunnel(UUID_H, UUID_A)]]))
    bus.publish({ scope: { kind: 'host', hostId: UUID_H }, at: 1, reason: 'compromised' })
    await new Promise((r) => setTimeout(r, 0))
    expect(drainHost).toHaveBeenCalledTimes(1)
    expect(drainHost).toHaveBeenCalledWith(UUID_H, DRAIN_REASON.revoked)
    expect(REVOCATION_PUSH_BUDGET_MS).toBe(2000)
  })

  test('a host NOT on this node → no-op (0 calls, no cross-tenant reach)', async () => {
    const { bus, drainHost } = setup(new Map([[UUID_H, tunnel(UUID_H, UUID_A)]]))
    bus.publish({ scope: { kind: 'host', hostId: UUID_H2 }, at: 1, reason: 'x' })
    await new Promise((r) => setTimeout(r, 0))
    expect(drainHost).not.toHaveBeenCalled()
  })

  test('account scope tears down only that account’s hosts (sibling account keeps flowing, INV1)', async () => {
    const { bus, drainHost } = setup(
      new Map([
        [UUID_H, tunnel(UUID_H, UUID_A)],
        [UUID_H2, tunnel(UUID_H2, 'other-acct')],
      ]),
    )
    bus.publish({ scope: { kind: 'account', accountId: UUID_A }, at: 1, reason: 'x' })
    await new Promise((r) => setTimeout(r, 0))
    expect(drainHost).toHaveBeenCalledTimes(1)
    expect(drainHost).toHaveBeenCalledWith(UUID_H, DRAIN_REASON.revoked)
  })

  test('global scope tears down every live tunnel', async () => {
    const { bus, drainHost } = setup(
      new Map([
        [UUID_H, tunnel(UUID_H, UUID_A)],
        [UUID_H2, tunnel(UUID_H2, 'other-acct')],
      ]),
    )
    bus.publish({ scope: { kind: 'global' }, at: 1, reason: 'x' })
    await new Promise((r) => setTimeout(r, 0))
    expect(drainHost).toHaveBeenCalledTimes(2)
  })

  test('malformed signal → dropped + counted, no teardown, subscriber stays alive', async () => {
    const { bus, drainHost, onDropped } = setup(new Map([[UUID_H, tunnel(UUID_H, UUID_A)]]))
    bus.publish('{ not valid json')
    bus.publish(JSON.stringify({ scope: { kind: 'nonsense' }, at: 1, reason: 'x' }))
    await new Promise((r) => setTimeout(r, 0))
    expect(onDropped).toHaveBeenCalledTimes(2)
    expect(drainHost).not.toHaveBeenCalled()
    // Still alive: a good signal after the poison ones is applied.
    bus.publish({ scope: { kind: 'host', hostId: UUID_H }, at: 2, reason: 'x' })
    await new Promise((r) => setTimeout(r, 0))
    expect(drainHost).toHaveBeenCalledTimes(1)
  })

  test('onApplied reports metadata only (count + elapsed), never the reason as payload (INV10)', async () => {
    const { bus, onApplied } = setup(new Map([[UUID_H, tunnel(UUID_H, UUID_A)]]))
    bus.publish({ scope: { kind: 'host', hostId: UUID_H }, at: 1, reason: 'super-secret-reason' })
    await new Promise((r) => setTimeout(r, 0))
    expect(onApplied).toHaveBeenCalledTimes(1)
    const [, hostsAffected, elapsed] = onApplied.mock.calls[0]!
    expect(hostsAffected).toBe(1)
    expect(typeof elapsed).toBe('number')
  })

  test('close() unsubscribes cleanly', () => {
    const bus = fakeBus()
    const handle = startRevocationSubscriber({
      subscribe: bus.subscribe,
      drainHost: vi.fn(async () => {}),
      tunnels: () => new Map(),
      now: () => 0,
    })
    handle.close()
    expect(bus.sub.close).toHaveBeenCalledTimes(1)
  })
})