From 4f02d30393cefd6b98f35806260f17d239270626 Mon Sep 17 00:00:00 2001 From: tiennm99 Date: Sun, 10 May 2026 03:05:56 +0700 Subject: [PATCH] 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 | 33 +++++++++++++++ src/client/components/CanvasRenderer.svelte | 34 ++++++++++++--- src/durable-objects/canvas-room.js | 26 +++++++++--- src/lib/constants.js | 4 ++ src/worker.js | 34 +++++++++++++-- test/worker-validation.test.js | 47 +++++++++++++++++++++ wrangler.json | 4 ++ 7 files changed, 167 insertions(+), 15 deletions(-) diff --git a/src/client/App.svelte b/src/client/App.svelte index bf5c8d7..8ed3733 100644 --- a/src/client/App.svelte +++ b/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); diff --git a/src/client/components/CanvasRenderer.svelte b/src/client/components/CanvasRenderer.svelte index 52c0cbd..e48dee3 100644 --- a/src/client/components/CanvasRenderer.svelte +++ b/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) { - 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); - } + if (!isLoadComplete) { + // Buffer until loadCanvas resolves; replay against the freshly-fetched + // committedColors there. Without this, the fetch overwrites these edits. + pendingWsEdits.push(...pixels); + 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); diff --git a/src/durable-objects/canvas-room.js b/src/durable-objects/canvas-room.js index 9edb5bd..de5554f 100644 --- a/src/durable-objects/canvas-room.js +++ b/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) { - // Protocol is broadcast-only; reject any inbound payload. + webSocketMessage(ws, message) { + // 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'); } diff --git a/src/lib/constants.js b/src/lib/constants.js index bec8ab9..a12a94a 100644 --- a/src/lib/constants.js +++ b/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- diff --git a/src/worker.js b/src/worker.js index 67ce9ec..4fb109b 100644 --- a/src/worker.js +++ b/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 - * preserving the original headers (including Upgrade) via the request init. */ + * Origin allowlist + identity resolution + per-identity cap (in DO). */ 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; diff --git a/test/worker-validation.test.js b/test/worker-validation.test.js index 5eaef55..5e0b562 100644 --- a/test/worker-validation.test.js +++ b/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); + }); }); diff --git a/wrangler.json b/wrangler.json index 84e802d..75b881c 100644 --- a/wrangler.json +++ b/wrangler.json @@ -27,5 +27,9 @@ "traces": { "enabled": true } + }, + "vars": { + "ALLOWED_ORIGINS": "", + "ENVIRONMENT": "development" } }