feat(canvas): WS hardening, client race fix, and ping/pong heartbeat

Server:
- Origin allowlist on /api/ws (env.ALLOWED_ORIGINS, comma-separated; empty
  = allow all for dev/preview)
- per-identity WS connection cap (MAX_WS_PER_IDENTITY = 5) using
  acceptWebSocket(socket, [identity]) tagging; 6th upgrade returns 429
- ws.send 'ping' triggers a {type:'pong'} reply so dead connections fire
  onclose promptly instead of waiting on TCP keepalive

Client:
- buffer WS pixels arriving during the initial canvas fetch and replay
  them after committedColors is replaced; previously the post-fetch
  Uint8Array assignment silently overwrote any pixels broadcast in the
  fetch window (the documented C2 race)
- 30s ping / 60s pong watchdog closes the socket if pong stops arriving,
  routing through the existing exponential-backoff reconnect path

Tests:
- four /api/ws cases: missing upgrade, disallowed origin, allowed origin,
  empty allowlist (dev default). Sentinel uses status 200 because undici
  rejects 101 in Node-side Response constructors.

Plan: plans/260510-0232-fix-do-migration-followups/phase-03-ws-hardening-client-race.md

src/client/App.svelte

@@ -108,6 +108,33 @@
   // Monotonic broadcast counter from the server. A gap means we missed at
   // least one frame (DO hibernation, network glitch) — refetch to resync.
   let lastSeq = null;
+  // App-level heartbeat — if no pong arrives within PONG_TIMEOUT_MS we close
+  // the socket so onclose triggers reconnect (avoids zombie-connection wait
+  // for OS-level TCP keepalive).
+  const PING_INTERVAL_MS = 30_000;
+  const PONG_TIMEOUT_MS = 60_000;
+  let pingInterval = null;
+  let pongWatchdog = null;
+  let lastPongAt = 0;
+
+  function startHeartbeat(ws) {
+    lastPongAt = Date.now();
+    pingInterval = setInterval(() => {
+      if (ws.readyState === 1) {
+        try { ws.send('ping'); } catch { /* socket closing */ }
+      }
+    }, PING_INTERVAL_MS);
+    pongWatchdog = setInterval(() => {
+      if (Date.now() - lastPongAt > PONG_TIMEOUT_MS) {
+        try { ws.close(); } catch { /* ignore */ }
+      }
+    }, PING_INTERVAL_MS / 3);
+  }
+
+  function stopHeartbeat() {
+    if (pingInterval) { clearInterval(pingInterval); pingInterval = null; }
+    if (pongWatchdog) { clearInterval(pongWatchdog); pongWatchdog = null; }
+  }
 
   function connectWebSocket() {
     const protocol = location.protocol === 'https:' ? 'wss:' : 'ws:';
@@ -116,6 +143,10 @@
     ws.onmessage = (event) => {
       try {
         const data = JSON.parse(event.data);
+        if (data.type === 'pong') {
+          lastPongAt = Date.now();
+          return;
+        }
         if (data.type === 'pixels' && canvasRenderer) {
           if (data.seq != null) {
             const expected = lastSeq == null ? data.seq : ((lastSeq + 1) >>> 0);
@@ -139,8 +170,10 @@
       wsRetryDelay = 1000;
       isReconnect = true;
       wsState = 'open';
+      startHeartbeat(ws);
     };
     ws.onclose = () => {
+      stopHeartbeat();
       wsState = 'reconnecting';
       setTimeout(connectWebSocket, wsRetryDelay);
       wsRetryDelay = Math.min(wsRetryDelay * 2, 30000);

src/client/components/CanvasRenderer.svelte

@@ -14,6 +14,12 @@
   /** Committed color index per pixel (server-confirmed state). Allocated upfront so WS
    *  updates that arrive during the initial canvas fetch don't null-deref. */
   let committedColors = new Uint8Array(CANVAS_WIDTH * CANVAS_HEIGHT);
+  /** WS pixel updates that arrive while loadCanvas is in flight. Buffered
+   *  here and replayed after committedColors is replaced — without this, the
+   *  fetch's `new Uint8Array(indices)` assignment would silently overwrite
+   *  any pixels broadcast during the fetch window. */
+  let pendingWsEdits = [];
+  let isLoadComplete = false;
   let pan = { x: 0, y: 0 };
   let dragging = $state(false);
   let lastMouse = { x: 0, y: 0 };
@@ -166,16 +172,24 @@
   // --- Public API (called by App.svelte) ---
 
   export function applyUpdates(pixels) {
-    for (const { x, y, color } of pixels) {
+    if (!isLoadComplete) {
-      committedColors[y * CANVAS_WIDTH + x] = color;
+      // Buffer until loadCanvas resolves; replay against the freshly-fetched
-      // Only update display if no pending or active stroke covers this pixel
+      // committedColors there. Without this, the fetch overwrites these edits.
-      if (buffer.getColorAt(x, y) < 0 && !currentStrokeKeys.has(y * 65536 + x)) {
+      pendingWsEdits.push(...pixels);
-        setPixelRgba(x, y, color);
+      return;
-      }
     }
+    for (const { x, y, color } of pixels) applySingleUpdate(x, y, color);
     render();
   }
 
+  function applySingleUpdate(x, y, color) {
+    committedColors[y * CANVAS_WIDTH + x] = color;
+    // Only update display if no pending or active stroke covers this pixel.
+    if (buffer.getColorAt(x, y) < 0 && !currentStrokeKeys.has(y * 65536 + x)) {
+      setPixelRgba(x, y, color);
+    }
+  }
+
   export function undo() {
     const stroke = buffer.undo();
     if (!stroke) return;
@@ -465,6 +479,7 @@
   async function loadCanvas() {
     loading = true;
     loadError = null;
+    isLoadComplete = false;
     try {
       const res = await fetch('/api/canvas');
       if (!res.ok) throw new Error(`HTTP ${res.status}`);
@@ -473,7 +488,14 @@
       committedColors = new Uint8Array(indices); // replace pre-allocated zero array
       const rgba = indicesToRgba(indices);
       imageData = new ImageData(rgba, CANVAS_WIDTH, CANVAS_HEIGHT);
+      // Replay any WS pixels that arrived during the fetch — they raced the
+      // assignment above and would otherwise be lost.
+      if (pendingWsEdits.length > 0) {
+        for (const { x, y, color } of pendingWsEdits) applySingleUpdate(x, y, color);
+        pendingWsEdits = [];
+      }
       imageDataDirty = true;
+      isLoadComplete = true;
       render();
     } catch (err) {
       console.error('Failed to load canvas:', err);

src/durable-objects/canvas-room.js

@@ -1,4 +1,4 @@
-import { CANVAS_WIDTH, CANVAS_HEIGHT, MAX_COLORS, MAX_BATCH_SIZE } from '../lib/constants.js';
+import { CANVAS_WIDTH, CANVAS_HEIGHT, MAX_COLORS, MAX_BATCH_SIZE, MAX_WS_PER_IDENTITY } from '../lib/constants.js';
 import { init as initSchema } from './lib/schema.js';
 import { readAllChunks, writePixels } from './lib/chunk-storage.js';
 import { tryAcquire, release } from './lib/cooldown-store.js';
@@ -32,7 +32,7 @@ export class CanvasRoom {
     switch (url.pathname) {
       case '/canvas': return this.#handleGetCanvas();
       case '/place':  return this.#handlePlace(request);
-      case '/ws':     return this.#handleWsUpgrade();
+      case '/ws':     return this.#handleWsUpgrade(url);
       default:        return new Response('not found', { status: 404 });
     }
   }
@@ -103,10 +103,17 @@ export class CanvasRoom {
     return Response.json({ ok: true });
   }
 
-  #handleWsUpgrade() {
+  #handleWsUpgrade(url) {
+    const identity = url.searchParams.get('identity') || 'anon:unknown';
+    const existing = this.state.getWebSockets(identity);
+    if (existing.length >= MAX_WS_PER_IDENTITY) {
+      return new Response('too_many_sockets', { status: 429 });
+    }
     const pair = new WebSocketPair();
     const [client, server] = Object.values(pair);
-    this.state.acceptWebSocket(server);
+    // Tag the socket with identity so getWebSockets(identity) can count it
+    // toward the per-identity cap on subsequent upgrades.
+    this.state.acceptWebSocket(server, [identity]);
     return new Response(null, { status: 101, webSocket: client });
   }
 
@@ -125,8 +132,15 @@ export class CanvasRoom {
 
   /** Hibernation-API callbacks. */
 
-  webSocketMessage(ws) {
+  webSocketMessage(ws, message) {
-    // Protocol is broadcast-only; reject any inbound payload.
+    // Heartbeat: client sends 'ping', server replies with {type:'pong'} so
+    // dead connections fire onclose promptly instead of waiting for OS-level
+    // TCP keepalive (potentially minutes).
+    if (typeof message === 'string' && message === 'ping') {
+      try { ws.send(JSON.stringify({ type: 'pong' })); } catch { /* socket closing */ }
+      return;
+    }
+    // Anything else is unexpected — close the socket.
     ws.close(1003, 'unexpected client message');
   }
 

src/lib/constants.js

@@ -11,6 +11,10 @@ export const MAX_COLORS = 256;
 export const REQUEST_COOLDOWN_SEC = 1;
 export const MAX_BATCH_SIZE = 2048;
 
+/** Per-identity concurrent WebSocket cap. Prevents broadcast amplification
+ *  from a single client opening thousands of sockets. */
+export const MAX_WS_PER_IDENTITY = 5;
+
 /** Canvas chunked storage layout (DO SQLite). 64 KB chunks → 256 chunks for
  *  the 16 MB / 4096² canvas. CHUNK_COUNT is derived: bumping CANVAS_WIDTH or
  *  CANVAS_HEIGHT and redeploying transparently allocates more chunks (lazy-

src/worker.js

@@ -105,14 +105,42 @@ app.post('/api/place', async (c) => {
 });
 
 /** GET /api/ws — WebSocket upgrade routed to the DO.
- *  The DO routes by url.pathname, so we rewrite the URL to `/ws` while
+ *  Origin allowlist + identity resolution + per-identity cap (in DO). */
- *  preserving the original headers (including Upgrade) via the request init. */
 app.get('/api/ws', async (c) => {
   const upgradeHeader = c.req.header('Upgrade');
   if (upgradeHeader !== 'websocket') {
     return c.text('Expected WebSocket', 426);
   }
-  return room(c.env).fetch('http://do/ws', c.req.raw);
+  // Origin allowlist. Empty ALLOWED_ORIGINS ⇒ allow all (dev / preview).
+  const origin = c.req.header('Origin');
+  const allowed = parseAllowedOrigins(c.env?.ALLOWED_ORIGINS);
+  if (origin && allowed.size > 0 && !allowed.has(origin)) {
+    return c.text('forbidden_origin', 403);
+  }
+  let identity;
+  try {
+    identity = await resolveIdentity(c.req.raw, c.env);
+  } catch (err) {
+    if (err instanceof NoIdentityError) {
+      return c.json({ error: 'no_identity' }, 500);
+    }
+    throw err;
+  }
+  // Identity goes via URL query so the original request (with Upgrade header)
+  // forwards unchanged. CF DO routing reads it from request.url.
+  const url = `http://do/ws?identity=${encodeURIComponent(identity.id)}`;
+  return room(c.env).fetch(url, c.req.raw);
 });
 
+/** Parse the comma-separated ALLOWED_ORIGINS env var into a Set. */
+function parseAllowedOrigins(raw) {
+  const out = new Set();
+  if (!raw) return out;
+  for (const o of String(raw).split(',')) {
+    const trimmed = o.trim();
+    if (trimmed) out.add(trimmed);
+  }
+  return out;
+}
+
 export default app;

test/worker-validation.test.js

@@ -189,4 +189,51 @@ describe('GET /api/ws', () => {
     const res = await app.fetch(req, env);
     expect(res.status).toBe(426);
   });
+
+  it('rejects WS upgrade from a disallowed Origin', async () => {
+    // undici disallows constructing a Response with status 101; use a 2xx
+    // sentinel to assert the upstream DO was reached.
+    doResponse = () => new Response(null, { status: 200 });
+    const req = new Request('http://localhost/api/ws', {
+      headers: {
+        Upgrade: 'websocket',
+        Origin: 'https://evil.example',
+      },
+    });
+    const restrictedEnv = {
+      ...env,
+      ALLOWED_ORIGINS: 'https://rplace.miti99.workers.dev',
+    };
+    const res = await app.fetch(req, restrictedEnv);
+    expect(res.status).toBe(403);
+  });
+
+  it('forwards WS upgrade when Origin is in the allowlist', async () => {
+    // undici disallows constructing a Response with status 101; use a 2xx
+    // sentinel to assert the upstream DO was reached.
+    doResponse = () => new Response(null, { status: 200 });
+    const req = new Request('http://localhost/api/ws', {
+      headers: {
+        Upgrade: 'websocket',
+        Origin: 'https://rplace.miti99.workers.dev',
+      },
+    });
+    const restrictedEnv = {
+      ...env,
+      ALLOWED_ORIGINS: 'https://rplace.miti99.workers.dev',
+    };
+    const res = await app.fetch(req, restrictedEnv);
+    expect(res.status).toBe(200);
+  });
+
+  it('forwards WS upgrade when ALLOWED_ORIGINS is empty (dev default)', async () => {
+    // undici disallows constructing a Response with status 101; use a 2xx
+    // sentinel to assert the upstream DO was reached.
+    doResponse = () => new Response(null, { status: 200 });
+    const req = new Request('http://localhost/api/ws', {
+      headers: { Upgrade: 'websocket', Origin: 'https://anything.example' },
+    });
+    const res = await app.fetch(req, env);
+    expect(res.status).toBe(200);
+  });
 });

wrangler.json

@@ -27,5 +27,9 @@
     "traces": {
       "enabled": true
     }
+  },
+  "vars": {
+    "ALLOWED_ORIGINS": "",
+    "ENVIRONMENT": "development"
   }
 }