server.js

const express = require('express');
const http = require('http');
const { Server } = require('socket.io');
const crypto = require('crypto');

const app = express();
const server = http.createServer(app);
const io = new Server(server, {
  maxHttpBufferSize: 1e7, // 10 MB limit for WebSocket messages
});

app.use(express.static('public'));

// In-Memory Data Stores
const usersBySocket = new Map();
const usersById = new Map();
const groups = new Map();

io.on('connection', (socket) => {
  // Rate limiting map per socket
  const rateLimit = new Map();

  const isRateLimited = (eventName) => {
    const now = Date.now();
    const lastTime = rateLimit.get(eventName) || 0;
    if (now - lastTime < 100) { // 100ms between same events
      return true;
    }
    rateLimit.set(eventName, now);
    return false;
  };

  socket.on('register', ({ displayName, avatarColor, publicKey }) => {
    if (isRateLimited('register')) return;

    const name = (displayName || '').trim().substring(0, 30).replace(/</g, "&lt;").replace(/>/g, "&gt;");
    const sanitizedColor = sanitizeColor(avatarColor);
    if (!name || !publicKey) return;

    const userId = crypto.randomUUID();
    const user = {
      userId,
      displayName: name,
      avatarColor: sanitizedColor,
      socketId: socket.id,
      publicKey,
      status: 'online',
      joinedAt: Date.now()
    };

    usersBySocket.set(socket.id, user);
    usersById.set(userId, user);

    socket.emit('registered', { userId, displayName: name });
    io.emit('user_joined', getSafeUserRecord(user));
    
    socket.emit('active_users', Array.from(usersById.values()).map(getSafeUserRecord));
  });

  socket.on('resume_session', ({ userId, displayName, publicKey }) => {
     if (isRateLimited('resume_session')) return;
     const name = (displayName || '').trim().substring(0, 30).replace(/</g, "&lt;").replace(/>/g, "&gt;");
     if (!name || !publicKey) return;

     // Security: Generate new userId on resume to prevent hijacking
     // Client-provided userId is ignored for security
     const newUserId = crypto.randomUUID();
     
     const user = {
        userId: newUserId,
        displayName: name,
        avatarColor: '#6366f1',
        socketId: socket.id,
        publicKey,
        status: 'online',
        joinedAt: Date.now()
      };
      
      usersBySocket.set(socket.id, user);
      usersById.set(newUserId, user);
      
      // Emit the NEW userId back to client
      socket.emit('registered', { userId: newUserId, displayName: name });
      
      io.emit('user_joined', getSafeUserRecord(user));
      socket.emit('active_users', Array.from(usersById.values()).map(getSafeUserRecord));
      
      // Note: Groups are lost on session resume due to ephemeral design
      // Client will need to be re-added to groups by other members
      const myGroups = [];
      for (const group of groups.values()) {
          if (group.memberIds.has(userId)) {
              myGroups.push({ ...group, memberIds: Array.from(group.memberIds), adminIds: Array.from(group.adminIds) });
          }
      }
      socket.emit('active_groups', myGroups);

      // Re-join socket rooms for groups
      for (const group of groups.values()) {
          if (group.memberIds.has(userId)) {
              socket.join(`group_${group.id}`);
          }
      }
  });

  socket.on('update_status', ({ status }) => {
    if (isRateLimited('update_status')) return;
    const user = usersBySocket.get(socket.id);
    if (!user) return;
    
    const validStatuses = ['online', 'away', 'dnd'];
    if (!validStatuses.includes(status)) return;
    
    user.status = status;
    io.emit('user_status_changed', { userId: user.userId, status });
  });

  socket.on('send_dm', ({ toUserId, message }) => {
    if (isRateLimited('send_dm')) return;
    if (Buffer.byteLength(JSON.stringify(message)) > 10 * 1024 * 1024) return; // double check size

    const sender = usersBySocket.get(socket.id);
    const recipient = usersById.get(toUserId);
    
    if (!sender || !recipient) return;
    
    io.to(recipient.socketId).emit('receive_dm', {
      fromUserId: sender.userId,
      message,
      timestamp: Date.now()
    });
    
    socket.emit('dm_sent_receipt', { toUserId, tempId: message.tempId, timestamp: Date.now() });
  });
  
  socket.on('delete_message', ({ targetId, isGroup, messageId }) => {
     if (isRateLimited('delete_message')) return;
     const sender = usersBySocket.get(socket.id);
     if (!sender) return;

     if (isGroup) {
         socket.to(`group_${targetId}`).emit('message_deleted', { targetId, isGroup: true, messageId });
     } else {
         const recipient = usersById.get(targetId);
         if (recipient) {
             io.to(recipient.socketId).emit('message_deleted', { targetId: sender.userId, isGroup: false, messageId });
         }
     }
  });

  socket.on('create_group', ({ name, avatar, memberIds, encryptedKeys }) => {
    if (isRateLimited('create_group')) return;
    const sender = usersBySocket.get(socket.id);
    if (!sender) return;

    const groupId = crypto.randomUUID();
    const sanitizedName = (name || '').trim().substring(0, 30).replace(/</g, "&lt;").replace(/>/g, "&gt;");
    const sanitizedAvatar = sanitizeColor(avatar);

    if (!Array.isArray(memberIds) || memberIds.length === 0 || memberIds.length > 100) return;
    
    const validMemberIds = new Set();
    for (const mId of memberIds) {
      if (typeof mId === 'string' && usersById.has(mId)) {
        validMemberIds.add(mId);
      }
    }
    validMemberIds.add(sender.userId);

    const group = {
       id: groupId,
       name: sanitizedName, 
       avatar: sanitizedAvatar,
       adminIds: new Set([sender.userId]),
       memberIds: validMemberIds
    };
    groups.set(groupId, group);

    for (const mId of validMemberIds) {
       const m = usersById.get(mId);
       if (m) {
          const encKey = encryptedKeys[mId];
          if (!encKey) {
             // Skip members without encrypted keys - they can't decrypt anyway
             continue;
          }
          const s = io.sockets.sockets.get(m.socketId);
          if (s) s.join(`group_${groupId}`);

          io.to(m.socketId).emit('group_created', {
             id: groupId,
             name: group.name,
             avatar: group.avatar,
             memberIds: Array.from(group.memberIds),
             adminIds: Array.from(group.adminIds),
             groupKeyEncrypted: encKey,
             createdBy: sender.userId
          });
       }
    }
  });

  socket.on('send_group_msg', ({ groupId, message }) => {
      if (isRateLimited('send_group_msg')) return;
      if (Buffer.byteLength(JSON.stringify(message)) > 10 * 1024 * 1024) return;

      const sender = usersBySocket.get(socket.id);
      if (!sender) return;
      const group = groups.get(groupId);
      if (!group || !group.memberIds.has(sender.userId)) return;

      socket.to(`group_${groupId}`).emit('receive_group_msg', {
          groupId,
          fromUserId: sender.userId,
          message,
          timestamp: Date.now()
      });
      
      socket.emit('group_msg_sent_receipt', { groupId, tempId: message.tempId, timestamp: Date.now() });
  });

  socket.on('typing', ({ targetId, isGroup }) => {
     if (isRateLimited('typing')) return;
     const sender = usersBySocket.get(socket.id);
     if (!sender) return;
     if (isGroup) {
        socket.to(`group_${targetId}`).emit('user_typing', { userId: sender.userId, targetId, isGroup: true });
     } else {
        const recipient = usersById.get(targetId);
        if (recipient) {
           io.to(recipient.socketId).emit('user_typing', { userId: sender.userId, targetId, isGroup: false });
        }
     }
  });

  socket.on('stop_typing', ({ targetId, isGroup }) => {
     if (isRateLimited('stop_typing')) return;
     const sender = usersBySocket.get(socket.id);
     if (!sender) return;
     if (isGroup) {
        socket.to(`group_${targetId}`).emit('user_stop_typing', { userId: sender.userId, targetId, isGroup: true });
     } else {
        const recipient = usersById.get(targetId);
        if (recipient) {
           io.to(recipient.socketId).emit('user_stop_typing', { userId: sender.userId, targetId, isGroup: false });
        }
     }
  });

  socket.on('disconnect', () => {
    const user = usersBySocket.get(socket.id);
    if (user) {
      const disconnectedSocketId = socket.id;
      user.status = 'offline';
      user.lastSeen = Date.now();
      usersBySocket.delete(socket.id);
      
      setTimeout(() => {
          // Check if user reconnected by verifying their userId still exists with a different socketId
          const currentUser = usersById.get(user.userId);
          if (!currentUser || currentUser.socketId === disconnectedSocketId) {
             // User didn't reconnect or is still on the same (now disconnected) socket
             usersById.delete(user.userId);
             io.emit('user_left', { userId: user.userId, lastSeen: user.lastSeen });
          }
      }, 5000);
    }
  });
});

function getSafeUserRecord(user) {
    return {
       userId: user.userId,
       displayName: user.displayName,
       avatarColor: user.avatarColor,
       publicKey: user.publicKey,
       status: user.status,
       lastSeen: user.lastSeen || null
    };
}

function sanitizeColor(color) {
    // Only allow valid hex colors
    if (!color || typeof color !== 'string') return '#6366f1';
    const hexPattern = /^#[0-9A-Fa-f]{6}$/;
    return hexPattern.test(color) ? color : '#6366f1';
}

const PORT = process.env.PORT || 3000;
server.listen(PORT, () => {
  console.log(`Server listening on port ${PORT}`);
});