scottlanoue.com / atmosphere-office
Full document, spreadsheet, slideshow, and diagram tooling
0
fork

Configure Feed

Select the types of activity you want to include in your feed.

Merge pull request 'test: add WebSocket relay integration tests (24 tests)' (#351) from test/websocket-relay-coverage into main

scott 0a6c2290 1b274939

+524
+524
tests/server/websocket-relay.test.ts
··· 1 + /** 2 + * Integration tests for the WebSocket relay in server/index.ts. 3 + * 4 + * Spins up a minimal HTTP + WebSocket server with in-memory SQLite 5 + * that mirrors production upgrade handling: room management, peer counting, 6 + * message relay, expiry checks, and cleanup. 7 + */ 8 + 9 + import { describe, it, expect, beforeAll, afterAll } from 'vitest'; 10 + import Database from 'better-sqlite3'; 11 + import { createServer, type Server } from 'http'; 12 + import { WebSocketServer, WebSocket } from 'ws'; 13 + import type { IncomingMessage } from 'http'; 14 + import type { Duplex } from 'stream'; 15 + 16 + let wsUrl: string; 17 + let server: Server; 18 + let db: ReturnType<typeof Database>; 19 + let rooms: Map<string, Set<WebSocket>>; 20 + let stmts: { 21 + getOne: ReturnType<ReturnType<typeof Database>['prepare']>; 22 + insert: ReturnType<ReturnType<typeof Database>['prepare']>; 23 + updateExpiry: ReturnType<ReturnType<typeof Database>['prepare']>; 24 + }; 25 + 26 + // Track all WebSockets opened during tests for cleanup 27 + const allOpenWs: WebSocket[] = []; 28 + 29 + beforeAll(async () => { 30 + db = new Database(':memory:'); 31 + db.pragma('journal_mode = WAL'); 32 + 33 + db.exec(` 34 + CREATE TABLE documents ( 35 + id TEXT PRIMARY KEY, 36 + type TEXT NOT NULL, 37 + name_encrypted TEXT, 38 + snapshot BLOB, 39 + share_mode TEXT DEFAULT 'edit', 40 + expires_at TEXT, 41 + deleted_at TEXT, 42 + created_at TEXT DEFAULT (datetime('now')), 43 + updated_at TEXT DEFAULT (datetime('now')) 44 + ) 45 + `); 46 + 47 + stmts = { 48 + getOne: db.prepare('SELECT id, type, expires_at FROM documents WHERE id = ?'), 49 + insert: db.prepare('INSERT INTO documents (id, type) VALUES (?, ?)'), 50 + updateExpiry: db.prepare('UPDATE documents SET expires_at = ? WHERE id = ?'), 51 + }; 52 + 53 + rooms = new Map(); 54 + 55 + const app = (_req: IncomingMessage, res: import('http').ServerResponse) => { 56 + res.writeHead(404); 57 + res.end(); 58 + }; 59 + 60 + server = createServer(app); 61 + 62 + const wss = new WebSocketServer({ noServer: true, maxPayload: 5 * 1024 * 1024 }); 63 + 64 + // Mirror the production handleUpgrade from server/index.ts 65 + server.on('upgrade', (request: IncomingMessage, socket: Duplex, head: Buffer) => { 66 + const url = new URL(request.url || '', 'http://localhost'); 67 + if (url.pathname !== '/ws') { 68 + socket.destroy(); 69 + return; 70 + } 71 + wss.handleUpgrade(request, socket, head, (ws: WebSocket) => { 72 + const room = url.searchParams.get('room'); 73 + if (!room || room.length > 200 || !/^[a-zA-Z0-9_-]+$/.test(room)) { ws.close(); return; } 74 + 75 + // Check share link expiry 76 + try { 77 + const doc = stmts.getOne.get(room) as { expires_at: string | null } | undefined; 78 + if (doc?.expires_at) { 79 + const expiresAt = new Date(doc.expires_at.endsWith('Z') ? doc.expires_at : doc.expires_at + 'Z'); 80 + if (expiresAt <= new Date()) { 81 + ws.close(4410, 'Document link has expired'); 82 + return; 83 + } 84 + } 85 + } catch { 86 + ws.close(4500, 'Internal error checking document access'); 87 + return; 88 + } 89 + 90 + // Extract Tailscale identity from upgrade request headers 91 + const wsUserLogin = request.headers['tailscale-user-login'] as string | undefined; 92 + const wsUserName = request.headers['tailscale-user-name'] as string | undefined; 93 + const wsUserPic = request.headers['tailscale-user-profile-pic'] as string | undefined; 94 + const wsUser = wsUserLogin 95 + ? { login: wsUserLogin, name: wsUserName || wsUserLogin, profilePic: wsUserPic || null } 96 + : null; 97 + (ws as WebSocket & { tsUser?: unknown }).tsUser = wsUser; 98 + 99 + if (!rooms.has(room)) rooms.set(room, new Set()); 100 + const peers = rooms.get(room)!; 101 + peers.add(ws); 102 + 103 + const msg = { type: 'peer-count', count: peers.size - 1 }; 104 + ws.send(JSON.stringify(msg)); 105 + 106 + for (const peer of peers) { 107 + if (peer !== ws && peer.readyState === 1) { 108 + const joinMsg = { type: 'peer-joined', user: wsUser }; 109 + peer.send(JSON.stringify(joinMsg)); 110 + } 111 + } 112 + 113 + ws.on('message', (data: Buffer | ArrayBuffer | Buffer[], isBinary: boolean) => { 114 + for (const peer of peers) { 115 + if (peer !== ws && peer.readyState === 1) { 116 + peer.send(data, { binary: isBinary }); 117 + } 118 + } 119 + }); 120 + 121 + ws.on('close', () => { 122 + peers.delete(ws); 123 + if (peers.size === 0) { 124 + rooms.delete(room); 125 + } else { 126 + for (const peer of peers) { 127 + if (peer.readyState === 1) { 128 + const leftMsg = { type: 'peer-left', count: peers.size }; 129 + peer.send(JSON.stringify(leftMsg)); 130 + } 131 + } 132 + } 133 + }); 134 + }); 135 + }); 136 + 137 + await new Promise<void>((resolve) => { 138 + server.listen(0, () => { 139 + const addr = server.address(); 140 + if (typeof addr === 'object' && addr) { 141 + wsUrl = `ws://localhost:${addr.port}`; 142 + } 143 + resolve(); 144 + }); 145 + }); 146 + }); 147 + 148 + afterAll(async () => { 149 + // Close all tracked WebSockets 150 + for (const ws of allOpenWs) { 151 + if (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING) { 152 + ws.close(); 153 + } 154 + } 155 + await new Promise<void>((resolve) => server.close(() => resolve())); 156 + db.close(); 157 + }); 158 + 159 + /** 160 + * Connect to a room and return the WS along with the initial peer-count message. 161 + * The message listener is attached BEFORE the connection opens to avoid the 162 + * microtask race between `open` event and peer-count delivery. 163 + */ 164 + function connectWs( 165 + roomId: string, 166 + headers?: Record<string, string>, 167 + ): Promise<{ ws: WebSocket; peerCount: number }> { 168 + return new Promise((resolve, reject) => { 169 + const ws = new WebSocket(`${wsUrl}/ws?room=${roomId}`, { headers }); 170 + allOpenWs.push(ws); 171 + // Attach message listener immediately — before open fires — to capture peer-count 172 + ws.once('message', (data) => { 173 + const msg = JSON.parse(data.toString()); 174 + resolve({ ws, peerCount: msg.count }); 175 + }); 176 + ws.on('error', reject); 177 + }); 178 + } 179 + 180 + /** Wait for the next message on a WebSocket */ 181 + function nextMessage(ws: WebSocket, timeout = 2000): Promise<string> { 182 + return new Promise((resolve, reject) => { 183 + const timer = setTimeout(() => reject(new Error('Timeout waiting for message')), timeout); 184 + ws.once('message', (data) => { 185 + clearTimeout(timer); 186 + resolve(data.toString()); 187 + }); 188 + }); 189 + } 190 + 191 + /** 192 + * Set up a message listener BEFORE an action that will trigger it. 193 + * Returns a promise that resolves with the message. This avoids the 194 + * microtask race that occurs when using `await action; await nextMessage()`. 195 + */ 196 + function pendingMessage(ws: WebSocket, timeout = 2000): Promise<string> { 197 + return new Promise((resolve, reject) => { 198 + const timer = setTimeout(() => reject(new Error('Timeout waiting for message')), timeout); 199 + ws.once('message', (data) => { 200 + clearTimeout(timer); 201 + resolve(data.toString()); 202 + }); 203 + }); 204 + } 205 + 206 + /** Wait for a close event (or error indicating destroyed socket) */ 207 + function waitForClose(ws: WebSocket, timeout = 2000): Promise<{ code: number; reason: string }> { 208 + return new Promise((resolve, reject) => { 209 + const timer = setTimeout(() => reject(new Error('Timeout waiting for close')), timeout); 210 + ws.on('close', (code, reason) => { 211 + clearTimeout(timer); 212 + resolve({ code, reason: reason.toString() }); 213 + }); 214 + // socket.destroy() for non-/ws paths causes error, not clean close 215 + ws.on('error', () => { 216 + clearTimeout(timer); 217 + resolve({ code: 1006, reason: 'socket destroyed' }); 218 + }); 219 + }); 220 + } 221 + 222 + describe('WebSocket relay', () => { 223 + describe('connection and room management', () => { 224 + it('connects to a valid room and receives peer-count 0', async () => { 225 + const { ws, peerCount } = await connectWs('test-room-1'); 226 + expect(peerCount).toBe(0); 227 + ws.close(); 228 + }); 229 + 230 + it('creates a room entry on first connection', async () => { 231 + const { ws } = await connectWs('room-create-test'); 232 + expect(rooms.has('room-create-test')).toBe(true); 233 + expect(rooms.get('room-create-test')!.size).toBe(1); 234 + ws.close(); 235 + }); 236 + 237 + it('cleans up room when last peer disconnects', async () => { 238 + const { ws } = await connectWs('room-cleanup'); 239 + expect(rooms.has('room-cleanup')).toBe(true); 240 + ws.close(); 241 + await new Promise(r => setTimeout(r, 100)); 242 + expect(rooms.has('room-cleanup')).toBe(false); 243 + }); 244 + }); 245 + 246 + describe('room name validation', () => { 247 + it('rejects missing room parameter', async () => { 248 + const ws = new WebSocket(`${wsUrl}/ws`); 249 + allOpenWs.push(ws); 250 + const { code } = await waitForClose(ws); 251 + expect([1005, 1006, 1000]).toContain(code); 252 + }); 253 + 254 + it('rejects room names with invalid characters', async () => { 255 + const ws = new WebSocket(`${wsUrl}/ws?room=invalid%20room!`); 256 + allOpenWs.push(ws); 257 + const { code } = await waitForClose(ws); 258 + expect([1005, 1006, 1000]).toContain(code); 259 + }); 260 + 261 + it('rejects room names exceeding 200 chars', async () => { 262 + const longRoom = 'a'.repeat(201); 263 + const ws = new WebSocket(`${wsUrl}/ws?room=${longRoom}`); 264 + allOpenWs.push(ws); 265 + const { code } = await waitForClose(ws); 266 + expect([1005, 1006, 1000]).toContain(code); 267 + }); 268 + 269 + it('accepts room names with valid chars (alphanumeric, dash, underscore)', async () => { 270 + const { ws, peerCount } = await connectWs('valid_room-123'); 271 + expect(peerCount).toBe(0); 272 + ws.close(); 273 + }); 274 + }); 275 + 276 + describe('non-/ws path rejection', () => { 277 + it('destroys socket for non-/ws upgrade paths', async () => { 278 + const ws = new WebSocket(`${wsUrl}/other-path?room=test`); 279 + allOpenWs.push(ws); 280 + const { code } = await waitForClose(ws); 281 + expect(code).toBe(1006); // abnormal close = socket.destroy() 282 + }); 283 + }); 284 + 285 + describe('expired document rejection', () => { 286 + it('closes with 4410 for expired document', async () => { 287 + const docId = 'expired-doc-1'; 288 + stmts.insert.run(docId, 'doc'); 289 + stmts.updateExpiry.run('2020-01-01T00:00:00Z', docId); 290 + 291 + const ws = new WebSocket(`${wsUrl}/ws?room=${docId}`); 292 + allOpenWs.push(ws); 293 + const { code, reason } = await waitForClose(ws); 294 + expect(code).toBe(4410); 295 + expect(reason).toBe('Document link has expired'); 296 + }); 297 + 298 + it('allows connection for non-expired document', async () => { 299 + const docId = 'active-doc-1'; 300 + stmts.insert.run(docId, 'doc'); 301 + stmts.updateExpiry.run('2099-12-31T23:59:59Z', docId); 302 + 303 + const { ws, peerCount } = await connectWs(docId); 304 + expect(peerCount).toBe(0); 305 + ws.close(); 306 + }); 307 + 308 + it('allows connection for document without expiry', async () => { 309 + const docId = 'no-expiry-doc'; 310 + stmts.insert.run(docId, 'doc'); 311 + 312 + const { ws, peerCount } = await connectWs(docId); 313 + expect(peerCount).toBe(0); 314 + ws.close(); 315 + }); 316 + 317 + it('allows connection for room with no matching document', async () => { 318 + const { ws, peerCount } = await connectWs('no-db-record'); 319 + expect(peerCount).toBe(0); 320 + ws.close(); 321 + }); 322 + 323 + it('handles expires_at without trailing Z', async () => { 324 + const docId = 'expiry-noz'; 325 + stmts.insert.run(docId, 'doc'); 326 + stmts.updateExpiry.run('2020-01-01T00:00:00', docId); 327 + 328 + const ws = new WebSocket(`${wsUrl}/ws?room=${docId}`); 329 + allOpenWs.push(ws); 330 + const { code } = await waitForClose(ws); 331 + expect(code).toBe(4410); 332 + }); 333 + }); 334 + 335 + describe('peer counting', () => { 336 + it('sends peer-count 0 to first client', async () => { 337 + const { ws, peerCount } = await connectWs('pc-room-1'); 338 + expect(peerCount).toBe(0); 339 + ws.close(); 340 + }); 341 + 342 + it('sends peer-count 1 to second client', async () => { 343 + const { ws: ws1 } = await connectWs('pc-room-2'); 344 + 345 + // Set up listener for peer-joined BEFORE connecting ws2 346 + const joinPromise = pendingMessage(ws1); 347 + const { ws: ws2, peerCount } = await connectWs('pc-room-2'); 348 + expect(peerCount).toBe(1); 349 + await joinPromise; // consume peer-joined 350 + 351 + ws1.close(); 352 + ws2.close(); 353 + }); 354 + 355 + it('sends peer-count 2 to third client', async () => { 356 + const { ws: ws1 } = await connectWs('pc-room-3'); 357 + 358 + const join1 = pendingMessage(ws1); 359 + const { ws: ws2 } = await connectWs('pc-room-3'); 360 + await join1; // ws2 joined notification on ws1 361 + 362 + const join2a = pendingMessage(ws1); 363 + const join2b = pendingMessage(ws2); 364 + const { ws: ws3, peerCount } = await connectWs('pc-room-3'); 365 + expect(peerCount).toBe(2); 366 + await join2a; 367 + await join2b; 368 + 369 + ws1.close(); 370 + ws2.close(); 371 + ws3.close(); 372 + }); 373 + }); 374 + 375 + describe('peer-joined notifications', () => { 376 + it('notifies existing peers when a new peer joins', async () => { 377 + const { ws: ws1 } = await connectWs('join-room'); 378 + 379 + const joinPromise = pendingMessage(ws1); 380 + const { ws: ws2 } = await connectWs('join-room'); 381 + 382 + const joinMsg = JSON.parse(await joinPromise); 383 + expect(joinMsg.type).toBe('peer-joined'); 384 + expect(joinMsg.user).toBeNull(); // no Tailscale headers 385 + 386 + ws1.close(); 387 + ws2.close(); 388 + }); 389 + 390 + it('includes Tailscale identity in peer-joined', async () => { 391 + const { ws: ws1 } = await connectWs('join-identity-room'); 392 + 393 + const joinPromise = pendingMessage(ws1); 394 + const { ws: ws2 } = await connectWs('join-identity-room', { 395 + 'tailscale-user-login': 'alice@example.com', 396 + 'tailscale-user-name': 'Alice', 397 + 'tailscale-user-profile-pic': 'https://example.com/pic.jpg', 398 + }); 399 + 400 + const joinMsg = JSON.parse(await joinPromise); 401 + expect(joinMsg.type).toBe('peer-joined'); 402 + expect(joinMsg.user).toEqual({ 403 + login: 'alice@example.com', 404 + name: 'Alice', 405 + profilePic: 'https://example.com/pic.jpg', 406 + }); 407 + 408 + ws1.close(); 409 + ws2.close(); 410 + }); 411 + }); 412 + 413 + describe('message relay', () => { 414 + it('relays messages between peers in the same room', async () => { 415 + const { ws: ws1 } = await connectWs('relay-room'); 416 + const join = pendingMessage(ws1); 417 + const { ws: ws2 } = await connectWs('relay-room'); 418 + await join; 419 + 420 + const recvPromise = pendingMessage(ws2); 421 + ws1.send('hello from ws1'); 422 + const received = await recvPromise; 423 + expect(received).toBe('hello from ws1'); 424 + 425 + ws1.close(); 426 + ws2.close(); 427 + }); 428 + 429 + it('does not echo messages back to sender', async () => { 430 + const { ws: ws1 } = await connectWs('no-echo-room'); 431 + const join = pendingMessage(ws1); 432 + const { ws: ws2 } = await connectWs('no-echo-room'); 433 + await join; 434 + 435 + const recvPromise = pendingMessage(ws2); 436 + ws1.send('test message'); 437 + await recvPromise; // ws2 gets it 438 + 439 + // ws1 should NOT get the message back 440 + let echoed = false; 441 + ws1.once('message', () => { echoed = true; }); 442 + await new Promise(r => setTimeout(r, 200)); 443 + expect(echoed).toBe(false); 444 + 445 + ws1.close(); 446 + ws2.close(); 447 + }); 448 + 449 + it('relays binary messages', async () => { 450 + const { ws: ws1 } = await connectWs('binary-room'); 451 + const join = pendingMessage(ws1); 452 + const { ws: ws2 } = await connectWs('binary-room'); 453 + await join; 454 + 455 + const binaryData = Buffer.from([0x01, 0x02, 0x03, 0xff]); 456 + const recvPromise = new Promise<Buffer>((resolve) => { 457 + ws2.once('message', (data) => resolve(data as Buffer)); 458 + }); 459 + 460 + ws1.send(binaryData); 461 + const received = await recvPromise; 462 + expect(Buffer.compare(received, binaryData)).toBe(0); 463 + 464 + ws1.close(); 465 + ws2.close(); 466 + }); 467 + 468 + it('does not relay between different rooms', async () => { 469 + const { ws: ws1 } = await connectWs('room-a'); 470 + const { ws: ws2 } = await connectWs('room-b'); 471 + 472 + ws1.send('room-a message'); 473 + 474 + let crossRoom = false; 475 + ws2.once('message', () => { crossRoom = true; }); 476 + await new Promise(r => setTimeout(r, 200)); 477 + expect(crossRoom).toBe(false); 478 + 479 + ws1.close(); 480 + ws2.close(); 481 + }); 482 + }); 483 + 484 + describe('peer-left notifications', () => { 485 + it('notifies remaining peers when one disconnects', async () => { 486 + const { ws: ws1 } = await connectWs('left-room'); 487 + const join = pendingMessage(ws1); 488 + const { ws: ws2 } = await connectWs('left-room'); 489 + await join; 490 + 491 + const leftPromise = pendingMessage(ws1); 492 + ws2.close(); 493 + const leftMsg = JSON.parse(await leftPromise); 494 + expect(leftMsg).toEqual({ type: 'peer-left', count: 1 }); 495 + }); 496 + 497 + it('updates count correctly when multiple peers leave', async () => { 498 + const { ws: ws1 } = await connectWs('multi-left'); 499 + 500 + const j1 = pendingMessage(ws1); 501 + const { ws: ws2 } = await connectWs('multi-left'); 502 + await j1; 503 + 504 + const j2a = pendingMessage(ws1); 505 + const j2b = pendingMessage(ws2); 506 + const { ws: ws3 } = await connectWs('multi-left'); 507 + await j2a; 508 + await j2b; 509 + 510 + // ws3 leaves — ws1 and ws2 should get count=2 511 + const left1 = pendingMessage(ws1); 512 + const left2 = pendingMessage(ws2); 513 + ws3.close(); 514 + 515 + const msg1 = JSON.parse(await left1); 516 + const msg2 = JSON.parse(await left2); 517 + expect(msg1).toEqual({ type: 'peer-left', count: 2 }); 518 + expect(msg2).toEqual({ type: 'peer-left', count: 2 }); 519 + 520 + ws1.close(); 521 + ws2.close(); 522 + }); 523 + }); 524 + });