company-tool/scripts/ws-smoke-test.mjs

#!/usr/bin/env node

import { webcrypto } from "node:crypto";

const wsUrl = process.argv[2] ?? process.env.WS_URL ?? "ws://localhost:8080/ws";
const protocolVersion = 1;

const encoder = new TextEncoder();
const decoder = new TextDecoder();

function bytesToBase64(bytes) {
  return Buffer.from(bytes).toString("base64");
}

function base64ToBytes(value) {
  return Buffer.from(value, "base64");
}

async function createSession() {
  const key = await webcrypto.subtle.generateKey(
    { name: "AES-GCM", length: 256 },
    true,
    ["encrypt", "decrypt"],
  );
  const rawKey = await webcrypto.subtle.exportKey("raw", key);

  return {
    key,
    keyId: webcrypto.randomUUID(),
    exportedKey: bytesToBase64(new Uint8Array(rawKey)),
  };
}

async function encryptMessage(session, message) {
  const nonce = webcrypto.getRandomValues(new Uint8Array(12));
  const plaintext = encoder.encode(JSON.stringify(message));
  const ciphertext = await webcrypto.subtle.encrypt(
    { name: "AES-GCM", iv: nonce },
    session.key,
    plaintext,
  );

  return {
    enc: `aes-256-gcm-v${protocolVersion}`,
    key_id: session.keyId,
    nonce: bytesToBase64(nonce),
    ciphertext: bytesToBase64(new Uint8Array(ciphertext)),
  };
}

async function decryptMessage(session, envelope) {
  const plaintext = await webcrypto.subtle.decrypt(
    { name: "AES-GCM", iv: base64ToBytes(envelope.nonce) },
    session.key,
    base64ToBytes(envelope.ciphertext),
  );

  return JSON.parse(decoder.decode(plaintext));
}

function waitForOpen(socket) {
  return new Promise((resolve, reject) => {
    socket.addEventListener("open", resolve, { once: true });
    socket.addEventListener("error", reject, { once: true });
  });
}

function waitForMessage(socket, timeoutMs = 5000) {
  return new Promise((resolve, reject) => {
    const timeout = setTimeout(() => {
      socket.removeEventListener("message", onMessage);
      reject(new Error(`timeout after ${timeoutMs}ms`));
    }, timeoutMs);

    function onMessage(event) {
      clearTimeout(timeout);
      resolve(JSON.parse(event.data));
    }

    socket.addEventListener("message", onMessage, { once: true });
  });
}

async function main() {
  const session = await createSession();
  const socket = new WebSocket(wsUrl);

  console.log(`connecting ${wsUrl}`);
  await waitForOpen(socket);
  console.log("socket open");

  socket.send(JSON.stringify({
    type: "hello",
    payload: {
      protocol_version: protocolVersion,
      key_id: session.keyId,
      session_key: session.exportedKey,
    },
  }));

  const ack = await waitForMessage(socket);
  if (ack.type !== "hello_ack") {
    throw new Error(`expected hello_ack, got ${JSON.stringify(ack)}`);
  }
  console.log(`hello_ack protocol=${ack.payload.protocol_version} key=${ack.payload.key_id}`);

  const firstEncrypted = await waitForMessage(socket);
  if (firstEncrypted.type !== "encrypted") {
    throw new Error(`expected encrypted snapshot, got ${JSON.stringify(firstEncrypted)}`);
  }
  const snapshot = await decryptMessage(session, firstEncrypted.payload);
  console.log(`decrypted first server message: ${snapshot.type}`);

  const subscribe = await encryptMessage(session, {
    type: "subscribe",
    payload: { topic: "records" },
  });
  socket.send(JSON.stringify({ type: "encrypted", payload: subscribe }));
  console.log("sent encrypted subscribe");

  const ping = await encryptMessage(session, { type: "ping" });
  socket.send(JSON.stringify({ type: "encrypted", payload: ping }));

  const pongEnvelope = await waitForMessage(socket);
  if (pongEnvelope.type !== "encrypted") {
    throw new Error(`expected encrypted pong, got ${JSON.stringify(pongEnvelope)}`);
  }
  const pong = await decryptMessage(session, pongEnvelope.payload);
  console.log(`decrypted ping response: ${pong.type}`);

  socket.close();
  console.log("communication smoke test ok");
}

main().catch((error) => {
  console.error(error);
  process.exit(1);
});