WebSockets

GET /chat HTTP/1.1
Host: api.example.com
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Version: 13
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Origin: https://app.example.com
Sec-WebSocket-Protocol: chat.v2, chat.v1
Sec-WebSocket-Extensions: permessage-deflate
 
HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=
Sec-WebSocket-Protocol: chat.v2

const GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
 
async function acceptFor(key: string): Promise<string> {
  const bytes = new TextEncoder().encode(key + GUID);
  const hash = await crypto.subtle.digest("SHA-1", bytes);
  return btoa(String.fromCharCode(...new Uint8Array(hash)));
}
 
// RFC 6455 example: resolves to
// "s3pPLMBiTxaQ9kYGzzhZRbK+xOo="
await acceptFor("dGhlIHNhbXBsZSBub25jZQ==");

const ws = new WebSocket("wss://api.example.com/chat", [
  "chat.v2",
]);
ws.binaryType = "arraybuffer"; // default "blob"
 
ws.addEventListener("open", () => {
  ws.send(JSON.stringify({ type: "join", room: "lobby" }));
});
 
ws.addEventListener("message", (e: MessageEvent) => {
  if (typeof e.data === "string") handleText(e.data);
  else handleBinary(new Uint8Array(e.data as ArrayBuffer));
});
 
ws.addEventListener("error", () => {
  // no details by design; a close event always follows
});
 
ws.addEventListener("close", (e) => {
  console.info(e.code, e.reason, e.wasClean);
});
 
// later
ws.close(1000, "bye");

import { z } from "zod";
 
export const ClientMsg = z.discriminatedUnion("type", [
  z.object({
    type: z.literal("join"),
    room: z.string().min(1).max(64),
  }),
  z.object({
    type: z.literal("chat"),
    room: z.string().min(1).max(64),
    text: z.string().min(1).max(2000),
  }),
  z.object({ type: z.literal("ping"), t: z.number() }),
]);
export type ClientMsg = z.infer<typeof ClientMsg>;
 
export type ServerMsg =
  | { type: "joined"; room: string; members: number }
  | {
      type: "chat";
      room: string;
      from: string;
      text: string;
    }
  | { type: "pong"; t: number }
  | { type: "error"; code: "bad_message" | "forbidden" };
 
export function parseClientMsg(
  raw: unknown,
): ClientMsg | null {
  if (typeof raw !== "string") return null;
  try {
    const result = ClientMsg.safeParse(JSON.parse(raw));
    return result.success ? result.data : null;
  } catch {
    return null; // not JSON
  }
}
 
interface Sink {
  send(data: string): unknown;
}
 
export function sendMsg(ws: Sink, msg: ServerMsg): void {
  ws.send(JSON.stringify(msg));
}

import {
  type ClientMsg,
  type ServerMsg,
  parseClientMsg,
  sendMsg,
} from "./protocol.js";
 
type Conn = { send(data: string): unknown; userId: string };
 
export function onMessage(conn: Conn, raw: unknown): void {
  const msg = parseClientMsg(raw);
  if (!msg) {
    return sendMsg(conn, {
      type: "error",
      code: "bad_message",
    });
  }
  switch (msg.type) {
    case "join":
      return sendMsg(conn, joinRoom(conn, msg.room));
    case "chat":
      return broadcast(msg.room, {
        type: "chat",
        room: msg.room,
        from: conn.userId,
        text: msg.text,
      });
    case "ping":
      return sendMsg(conn, { type: "pong", t: msg.t });
    default:
      msg satisfies never;
  }
}

type WsData = { userId: string; rooms: Set<string> };
 
const server = Bun.serve({
  port: 3000,
  fetch(req, server) {
    const url = new URL(req.url);
    if (url.pathname !== "/ws") {
      return new Response("Not found", { status: 404 });
    }
    const userId = userFromCookie(req);
    if (!userId) {
      return new Response("Unauthorized", { status: 401 });
    }
    const ok = server.upgrade(req, {
      data: { userId, rooms: new Set() },
    });
    if (ok) return; // Bun sends the 101 itself
    return new Response("Upgrade required", { status: 426 });
  },
  websocket: {
    data: {} as WsData,
    maxPayloadLength: 64 * 1024, // bytes, default 16 MB
    idleTimeout: 60, // seconds, default 120
    open(ws) {
      ws.subscribe(`user:${ws.data.userId}`);
    },
    message(ws, message) {
      onMessage(ws, message); // string | Buffer
    },
    close(ws, code, reason) {
      console.info(ws.data.userId, "left", code, reason);
    },
  },
});
 
console.info(`listening on ${server.url}`);

import { createServer } from "node:http";
import { WebSocketServer } from "ws";
 
const http = createServer((_req, res) => {
  res.writeHead(404).end();
});
const wss = new WebSocketServer({
  noServer: true,
  maxPayload: 64 * 1024,
});
 
http.on("upgrade", async (req, socket, head) => {
  const userId = await verifyUpgrade(req); // origin + cookie
  if (!userId) {
    socket.end("HTTP/1.1 401 Unauthorized\r\n\r\n");
    return;
  }
  wss.handleUpgrade(req, socket, head, (ws) => {
    ws.on("error", (err) => console.error(err));
    ws.on("message", (data, isBinary) => {
      if (isBinary) return ws.close(1003, "text only");
      const conn = {
        send: (s: string) => ws.send(s),
        userId,
      };
      onMessage(conn, String(data));
    });
  });
});
 
http.listen(3000);

import { WebSocketServer, type WebSocket } from "ws";
 
const wss = new WebSocketServer({ port: 8080 });
const alive = new WeakSet<WebSocket>();
 
wss.on("connection", (ws) => {
  alive.add(ws);
  ws.on("pong", () => alive.add(ws));
});
 
const timer = setInterval(() => {
  for (const ws of wss.clients) {
    if (!alive.has(ws)) {
      ws.terminate();
      continue;
    }
    alive.delete(ws);
    ws.ping();
  }
}, 30_000);
 
wss.on("close", () => clearInterval(timer));

export function backoff(
  attempt: number,
  baseMs = 500,
  capMs = 30_000,
): number {
  return Math.random() * Math.min(
    capMs,
    baseMs * 2 ** attempt,
  );
}
 
type Options = {
  url: string;
  onOpen: (ws: WebSocket) => void; // resubscribe here
  onMessage: (data: unknown) => void;
};
 
export function connect({
  url,
  onOpen,
  onMessage,
}: Options) {
  let attempt = 0;
  let stopped = false;
  let timer: ReturnType<typeof setTimeout> | undefined;
  let ws: WebSocket;
 
  const open = () => {
    ws = new WebSocket(url);
    ws.onopen = () => {
      attempt = 0;
      onOpen(ws);
    };
    ws.onmessage = (e) => onMessage(e.data);
    ws.onclose = (e) => {
      // 1008: auth/policy failure, retrying will not help
      if (stopped || e.code === 1008) return;
      timer = setTimeout(open, backoff(attempt++));
    };
  };
  open();
 
  return {
    send: (data: string) => ws.send(data),
    close() {
      stopped = true;
      clearTimeout(timer);
      ws.close(1000);
    },
  };
}

const HIGH_WATER = 1024 * 1024; // 1 MiB
 
async function upload(
  ws: WebSocket,
  chunks: AsyncIterable<Uint8Array<ArrayBuffer>>,
): Promise<void> {
  for await (const chunk of chunks) {
    while (ws.bufferedAmount > HIGH_WATER) {
      await new Promise((r) => setTimeout(r, 50));
    }
    ws.send(chunk);
  }
}

type Data = { queue: string[] };
 
function flush(ws: Bun.ServerWebSocket<Data>): void {
  while (ws.data.queue.length > 0) {
    const status = ws.send(ws.data.queue[0]!);
    if (status === -1) return; // queued; wait for drain
    ws.data.queue.shift(); // sent (>0) or dropped (0)
  }
}
 
Bun.serve({
  fetch: (req, server) =>
    server.upgrade(req, { data: { queue: [] } })
      ? undefined
      : new Response("Expected WebSocket", { status: 426 }),
  websocket: {
    data: {} as Data,
    backpressureLimit: 1024 * 1024,
    closeOnBackpressureLimit: true, // shed slow consumers
    message(ws, msg) {
      ws.data.queue.push(String(msg));
      flush(ws);
    },
    drain: flush,
  },
});

const ALLOWED = new Set(["https://app.example.com"]);
 
Bun.serve({
  async fetch(req, server) {
    const origin = req.headers.get("origin");
    if (!origin || !ALLOWED.has(origin)) {
      return new Response("Forbidden", { status: 403 });
    }
    const cookies = new Bun.CookieMap(
      req.headers.get("cookie") ?? "",
    );
    const session = await sessions.get(
      cookies.get("__Host-sid") ?? "",
    );
    if (!session) return new Response(null, { status: 401 });
    if (server.upgrade(req, { data: session })) return;
    return new Response(null, { status: 426 });
  },
  websocket: {
    data: {} as Session,
    message(ws, msg) {
      // ws.data.userId is trusted from here on
    },
  },
});

import { RedisClient } from "bun";
 
const pub = new RedisClient(process.env.REDIS_URL);
const sub = await pub.duplicate(); // subscriber conn
 
type Fanout = { topic: string; payload: string };
 
// every node relays broker messages to local sockets
await sub.subscribe("fanout", (message) => {
  const { topic, payload } = JSON.parse(message) as Fanout;
  server.publish(topic, payload);
});
 
export function broadcast(topic: string, payload: string) {
  const msg: Fanout = { topic, payload };
  return pub.publish("fanout", JSON.stringify(msg));
}

# nginx: proxy WebSockets
location /ws {
  proxy_pass http://app;
  proxy_http_version 1.1;
  proxy_set_header Upgrade $http_upgrade;
  proxy_set_header Connection "upgrade";
  proxy_read_timeout 75s;
}

// token bucket: `rate` tokens per second,
// bursts up to `burst`
export function tokenBucket(rate: number, burst: number) {
  let tokens = burst;
  let last = performance.now();
  return function take(): boolean {
    const now = performance.now();
    const refill = ((now - last) / 1000) * rate;
    tokens = Math.min(burst, tokens + refill);
    last = now;
    if (tokens < 1) return false;
    tokens -= 1;
    return true;
  };
}
 
// per connection: const allow = tokenBucket(10, 20);
// if (!allow()) ws.close(1008, "rate limit");

export function opened(ws: WebSocket): Promise<void> {
  if (ws.readyState === ws.OPEN) return Promise.resolve();
  return new Promise((resolve, reject) => {
    const off = new AbortController();
    const opts = { once: true, signal: off.signal };
    ws.addEventListener("open", () => {
      off.abort();
      resolve();
    }, opts);
    ws.addEventListener("close", (e) => {
      off.abort();
      reject(new Error(`closed before open: ${e.code}`));
    }, opts);
  });
}
 
const ws = new WebSocket("wss://api.example.com/ws");
await opened(ws);
ws.send(JSON.stringify({ type: "join", room: "lobby" }));

type Call = PromiseWithResolvers<unknown>;
type Reply = { id?: string; result?: unknown; err?: string };
 
export function rpc(ws: WebSocket, timeoutMs = 10_000) {
  const pending = new Map<string, Call>();
  ws.addEventListener("message", (e) => {
    const msg = JSON.parse(String(e.data)) as Reply;
    const call = msg.id ? pending.get(msg.id) : undefined;
    if (!call || !msg.id) return; // a push, not a reply
    pending.delete(msg.id);
    if (msg.err) call.reject(new Error(msg.err));
    else call.resolve(msg.result);
  });
  return (method: string, params?: unknown) => {
    const id = crypto.randomUUID();
    const call = Promise.withResolvers<unknown>();
    pending.set(id, call);
    ws.send(JSON.stringify({ id, method, params }));
    const timer = setTimeout(() => {
      pending.delete(id);
      call.reject(new Error(`${method} timed out`));
    }, timeoutMs);
    return call.promise.finally(() => clearTimeout(timer));
  };
}

export function outbox(max = 100) {
  let queue: string[] = [];
  let socket: WebSocket | undefined;
  return {
    send(data: string): void {
      if (socket?.readyState === WebSocket.OPEN) {
        socket.send(data);
      } else {
        queue = [...queue, data].slice(-max); // keep newest
      }
    },
    attach(ws: WebSocket): void { // call from onOpen
      socket = ws;
      for (const data of queue) ws.send(data);
      queue = [];
    },
  };
}
 
const box = outbox();
box.send(JSON.stringify({ type: "typing" })); // queued
box.attach(new WebSocket("wss://api.example.com/ws"));

type Data = { userId: string };
type Msg =
  | { type: "join" | "leave"; room: string }
  | { type: "say"; room: string; text: string };
 
Bun.serve({
  fetch(req, server) {
    const userId = crypto.randomUUID(); // real app: session
    if (server.upgrade(req, { data: { userId } })) return;
    return new Response("Upgrade required", { status: 426 });
  },
  websocket: {
    data: {} as Data,
    message(ws, raw) {
      const msg = JSON.parse(String(raw)) as Msg; // Zod!
      const room = `room:${msg.room}`;
      if (msg.type === "join") ws.subscribe(room);
      else if (msg.type === "leave") ws.unsubscribe(room);
      else if (msg.type === "say" && ws.isSubscribed(room)) {
        const out = { from: ws.data.userId, text: msg.text };
        ws.publish(room, JSON.stringify(out)); // not to self
      }
    },
  },
});

const sockets = new Set<Bun.ServerWebSocket<undefined>>();
 
const server = Bun.serve({
  fetch(req, server) {
    if (server.upgrade(req)) return;
    return new Response("Upgrade required", { status: 426 });
  },
  websocket: {
    open: (ws) => void sockets.add(ws),
    close: (ws) => void sockets.delete(ws),
    message(ws, msg) { ws.send(msg); },
  },
});
 
async function shutdown(): Promise<void> {
  // 1012 "service restart": clients reconnect with jitter
  for (const ws of sockets) ws.close(1012, "restarting");
  await server.stop(); // stop accepting, finish in-flight
  process.exit(0);
}
process.once("SIGTERM", shutdown);
process.once("SIGINT", shutdown);