web-terminal/relay-web/test/e2e-socket.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest'
import { importAeadKey, createE2ESession } from 'relay-e2e'
import { randomBytes } from 'relay-e2e'
import type { AeadKey, E2ESession, HandshakeResult } from 'relay-contracts'
import { readConfig } from '../src/config'
import { createHostPinStore } from '../src/host-pin-store'
import {
  createE2ETransport,
  type E2EHandshakeIo,
  type E2ETransportDeps,
  type FprDecision,
  type RunHandshakeArgs,
} from '../src/e2e-socket'
import type { WebSocketLike } from '../src/ws-transport'
import { APP_SUBPROTOCOL } from 'relay-contracts'

const cfg = readConfig({
  protocol: 'https:',
  host: 'alice.term.example.com',
  hostname: 'alice.term.example.com',
})

/** A real relay-e2e loopback: client + host sessions sharing DirectionalKeys (genuine crypto). */
function loopbackPair(): { client: E2ESession; host: E2ESession } {
  const c2h: AeadKey = importAeadKey(randomBytes(32), 'aes-256-gcm', 'c2h')
  const h2c: AeadKey = importAeadKey(randomBytes(32), 'aes-256-gcm', 'h2c')
  const result: HandshakeResult = { keys: { c2h, h2c }, aead: 'aes-256-gcm', transcript: new Uint8Array() }
  return { client: createE2ESession('client', result), host: createE2ESession('host', result) }
}

class MockWs implements WebSocketLike {
  static last: MockWs | null = null
  binaryType = ''
  protocol = ''
  readonly sent: Array<ArrayBufferView | ArrayBufferLike | string> = []
  closed = false
  onopen: ((ev: unknown) => void) | null = null
  onmessage: ((ev: { data: unknown }) => void) | null = null
  onclose: ((ev: { code?: number; reason?: string }) => void) | null = null
  onerror: ((ev: unknown) => void) | null = null
  constructor(
    readonly url: string,
    readonly protocols?: string | readonly string[],
  ) {
    MockWs.last = this
  }
  send(data: ArrayBufferView | ArrayBufferLike | string): void {
    this.sent.push(data)
  }
  close(): void {
    this.closed = true
    this.onclose?.({ code: 1000 })
  }
  fireOpen(): void {
    this.protocol = APP_SUBPROTOCOL
    this.onopen?.({})
  }
  fireMessage(data: unknown): void {
    this.onmessage?.({ data })
  }
}

const ctx = {
  deviceAuthProofProvider: { proofFor: async () => 'proof' },
  pinStore: { get: async () => null, pin: async () => {} },
  hostContentSecret: new Uint8Array(32),
}

/**
 * Build a transport whose injected `runHandshake` honours the fingerprint gate (calling
 * verifyOfferedFpr with `offeredFpr`) and, on 'trust', returns the real relay-e2e client session.
 */
function makeTransport(opts: {
  expectedFpr: string
  offeredFpr: string
  onPinMismatch?: (e: string, o: string, s: 'api' | 'cache') => Promise<FprDecision>
  pins?: ReturnType<typeof createHostPinStore>
}) {
  const { client, host } = loopbackPair()
  const pins = opts.pins ?? createHostPinStore(localStorage)
  const runHandshake = vi.fn(
    async (_io: E2EHandshakeIo, args: RunHandshakeArgs): Promise<E2ESession> => {
      const decision = await args.verifyOfferedFpr(opts.offeredFpr)
      if (decision === 'abort') throw new Error('handshake aborted: fingerprint mismatch')
      return client
    },
  )
  const deps: E2ETransportDeps = {
    runHandshake,
    context: ctx,
    agentPubkey: new Uint8Array([1, 2, 3]),
  }
  const onPinMismatch = opts.onPinMismatch ?? (async () => 'abort' as FprDecision)
  const transport = createE2ETransport(
    cfg,
    'RAWTOK',
    'h1',
    opts.expectedFpr,
    pins,
    onPinMismatch,
    deps,
  )
  return { transport, host, pins, runHandshake }
}

async function openTransport(t: ReturnType<typeof makeTransport>['transport']): Promise<void> {
  const opened = t.open()
  MockWs.last!.fireOpen()
  await opened
}

beforeEach(() => {
  localStorage.clear()
  MockWs.last = null
})

// Patch the ctor into deps via globalThis for these tests.
function withMockWs<T>(fn: () => T): T {
  const original = (globalThis as { WebSocket?: unknown }).WebSocket
  ;(globalThis as { WebSocket?: unknown }).WebSocket = MockWs as unknown as typeof WebSocket
  try {
    return fn()
  } finally {
    ;(globalThis as { WebSocket?: unknown }).WebSocket = original
  }
}

describe('createE2ETransport (T8) — INV2 / INV13 / TOFU trust root', () => {
  it('first-connect MITM: relay fpr ≠ API expectedFpr → abort, no session, nothing recorded', async () => {
    await withMockWs(async () => {
      const onPinMismatch = vi.fn(async () => 'abort' as FprDecision)
      const { transport, pins } = makeTransport({
        expectedFpr: 'fpr-API',
        offeredFpr: 'fpr-RELAY-FORGED',
        onPinMismatch,
      })
      await expect(openTransport(transport)).rejects.toThrow()
      expect(onPinMismatch).toHaveBeenCalledWith('fpr-API', 'fpr-RELAY-FORGED', 'api')
      expect(pins.get('h1')).toBeNull() // never recorded a relay-forwarded value
      expect(MockWs.last!.closed).toBe(true)
    })
  })

  it('happy path: fpr == expectedFpr → session established, API value cached (audit trail)', async () => {
    await withMockWs(async () => {
      const { transport, pins } = makeTransport({ expectedFpr: 'fpr-API', offeredFpr: 'fpr-API' })
      await openTransport(transport)
      expect(pins.get('h1')).toBe('fpr-API') // the API-sourced value, not a handshake-derived one
    })
  })

  it('INV2: a known plaintext marker never appears un-AEAD\'d on the wire', async () => {
    await withMockWs(async () => {
      const { transport } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      await openTransport(transport)
      const marker = new TextEncoder().encode('SECRET-MARKER-12345')
      transport.send(marker)
      const framed = MockWs.last!.sent.at(-1) as Uint8Array
      expect(new TextDecoder().decode(framed)).not.toContain('SECRET-MARKER-12345')
    })
  })

  it('round-trips: a host-sealed frame decrypts to the original plaintext via onMessage', async () => {
    await withMockWs(async () => {
      const { transport, host } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      const received: Uint8Array[] = []
      transport.onMessage((b) => received.push(b))
      await openTransport(transport)
      MockWs.last!.fireMessage(host.seal(new TextEncoder().encode('hello')))
      expect(new TextDecoder().decode(received[0]!)).toBe('hello')
    })
  })

  it('INV13 tamper: flipping a ciphertext byte → AEAD tag fails → socket torn down', async () => {
    await withMockWs(async () => {
      const { transport, host } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      let closeReason = ''
      transport.onClose((r) => (closeReason = r))
      await openTransport(transport)
      const frame = Uint8Array.from(host.seal(new TextEncoder().encode('data')))
      const last = frame.length - 1
      frame[last] = (frame[last] ?? 0) ^ 0xff // flip a tag/ciphertext byte
      MockWs.last!.fireMessage(frame)
      expect(closeReason).toBe('e2e-verify-failed')
      expect(MockWs.last!.closed).toBe(true)
    })
  })

  it('INV13 replay: re-injecting a captured frame → rejected → tear-down', async () => {
    await withMockWs(async () => {
      const { transport, host } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      let closeReason = ''
      transport.onClose((r) => (closeReason = r))
      await openTransport(transport)
      const frame = host.seal(new TextEncoder().encode('data'))
      MockWs.last!.fireMessage(frame) // seq 0 accepted
      MockWs.last!.fireMessage(frame) // replay of seq 0 → strict-successor guard rejects
      expect(closeReason).toBe('e2e-verify-failed')
    })
  })

  it('INV13 reorder: seq=1 before seq=0 → rejected', async () => {
    await withMockWs(async () => {
      const { transport, host } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      let closeReason = ''
      transport.onClose((r) => (closeReason = r))
      await openTransport(transport)
      host.seal(new TextEncoder().encode('zero')) // consume seq 0 on the host sender
      const frameOne = host.seal(new TextEncoder().encode('one')) // seq 1
      MockWs.last!.fireMessage(frameOne) // client expects seq 0 first
      expect(closeReason).toBe('e2e-verify-failed')
    })
  })

  it('INV13 reflection: a client→host frame re-injected as host→client → rejected (direction-bound)', async () => {
    await withMockWs(async () => {
      const { transport } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      let closeReason = ''
      transport.onClose((r) => (closeReason = r))
      await openTransport(transport)
      // Capture a genuine outbound (client→host) frame and reflect it into the inbound path.
      transport.send(new TextEncoder().encode('outbound'))
      const reflected = MockWs.last!.sent.at(-1) as Uint8Array
      MockWs.last!.fireMessage(reflected)
      expect(closeReason).toBe('e2e-verify-failed') // c2h frame cannot verify under h2c
    })
  })

  it('key non-exposure: no key bytes reach localStorage; only the public fpr is cached', async () => {
    await withMockWs(async () => {
      const { transport } = makeTransport({ expectedFpr: 'fpr-API', offeredFpr: 'fpr-API' })
      await openTransport(transport)
      transport.send(new TextEncoder().encode('x'))
      const dump = JSON.stringify(localStorage)
      expect(dump).toContain('fpr-API')
      expect(dump).not.toMatch(/__aeadKey|CryptoKey/)
    })
  })

  it('a socket error during the E2E open rejects (no partial session)', async () => {
    await withMockWs(async () => {
      const { transport } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      const opened = transport.open()
      MockWs.last!.onerror?.({})
      await expect(opened).rejects.toThrow('websocket error')
    })
  })

  it('a plain server close during DATA phase surfaces its reason string', async () => {
    await withMockWs(async () => {
      const { transport } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      let reason = ''
      transport.onClose((r) => (reason = r))
      await openTransport(transport)
      MockWs.last!.onclose?.({ code: 1000, reason: 'agent shutdown' })
      expect(reason).toBe('agent shutdown')
    })
  })

  it('a bare close (no code/reason) surfaces the default "closed" reason', async () => {
    await withMockWs(async () => {
      const { transport } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      let reason = ''
      transport.onClose((r) => (reason = r))
      await openTransport(transport)
      MockWs.last!.onclose?.({})
      expect(reason).toBe('closed')
    })
  })

  it('operator override: onPinMismatch("api") → "trust" lets the handshake proceed (branch)', async () => {
    await withMockWs(async () => {
      const onPinMismatch = vi.fn(async () => 'trust' as FprDecision)
      const { transport, pins } = makeTransport({
        expectedFpr: 'fpr-API',
        offeredFpr: 'fpr-DIFFERENT',
        onPinMismatch,
      })
      await openTransport(transport)
      expect(onPinMismatch).toHaveBeenCalledWith('fpr-API', 'fpr-DIFFERENT', 'api')
      expect(pins.get('h1')).toBe('fpr-API') // still records the authoritative API value
    })
  })

  it('send() before the session is established throws (never sends plaintext)', () => {
    withMockWs(() => {
      const { transport } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      expect(() => transport.send(new Uint8Array([1]))).toThrow('e2e transport not open')
    })
  })

  it('a non-abort handshake failure tears the socket down and rejects open()', async () => {
    await withMockWs(async () => {
      const pins = createHostPinStore(localStorage)
      const deps: E2ETransportDeps = {
        runHandshake: vi.fn(async () => {
          throw new Error('agent unreachable')
        }),
        context: ctx,
        agentPubkey: new Uint8Array([1]),
      }
      const transport = createE2ETransport(cfg, 'RAWTOK', 'h1', 'f', pins, async () => 'trust', deps)
      let closeReason = ''
      transport.onClose((r) => (closeReason = r))
      const opened = transport.open()
      MockWs.last!.fireOpen()
      await expect(opened).rejects.toThrow('agent unreachable')
      expect(closeReason).toBe('e2e-handshake-failed')
      expect(pins.get('h1')).toBeNull()
    })
  })

  it('a 4403 server close during the DATA phase surfaces onClose("forbidden")', async () => {
    await withMockWs(async () => {
      const { transport } = makeTransport({ expectedFpr: 'f', offeredFpr: 'f' })
      let closeReason = ''
      transport.onClose((r) => (closeReason = r))
      await openTransport(transport)
      MockWs.last!.onclose?.({ code: 4403 })
      expect(closeReason).toBe('forbidden')
    })
  })

  it('cache drift: a later expectedFpr differing from the cache surfaces onPinMismatch(...,"cache")', async () => {
    await withMockWs(async () => {
      const pins = createHostPinStore(localStorage)
      pins.record('h1', 'fpr-OLD')
      const onPinMismatch = vi.fn(async () => 'trust' as FprDecision)
      const { transport } = makeTransport({
        expectedFpr: 'fpr-NEW',
        offeredFpr: 'fpr-NEW',
        onPinMismatch,
        pins,
      })
      await openTransport(transport)
      expect(onPinMismatch).toHaveBeenCalledWith('fpr-NEW', 'fpr-OLD', 'cache')
      expect(pins.get('h1')).toBe('fpr-NEW') // API value is authoritative — cache updated
    })
  })
})