WebSocket Chat Example #

// server.ts
import { createServer, ServerProtocol } from "verb";
import { cors, json, staticFiles } from "verb/middleware";
import { ChatServer } from "./src/ChatServer";
import { authRouter } from "./src/routes/auth";
import { chatRouter } from "./src/routes/chat";

const app = createServer(ServerProtocol.WEBSOCKET);

// Initialize chat server
const chatServer = new ChatServer();

// Middleware
app.use(cors({
  origin: ["http://localhost:3000"],
  credentials: true
}));
app.use(json());
app.use(staticFiles({ root: "./public" }));

// HTTP routes
app.use("/api/auth", authRouter);
app.use("/api/chat", chatRouter);

// Serve main chat page
app.get("/", (req, res) => {
  res.sendFile("./public/index.html");
});

// WebSocket configuration
app.withOptions({
  port: 3000,
  websocket: {
    message: (ws, message) => chatServer.handleMessage(ws, message),
    open: (ws) => chatServer.handleConnection(ws),
    close: (ws, code, reason) => chatServer.handleDisconnection(ws, code, reason),
    drain: (ws) => console.log("WebSocket buffer drained"),
    ping: (ws, data) => ws.pong(data),
    
    // Performance settings
    maxCompressedSize: 64 * 1024 * 1024,
    maxBackpressure: 64 * 1024 * 1024,
    compression: "shared"
  }
});

app.listen(3000);
console.log("🚀 Chat server running on http://localhost:3000");

// src/ChatServer.ts
import { Database } from "bun:sqlite";
import { verify } from "jsonwebtoken";
import { ChatRoom } from "./ChatRoom";
import { User, Message, ChatEvent } from "./types";

export class ChatServer {
  private clients = new Map<WebSocket, User>();
  private rooms = new Map<string, ChatRoom>();
  private db: Database;
  
  constructor() {
    this.db = new Database("chat.db");
    this.initializeDatabase();
    this.createDefaultRooms();
  }
  
  private initializeDatabase() {
    // Create tables
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS users (
        id TEXT PRIMARY KEY,
        username TEXT UNIQUE NOT NULL,
        email TEXT UNIQUE NOT NULL,
        password_hash TEXT NOT NULL,
        avatar_url TEXT,
        role TEXT DEFAULT 'user',
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        last_seen DATETIME DEFAULT CURRENT_TIMESTAMP
      );
      
      CREATE TABLE IF NOT EXISTS rooms (
        id TEXT PRIMARY KEY,
        name TEXT NOT NULL,
        description TEXT,
        is_private BOOLEAN DEFAULT FALSE,
        created_by TEXT,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY (created_by) REFERENCES users (id)
      );
      
      CREATE TABLE IF NOT EXISTS messages (
        id TEXT PRIMARY KEY,
        room_id TEXT NOT NULL,
        user_id TEXT NOT NULL,
        content TEXT NOT NULL,
        message_type TEXT DEFAULT 'text',
        reply_to TEXT,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        edited_at DATETIME,
        FOREIGN KEY (room_id) REFERENCES rooms (id),
        FOREIGN KEY (user_id) REFERENCES users (id),
        FOREIGN KEY (reply_to) REFERENCES messages (id)
      );
      
      CREATE TABLE IF NOT EXISTS room_members (
        room_id TEXT,
        user_id TEXT,
        role TEXT DEFAULT 'member',
        joined_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (room_id, user_id),
        FOREIGN KEY (room_id) REFERENCES rooms (id),
        FOREIGN KEY (user_id) REFERENCES users (id)
      );
    `);
  }
  
  private createDefaultRooms() {
    const defaultRooms = [
      { id: "general", name: "General", description: "General discussion" },
      { id: "random", name: "Random", description: "Random topics" },
      { id: "tech", name: "Tech Talk", description: "Technology discussions" }
    ];
    
    for (const room of defaultRooms) {
      this.db.query(`
        INSERT OR IGNORE INTO rooms (id, name, description, is_private)
        VALUES (?, ?, ?, ?)
      `).run(room.id, room.name, room.description, false);
      
      if (!this.rooms.has(room.id)) {
        this.rooms.set(room.id, new ChatRoom(room.id, room.name, this.db));
      }
    }
  }
  
  async handleConnection(ws: WebSocket) {
    console.log("New WebSocket connection");
    
    // Send welcome message
    this.sendToClient(ws, {
      type: "connection_established",
      data: {
        message: "Connected to chat server",
        timestamp: new Date().toISOString()
      }
    });
  }
  
  async handleMessage(ws: WebSocket, message: string) {
    try {
      const data = JSON.parse(message);
      
      switch (data.type) {
        case "authenticate":
          await this.handleAuthentication(ws, data);
          break;
          
        case "join_room":
          await this.handleJoinRoom(ws, data);
          break;
          
        case "leave_room":
          await this.handleLeaveRoom(ws, data);
          break;
          
        case "send_message":
          await this.handleSendMessage(ws, data);
          break;
          
        case "typing_start":
          await this.handleTypingStart(ws, data);
          break;
          
        case "typing_stop":
          await this.handleTypingStop(ws, data);
          break;
          
        case "private_message":
          await this.handlePrivateMessage(ws, data);
          break;
          
        case "get_room_history":
          await this.handleGetRoomHistory(ws, data);
          break;
          
        case "edit_message":
          await this.handleEditMessage(ws, data);
          break;
          
        case "delete_message":
          await this.handleDeleteMessage(ws, data);
          break;
          
        default:
          this.sendError(ws, "Unknown message type", data.type);
      }
    } catch (error) {
      console.error("Error handling message:", error);
      this.sendError(ws, "Invalid message format");
    }
  }
  
  async handleAuthentication(ws: WebSocket, data: any) {
    try {
      const { token } = data;
      
      if (!token) {
        return this.sendError(ws, "Authentication token required");
      }
      
      const decoded = verify(token, process.env.JWT_SECRET!) as any;
      const user = this.db.query(`
        SELECT id, username, email, avatar_url, role 
        FROM users 
        WHERE id = ?
      `).get(decoded.userId) as User;
      
      if (!user) {
        return this.sendError(ws, "User not found");
      }
      
      // Update last seen
      this.db.query(`
        UPDATE users SET last_seen = CURRENT_TIMESTAMP WHERE id = ?
      `).run(user.id);
      
      // Store user connection
      this.clients.set(ws, user);
      
      // Send authentication success
      this.sendToClient(ws, {
        type: "authenticated",
        data: {
          user: {
            id: user.id,
            username: user.username,
            avatar_url: user.avatar_url,
            role: user.role
          },
          rooms: this.getAvailableRooms(user.id)
        }
      });
      
      // Broadcast user online status
      this.broadcastUserStatus(user, "online");
      
      console.log(`User ${user.username} authenticated`);
    } catch (error) {
      this.sendError(ws, "Invalid authentication token");
    }
  }
  
  async handleJoinRoom(ws: WebSocket, data: any) {
    const user = this.clients.get(ws);
    if (!user) {
      return this.sendError(ws, "Authentication required");
    }
    
    const { roomId } = data;
    const room = this.rooms.get(roomId);
    
    if (!room) {
      return this.sendError(ws, "Room not found");
    }
    
    // Check if user can join room
    if (!await this.canUserJoinRoom(user.id, roomId)) {
      return this.sendError(ws, "Access denied to room");
    }
    
    // Add user to room
    room.addUser(ws, user);
    
    // Add to database
    this.db.query(`
      INSERT OR REPLACE INTO room_members (room_id, user_id, joined_at)
      VALUES (?, ?, CURRENT_TIMESTAMP)
    `).run(roomId, user.id);
    
    // Send join confirmation
    this.sendToClient(ws, {
      type: "room_joined",
      data: {
        roomId,
        roomName: room.name,
        members: room.getMembers()
      }
    });
    
    // Broadcast to room that user joined
    room.broadcast({
      type: "user_joined",
      data: {
        user: {
          id: user.id,
          username: user.username,
          avatar_url: user.avatar_url
        },
        roomId,
        timestamp: new Date().toISOString()
      }
    }, ws);
    
    console.log(`User ${user.username} joined room ${roomId}`);
  }
  
  async handleSendMessage(ws: WebSocket, data: any) {
    const user = this.clients.get(ws);
    if (!user) {
      return this.sendError(ws, "Authentication required");
    }
    
    const { roomId, content, messageType = "text", replyTo } = data;
    const room = this.rooms.get(roomId);
    
    if (!room) {
      return this.sendError(ws, "Room not found");
    }
    
    if (!room.hasUser(ws)) {
      return this.sendError(ws, "You are not a member of this room");
    }
    
    // Validate message content
    if (!content || content.trim().length === 0) {
      return this.sendError(ws, "Message content cannot be empty");
    }
    
    if (content.length > 2000) {
      return this.sendError(ws, "Message too long (max 2000 characters)");
    }
    
    // Create message
    const messageId = crypto.randomUUID();
    const timestamp = new Date().toISOString();
    
    // Save to database
    this.db.query(`
      INSERT INTO messages (id, room_id, user_id, content, message_type, reply_to, created_at)
      VALUES (?, ?, ?, ?, ?, ?, ?)
    `).run(messageId, roomId, user.id, content, messageType, replyTo, timestamp);
    
    const message: Message = {
      id: messageId,
      roomId,
      userId: user.id,
      username: user.username,
      avatar_url: user.avatar_url,
      content,
      messageType,
      replyTo,
      timestamp,
      edited: false
    };
    
    // Broadcast message to room
    room.broadcast({
      type: "new_message",
      data: { message }
    });
    
    console.log(`Message sent in room ${roomId} by ${user.username}`);
  }
  
  async handleTypingStart(ws: WebSocket, data: any) {
    const user = this.clients.get(ws);
    if (!user) return;
    
    const { roomId } = data;
    const room = this.rooms.get(roomId);
    
    if (room && room.hasUser(ws)) {
      room.broadcast({
        type: "user_typing",
        data: {
          userId: user.id,
          username: user.username,
          roomId,
          typing: true
        }
      }, ws);
    }
  }
  
  async handleTypingStop(ws: WebSocket, data: any) {
    const user = this.clients.get(ws);
    if (!user) return;
    
    const { roomId } = data;
    const room = this.rooms.get(roomId);
    
    if (room && room.hasUser(ws)) {
      room.broadcast({
        type: "user_typing",
        data: {
          userId: user.id,
          username: user.username,
          roomId,
          typing: false
        }
      }, ws);
    }
  }
  
  async handlePrivateMessage(ws: WebSocket, data: any) {
    const sender = this.clients.get(ws);
    if (!sender) {
      return this.sendError(ws, "Authentication required");
    }
    
    const { targetUserId, content } = data;
    
    // Find target user's websocket
    let targetWs: WebSocket | null = null;
    for (const [socket, user] of this.clients) {
      if (user.id === targetUserId) {
        targetWs = socket;
        break;
      }
    }
    
    if (!targetWs) {
      return this.sendError(ws, "Target user not online");
    }
    
    const message = {
      type: "private_message",
      data: {
        from: {
          id: sender.id,
          username: sender.username,
          avatar_url: sender.avatar_url
        },
        content,
        timestamp: new Date().toISOString()
      }
    };
    
    // Send to target user
    this.sendToClient(targetWs, message);
    
    // Send confirmation to sender
    this.sendToClient(ws, {
      type: "private_message_sent",
      data: {
        targetUserId,
        content,
        timestamp: message.data.timestamp
      }
    });
  }
  
  async handleGetRoomHistory(ws: WebSocket, data: any) {
    const user = this.clients.get(ws);
    if (!user) {
      return this.sendError(ws, "Authentication required");
    }
    
    const { roomId, limit = 50, before } = data;
    
    let query = `
      SELECT m.*, u.username, u.avatar_url
      FROM messages m
      JOIN users u ON m.user_id = u.id
      WHERE m.room_id = ?
    `;
    
    const params = [roomId];
    
    if (before) {
      query += ` AND m.created_at < ?`;
      params.push(before);
    }
    
    query += ` ORDER BY m.created_at DESC LIMIT ?`;
    params.push(limit);
    
    const messages = this.db.query(query).all(...params);
    
    this.sendToClient(ws, {
      type: "room_history",
      data: {
        roomId,
        messages: messages.reverse().map(msg => ({
          id: msg.id,
          roomId: msg.room_id,
          userId: msg.user_id,
          username: msg.username,
          avatar_url: msg.avatar_url,
          content: msg.content,
          messageType: msg.message_type,
          replyTo: msg.reply_to,
          timestamp: msg.created_at,
          edited: !!msg.edited_at
        }))
      }
    });
  }
  
  handleDisconnection(ws: WebSocket, code: number, reason: string) {
    const user = this.clients.get(ws);
    if (user) {
      console.log(`User ${user.username} disconnected`);
      
      // Remove from all rooms
      for (const room of this.rooms.values()) {
        if (room.hasUser(ws)) {
          room.removeUser(ws);
          
          // Broadcast user left
          room.broadcast({
            type: "user_left",
            data: {
              user: {
                id: user.id,
                username: user.username
              },
              timestamp: new Date().toISOString()
            }
          });
        }
      }
      
      // Broadcast user offline status
      this.broadcastUserStatus(user, "offline");
      
      // Remove from clients
      this.clients.delete(ws);
      
      // Update last seen in database
      this.db.query(`
        UPDATE users SET last_seen = CURRENT_TIMESTAMP WHERE id = ?
      `).run(user.id);
    }
  }
  
  private sendToClient(ws: WebSocket, message: ChatEvent) {
    if (ws.readyState === 1) {
      ws.send(JSON.stringify(message));
    }
  }
  
  private sendError(ws: WebSocket, message: string, code?: string) {
    this.sendToClient(ws, {
      type: "error",
      data: { message, code }
    });
  }
  
  private broadcastUserStatus(user: User, status: "online" | "offline") {
    const statusMessage = {
      type: "user_status",
      data: {
        userId: user.id,
        username: user.username,
        status,
        timestamp: new Date().toISOString()
      }
    };
    
    // Broadcast to all connected clients
    for (const [ws, client] of this.clients) {
      if (client.id !== user.id) {
        this.sendToClient(ws, statusMessage);
      }
    }
  }
  
  private getAvailableRooms(userId: string) {
    const rooms = this.db.query(`
      SELECT r.id, r.name, r.description, r.is_private,
             (SELECT COUNT(*) FROM room_members rm WHERE rm.room_id = r.id) as member_count
      FROM rooms r
      LEFT JOIN room_members rm ON r.id = rm.room_id AND rm.user_id = ?
      WHERE r.is_private = FALSE OR rm.user_id IS NOT NULL
      ORDER BY r.name
    `).all(userId);
    
    return rooms;
  }
  
  private async canUserJoinRoom(userId: string, roomId: string): Promise<boolean> {
    const room = this.db.query(`
      SELECT is_private FROM rooms WHERE id = ?
    `).get(roomId) as any;
    
    if (!room) return false;
    
    // Public rooms can be joined by anyone
    if (!room.is_private) return true;
    
    // Private rooms require membership
    const membership = this.db.query(`
      SELECT user_id FROM room_members WHERE room_id = ? AND user_id = ?
    `).get(roomId, userId);
    
    return !!membership;
  }
}

// src/ChatRoom.ts
import { Database } from "bun:sqlite";
import { User, ChatEvent } from "./types";

export class ChatRoom {
  private users = new Map<WebSocket, User>();
  private typingUsers = new Set<string>();
  
  constructor(
    public id: string,
    public name: string,
    private db: Database
  ) {}
  
  addUser(ws: WebSocket, user: User) {
    this.users.set(ws, user);
  }
  
  removeUser(ws: WebSocket) {
    const user = this.users.get(ws);
    if (user) {
      this.users.delete(ws);
      this.typingUsers.delete(user.id);
    }
  }
  
  hasUser(ws: WebSocket): boolean {
    return this.users.has(ws);
  }
  
  getMembers() {
    return Array.from(this.users.values()).map(user => ({
      id: user.id,
      username: user.username,
      avatar_url: user.avatar_url,
      role: user.role
    }));
  }
  
  broadcast(message: ChatEvent, excludeWs?: WebSocket) {
    for (const [ws, user] of this.users) {
      if (ws !== excludeWs && ws.readyState === 1) {
        ws.send(JSON.stringify(message));
      }
    }
  }
  
  getUserCount(): number {
    return this.users.size;
  }
  
  getUser(ws: WebSocket): User | undefined {
    return this.users.get(ws);
  }
}

// src/types.ts
export interface User {
  id: string;
  username: string;
  email: string;
  avatar_url?: string;
  role: "user" | "moderator" | "admin";
}

export interface Message {
  id: string;
  roomId: string;
  userId: string;
  username: string;
  avatar_url?: string;
  content: string;
  messageType: "text" | "image" | "file" | "system";
  replyTo?: string;
  timestamp: string;
  edited: boolean;
}

export interface ChatEvent {
  type: string;
  data: any;
}

export interface Room {
  id: string;
  name: string;
  description?: string;
  isPrivate: boolean;
  memberCount: number;
}