import { createHash } from "node:crypto"; import type { IncomingMessage, Server as HttpServer } from "node:http"; import type { Duplex } from "node:stream"; import { and, eq, isNull } from "drizzle-orm"; import type { Db } from "@paperclip/db"; import { agentApiKeys } from "@paperclip/db"; import { WebSocket, WebSocketServer } from "ws"; import { logger } from "../middleware/logger.js"; import { subscribeCompanyLiveEvents } from "../services/live-events.js"; interface UpgradeContext { companyId: string; actorType: "board" | "agent"; actorId: string; } interface IncomingMessageWithContext extends IncomingMessage { paperclipUpgradeContext?: UpgradeContext; } function hashToken(token: string) { return createHash("sha256").update(token).digest("hex"); } function rejectUpgrade(socket: Duplex, statusLine: string, message: string) { const safe = message.replace(/[\r\n]+/g, " ").trim(); socket.write(`HTTP/1.1 ${statusLine}\r\nConnection: close\r\nContent-Type: text/plain\r\n\r\n${safe}`); socket.destroy(); } function parseCompanyId(pathname: string) { const match = pathname.match(/^\/api\/companies\/([^/]+)\/events\/ws$/); if (!match) return null; try { return decodeURIComponent(match[1] ?? ""); } catch { return null; } } function parseBearerToken(rawAuth: string | string[] | undefined) { const auth = Array.isArray(rawAuth) ? rawAuth[0] : rawAuth; if (!auth) return null; if (!auth.toLowerCase().startsWith("bearer ")) return null; const token = auth.slice("bearer ".length).trim(); return token.length > 0 ? token : null; } async function authorizeUpgrade( db: Db, req: IncomingMessage, companyId: string, url: URL, ): Promise { const queryToken = url.searchParams.get("token")?.trim() ?? ""; const authToken = parseBearerToken(req.headers.authorization); const token = authToken ?? (queryToken.length > 0 ? queryToken : null); // Browser board context has no bearer token in V1. if (!token) { return { companyId, actorType: "board", actorId: "board", }; } const tokenHash = hashToken(token); const key = await db .select() .from(agentApiKeys) .where(and(eq(agentApiKeys.keyHash, tokenHash), isNull(agentApiKeys.revokedAt))) .then((rows) => rows[0] ?? null); if (!key || key.companyId !== companyId) { return null; } await db .update(agentApiKeys) .set({ lastUsedAt: new Date() }) .where(eq(agentApiKeys.id, key.id)); return { companyId, actorType: "agent", actorId: key.agentId, }; } export function setupLiveEventsWebSocketServer(server: HttpServer, db: Db) { const wss = new WebSocketServer({ noServer: true }); const cleanupByClient = new Map void>(); const aliveByClient = new Map(); const pingInterval = setInterval(() => { for (const socket of wss.clients) { if (!aliveByClient.get(socket)) { socket.terminate(); continue; } aliveByClient.set(socket, false); socket.ping(); } }, 30000); wss.on("connection", (socket, req) => { const context = (req as IncomingMessageWithContext).paperclipUpgradeContext; if (!context) { socket.close(1008, "missing context"); return; } const unsubscribe = subscribeCompanyLiveEvents(context.companyId, (event) => { if (socket.readyState !== WebSocket.OPEN) return; socket.send(JSON.stringify(event)); }); cleanupByClient.set(socket, unsubscribe); aliveByClient.set(socket, true); socket.on("pong", () => { aliveByClient.set(socket, true); }); socket.on("close", () => { const cleanup = cleanupByClient.get(socket); if (cleanup) cleanup(); cleanupByClient.delete(socket); aliveByClient.delete(socket); }); socket.on("error", (err) => { logger.warn({ err, companyId: context.companyId }, "live websocket client error"); }); }); wss.on("close", () => { clearInterval(pingInterval); }); server.on("upgrade", (req, socket, head) => { if (!req.url) { rejectUpgrade(socket, "400 Bad Request", "missing url"); return; } const url = new URL(req.url, "http://localhost"); const companyId = parseCompanyId(url.pathname); if (!companyId) { socket.destroy(); return; } void authorizeUpgrade(db, req, companyId, url) .then((context) => { if (!context) { rejectUpgrade(socket, "403 Forbidden", "forbidden"); return; } const reqWithContext = req as IncomingMessageWithContext; reqWithContext.paperclipUpgradeContext = context; wss.handleUpgrade(req, socket, head, (ws) => { wss.emit("connection", ws, reqWithContext); }); }) .catch((err) => { logger.error({ err, path: req.url }, "failed websocket upgrade authorization"); rejectUpgrade(socket, "500 Internal Server Error", "upgrade failed"); }); }); return wss; }