Building Real-Time Applications with WebSockets

// server.ts
import { WebSocketServer, WebSocket } from 'ws';
import { createServer } from 'http';
import { v4 as uuid } from 'uuid';
 
interface Client {
  id: string;
  socket: WebSocket;
  userId?: string;
  rooms: Set<string>;
  lastPing: number;
}
 
class WebSocketManager {
  private wss: WebSocketServer;
  private clients: Map<string, Client> = new Map();
  private rooms: Map<string, Set<string>> = new Map();
 
  constructor(server: ReturnType<typeof createServer>) {
    this.wss = new WebSocketServer({ server });
    this.setupConnectionHandler();
    this.startHeartbeat();
  }
 
  private setupConnectionHandler() {
    this.wss.on('connection', (socket, request) => {
      const clientId = uuid();
      const client: Client = {
        id: clientId,
        socket,
        rooms: new Set(),
        lastPing: Date.now(),
      };
 
      this.clients.set(clientId, client);
      console.log(`Client connected: ${clientId}`);
 
      // Send connection acknowledgment
      this.send(socket, {
        type: 'connected',
        clientId,
      });
 
      // Handle incoming messages
      socket.on('message', (data) => {
        try {
          const message = JSON.parse(data.toString());
          this.handleMessage(client, message);
        } catch (error) {
          console.error('Invalid message format:', error);
        }
      });
 
      // Handle disconnection
      socket.on('close', () => {
        this.handleDisconnect(client);
      });
 
      // Handle errors
      socket.on('error', (error) => {
        console.error(`Client ${clientId} error:`, error);
      });
 
      // Respond to pings
      socket.on('pong', () => {
        client.lastPing = Date.now();
      });
    });
  }
 
  private handleMessage(client: Client, message: any) {
    switch (message.type) {
      case 'authenticate':
        this.authenticateClient(client, message.token);
        break;
 
      case 'join_room':
        this.joinRoom(client, message.room);
        break;
 
      case 'leave_room':
        this.leaveRoom(client, message.room);
        break;
 
      case 'broadcast_room':
        this.broadcastToRoom(message.room, {
          type: 'room_message',
          room: message.room,
          from: client.userId || client.id,
          content: message.content,
          timestamp: Date.now(),
        }, client.id);
        break;
 
      case 'ping':
        this.send(client.socket, { type: 'pong' });
        break;
 
      default:
        console.log('Unknown message type:', message.type);
    }
  }
 
  private authenticateClient(client: Client, token: string) {
    // Verify token and extract user info
    try {
      const user = this.verifyToken(token);
      client.userId = user.id;
      this.send(client.socket, {
        type: 'authenticated',
        userId: user.id,
      });
    } catch (error) {
      this.send(client.socket, {
        type: 'auth_error',
        message: 'Invalid token',
      });
    }
  }
 
  private joinRoom(client: Client, roomId: string) {
    // Add client to room
    if (!this.rooms.has(roomId)) {
      this.rooms.set(roomId, new Set());
    }
    this.rooms.get(roomId)!.add(client.id);
    client.rooms.add(roomId);
 
    // Notify client
    this.send(client.socket, {
      type: 'joined_room',
      room: roomId,
    });
 
    // Notify others in room
    this.broadcastToRoom(roomId, {
      type: 'user_joined',
      room: roomId,
      userId: client.userId || client.id,
    }, client.id);
  }
 
  private leaveRoom(client: Client, roomId: string) {
    const room = this.rooms.get(roomId);
    if (room) {
      room.delete(client.id);
      client.rooms.delete(roomId);
 
      // Notify others
      this.broadcastToRoom(roomId, {
        type: 'user_left',
        room: roomId,
        userId: client.userId || client.id,
      });
 
      // Clean up empty rooms
      if (room.size === 0) {
        this.rooms.delete(roomId);
      }
    }
  }
 
  private handleDisconnect(client: Client) {
    // Leave all rooms
    for (const roomId of client.rooms) {
      this.leaveRoom(client, roomId);
    }
 
    // Remove from clients
    this.clients.delete(client.id);
    console.log(`Client disconnected: ${client.id}`);
  }
 
  private broadcastToRoom(roomId: string, message: any, excludeClientId?: string) {
    const room = this.rooms.get(roomId);
    if (!room) return;
 
    for (const clientId of room) {
      if (clientId === excludeClientId) continue;
 
      const client = this.clients.get(clientId);
      if (client && client.socket.readyState === WebSocket.OPEN) {
        this.send(client.socket, message);
      }
    }
  }
 
  private send(socket: WebSocket, message: any) {
    if (socket.readyState === WebSocket.OPEN) {
      socket.send(JSON.stringify(message));
    }
  }
 
  private startHeartbeat() {
    setInterval(() => {
      const now = Date.now();
      const timeout = 30000; // 30 seconds
 
      for (const [clientId, client] of this.clients) {
        if (now - client.lastPing > timeout) {
          // Connection is dead
          console.log(`Client ${clientId} timed out`);
          client.socket.terminate();
          this.handleDisconnect(client);
        } else if (client.socket.readyState === WebSocket.OPEN) {
          // Send ping
          client.socket.ping();
        }
      }
    }, 10000); // Check every 10 seconds
  }
 
  private verifyToken(token: string): { id: string } {
    // Implement your JWT verification here
    return { id: 'user-123' };
  }
}
 
// Start server
const server = createServer();
const wsManager = new WebSocketManager(server);
server.listen(8080, () => {
  console.log('WebSocket server running on port 8080');
});

// useWebSocket.ts
import { useEffect, useRef, useCallback, useState } from 'react';
 
interface UseWebSocketOptions {
  url: string;
  onMessage?: (message: any) => void;
  onOpen?: () => void;
  onClose?: () => void;
  onError?: (error: Event) => void;
  reconnect?: boolean;
  reconnectAttempts?: number;
  reconnectInterval?: number;
}
 
interface UseWebSocketReturn {
  send: (message: any) => void;
  isConnected: boolean;
  lastMessage: any;
}
 
export function useWebSocket({
  url,
  onMessage,
  onOpen,
  onClose,
  onError,
  reconnect = true,
  reconnectAttempts = 5,
  reconnectInterval = 3000,
}: UseWebSocketOptions): UseWebSocketReturn {
  const [isConnected, setIsConnected] = useState(false);
  const [lastMessage, setLastMessage] = useState<any>(null);
 
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectCountRef = useRef(0);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout>();
 
  const connect = useCallback(() => {
    try {
      const ws = new WebSocket(url);
 
      ws.onopen = () => {
        setIsConnected(true);
        reconnectCountRef.current = 0;
        onOpen?.();
 
        // Start client-side heartbeat
        const heartbeat = setInterval(() => {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: 'ping' }));
          }
        }, 25000);
 
        ws.addEventListener('close', () => clearInterval(heartbeat));
      };
 
      ws.onmessage = (event) => {
        try {
          const message = JSON.parse(event.data);
          setLastMessage(message);
          onMessage?.(message);
        } catch (error) {
          console.error('Failed to parse message:', error);
        }
      };
 
      ws.onclose = () => {
        setIsConnected(false);
        onClose?.();
 
        // Attempt reconnection
        if (reconnect && reconnectCountRef.current < reconnectAttempts) {
          reconnectCountRef.current++;
          const delay = reconnectInterval * Math.pow(2, reconnectCountRef.current - 1);
 
          console.log(`Reconnecting in ${delay}ms (attempt ${reconnectCountRef.current})`);
 
          reconnectTimeoutRef.current = setTimeout(connect, delay);
        }
      };
 
      ws.onerror = (error) => {
        onError?.(error);
      };
 
      wsRef.current = ws;
    } catch (error) {
      console.error('WebSocket connection error:', error);
    }
  }, [url, onMessage, onOpen, onClose, onError, reconnect, reconnectAttempts, reconnectInterval]);
 
  useEffect(() => {
    connect();
 
    return () => {
      clearTimeout(reconnectTimeoutRef.current);
      wsRef.current?.close();
    };
  }, [connect]);
 
  const send = useCallback((message: any) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    } else {
      console.warn('WebSocket not connected');
    }
  }, []);
 
  return { send, isConnected, lastMessage };
}
 
// Usage example
function ChatRoom({ roomId }: { roomId: string }) {
  const [messages, setMessages] = useState<Message[]>([]);
 
  const { send, isConnected } = useWebSocket({
    url: 'wss://api.example.com/ws',
    onMessage: (message) => {
      if (message.type === 'room_message') {
        setMessages(prev => [...prev, message]);
      }
    },
    onOpen: () => {
      // Join room on connect
      send({ type: 'join_room', room: roomId });
    },
  });
 
  const sendMessage = (content: string) => {
    send({
      type: 'broadcast_room',
      room: roomId,
      content,
    });
  };
 
  return (
    <div>
      <div className={`status ${isConnected ? 'connected' : 'disconnected'}`}>
        {isConnected ? 'Connected' : 'Reconnecting...'}
      </div>
      <MessageList messages={messages} />
      <MessageInput onSend={sendMessage} disabled={!isConnected} />
    </div>
  );
}

// Scaling WebSockets across multiple servers
import Redis from 'ioredis';
 
class ScaledWebSocketManager extends WebSocketManager {
  private publisher: Redis;
  private subscriber: Redis;
 
  constructor(server: ReturnType<typeof createServer>) {
    super(server);
 
    this.publisher = new Redis(process.env.REDIS_URL);
    this.subscriber = new Redis(process.env.REDIS_URL);
 
    this.setupPubSub();
  }
 
  private setupPubSub() {
    // Subscribe to all room channels
    this.subscriber.psubscribe('room:*');
 
    this.subscriber.on('pmessage', (pattern, channel, message) => {
      const roomId = channel.replace('room:', '');
      const parsed = JSON.parse(message);
 
      // Only broadcast if the message didn't originate from this server
      if (parsed.serverId !== this.serverId) {
        this.localBroadcastToRoom(roomId, parsed.message);
      }
    });
  }
 
  // Override to publish to Redis instead of local-only broadcast
  protected broadcastToRoom(roomId: string, message: any, excludeClientId?: string) {
    // Publish to Redis for other servers
    this.publisher.publish(`room:${roomId}`, JSON.stringify({
      serverId: this.serverId,
      message,
      excludeClientId,
    }));
 
    // Also broadcast locally
    this.localBroadcastToRoom(roomId, message, excludeClientId);
  }
}

// Secure WebSocket server with JWT authentication
import { WebSocketServer } from 'ws';
import { verifyJWT } from './auth';
 
const wss = new WebSocketServer({
  server,
  verifyClient: async (info, callback) => {
    try {
      // Extract token from query string or header
      const url = new URL(info.req.url!, `http://${info.req.headers.host}`);
      const token = url.searchParams.get('token');
 
      if (!token) {
        callback(false, 401, 'Unauthorized');
        return;
      }
 
      // Verify JWT
      const user = await verifyJWT(token);
 
      // Attach user to request for later use
      (info.req as any).user = user;
 
      callback(true);
    } catch (error) {
      callback(false, 401, 'Invalid token');
    }
  },
});

import { z } from 'zod';
 
// Define message schemas
const MessageSchema = z.discriminatedUnion('type', [
  z.object({
    type: z.literal('join_room'),
    room: z.string().max(100),
  }),
  z.object({
    type: z.literal('broadcast_room'),
    room: z.string().max(100),
    content: z.string().max(10000),
  }),
  z.object({
    type: z.literal('ping'),
  }),
]);
 
function handleMessage(client: Client, rawMessage: unknown) {
  const result = MessageSchema.safeParse(rawMessage);
 
  if (!result.success) {
    send(client.socket, {
      type: 'error',
      message: 'Invalid message format',
    });
    return;
  }
 
  const message = result.data;
  // Now message is properly typed and validated
  processMessage(client, message);
}

class RateLimitedWebSocketManager extends WebSocketManager {
  private messageRates: Map<string, number[]> = new Map();
  private readonly maxMessagesPerMinute = 60;
  private readonly maxConnectionsPerIP = 10;
  private connectionsByIP: Map<string, number> = new Map();
 
  protected handleConnection(socket: WebSocket, request: IncomingMessage) {
    const ip = request.socket.remoteAddress || 'unknown';
 
    // Check connection limit per IP
    const currentConnections = this.connectionsByIP.get(ip) || 0;
    if (currentConnections >= this.maxConnectionsPerIP) {
      socket.close(1008, 'Too many connections');
      return;
    }
 
    this.connectionsByIP.set(ip, currentConnections + 1);
 
    socket.on('close', () => {
      const count = this.connectionsByIP.get(ip) || 1;
      this.connectionsByIP.set(ip, count - 1);
    });
 
    super.handleConnection(socket, request);
  }
 
  protected handleMessage(client: Client, message: any) {
    // Rate limiting
    const now = Date.now();
    const rates = this.messageRates.get(client.id) || [];
 
    // Remove old timestamps
    const recentRates = rates.filter(t => now - t < 60000);
 
    if (recentRates.length >= this.maxMessagesPerMinute) {
      this.send(client.socket, {
        type: 'error',
        message: 'Rate limit exceeded',
      });
      return;
    }
 
    recentRates.push(now);
    this.messageRates.set(client.id, recentRates);
 
    super.handleMessage(client, message);
  }
}

// websocket.test.ts
import { WebSocket } from 'ws';
 
describe('WebSocket Server', () => {
  let client: WebSocket;
 
  beforeEach((done) => {
    client = new WebSocket('ws://localhost:8080');
    client.on('open', done);
  });
 
  afterEach(() => {
    client.close();
  });
 
  test('receives connected message on connection', (done) => {
    client.on('message', (data) => {
      const message = JSON.parse(data.toString());
      expect(message.type).toBe('connected');
      expect(message.clientId).toBeDefined();
      done();
    });
  });
 
  test('can join and leave rooms', async () => {
    const messages: any[] = [];
 
    client.on('message', (data) => {
      messages.push(JSON.parse(data.toString()));
    });
 
    // Wait for connected message
    await new Promise(r => setTimeout(r, 100));
 
    // Join room
    client.send(JSON.stringify({ type: 'join_room', room: 'test-room' }));
 
    await new Promise(r => setTimeout(r, 100));
 
    expect(messages.some(m => m.type === 'joined_room')).toBe(true);
  });
});