yaoyao/tests/server/websocket.test.ts

import { createServer, type Server } from 'node:http'
import type { AddressInfo } from 'node:net'
import WebSocket, { WebSocketServer } from 'ws'
import { afterEach, describe, expect, it } from 'vitest'
import { createNodeServer, type NodeServerRuntime } from '../../src/server/app.js'
import type { ServerConfig } from '../../src/server/config.js'
import { mkdtempSync, rmSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { createAuthenticatedApplication as createApplication } from './authenticatedApplication.js'

const closers: Array<() => Promise<void>> = []
const homes: string[] = []

afterEach(async () => {
  for (const close of closers.splice(0).reverse()) await close()
  for (const home of homes.splice(0)) rmSync(home, { recursive: true, force: true })
})

function listen(server: Server): Promise<number> {
  return new Promise((resolve, reject) => {
    server.once('error', reject)
    server.listen(0, '127.0.0.1', () => resolve((server.address() as AddressInfo).port))
  })
}

function closeServer(server: Server): Promise<void> {
  return new Promise((resolve, reject) => server.close((error) => error ? reject(error) : resolve()))
}

function cookies(response: Response): string {
  return response.headers.getSetCookie().map((value) => value.split(';', 1)[0]).join('; ')
}

function wsOpen(url: string, origin: string, cookie: string): Promise<WebSocket> {
  return new Promise((resolve, reject) => {
    const socket = new WebSocket(url, { origin, headers: { Cookie: cookie } })
    socket.once('open', () => resolve(socket))
    socket.once('error', reject)
  })
}

function wsMessage(socket: WebSocket): Promise<Record<string, unknown>> {
  return new Promise((resolve, reject) => {
    socket.once('message', (data) => {
      try {
        resolve(JSON.parse(data.toString()) as Record<string, unknown>)
      } catch (error) {
        reject(error)
      }
    })
    socket.once('error', reject)
  })
}

function closeRuntime(runtime: NodeServerRuntime): () => Promise<void> {
  return () => runtime.close()
}

describe('WebSocket relay', () => {
  it('lets a QR-paired iOS node use Hermes-compatible REST and WebSocket paths', async () => {
    let pairedCookie = ''
    const pairedFrames: Array<Record<string, unknown>> = []
    const upstream = createServer((request, response) => {
      response.setHeader('Content-Type', 'application/json')
      if (request.url === '/api/status') {
        pairedCookie = request.headers.cookie ?? pairedCookie
        response.setHeader('Set-Cookie', 'hermes_session_at=paired-user; Path=/; HttpOnly')
        response.end(JSON.stringify({ auth_required: true }))
      } else if (request.url === '/api/auth/me') {
        pairedCookie = request.headers.cookie ?? pairedCookie
        response.end(JSON.stringify({ user_id: 'paired-user', display_name: 'Paired User' }))
      } else if (request.url === '/api/profiles') {
        pairedCookie = request.headers.cookie ?? pairedCookie
        response.end(JSON.stringify({ profiles: [{ name: 'default', is_default: true }] }))
      } else if (request.url === '/api/plugins/yaoyao/profiles') {
        response.end(JSON.stringify({ profiles: [{ name: 'default', botName: '竹儿', agentName: '旧插件名称' }] }))
      } else if (request.url === '/api/auth/providers') {
        response.end(JSON.stringify({ providers: [{ name: 'basic', supports_password: true }] }))
      } else if (request.url === '/auth/password-login' && request.method === 'POST') {
        response.setHeader('Set-Cookie', 'hermes_session_rt=paired-refresh; Path=/; HttpOnly')
        response.end(JSON.stringify({ ok: true }))
      } else if (request.url === '/api/auth/ws-ticket' && request.method === 'POST') {
        pairedCookie = request.headers.cookie ?? pairedCookie
        response.end(JSON.stringify({ ticket: 'paired-upstream-ticket' }))
      } else {
        response.statusCode = 404
        response.end(JSON.stringify({ error: 'not found' }))
      }
    })
    const upstreamSockets = new WebSocketServer({ noServer: true })
    upstream.on('upgrade', (request, socket, head) => {
      const url = new URL(request.url ?? '/', 'http://127.0.0.1')
      if (url.pathname !== '/api/ws' || url.searchParams.get('ticket') !== 'paired-upstream-ticket') {
        socket.destroy()
        return
      }
      upstreamSockets.handleUpgrade(request, socket, head, client => {
        client.send(JSON.stringify({ method: 'gateway.ready', params: { paired: true } }))
        client.on('message', data => {
          const frame = JSON.parse(data.toString()) as Record<string, unknown>
          pairedFrames.push(frame)
          client.send(JSON.stringify({ id: frame.id, result: {} }))
        })
      })
    })
    const upstreamPort = await listen(upstream)
    closers.push(async () => {
      for (const client of upstreamSockets.clients) client.terminate()
      upstreamSockets.close()
      await closeServer(upstream)
    })

    const home = mkdtempSync(join(tmpdir(), 'hermes-yaoyao-paired-ws-'))
    homes.push(home)
    const config: ServerConfig = {
      host: '127.0.0.1', port: 8800,
      upstream: new URL(`http://127.0.0.1:${upstreamPort}`),
      allowedHosts: new Set(), home,
      mediaRoot: join(home, 'media'), attachmentsRoot: join(home, 'attachments'), imagesRoot: join(home, 'images'), mediaOwner: 'tester',
      allowInsecureLan: false, insecureLan: false, production: false,
    }
    const application = createApplication({ config })
    const node = createNodeServer(application)
    await listen(node.server)
    closers.push(closeRuntime(node))
    const port = (node.server.address() as AddressInfo).port
    config.port = port
    const origin = `http://127.0.0.1:${port}`

    const bootstrap = await fetch(`${origin}/api/app/bootstrap`)
    const bootstrapBody = await bootstrap.json() as { csrfToken: string }
    const cookie = cookies(bootstrap)
    const pairingResponse = await fetch(`${origin}/api/app/pairings`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json', Cookie: cookie, Origin: origin,
        'X-CSRF-Token': bootstrapBody.csrfToken,
      },
      body: JSON.stringify({
        username: 'paired-user', password: 'paired-password',
      }),
    })
    expect(pairingResponse.status).toBe(201)
    const pairing = await pairingResponse.json() as { pairingId: string; qrPayload: string }
    const qr = new URL(pairing.qrPayload)
    const claimResponse = await fetch(`${origin}/api/pair/v1/claim`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        pairingId: pairing.pairingId,
        secret: qr.searchParams.get('secret'),
        deviceName: 'Test iPhone',
      }),
    })
    expect(claimResponse.status).toBe(201)
    const claim = await claimResponse.json() as { serverUrl: string; token: string; deviceId: string }
    const authorization = `Bearer ${claim.token}`

    const status = await fetch(`${claim.serverUrl}/api/status`, {
      headers: { Authorization: authorization },
    })
    expect(status.status).toBe(200)
    expect((await status.json() as { auth_required: boolean }).auth_required).toBe(true)
    const ticketResponse = await fetch(`${claim.serverUrl}/api/auth/ws-ticket`, {
      method: 'POST', headers: { Authorization: authorization },
    })
    expect(ticketResponse.status).toBe(200)
    const ticket = (await ticketResponse.json() as { ticket: string }).ticket
    expect(pairedCookie).toContain('hermes_session_at=paired-user')

    const socket = await wsOpen(
      `${claim.serverUrl.replace(/^http/, 'ws')}/api/ws?ticket=${encodeURIComponent(ticket)}`,
      origin,
      '',
    )
    expect(await wsMessage(socket)).toMatchObject({
      method: 'gateway.ready', params: { paired: true },
    })
    socket.send(JSON.stringify({
      id: 'profiles', method: 'profiles.list',
      params: { include_sessions: true },
    }))
    await wsMessage(socket)
    socket.send(JSON.stringify({
      id: 'profile-meta', method: 'profiles.configure',
      params: {
        name: 'default',
        ui_meta: { 'hermes-bots': { chat: 'stored-session' } },
        ui_meta_expected_revisions: { 'hermes-bots': 2 },
      },
    }))
    await wsMessage(socket)
    socket.send(JSON.stringify({
      id: 'profile-avatar', method: 'profiles.set_asset',
      params: { name: 'default', asset: 'avatar', data: 'data:image/png;base64,AA==' },
    }))
    await wsMessage(socket)
    socket.send(JSON.stringify({
      id: 'hidden-bot', method: 'session.create',
      params: {
        profile: 'default', title: 'Bot Chat', source: 'ios_bot_group',
        hidden: true, close_on_disconnect: false,
      },
    }))
    await wsMessage(socket)
    expect(pairedFrames[0]).toEqual({
      id: 'profiles', method: 'profiles.list', params: { include_sessions: true },
    })
    expect(pairedFrames[1]).toMatchObject({
      id: 'profile-meta', method: 'profiles.configure',
      params: { name: 'default' },
    })
    expect(pairedFrames[2]).toEqual({
      id: 'profile-avatar', method: 'profiles.set_asset',
      params: { name: 'default', asset: 'avatar', data: 'data:image/png;base64,AA==' },
    })
    expect(pairedFrames[3]).toMatchObject({
      id: 'hidden-bot', method: 'session.create',
      params: { profile: 'default', source: 'ios_bot_group', hidden: true },
    })
    socket.close()

    const revoke = await fetch(`${origin}/api/pair/v1/devices/${claim.deviceId}`, {
      method: 'DELETE', headers: { Authorization: authorization },
    })
    expect(revoke.status).toBe(200)
    expect(await fetch(`${claim.serverUrl}/api/status`, {
      headers: { Authorization: authorization },
    })).toHaveProperty('status', 401)
  })

  it('keeps the upstream ticket server-side and filters chat RPC methods', async () => {
    const received: Array<Record<string, unknown>> = []
    const upstream = createServer((request, response) => {
      response.setHeader('Content-Type', 'application/json')
      if (request.url === '/api/status') {
        response.setHeader('Set-Cookie', 'hermes_session_at=socket-user; Path=/; HttpOnly')
        response.end(JSON.stringify({ auth_required: true }))
      } else if (request.url === '/api/auth/me') {
        response.end(JSON.stringify({ user_id: 'socket-user', display_name: 'Socket User' }))
      } else if (request.url === '/api/profiles') {
        response.end(JSON.stringify({ profiles: [{ name: 'default', is_default: true }] }))
      } else if (request.url === '/api/auth/ws-ticket' && request.method === 'POST') {
        response.end(JSON.stringify({ ticket: 'one-time-upstream-secret' }))
      } else {
        response.statusCode = 404
        response.end(JSON.stringify({ error: 'not found' }))
      }
    })
    const upstreamSockets = new WebSocketServer({ noServer: true })
    upstream.on('upgrade', (request, socket, head) => {
      const url = new URL(request.url ?? '/', 'http://127.0.0.1')
      if (url.pathname !== '/api/ws' || url.searchParams.get('ticket') !== 'one-time-upstream-secret') {
        socket.destroy()
        return
      }
      upstreamSockets.handleUpgrade(request, socket, head, (client) => {
        client.send(JSON.stringify({ method: 'gateway.ready', params: {} }))
        client.on('message', (data) => {
          received.push(JSON.parse(data.toString()) as Record<string, unknown>)
          client.send(JSON.stringify({ id: 'ok', result: {} }))
        })
      })
    })
    const upstreamPort = await listen(upstream)
    closers.push(async () => {
      for (const client of upstreamSockets.clients) client.terminate()
      upstreamSockets.close()
      await closeServer(upstream)
    })

    const home = mkdtempSync(join(tmpdir(), 'hermes-yaoyao-ws-'))
    homes.push(home)
    const config: ServerConfig = {
      host: '127.0.0.1',
      port: 8800,
      upstream: new URL(`http://127.0.0.1:${upstreamPort}`),
      allowedHosts: new Set(),
      home,
      allowInsecureLan: false,
      insecureLan: false,
      production: false,
    }
    const application = createApplication({ config })
    const node = createNodeServer(application)
    await listen(node.server)
    closers.push(closeRuntime(node))
    const port = (node.server.address() as AddressInfo).port
    config.port = port
    const origin = `http://127.0.0.1:${port}`

    const bootstrap = await fetch(`${origin}/api/app/bootstrap`)
    expect(bootstrap.status).toBe(200)
    const bootstrapBody = await bootstrap.json() as { csrfToken: string }
    const cookie = cookies(bootstrap)
    const leaseResponse = await fetch(`${origin}/api/app/realtime-leases`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        Cookie: cookie,
        Origin: origin,
        'X-CSRF-Token': bootstrapBody.csrfToken,
      },
      body: JSON.stringify({ channel: 'chat' }),
    })
    expect(leaseResponse.status).toBe(201)
    const leaseBody = await leaseResponse.json() as { lease: string }
    expect(JSON.stringify(leaseBody)).not.toContain('upstream-secret')
    const rotatedCookie = [cookie, cookies(leaseResponse)].filter(Boolean).join('; ')

    const socket = await wsOpen(
      `ws://127.0.0.1:${port}/ws/chat?lease=${encodeURIComponent(leaseBody.lease)}`,
      origin,
      rotatedCookie,
    )
    expect(await wsMessage(socket)).toMatchObject({ method: 'gateway.ready' })
    socket.send(JSON.stringify({
      id: 'resume-1',
      method: 'session.resume',
      params: { session_id: 'stored-session', profile: 'default', source: 'ios' },
    }))
    await wsMessage(socket)
    expect(received[0]).toMatchObject({
      method: 'session.resume',
      params: { session_id: 'stored-session', profile: 'default', source: 'web' },
    })

    socket.send(JSON.stringify({
      id: 'clarify-1',
      method: 'clarify.respond',
      params: { request_id: 'question-1', answer: '继续', },
    }))
    await wsMessage(socket)
    expect(received[1]).toMatchObject({
      method: 'clarify.respond',
      params: { request_id: 'question-1', answer: '继续' },
    })

    socket.send(JSON.stringify({
      id: 'queue-1',
      method: 'prompt.submit',
      params: { session_id: 'stored-session', text: '下一条', queued: true, ignored: 'drop-me' },
    }))
    await wsMessage(socket)
    expect(received[2]).toEqual({
      id: 'queue-1',
      method: 'prompt.submit',
      params: { session_id: 'stored-session', text: '下一条', queued: true },
    })

    socket.send(JSON.stringify({
      id: 'model-with-scoped-provider',
      method: 'config.set',
      params: { session_id: 'stored-session', key: 'model', value: 'omni --provider Custom:Tingly --session' },
    }))
    await wsMessage(socket)
    expect(received[3]).toEqual({
      id: 'model-with-scoped-provider',
      method: 'config.set',
      params: { session_id: 'stored-session', key: 'model', value: 'omni --provider Custom:Tingly --session' },
    })

    socket.send(JSON.stringify({
      id: 'forbidden',
      method: 'config.set',
      params: { session_id: 'stored-session', key: 'yolo', value: 'true' },
    }))
    const closeCode = await new Promise<number>((resolve) => socket.once('close', resolve))
    expect(closeCode).toBe(1008)
  })

  it('rejects unknown production Upgrade requests instead of leaking a hanging socket', async () => {
    const home = mkdtempSync(join(tmpdir(), 'hermes-yaoyao-ws-'))
    homes.push(home)
    const config: ServerConfig = {
      host: '127.0.0.1',
      port: 8800,
      upstream: new URL('http://127.0.0.1:9119'),
      allowedHosts: new Set(),
      home,
      allowInsecureLan: false,
      insecureLan: false,
      production: true,
    }
    const application = createApplication({ config })
    const node = createNodeServer(application)
    await listen(node.server)
    closers.push(closeRuntime(node))
    const port = (node.server.address() as AddressInfo).port
    config.port = port
    const status = await new Promise<number>((resolve, reject) => {
      const socket = new WebSocket(`ws://127.0.0.1:${port}/vite-hmr`, {
        origin: `http://127.0.0.1:${port}`,
      })
      socket.once('unexpected-response', (_request, response) => {
        resolve(response.statusCode ?? 0)
        response.resume()
      })
      socket.once('error', reject)
    })
    expect(status).toBe(404)
  })
})