vdm/tests/conformance/ts/client.ts

/**
 * A minimal client for each transport, built on the generated types.
 *
 * Deliberately not the extension's transport implementation: the conformance suite must
 * fail when the *contract* is broken, not when the extension's reconnect logic is. It
 * speaks the two framings and nothing else.
 */

import net from 'node:net';
import { WebSocket } from 'ws';
import type { MethodName, Params, Result } from '../../../extension/src/shared/protocol/methods.js';

export type TransportName = 'uds' | 'ws';

export interface RpcFrame {
  jsonrpc: '2.0';
  id?: number | string;
  method?: string;
  params?: unknown;
  result?: unknown;
  error?: { code: number; message: string; data?: unknown };
}

export interface Conn {
  readonly transport: TransportName;
  /** Send and wait. Resolves to null when nothing arrives inside `timeoutMs`. */
  request(method: string, params: unknown, timeoutMs: number): Promise<RpcFrame | null>;
  /** Typed convenience wrapper, so the suite itself is checked against the contract. */
  call<M extends MethodName>(method: M, params: Params<M>): Promise<Result<M>>;
  notifications(): RpcFrame[];
  close(): void;
}

abstract class BaseConn implements Conn {
  abstract readonly transport: TransportName;
  protected nextId = 1;
  protected readonly pending = new Map<number | string, (f: RpcFrame) => void>();
  private readonly received: RpcFrame[] = [];

  protected abstract write(text: string): void;
  abstract close(): void;

  protected onFrame(frame: RpcFrame): void {
    if (frame.id !== undefined && this.pending.has(frame.id)) {
      const resolve = this.pending.get(frame.id);
      this.pending.delete(frame.id);
      resolve?.(frame);
      return;
    }
    if (frame.method !== undefined) this.received.push(frame);
  }

  notifications(): RpcFrame[] {
    return [...this.received];
  }

  request(method: string, params: unknown, timeoutMs: number): Promise<RpcFrame | null> {
    const id = this.nextId++;
    return new Promise((resolve) => {
      const timer = setTimeout(() => {
        this.pending.delete(id);
        resolve(null); // the fail-open case: no answer inside the deadline
      }, timeoutMs);
      this.pending.set(id, (frame) => {
        clearTimeout(timer);
        resolve(frame);
      });
      this.write(JSON.stringify({ jsonrpc: '2.0', id, method, params }));
    });
  }

  async call<M extends MethodName>(method: M, params: Params<M>): Promise<Result<M>> {
    const frame = await this.request(method, params, 10_000);
    if (frame === null) throw new Error(`${method}: no response`);
    if (frame.error) throw new Error(`${method}: error ${frame.error.code}: ${frame.error.message}`);
    return frame.result as Result<M>;
  }
}

class UdsConn extends BaseConn {
  readonly transport = 'uds' as const;
  private buffer = '';

  constructor(private readonly socket: net.Socket) {
    super();
    socket.on('data', (chunk) => {
      this.buffer += chunk.toString('utf8');
      let nl = this.buffer.indexOf('\n');
      while (nl !== -1) {
        const line = this.buffer.slice(0, nl).trim();
        this.buffer = this.buffer.slice(nl + 1);
        nl = this.buffer.indexOf('\n');
        if (line) this.onFrame(JSON.parse(line) as RpcFrame);
      }
    });
  }

  protected write(text: string): void {
    this.socket.write(text + '\n');
  }

  close(): void {
    this.socket.destroy();
  }
}

class WsConn extends BaseConn {
  readonly transport = 'ws' as const;

  constructor(private readonly socket: WebSocket) {
    super();
    socket.on('message', (data) => this.onFrame(JSON.parse(data.toString()) as RpcFrame));
  }

  protected write(text: string): void {
    this.socket.send(text);
  }

  close(): void {
    this.socket.close();
  }
}

export async function connectUds(path: string): Promise<Conn> {
  const socket = net.connect(path);
  await new Promise<void>((resolve, reject) => {
    socket.once('connect', () => resolve());
    socket.once('error', reject);
  });
  return new UdsConn(socket);
}

export async function connectWs(port: number): Promise<Conn> {
  // The daemon verifies this Origin on the upgrade, so the suite must present a real one.
  const socket = new WebSocket(`ws://127.0.0.1:${port}`, {
    headers: { Origin: 'moz-extension://11111111-2222-3333-4444-555555555555' },
  });
  await new Promise<void>((resolve, reject) => {
    socket.once('open', () => resolve());
    socket.once('error', reject);
  });
  return new WsConn(socket);
}