web-terminal/relay-run/tests/wiring/stores-pg.test.ts

import { describe, it, expect } from 'vitest'
import type { AuditEvent } from 'relay-auth'
import { NO_STEPUP_POLICY } from 'relay-auth/src/human/stepup/stepup.js'
import type { QueryFn } from 'control-plane/src/db/pool.js'
import { createRelayEnforceDeps, type RedisLike } from '../../src/wiring/stores-pg.js'

// ── Fakes ─────────────────────────────────────────────────────────────────────────────────────

interface QueryCall {
  readonly sql: string
  readonly params: readonly unknown[]
}

/** A `QueryFn` that records every call and returns rows chosen by `rowsFor`. */
function makeFakeQuery(
  rowsFor: (sql: string, params: readonly unknown[]) => unknown[],
): { query: QueryFn; calls: QueryCall[] } {
  const calls: QueryCall[] = []
  const query: QueryFn = async <T>(sql: string, params: readonly unknown[]): Promise<readonly T[]> => {
    calls.push({ sql, params })
    return rowsFor(sql, params) as T[]
  }
  return { query, calls }
}

interface RedisCall {
  readonly method: string
  readonly args: readonly unknown[]
}

/** A `RedisLike` mock whose methods can be overridden per test; every call is recorded. */
function makeMockRedis(over: Partial<RedisLike> = {}): RedisLike & { calls: RedisCall[] } {
  const calls: RedisCall[] = []
  const record = (method: string, args: unknown[]): void => void calls.push({ method, args })
  const redis: RedisLike = {
    exists: over.exists ?? (async (key) => (record('exists', [key]), 0)),
    set: over.set ?? (async (key, value, mode) => (record('set', [key, value, mode]), 'OK')),
    expireat: over.expireat ?? (async (key, ts) => (record('expireat', [key, ts]), 1)),
    eval: over.eval ?? (async (...args) => (record('eval', args), 1)),
  }
  return Object.assign(redis, { calls })
}

const NOOP_QUERY: QueryFn = async () => []

// ── hosts (Postgres pass-through) ───────────────────────────────────────────────────────────────

describe('createRelayEnforceDeps.hosts', () => {
  it('maps a CP host row to a relay-auth HostRecord', async () => {
    const { query, calls } = makeFakeQuery(() => [
      {
        host_id: 'host-1',
        account_id: 'acct-1',
        subdomain: 'demo',
        agent_pubkey: Buffer.from([1, 2, 3]),
        enroll_fpr: 'fpr-1',
        status: 'online',
        last_seen: '2026-07-06T00:00:00.000Z',
        created_at: '2026-07-05T00:00:00.000Z',
        revoked_at: null,
      },
    ])
    const deps = createRelayEnforceDeps({ query, redis: makeMockRedis() })

    const host = await deps.hosts.getById('host-1')

    expect(host).not.toBeNull()
    expect(host?.hostId).toBe('host-1')
    expect(host?.accountId).toBe('acct-1')
    expect(host?.subdomain).toBe('demo')
    expect(host?.status).toBe('online')
    expect(host?.revokedAt).toBeNull()
    expect(Array.from(host?.agentPubkey ?? [])).toEqual([1, 2, 3])
    // accountId came from the stored row, never fabricated (INV3).
    expect(calls[0]?.params).toEqual(['host-1'])
  })

  it('returns null for an unknown host', async () => {
    const { query } = makeFakeQuery(() => [])
    const deps = createRelayEnforceDeps({ query, redis: makeMockRedis() })
    expect(await deps.hosts.getById('nope')).toBeNull()
  })
})

// ── sessions (Postgres, remapped to the port shape) ──────────────────────────────────────────────

describe('createRelayEnforceDeps.sessions', () => {
  it('maps a session row to { hostId, accountId }', async () => {
    const { query } = makeFakeQuery(() => [
      {
        session_id: 'sess-1',
        host_id: 'host-1',
        account_id: 'acct-1',
        created_at: '2026-07-06T00:00:00.000Z',
        last_attach_at: '2026-07-06T00:00:00.000Z',
      },
    ])
    const deps = createRelayEnforceDeps({ query, redis: makeMockRedis() })

    expect(await deps.sessions.getById('sess-1')).toEqual({ hostId: 'host-1', accountId: 'acct-1' })
  })

  it('returns null for an unknown session', async () => {
    const { query } = makeFakeQuery(() => [])
    const deps = createRelayEnforceDeps({ query, redis: makeMockRedis() })
    expect(await deps.sessions.getById('nope')).toBeNull()
  })
})

// ── revocation (Redis) ───────────────────────────────────────────────────────────────────────────

describe('createRelayEnforceDeps.revocation', () => {
  it('isRevoked is true iff the revoked:<jti> key exists', async () => {
    const present = createRelayEnforceDeps({
      query: NOOP_QUERY,
      redis: makeMockRedis({ exists: async () => 1 }),
    })
    const absent = createRelayEnforceDeps({
      query: NOOP_QUERY,
      redis: makeMockRedis({ exists: async () => 0 }),
    })
    expect(await present.revocation.isRevoked('j1')).toBe(true)
    expect(await absent.revocation.isRevoked('j1')).toBe(false)
  })

  it('revokeJti sets revoked:<jti> and pins its expiry to exp', async () => {
    const redis = makeMockRedis()
    const deps = createRelayEnforceDeps({ query: NOOP_QUERY, redis })

    await deps.revocation.revokeJti('j1', 1_800_000_000)

    expect(redis.calls).toEqual([
      { method: 'set', args: ['revoked:j1', '1', undefined] },
      { method: 'expireat', args: ['revoked:j1', 1_800_000_000] },
    ])
  })

  it('consumeOnce burns the jti: first use wins (SET NX), replays lose', async () => {
    let existing = false
    const redis = makeMockRedis({
      set: async (_k, _v, mode) => {
        if (mode !== 'NX') return 'OK'
        if (existing) return null
        existing = true
        return 'OK'
      },
    })
    const deps = createRelayEnforceDeps({ query: NOOP_QUERY, redis })

    expect(await deps.revocation.consumeOnce('j1', 1_800_000_000)).toBe(true)
    expect(await deps.revocation.consumeOnce('j1', 1_800_000_000)).toBe(false)
    // expiry is only pinned on the winning first use.
    expect(redis.calls.filter((c) => c.method === 'expireat')).toEqual([
      { method: 'expireat', args: ['used:j1', 1_800_000_000] },
    ])
  })
})

// ── buckets (Redis token bucket) ─────────────────────────────────────────────────────────────────

describe('createRelayEnforceDeps.buckets', () => {
  it('namespaces the key and forwards refill/burst/now + a computed ttl to the Lua script', async () => {
    let evalArgs: readonly unknown[] = []
    const redis = makeMockRedis({
      eval: async (...args) => {
        evalArgs = args
        return 1
      },
    })
    const deps = createRelayEnforceDeps({ query: NOOP_QUERY, redis })

    const ok = await deps.buckets.take('connect:acct:a1', 1, 60, 1000)

    expect(ok).toBe(true)
    // eval(script, numKeys, key, refillPerSec, burst, now, ttl)
    expect(evalArgs[1]).toBe(1) // numKeys
    expect(evalArgs[2]).toBe('bucket:connect:acct:a1')
    expect(evalArgs[3]).toBe(1) // refillPerSec
    expect(evalArgs[4]).toBe(60) // burst
    expect(evalArgs[5]).toBe(1000) // now
    expect(evalArgs[6]).toBe(61) // ttl = ceil(60/1) + 1
  })

  it('maps a 0 result to throttled (false)', async () => {
    const deps = createRelayEnforceDeps({
      query: NOOP_QUERY,
      redis: makeMockRedis({ eval: async () => 0 }),
    })
    expect(await deps.buckets.take('preauth:ip:x', 1, 60, 1000)).toBe(false)
  })
})

// ── audit (Postgres audit_log, metadata only) ────────────────────────────────────────────────────

describe('createRelayEnforceDeps.audit', () => {
  it('maps an AuditEvent onto the audit_log row, folding non-column fields into meta (INV10)', async () => {
    const { query, calls } = makeFakeQuery(() => [])
    const deps = createRelayEnforceDeps({ query, redis: makeMockRedis() })
    const event: AuditEvent = {
      ts: '2026-07-06T00:00:00.000Z',
      action: 'attach',
      principalId: 'cred-1',
      accountId: 'acct-1',
      hostId: 'host-1',
      sessionId: 'sess-1',
      jti: 'jti-1',
      outcome: 'allow',
      reason: 'ok',
      remoteAddrHash: 'iphash',
    }

    await deps.audit.append(event)

    const params = calls[0]?.params ?? []
    expect(params[0]).toBe('attach') // action
    expect(params[1]).toBe('cred-1') // principal_id
    expect(params[2]).toBe('acct-1') // account_id
    expect(params[3]).toBe('host-1') // host_id
    expect(params[4]).toBe('2026-07-06T00:00:00.000Z') // ts
    expect(JSON.parse(params[5] as string)).toEqual({
      outcome: 'allow',
      reason: 'ok',
      remoteAddrHash: 'iphash',
      sessionId: 'sess-1',
      jti: 'jti-1',
    })
  })

  it('omits null sessionId/jti from meta', async () => {
    const { query, calls } = makeFakeQuery(() => [])
    const deps = createRelayEnforceDeps({ query, redis: makeMockRedis() })
    const event: AuditEvent = {
      ts: '2026-07-06T00:00:00.000Z',
      action: 'attach',
      principalId: '',
      accountId: '',
      hostId: null,
      sessionId: null,
      jti: null,
      outcome: 'deny',
      reason: 'bad_origin',
      remoteAddrHash: 'iphash',
    }

    await deps.audit.append(event)

    expect(calls[0]?.params[3]).toBeNull() // host_id
    expect(JSON.parse((calls[0]?.params[5] as string) ?? '{}')).toEqual({
      outcome: 'deny',
      reason: 'bad_origin',
      remoteAddrHash: 'iphash',
    })
  })
})

// ── stepUpPolicyFor (staging) ────────────────────────────────────────────────────────────────────

describe('createRelayEnforceDeps.stepUpPolicyFor', () => {
  it('returns NO_STEPUP_POLICY (staging single-operator)', () => {
    const deps = createRelayEnforceDeps({ query: NOOP_QUERY, redis: makeMockRedis() })
    // host argument is irrelevant in staging; the policy is never-required.
    expect(deps.stepUpPolicyFor({} as never)).toBe(NO_STEPUP_POLICY)
    expect(deps.stepUpPolicyFor({} as never).required).toBe(false)
  })
})