@mooncompany/uplink-chat 0.5.0

package/server/websocket/connections.js

@@ -0,0 +1,339 @@
+/**
+ * WebSocket Connections Module
+ * Client tracking, rate limiting, heartbeats, IP-based limits, cleanup
+ */
+
+import { WebSocket } from 'ws';
+import { randomUUID } from 'crypto';
+import { log } from '../utils.js';
+import { WEBSOCKET } from '../config.js';
+
+// ============================================
+// Shared State
+// ============================================
+
+/** Track connected clients: clientId -> { ws, lastPing, lastPong, missedPongs, sessionUser, connectedAt, clientIp } */
+export const wsClients = new Map();
+
+/** Rate limiting per client: clientId -> { count, resetAt } */
+const wsRateLimits = new Map();
+
+/** IP-based connection tracking: ip -> Set<clientId> */
+const wsConnectionsByIp = new Map();
+
+/** Active streaming requests: requestId -> { abortController, clientId, startedAt } */
+export const activeStreamingRequests = new Map();
+
+// ============================================
+// Config Constants
+// ============================================
+
+const WS_RATE_LIMIT = {
+  windowMs: WEBSOCKET.messageRateLimitWindow,
+  maxMessages: WEBSOCKET.maxMessagesPerWindow,
+};
+const MAX_WS_RATE_LIMIT_ENTRIES = WEBSOCKET.maxRateLimitEntries;
+const MAX_ACTIVE_STREAMING_REQUESTS = WEBSOCKET.maxStreamingRequests;
+export const WS_CONNECTION_LIMITS = WEBSOCKET.connectionLimits;
+export const MAX_MESSAGE_SIZE = WEBSOCKET.maxMessageSize;
+const HEARTBEAT_INTERVAL_MS = WEBSOCKET.heartbeatIntervalMs;
+const MAX_MISSED_PONGS = WEBSOCKET.maxMissedPongs;
+
+// ============================================
+// Client ID Generation
+// ============================================
+
+export function generateClientId() {
+  return `ws-${Date.now()}-${randomUUID().slice(0, 8)}`;
+}
+
+// ============================================
+// Safe Send Helper
+// ============================================
+
+/**
+ * Safely send a message to a WebSocket, checking readyState first
+ * Returns true if message was sent, false if connection was not open
+ */
+export function safeSend(ws, data) {
+  if (ws.readyState === WebSocket.OPEN) {
+    try {
+      ws.send(typeof data === 'string' ? data : JSON.stringify(data));
+      return true;
+    } catch (sendError) {
+      log('warn', `[WS] Failed to send message:`, sendError.message);
+      return false;
+    }
+  }
+  return false;
+}
+
+// ============================================
+// Rate Limiting
+// ============================================
+
+/**
+ * Trim oldest rate limit entries when map exceeds size limit
+ * Removes entries that are expired first, then oldest by resetAt
+ */
+function trimRateLimitEntries() {
+  if (wsRateLimits.size <= MAX_WS_RATE_LIMIT_ENTRIES) return;
+
+  const now = Date.now();
+  const entries = Array.from(wsRateLimits.entries());
+
+  // Sort by resetAt (oldest first)
+  entries.sort((a, b) => a[1].resetAt - b[1].resetAt);
+
+  const toRemove = wsRateLimits.size - MAX_WS_RATE_LIMIT_ENTRIES;
+  let removed = 0;
+
+  for (let i = 0; i < entries.length && removed < toRemove; i++) {
+    const [clientId] = entries[i];
+    // Only remove if client is disconnected OR entry is expired
+    if (!wsClients.has(clientId) || entries[i][1].resetAt < now) {
+      wsRateLimits.delete(clientId);
+      removed++;
+    }
+  }
+
+  // If still over limit, force remove oldest
+  if (wsRateLimits.size > MAX_WS_RATE_LIMIT_ENTRIES) {
+    const remaining = Array.from(wsRateLimits.entries())
+      .sort((a, b) => a[1].resetAt - b[1].resetAt);
+    const forceRemove = wsRateLimits.size - MAX_WS_RATE_LIMIT_ENTRIES;
+    for (let i = 0; i < forceRemove; i++) {
+      wsRateLimits.delete(remaining[i][0]);
+    }
+  }
+}
+
+export function checkWsRateLimit(clientId) {
+  const now = Date.now();
+  let limit = wsRateLimits.get(clientId);
+
+  if (!limit || now > limit.resetAt) {
+    limit = { count: 0, resetAt: now + WS_RATE_LIMIT.windowMs };
+    wsRateLimits.set(clientId, limit);
+  }
+
+  limit.count++;
+
+  // Ensure rate limit map doesn't grow unbounded
+  trimRateLimitEntries();
+
+  return limit.count <= WS_RATE_LIMIT.maxMessages;
+}
+
+/**
+ * Periodic cleanup of stale rate limit entries
+ * Removes entries that have expired and aren't associated with active clients
+ */
+export function cleanupStaleRateLimits() {
+  const now = Date.now();
+  let cleaned = 0;
+
+  for (const [clientId, limit] of wsRateLimits) {
+    // Remove if: entry expired AND client no longer connected
+    if (now > limit.resetAt && !wsClients.has(clientId)) {
+      wsRateLimits.delete(clientId);
+      cleaned++;
+    }
+  }
+
+  if (cleaned > 0) {
+    log('debug', `[WS] Cleaned ${cleaned} stale rate limit entries`);
+  }
+}
+
+// ============================================
+// Streaming Request Management
+// ============================================
+
+/**
+ * Trim oldest streaming requests when map exceeds size limit
+ * Aborts and removes oldest requests
+ */
+export function trimStreamingRequests() {
+  if (activeStreamingRequests.size <= MAX_ACTIVE_STREAMING_REQUESTS) return;
+
+  const entries = Array.from(activeStreamingRequests.entries());
+  // Sort by startedAt (oldest first)
+  entries.sort((a, b) => (a[1].startedAt || 0) - (b[1].startedAt || 0));
+
+  const toRemove = activeStreamingRequests.size - MAX_ACTIVE_STREAMING_REQUESTS;
+
+  for (let i = 0; i < toRemove; i++) {
+    const [requestId, request] = entries[i];
+    try {
+      request.abortController.abort();
+    } catch (abortError) {
+      // Ignore abort errors
+    }
+    activeStreamingRequests.delete(requestId);
+    log('warn', `[WS] Evicted oldest streaming request ${requestId} due to size limit`);
+  }
+}
+
+// ============================================
+// Connection Limits & IP Tracking
+// ============================================
+
+/**
+ * Check connection limits and return whether a new connection should be allowed
+ */
+export function checkConnectionLimits(clientIp) {
+  // Check total connection limit
+  if (wsClients.size >= WS_CONNECTION_LIMITS.maxTotal) {
+    return { allowed: false, reason: 'Server at capacity' };
+  }
+
+  // Check per-IP limit
+  const existingConnections = wsConnectionsByIp.get(clientIp);
+  if (existingConnections && existingConnections.size >= WS_CONNECTION_LIMITS.maxPerIp) {
+    return { allowed: false, reason: 'Connection limit exceeded for this IP' };
+  }
+
+  return { allowed: true };
+}
+
+/**
+ * Track a new connection for an IP address
+ */
+export function trackConnection(clientId, clientIp) {
+  if (!wsConnectionsByIp.has(clientIp)) {
+    wsConnectionsByIp.set(clientIp, new Set());
+  }
+  wsConnectionsByIp.get(clientIp).add(clientId);
+}
+
+/**
+ * Remove a connection from IP tracking
+ */
+function untrackConnection(clientId, clientIp) {
+  const ipConnections = wsConnectionsByIp.get(clientIp);
+  if (ipConnections) {
+    ipConnections.delete(clientId);
+    if (ipConnections.size === 0) {
+      wsConnectionsByIp.delete(clientIp);
+    }
+  }
+}
+
+// ============================================
+// Client Cleanup
+// ============================================
+
+/**
+ * Clean up all resources associated with a client
+ * Called on disconnect (close/error) to prevent memory leaks
+ * Safe to call multiple times — early-returns if already cleaned up (H-24)
+ *
+ * @param {string} clientId
+ * @param {string} clientIp
+ * @param {Function} cleanupSyncForClient - callback to clean sync delta throttles for this client
+ */
+export function cleanupClient(clientId, clientIp, cleanupSyncForClient) {
+  // 1. Remove from clients map (guard against double-cleanup race condition)
+  const client = wsClients.get(clientId);
+  if (!client) return; // Already cleaned up
+  wsClients.delete(clientId);
+
+  // 2. Remove rate limit entry
+  wsRateLimits.delete(clientId);
+
+  // 3. Remove from IP tracking
+  if (clientIp) {
+    untrackConnection(clientId, clientIp);
+  }
+
+  // 4. Abort any active streaming requests for this client
+  let abortedCount = 0;
+  for (const [requestId, request] of activeStreamingRequests) {
+    if (request.clientId === clientId) {
+      try {
+        request.abortController.abort();
+      } catch (abortError) {
+        // Ignore abort errors
+      }
+      activeStreamingRequests.delete(requestId);
+      abortedCount++;
+    }
+  }
+
+  // 5. Clean up any sync delta throttles for this client
+  if (cleanupSyncForClient) {
+    cleanupSyncForClient(clientId);
+  }
+
+  // 6. Trim streaming requests map after cleanup
+  trimStreamingRequests();
+
+  log('debug', `[WS] Cleaned up client ${clientId}: rate limit removed, ${abortedCount} active requests aborted`);
+}
+
+// ============================================
+// Heartbeat
+// ============================================
+
+/**
+ * Start server-initiated heartbeat interval
+ * Pings all connected clients and terminates those that miss too many pongs
+ * @param {Function} cleanupSyncForClient - callback to clean sync delta throttles
+ * @returns {NodeJS.Timeout} interval handle
+ */
+export function startHeartbeat(cleanupSyncForClient) {
+  return setInterval(() => {
+    for (const [clientId, client] of wsClients) {
+      if (client.ws.readyState === WebSocket.OPEN) {
+        // Check if client missed too many pongs
+        if (client.missedPongs >= MAX_MISSED_PONGS) {
+          log('warn', `[WS] Client ${clientId} missed ${client.missedPongs} pongs, closing connection`);
+          client.ws.terminate();
+          cleanupClient(clientId, client.clientIp, cleanupSyncForClient);
+          continue;
+        }
+
+        // Increment missed pongs counter (will be reset when pong is received)
+        client.missedPongs++;
+
+        // Send ping
+        try {
+          client.ws.ping();
+        } catch (err) {
+          log('warn', `[WS] Failed to ping client ${clientId}:`, err.message);
+        }
+      }
+    }
+  }, HEARTBEAT_INTERVAL_MS);
+}
+
+/**
+ * Start periodic cleanup of dead connections (every 1 minute)
+ * @param {Function} cleanupSyncForClient - callback to clean sync delta throttles
+ * @returns {NodeJS.Timeout} interval handle
+ */
+export function startDeadConnectionCleanup(cleanupSyncForClient) {
+  return setInterval(() => {
+    const now = Date.now();
+    const timeout = 5 * 60 * 1000;
+
+    for (const [clientId, client] of wsClients) {
+      if (now - client.lastPing > timeout) {
+        log('info', `[WS] Cleaning up stale connection: ${clientId}`);
+        client.ws.terminate();
+        cleanupClient(clientId, client.clientIp, cleanupSyncForClient);
+      }
+    }
+  }, 60 * 1000);
+}
+
+/**
+ * Start periodic cleanup of stale rate limit entries (every 5 minutes)
+ * @returns {NodeJS.Timeout} interval handle
+ */
+export function startRateLimitCleanup() {
+  return setInterval(() => {
+    cleanupStaleRateLimits();
+  }, 5 * 60 * 1000);
+}

package/server/websocket/index.js

@@ -0,0 +1,215 @@
+/**
+ * WebSocket Module - Real-time bidirectional communication
+ *
+ * Thin orchestrator: creates the WebSocket.Server, delegates to submodules
+ * for connection management, message routing, broadcasting, and sync.
+ */
+
+import { WebSocketServer } from 'ws';
+import { log } from '../utils.js';
+import { ALLOWED_ORIGINS, SESSION_USER } from '../config.js';
+import { verifyWebSocketToken, isAuthEnabled } from '../middleware/auth.js';
+
+// Connections
+import {
+  wsClients,
+  generateClientId,
+  safeSend,
+  checkWsRateLimit,
+  checkConnectionLimits,
+  trackConnection,
+  cleanupClient,
+  startHeartbeat,
+  startDeadConnectionCleanup,
+  startRateLimitCleanup,
+  MAX_MESSAGE_SIZE,
+} from './connections.js';
+
+// Routing
+import { handleWebSocketMessage } from './routing.js';
+
+// Sync
+import { cleanupSyncForClient } from './sync.js';
+
+// Re-export public API (same interface as the original websocket.js)
+export { wsClients } from './connections.js';
+export {
+  broadcast,
+  broadcastToAll,
+  sendToClient,
+  generateMessageId,
+  broadcastSyncMessage,
+  broadcastSyncThinking,
+  broadcastSyncTool,
+  broadcastSyncComplete,
+  broadcastOpenClawPush,
+  trackBroadcast,
+} from './broadcast.js';
+export {
+  broadcastSyncDelta,
+  cleanupSyncDeltaThrottle,
+} from './sync.js';
+
+// ============================================
+// Origin Verification
+// ============================================
+
+function verifyWebSocketOrigin(origin) {
+  // Reject null/undefined origins (L-02: WebSocket origin validation)
+  if (!origin) return false;
+  try {
+    const originUrl = new URL(origin);
+    // Allow localhost, configured origins, and Tailscale domains
+    return ALLOWED_ORIGINS.some(allowed => {
+      const allowedUrl = new URL(allowed);
+      return originUrl.hostname === allowedUrl.hostname;
+    }) || originUrl.hostname === 'localhost'
+      || originUrl.hostname === '127.0.0.1'
+      || originUrl.hostname.endsWith('.ts.net');
+  } catch {
+    return false;
+  }
+}
+
+// ============================================
+// Setup WebSocket Server
+// ============================================
+
+export function setupWebSocket(server, requestHelpers, saveMessageToSync) {
+  const wss = new WebSocketServer({
+    noServer: true,
+    maxPayload: MAX_MESSAGE_SIZE // 1MB limit
+  });
+
+  // Handle upgrade requests for /ws path (noServer mode to avoid
+  // aborting other WebSocket paths like /api/realtime and /gateway)
+  server.on('upgrade', (request, socket, head) => {
+    const { pathname } = new URL(request.url, `http://${request.headers.host}`);
+
+    if (pathname !== '/ws') return; // Let other handlers take it
+
+    const origin = request.headers.origin;
+    const isValidOrigin = verifyWebSocketOrigin(origin);
+    if (!isValidOrigin) {
+      log('warn', `[WS] Rejected connection from invalid origin: ${origin}`);
+      socket.write('HTTP/1.1 403 Forbidden\r\n\r\n');
+      socket.destroy();
+      return;
+    }
+
+    const clientIp = request.socket.remoteAddress;
+    const limitCheck = checkConnectionLimits(clientIp);
+    if (!limitCheck.allowed) {
+      log('warn', `[WS] Rejected connection from ${clientIp}: ${limitCheck.reason}`);
+      socket.write('HTTP/1.1 429 Too Many Requests\r\n\r\n');
+      socket.destroy();
+      return;
+    }
+
+    if (isAuthEnabled()) {
+      const url = new URL(request.url, `http://${request.headers.host}`);
+      const queryToken = url.searchParams.get('token');
+      const authHeader = request.headers.authorization;
+      const token = queryToken || (authHeader?.startsWith('Bearer ') ? authHeader.slice(7) : authHeader);
+
+      if (!verifyWebSocketToken(token)) {
+        log('warn', `[WS] Rejected connection from ${clientIp}: invalid or missing auth token`);
+        socket.write('HTTP/1.1 401 Unauthorized\r\n\r\n');
+        socket.destroy();
+        return;
+      }
+    }
+
+    wss.handleUpgrade(request, socket, head, (ws) => {
+      wss.emit('connection', ws, request);
+    });
+  });
+
+  wss.on('connection', (ws, req) => {
+    const clientId = generateClientId();
+    const clientIp = req.socket.remoteAddress;
+
+    // Track connection for IP-based limits
+    trackConnection(clientId, clientIp);
+
+    log('info', `[WS] Client connected: ${clientId} from ${clientIp}`);
+
+    wsClients.set(clientId, {
+      ws,
+      lastPing: Date.now(),
+      lastPong: Date.now(), // Track pong responses for server-initiated heartbeat
+      missedPongs: 0,       // Count of missed pong responses
+      sessionUser: SESSION_USER,
+      connectedAt: Date.now(),
+      clientIp
+    });
+
+    // Handle pong responses from server-initiated pings
+    ws.on('pong', () => {
+      const client = wsClients.get(clientId);
+      if (client) {
+        client.lastPong = Date.now();
+        client.missedPongs = 0;
+      }
+    });
+
+    safeSend(ws, {
+      type: 'connected',
+      clientId,
+      timestamp: new Date().toISOString()
+    });
+
+    ws.on('message', async (data) => {
+      // Manual message size validation (defense in depth with maxPayload)
+      if (data.length > MAX_MESSAGE_SIZE) {
+        log('warn', `[WS] Message too large from ${clientId}: ${data.length} bytes`);
+        safeSend(ws, {
+          type: 'error',
+          error: `Message too large. Maximum size is ${MAX_MESSAGE_SIZE / 1024 / 1024}MB.`
+        });
+        return;
+      }
+
+      if (!checkWsRateLimit(clientId)) {
+        log('warn', `[WS] Rate limit exceeded for ${clientId}`);
+        safeSend(ws, {
+          type: 'error',
+          error: 'Rate limit exceeded. Please slow down.'
+        });
+        return;
+      }
+
+      try {
+        const message = JSON.parse(data.toString());
+        await handleWebSocketMessage(clientId, message, requestHelpers, saveMessageToSync);
+      } catch (err) {
+        log('error', `[WS] Error handling message from ${clientId}:`, err.message);
+        safeSend(ws, {
+          type: 'error',
+          error: 'Invalid message format'
+        });
+      }
+    });
+
+    ws.on('close', (code) => {
+      log('info', `[WS] Client disconnected: ${clientId} (code: ${code})`);
+      cleanupClient(clientId, clientIp, cleanupSyncForClient);
+    });
+
+    ws.on('error', (err) => {
+      log('error', `[WS] Error for ${clientId}:`, err.message);
+      cleanupClient(clientId, clientIp, cleanupSyncForClient);
+    });
+  });
+
+  // Server-initiated heartbeat (ping/pong)
+  startHeartbeat(cleanupSyncForClient);
+
+  // Periodic cleanup of dead connections (every 1 minute)
+  startDeadConnectionCleanup(cleanupSyncForClient);
+
+  // Periodic cleanup of stale rate limit entries (every 5 minutes)
+  startRateLimitCleanup();
+
+  return wss;
+}