Building a WebSocket message router in TypeScript #

Your WebSocket server started with a switch (msg.type) over four cases. Two years later it is 900 lines long, every case parses its own payload slightly differently, an exception thrown in one handler closes the whole socket, adding a rate limit means touching thirty cases, and nobody is sure which message types the client still sends. HTTP frameworks solved this decades ago with routers and middleware, but a WebSocket message stream gets none of that by default. This page builds a small, typed router: handlers registered by message type, payloads validated before the handler runs, cross-cutting concerns as middleware, and failures isolated to the message that caused them.

Root cause #

A switch statement couples four concerns that change at different rates: dispatch (which code runs for which type), validation (is this payload well formed), policy (is this client allowed and within its rate), and business logic. Because they share one function, every change touches all of them, and every bug in one can break the others. The most damaging example is error handling. An uncaught exception in an async handler inside ws.on('message') becomes an unhandled rejection, and depending on your Node version and settings that can crash the process — taking every connection on the node with it, not just the one that sent the bad message.

Types make the problem worse rather than better when the switch is untyped: JSON.parse returns any, each case casts it, and the compiler cannot tell you that the client’s cursor.move payload has had a new required field for six months.

Router pipeline per inbound frame Each inbound frame passes through decode, handler lookup, schema validation, middleware and the handler, with each stage able to reject the message with its own error code. Router pipeline per inbound frame Decode size check, JSON.parse, envelope shape reject: bad_frame Lookup handler registry by message type reject: unknown_type Validate schema for this type, typed payload out reject: invalid Middleware auth, rate limit, metrics, tracing reject: policy Handler business logic with a typed context reply / emit Every stage fails the message, never the connection or the process
Five stages, each with one job and one error code.

Resolution #

The router below keeps a map from message type to a { schema, handler } pair. Handlers are registered with a Zod schema, and TypeScript infers the handler’s payload type from it, so the handler never sees unvalidated data. Middleware wraps every handler in registration order. The dispatch function catches every error and converts it into an error reply correlated with the request id, so one bad message cannot close the socket or the process.

import { z, ZodTypeAny } from 'zod';
import type { WebSocket } from 'ws';

const MAX_FRAME_BYTES = 64 * 1024;

const Envelope = z.object({
  type: z.string().min(1).max(64),
  id: z.string().max(64).optional(),     // request id, echoed in replies
  data: z.unknown(),
});

export interface Ctx {
  ws: WebSocket;
  userId: string;
  reqId?: string;
  reply: (data: unknown) => void;
}

type Handler<S extends ZodTypeAny> = (data: z.infer<S>, ctx: Ctx) => Promise<void> | void;
type Middleware = (type: string, ctx: Ctx, next: () => Promise<void>) => Promise<void>;

export class RouterError extends Error {
  constructor(public code: string, message: string) { super(message); }
}

export class MessageRouter {
  private routes = new Map<string, { schema: ZodTypeAny; handler: Handler<any> }>();
  private middleware: Middleware[] = [];

  on<S extends ZodTypeAny>(type: string, schema: S, handler: Handler<S>) {
    if (this.routes.has(type)) throw new Error(`duplicate route: ${type}`);
    this.routes.set(type, { schema, handler });
    return this;
  }

  use(mw: Middleware) { this.middleware.push(mw); return this; }

  async dispatch(raw: Buffer | string, base: Omit<Ctx, 'reply' | 'reqId'>) {
    let reqId: string | undefined;
    const send = (obj: object) => base.ws.send(JSON.stringify(obj));
    try {
      if (raw.length > MAX_FRAME_BYTES) throw new RouterError('too_large', 'frame exceeds limit');
      let parsed: unknown;
      try { parsed = JSON.parse(raw.toString()); }
      catch { throw new RouterError('bad_frame', 'not JSON'); }
      const env = Envelope.safeParse(parsed);
      if (!env.success) throw new RouterError('bad_frame', 'invalid envelope');
      reqId = env.data.id;

      const route = this.routes.get(env.data.type);
      if (!route) throw new RouterError('unknown_type', env.data.type);
      const data = route.schema.safeParse(env.data.data);
      if (!data.success) throw new RouterError('invalid', data.error.issues[0]?.message ?? 'invalid');

      const ctx: Ctx = { ...base, reqId, reply: (d) => send({ type: `${env.data.type}.ok`, id: reqId, data: d }) };
      // Compose middleware around the handler: mw[0](mw[1](… handler)).
      const run = this.middleware.reduceRight<() => Promise<void>>(
        (next, mw) => () => mw(env.data.type, ctx, next),
        async () => { await route.handler(data.data, ctx); },
      );
      await run();
    } catch (err) {
      const code = err instanceof RouterError ? err.code : 'internal';
      const message = err instanceof RouterError ? err.message : 'internal error';
      if (code === 'internal') console.error(err);       // unexpected: log with stack
      send({ type: 'error', id: reqId, code, message });  // the connection stays open
    }
  }
}

// Usage
const router = new MessageRouter()
  .use(async (type, ctx, next) => {                 // metrics middleware
    const t0 = performance.now();
    try { await next(); } finally { observeHandler(type, performance.now() - t0); }
  })
  .on('cursor.move', z.object({ docId: z.string(), x: z.number(), y: z.number() }), (d, ctx) => {
    broadcastCursor(d.docId, ctx.userId, d.x, d.y);  // d is fully typed here
  })
  .on('doc.rename', z.object({ docId: z.string(), title: z.string().min(1).max(200) }), async (d, ctx) => {
    await renameDoc(d.docId, d.title, ctx.userId);
    ctx.reply({ docId: d.docId });
  });

// ws.on('message', (raw) => void router.dispatch(raw as Buffer, { ws, userId }));
declare function observeHandler(type: string, ms: number): void;
declare function broadcastCursor(docId: string, userId: string, x: number, y: number): void;
declare function renameDoc(docId: string, title: string, userId: string): Promise<void>;

Schema validation here reuses the approach in validating WebSocket messages with Zod; the router simply makes it impossible to register a handler without one. The id echoed in *.ok and error replies is what lets the client match responses to requests — the correlation pattern described in request/response correlation over WebSockets.

Middleware is where cross-cutting policy lives. A rate limiter becomes one function that throws RouterError('rate_limited', …) instead of thirty copy-pasted checks; per-type authorization, tracing spans and metrics follow the same shape.

A failing handler, isolated A handler throws a database timeout; the router catches it, replies with an error correlated to request 17, and the next message on the same socket is processed normally. A failing handler, isolated Client Router Middleware Handler doc.rename id=17 decode + validate ok metrics, rate limit run handler throws: DB timeout error id=17 internal cursor.move (still works) The socket, the other handlers and the process are unaffected
One bad message costs one error reply — nothing more.

The difference from the switch statement is less about lines of code than about where each concern lives and what a mistake in one of them can break.

Switch statement vs typed router Comparison of a switch statement and a typed router on payload typing, validation, cross-cutting policy, handler exceptions and unknown message types. Switch statement vs typed router switch (msg.type) Typed router Payload types any, cast per case inferred from schema Validation ad hoc, optional required to register Rate limit, auth copied into cases one middleware Handler throws socket or process dies error reply only Unknown type silently ignored unknown_type reply The router is roughly a hundred lines; the switch it replaces was nine hundred
Separate dispatch, validation, policy and logic, and each can change without touching the others.

Edge cases #

Ordering. dispatch is async, and if you call it without awaiting in the message listener, two messages from the same client can run concurrently and finish out of order. For types where order matters (edits to the same document), chain dispatches per connection with a promise queue, or partition by key so unrelated messages still run in parallel.

Unknown types from newer clients. A client deployed ahead of the server sends types the server does not know. Replying with unknown_type and keeping the connection open lets the client degrade gracefully; closing the socket turns a staggered rollout into an outage. Pair it with the versioning scheme in versioning WebSocket message schemas.

Error detail leakage. Only RouterError messages reach the client. Unexpected exceptions become a generic internal error, so stack traces and SQL fragments stay in your logs.

Verification #

Test the router in isolation with a fake socket that records sends — no network needed:

it('isolates handler failures and keeps routing', async () => {
  const sent: any[] = [];
  const ws = { send: (s: string) => sent.push(JSON.parse(s)) } as any;
  const r = new MessageRouter()
    .on('boom', z.object({}), () => { throw new Error('db down'); })
    .on('ping', z.object({}), (_d, ctx) => ctx.reply('pong'));
  await r.dispatch(JSON.stringify({ type: 'boom', id: 'a', data: {} }), { ws, userId: 'u' });
  await r.dispatch(JSON.stringify({ type: 'ping', id: 'b', data: {} }), { ws, userId: 'u' });
  expect(sent).toEqual([
    { type: 'error', id: 'a', code: 'internal', message: 'internal error' },
    { type: 'ping.ok', id: 'b', data: 'pong' },
  ]);
});

In production, export error replies by code and handler latency by type. A sudden rise in invalid for one type almost always means a client release changed a payload; a rise in unknown_type means clients and server are out of step.

Operational checklist #

FAQ #

Why not use Socket.IO’s event handlers instead? #

Socket.IO gives you per-event dispatch and acknowledgements, which covers part of this. You still need validation, middleware for per-event policy and error isolation, and you take on Socket.IO’s protocol. With raw ws, the router above is about a hundred lines and fully under your control; see Socket.IO vs raw WebSockets.

Should an invalid message close the connection? #

Usually not. Reply with an error and keep going. Close only for protocol-level abuse — oversized frames, binary garbage on a text protocol, repeated violations — with 1008 or 1009.

How do I route binary frames? #

Give binary messages their own small header (a type byte or varint) and a separate registry, or decode them into the same envelope shape with MessagePack before routing. See binary WebSocket frames with MessagePack.

Where should rooms and subscriptions live? #

In a separate registry that handlers call into. The router decides which code runs; the room registry decides who receives output. Keeping them apart is covered in WebSocket rooms and channel subscriptions.

Back to Server-Side Routing Patterns.