codex-app-server-bridge/test/websocket-peer.ts

import { createHash } from "node:crypto";
import { EventEmitter, once } from "node:events";
import { createServer } from "node:net";
import type { Server, Socket } from "node:net";

type Frame = Readonly<{
  opcode: number;
  payload: Buffer<ArrayBufferLike>;
}>;

function encodeFrame(opcode: number, payload: Buffer<ArrayBufferLike>): Buffer<ArrayBufferLike> {
  const length = payload.length;
  if (length < 126) {
    return Buffer.concat([Buffer.from([0x80 | opcode, length]), payload]);
  }
  if (length <= 65_535) {
    const header = Buffer.alloc(4);
    header[0] = 0x80 | opcode;
    header[1] = 126;
    header.writeUInt16BE(length, 2);
    return Buffer.concat([header, payload]);
  }
  const header = Buffer.alloc(10);
  header[0] = 0x80 | opcode;
  header[1] = 127;
  header.writeBigUInt64BE(BigInt(length), 2);
  return Buffer.concat([header, payload]);
}

function decodeFrame(
  buffer: Buffer<ArrayBufferLike>,
): { frame?: Frame; remaining: Buffer<ArrayBufferLike> } {
  if (buffer.length < 2) {
    return { remaining: buffer };
  }

  const first = buffer[0];
  const second = buffer[1];
  if (first === undefined || second === undefined) {
    return { remaining: buffer };
  }
  const opcode = first & 0x0f;
  const masked = (second & 0x80) !== 0;
  let payloadLength = second & 0x7f;
  let offset = 2;

  if (payloadLength === 126) {
    if (buffer.length < 4) {
      return { remaining: buffer };
    }
    payloadLength = buffer.readUInt16BE(2);
    offset = 4;
  } else if (payloadLength === 127) {
    if (buffer.length < 10) {
      return { remaining: buffer };
    }
    const largeLength = buffer.readBigUInt64BE(2);
    if (largeLength > BigInt(Number.MAX_SAFE_INTEGER)) {
      throw new Error("test frame is too large");
    }
    payloadLength = Number(largeLength);
    offset = 10;
  }

  const maskLength = masked ? 4 : 0;
  if (buffer.length < offset + maskLength + payloadLength) {
    return { remaining: buffer };
  }

  const mask = masked ? buffer.subarray(offset, offset + 4) : undefined;
  offset += maskLength;
  const payload = Buffer.from(buffer.subarray(offset, offset + payloadLength));
  if (mask !== undefined) {
    for (let index = 0; index < payload.length; index += 1) {
      const maskByte = mask[index % 4];
      const payloadByte = payload[index];
      if (maskByte !== undefined && payloadByte !== undefined) {
        payload[index] = payloadByte ^ maskByte;
      }
    }
  }

  return {
    frame: { opcode, payload },
    remaining: buffer.subarray(offset + payloadLength),
  };
}

export class TestWebSocketConnection extends EventEmitter {
  readonly socket: Socket;
  readonly textFrames: string[] = [];
  abortOnCloseFrame = false;
  private frameBuffer: Buffer<ArrayBufferLike> = Buffer.alloc(0);

  constructor(socket: Socket, initialFrameBytes: Buffer<ArrayBufferLike>) {
    super();
    this.socket = socket;
    socket.on("data", (chunk: Buffer) => this.consume(chunk));
    socket.on("close", () => this.emit("socketClose"));
    socket.on("error", (error: Error) => this.emit("socketError", error));
    if (initialFrameBytes.length > 0) {
      this.consume(initialFrameBytes);
    }
  }

  private consume(chunk: Buffer<ArrayBufferLike>): void {
    this.frameBuffer = Buffer.concat([this.frameBuffer, chunk]);
    while (true) {
      const decoded = decodeFrame(this.frameBuffer);
      if (decoded.frame === undefined) {
        return;
      }
      this.frameBuffer = decoded.remaining;
      if (decoded.frame.opcode === 0x1) {
        const text = decoded.frame.payload.toString("utf8");
        this.textFrames.push(text);
        this.emit("text", text);
      } else if (decoded.frame.opcode === 0x8) {
        this.emit("closeFrame");
        if (this.abortOnCloseFrame) {
          this.socket.destroy(new Error("test peer aborted during local shutdown"));
        } else {
          this.socket.write(encodeFrame(0x8, decoded.frame.payload));
          this.socket.end();
        }
      } else if (decoded.frame.opcode === 0x9) {
        this.socket.write(encodeFrame(0xa, decoded.frame.payload));
      }
    }
  }

  sendText(value: string): void {
    this.socket.write(encodeFrame(0x1, Buffer.from(value)));
  }

  sendBinary(value: Buffer<ArrayBufferLike>): void {
    this.socket.write(encodeFrame(0x2, value));
  }

  close(code = 1000): void {
    const payload = Buffer.alloc(2);
    payload.writeUInt16BE(code);
    this.socket.write(encodeFrame(0x8, payload));
    this.socket.end();
  }
}

export class TestWebSocketPeer extends EventEmitter {
  readonly server: Server;
  readonly connections: TestWebSocketConnection[] = [];
  readonly completeHandshake: boolean;
  readonly handshakeDelayMs: number;

  private constructor(options: { completeHandshake: boolean; handshakeDelayMs: number }) {
    super();
    this.completeHandshake = options.completeHandshake;
    this.handshakeDelayMs = options.handshakeDelayMs;
    this.server = createServer((socket) => this.acceptSocket(socket));
  }

  static async start(
    options: { completeHandshake?: boolean; handshakeDelayMs?: number } = {},
  ): Promise<TestWebSocketPeer> {
    const peer = new TestWebSocketPeer({
      completeHandshake: options.completeHandshake ?? true,
      handshakeDelayMs: options.handshakeDelayMs ?? 0,
    });
    peer.server.listen(0, "127.0.0.1");
    await once(peer.server, "listening");
    return peer;
  }

  get port(): number {
    const address = this.server.address();
    if (address === null || typeof address === "string") {
      throw new Error("test server is not listening on TCP");
    }
    return address.port;
  }

  private acceptSocket(socket: Socket): void {
    let handshakeBuffer = Buffer.alloc(0);
    const onHandshakeData = (chunk: Buffer): void => {
      handshakeBuffer = Buffer.concat([handshakeBuffer, chunk]);
      const boundary = handshakeBuffer.indexOf("\r\n\r\n");
      if (boundary === -1 || !this.completeHandshake) {
        return;
      }

      socket.removeListener("data", onHandshakeData);
      const request = handshakeBuffer.subarray(0, boundary + 4).toString("utf8");
      const remaining = handshakeBuffer.subarray(boundary + 4);
      const key = /^Sec-WebSocket-Key:\s*(.+)$/im.exec(request)?.[1]?.trim();
      if (key === undefined) {
        socket.destroy(new Error("missing WebSocket key"));
        return;
      }

      const accept = createHash("sha1")
        .update(`${key}258EAFA5-E914-47DA-95CA-C5AB0DC85B11`)
        .digest("base64");
      setTimeout(() => {
        socket.write(
          "HTTP/1.1 101 Switching Protocols\r\n" +
            "Upgrade: websocket\r\n" +
            "Connection: Upgrade\r\n" +
            `Sec-WebSocket-Accept: ${accept}\r\n\r\n`,
        );
        const connection = new TestWebSocketConnection(socket, remaining);
        this.connections.push(connection);
        this.emit("connection", connection);
      }, this.handshakeDelayMs);
    };
    socket.on("data", onHandshakeData);
    socket.on("error", () => undefined);
  }

  async nextConnection(): Promise<TestWebSocketConnection> {
    const existing = this.connections[0];
    if (existing !== undefined) {
      return existing;
    }
    const [connection] = await once(this, "connection");
    return connection as TestWebSocketConnection;
  }

  async stop(): Promise<void> {
    for (const connection of this.connections) {
      connection.socket.destroy();
    }
    this.server.close();
    await once(this.server, "close");
  }
}