web-terminal/relay-web/test/ws-transport.test.ts

import { describe, expect, it } from 'vitest'
import { APP_SUBPROTOCOL, encodeTokenSubprotocol } from 'relay-contracts'
import { readConfig } from '../src/config'
import { createPassthroughTransport, type WebSocketLike } from '../src/ws-transport'

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

/** Controllable mock WebSocket capturing constructor args and driving lifecycle events. */
class MockWs implements WebSocketLike {
  static last: MockWs | null = null
  binaryType = ''
  protocol = ''
  readonly url: string
  readonly protocols: readonly string[]
  readonly sent: Array<ArrayBufferView | ArrayBufferLike | string> = []
  closed: { code?: number; reason?: string } | null = null
  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(url: string, protocols?: string | readonly string[]) {
    this.url = url
    this.protocols = protocols === undefined ? [] : typeof protocols === 'string' ? [protocols] : protocols
    MockWs.last = this
  }
  send(data: ArrayBufferView | ArrayBufferLike | string): void {
    this.sent.push(data)
  }
  close(code?: number, reason?: string): void {
    const ev: { code?: number; reason?: string } = {}
    if (code !== undefined) ev.code = code
    if (reason !== undefined) ev.reason = reason
    this.closed = ev
    this.onclose?.(ev)
  }
  fireOpen(echoedProtocol: string): void {
    this.protocol = echoedProtocol
    this.onopen?.({})
  }
  fireMessage(data: unknown): void {
    this.onmessage?.({ data })
  }
}

describe('createPassthroughTransport (T4)', () => {
  it('v0.8: opens same-origin wss with NO query-string token (cookie-only auth)', async () => {
    const t = createPassthroughTransport(cfg, { wsCtor: MockWs })
    const opened = t.open()
    MockWs.last!.fireOpen(APP_SUBPROTOCOL)
    await opened
    expect(MockWs.last!.url).toBe('wss://alice.term.example.com/term')
    expect(MockWs.last!.url).not.toMatch(/token|\?/)
    expect(MockWs.last!.protocols).toEqual([APP_SUBPROTOCOL]) // no token entry in v0.8
    expect(MockWs.last!.binaryType).toBe('arraybuffer')
  })

  it('round-trips bytes: send() reaches the socket, onMessage() delivers incoming bytes', async () => {
    const t = createPassthroughTransport(cfg, { wsCtor: MockWs })
    const received: Uint8Array[] = []
    t.onMessage((b) => received.push(b))
    const opened = t.open()
    MockWs.last!.fireOpen(APP_SUBPROTOCOL)
    await opened

    t.send(new Uint8Array([9, 8, 7]))
    expect(MockWs.last!.sent[0]).toEqual(new Uint8Array([9, 8, 7]))

    MockWs.last!.fireMessage(new Uint8Array([1, 2]).buffer)
    expect(received[0]).toEqual(new Uint8Array([1, 2]))
  })

  it('v0.9+ token attachment: token rides Sec-WebSocket-Protocol, NEVER the URL (§4.3)', async () => {
    const t = createPassthroughTransport(cfg, { wsCtor: MockWs, capabilityToken: 'RAWTOK' })
    const opened = t.open()
    MockWs.last!.fireOpen(APP_SUBPROTOCOL)
    await opened
    expect(MockWs.last!.protocols).toEqual([APP_SUBPROTOCOL, encodeTokenSubprotocol('RAWTOK')])
    expect(MockWs.last!.url).not.toContain('RAWTOK') // never in the URL/query
    expect(MockWs.last!.url).not.toContain(encodeTokenSubprotocol('RAWTOK'))
  })

  it('v0.9+ echo rule: a token echo (or foreign subprotocol) tears the socket down', async () => {
    const t = createPassthroughTransport(cfg, { wsCtor: MockWs, capabilityToken: 'RAWTOK' })
    const opened = t.open()
    MockWs.last!.fireOpen(encodeTokenSubprotocol('RAWTOK')) // relay wrongly echoes the token entry
    await expect(opened).rejects.toThrow()
    expect(MockWs.last!.closed).not.toBeNull()
  })

  it('server close reason surfaces via onClose (no silent swallow)', async () => {
    const t = createPassthroughTransport(cfg, { wsCtor: MockWs })
    let reason = ''
    t.onClose((r) => (reason = r))
    const opened = t.open()
    MockWs.last!.fireOpen(APP_SUBPROTOCOL)
    await opened
    MockWs.last!.close(1000, 'shell exited')
    expect(reason).toBe('shell exited')
  })

  it('a socket error during open rejects the open() promise', async () => {
    const t = createPassthroughTransport(cfg, { wsCtor: MockWs })
    const opened = t.open()
    MockWs.last!.onerror?.({})
    await expect(opened).rejects.toThrow('websocket error')
  })

  it('close() before open() is a safe no-op (no throw)', () => {
    const t = createPassthroughTransport(cfg, { wsCtor: MockWs })
    expect(() => t.close()).not.toThrow()
  })

  it('a 4403 upgrade denial maps to onClose("forbidden") (INV1/INV6 client mirror)', async () => {
    const t = createPassthroughTransport(cfg, { wsCtor: MockWs })
    let reason = ''
    t.onClose((r) => (reason = r))
    const opened = t.open()
    MockWs.last!.fireOpen(APP_SUBPROTOCOL)
    await opened
    MockWs.last!.close(4403, '')
    expect(reason).toBe('forbidden')
  })
})