ws WebSocket server with a Zod-validated message protocol: working versions (tested 2026-10-15)

{
  "name": "ws-zod-message-protocol",
  "private": true,
  "type": "module",
  "scripts": {
    "test": "node --import tsx --test test/*.test.ts"
  },
  "dependencies": {
    "ws": "8.22.0",
    "zod": "4.6.7"
  },
  "devDependencies": {
    "tsx": "4.23.15",
    "@types/node": "26.6.5",
    "@types/ws": "8.18.2"
  }
}

import { WebSocketServer, WebSocket } from "ws";
import type { IncomingMessage } from "node:http";
import { z } from "zod";

/** Client to server messages, as a discriminated union on "type". */
export const ClientMsg = z.discriminatedUnion("type", [
  z.object({ type: z.literal("join"), room: z.string().min(1).max(32) }),
  z.object({ type: z.literal("say"), room: z.string(), text: z.string().min(1).max(500) }),
  z.object({ type: z.literal("ping"), id: z.number().int() }),
]);
export type ServerMsg =
  | { type: "welcome"; user: string }
  | { type: "joined"; room: string; members: number }
  | { type: "message"; room: string; from: string; text: string }
  | { type: "pong"; id: number }
  | { type: "error"; code: "bad_json" | "invalid" | "not_in_room"; detail?: string };

const send = (ws: WebSocket, msg: ServerMsg) => ws.send(JSON.stringify(msg));

/**
 * A room chat server on ws. Auth is a token in the upgrade URL, checked in verifyClient-style
 * handling of the "connection" request; every inbound frame is validated with Zod.
 */
export function startServer(tokens: Record<string, string>, port = 0) {
  const wss = new WebSocketServer({ port });
  const rooms = new Map<string, Set<WebSocket>>();
  const userOf = new WeakMap<WebSocket, string>();

  wss.on("connection", (ws: WebSocket, req: IncomingMessage) => {
    const token = new URL(req.url ?? "/", "http://localhost").searchParams.get("token") ?? "";
    const user = tokens[token];
    if (!user) return ws.close(4401, "unauthorized");
    userOf.set(ws, user);
    send(ws, { type: "welcome", user });

    ws.on("message", (raw) => {
      let data: unknown;
      try { data = JSON.parse(raw.toString()); } catch { return send(ws, { type: "error", code: "bad_json" }); }
      const parsed = ClientMsg.safeParse(data);
      if (!parsed.success) return send(ws, { type: "error", code: "invalid", detail: parsed.error.issues[0].path.join(".") || parsed.error.issues[0].message });
      const msg = parsed.data;
      switch (msg.type) {
        case "ping":
          return send(ws, { type: "pong", id: msg.id });
        case "join": {
          const set = rooms.get(msg.room) ?? new Set();
          set.add(ws);
          rooms.set(msg.room, set);
          return send(ws, { type: "joined", room: msg.room, members: set.size });
        }
        case "say": {
          const set = rooms.get(msg.room);
          if (!set?.has(ws)) return send(ws, { type: "error", code: "not_in_room" });
          for (const peer of set) if (peer.readyState === WebSocket.OPEN) send(peer, { type: "message", room: msg.room, from: user, text: msg.text });
        }
      }
    });

    ws.on("close", () => { for (const set of rooms.values()) set.delete(ws); });
  });

  const ready = new Promise<number>((resolve) => wss.on("listening", () => resolve((wss.address() as { port: number }).port)));
  return { wss, ready, close: () => new Promise<void>((r) => { for (const c of wss.clients) c.terminate(); wss.close(() => r()); }) };
}

import { test, before, after } from "node:test";
import assert from "node:assert/strict";
import { WebSocket } from "ws";
import { startServer, type ServerMsg } from "../src/server.ts";

let srv: ReturnType<typeof startServer>;
let port: number;

before(async () => {
  srv = startServer({ "tok-ada": "ada", "tok-bob": "bob" });
  port = await srv.ready;
});
after(() => srv.close());

/** Connects and buffers incoming messages so tests can await the next one in order. */
async function connect(token: string) {
  const ws = new WebSocket(`ws://127.0.0.1:${port}/?token=${token}`);
  const queue: ServerMsg[] = [];
  const waiters: ((m: ServerMsg) => void)[] = [];
  ws.on("message", (raw) => {
    const m = JSON.parse(raw.toString()) as ServerMsg;
    const w = waiters.shift();
    w ? w(m) : queue.push(m);
  });
  const next = () => new Promise<ServerMsg>((resolve) => { const m = queue.shift(); m ? resolve(m) : waiters.push(resolve); });
  await new Promise((r, j) => { ws.once("open", r); ws.once("error", j); });
  return { ws, next, send: (v: unknown) => ws.send(typeof v === "string" ? v : JSON.stringify(v)) };
}

test("rejects connections without a valid token", async () => {
  const ws = new WebSocket(`ws://127.0.0.1:${port}/?token=nope`);
  const [code, reason] = await new Promise<[number, string]>((r) => ws.on("close", (c, why) => r([c, why.toString()])));
  assert.equal(code, 4401);
  assert.equal(reason, "unauthorized");
});

test("welcomes and answers ping", async () => {
  const c = await connect("tok-ada");
  assert.deepEqual(await c.next(), { type: "welcome", user: "ada" });
  c.send({ type: "ping", id: 7 });
  assert.deepEqual(await c.next(), { type: "pong", id: 7 });
  c.ws.close();
});

test("invalid frames get typed errors and the socket stays open", async () => {
  const c = await connect("tok-ada");
  await c.next();
  c.send("{not json");
  assert.deepEqual(await c.next(), { type: "error", code: "bad_json" });
  c.send({ type: "say", room: "general", text: "" });
  assert.deepEqual(await c.next(), { type: "error", code: "invalid", detail: "text" });
  c.send({ type: "shout" });
  assert.equal((await c.next() as { code: string }).code, "invalid");
  c.send({ type: "say", room: "general", text: "hi" });
  assert.deepEqual(await c.next(), { type: "error", code: "not_in_room" });
  assert.equal(c.ws.readyState, WebSocket.OPEN);
  c.ws.close();
});

test("broadcasts room messages to every member", async () => {
  const ada = await connect("tok-ada");
  const bob = await connect("tok-bob");
  await ada.next(); await bob.next();
  ada.send({ type: "join", room: "ops" });
  assert.deepEqual(await ada.next(), { type: "joined", room: "ops", members: 1 });
  bob.send({ type: "join", room: "ops" });
  assert.deepEqual(await bob.next(), { type: "joined", room: "ops", members: 2 });
  bob.send({ type: "say", room: "ops", text: "deploying" });
  const expected = { type: "message", room: "ops", from: "bob", text: "deploying" };
  assert.deepEqual(await ada.next(), expected);
  assert.deepEqual(await bob.next(), expected);
  ada.ws.close(); bob.ws.close();
});