Learn_System/backend/src/ws-server.js

/**
 * WebSocket server — full real-time classroom channel.
 *
 * Client → server messages:
 *   classroom_join  { sessionId }
 *   cursor          { sessionId, x, y, pageNum }
 *   preview         { sessionId, liveId, tool, data, pageNum, cancel }
 *   hand_raise      { sessionId }
 *   hand_lower      { sessionId, targetUserId? }   teacher can lower anyone's hand
 *   page_change     { sessionId, pageNum }          teacher only
 *   page_clear      { sessionId, pageNum }          teacher only
 *   template_change { sessionId, pageNum, template } teacher only
 *   allow_draw      { sessionId, targetUserId }     teacher only
 *   revoke_draw     { sessionId, targetUserId }     teacher only
 *   mute_peer       { sessionId, targetUserId }     teacher only
 *   screen_start    { sessionId }                   teacher only
 *   screen_stop     { sessionId }                   teacher only
 *
 * Server → client delivery:
 *   All classroom events are sent via WS to connected users,
 *   with automatic SSE fallback for users not on WS.
 *   Use broadcastToSession() / emitToUser() from other modules.
 */
const { WebSocketServer } = require('ws');
const jwt = require('jsonwebtoken');
const db  = require('./db/db');
const { emit, emitToGuests } = require('./sse');

/* ── Classroom connections: userId → Set<ws> ──────────────────────────── */
const _classroomConns = new Map();

function _registerUser(ws) {
  if (!_classroomConns.has(ws.userId)) _classroomConns.set(ws.userId, new Set());
  _classroomConns.get(ws.userId).add(ws);
}

function _unregisterUser(ws) {
  const set = _classroomConns.get(ws.userId);
  if (!set) return;
  set.delete(ws);
  if (set.size === 0) _classroomConns.delete(ws.userId);
}

/* ── Delivery: WS first, SSE fallback ────────────────────────────────── */
function emitToUser(userId, data) {
  const sockets = _classroomConns.get(userId);
  if (sockets?.size) {
    const raw = JSON.stringify(data);
    for (const ws of sockets) {
      if (ws.readyState === 1) try { ws.send(raw); } catch {}
    }
    return; // delivered via WS
  }
  emit(userId, data); // SSE fallback
}

function broadcastToSession(sessionId, data, includeGuests = false) {
  const members = _getMembers(sessionId);
  if (!members) return;
  for (const uid of members.userIds) emitToUser(uid, data);
  if (includeGuests) emitToGuests(sessionId, data);
}

/* ── Session member cache (30s TTL — avoids DB per message) ───────────── */
const _cache = new Map();
const CACHE_TTL = 30_000;

function _getMembers(sessionId) {
  const c = _cache.get(sessionId);
  if (c && Date.now() - c.ts < CACHE_TTL) return c;

  const session = db.prepare(
    "SELECT class_id, teacher_id FROM classroom_sessions WHERE id=? AND status='active'"
  ).get(sessionId);
  if (!session) { _cache.delete(sessionId); return null; }

  let userIds;
  if (session.class_id) {
    const rows = db.prepare('SELECT user_id FROM class_members WHERE class_id=?').all(session.class_id);
    userIds = [session.teacher_id, ...rows.map(r => r.user_id)];
  } else {
    const rows = db.prepare('SELECT user_id FROM classroom_invites WHERE session_id=?').all(sessionId);
    userIds = [session.teacher_id, ...rows.map(r => r.user_id)];
  }

  const entry = { teacherId: session.teacher_id, classId: session.class_id, userIds, ts: Date.now() };
  _cache.set(sessionId, entry);
  return entry;
}

function _invalidateSession(sessionId) {
  _cache.delete(sessionId);
  // Cleanup draw cache entries for this session
  for (const key of _drawCache.keys()) {
    if (key.startsWith(sessionId + ':')) _drawCache.delete(key);
  }
}

/* ── Draw permission cache (10s TTL) ─────────────────────────────────── */
const _drawCache = new Map();
const DRAW_TTL   = 10_000;

function _canDraw(sessionId, userId, members) {
  if (!members) return false;
  if (members.teacherId === userId) return true;
  const key = `${sessionId}:${userId}`;
  const c = _drawCache.get(key);
  if (c && Date.now() - c.ts < DRAW_TTL) return c.allowed;
  const allowed = !!db.prepare(
    'SELECT 1 FROM classroom_draw_permissions WHERE session_id=? AND user_id=?'
  ).get(sessionId, userId);
  _drawCache.set(key, { allowed, ts: Date.now() });
  return allowed;
}

function _invalidateDrawCache(sessionId, userId) {
  _drawCache.delete(`${sessionId}:${userId}`);
}

/* ── Message handler ─────────────────────────────────────────────────── */
function _handleMessage(ws, msg) {
  const { type, sessionId } = msg;
  if (!sessionId || typeof sessionId !== 'number') return;

  const members = _getMembers(sessionId);
  if (!members) return;
  const isMember  = members.userIds.includes(ws.userId);
  const isTeacher = members.teacherId === ws.userId;

  switch (type) {

    /* ── Register in session (sends WS events instead of SSE) ── */
    case 'classroom_join':
      ws.classroomSessionId = sessionId;
      _registerUser(ws);
      break;

    /* ── Cursor position ── */
    case 'cursor': {
      if (!isMember) return;
      broadcastToSession(sessionId, {
        type: 'classroom_cursor', sessionId,
        x: msg.x, y: msg.y, pageNum: msg.pageNum || 1,
        userId: ws.userId, userName: ws.userName,
      }, true);
      break;
    }

    /* ── Stroke live preview ── */
    case 'preview': {
      if (!_canDraw(sessionId, ws.userId, members)) return;
      if (!msg.liveId && !msg.cancel) return;
      broadcastToSession(sessionId, {
        type: 'classroom_stroke_preview', sessionId,
        pageNum: msg.pageNum || 1, liveId: msg.liveId,
        tool: msg.tool, data: msg.data, cancel: msg.cancel || false,
        userId: ws.userId, userName: ws.userName,
      }, true);
      break;
    }

    /* ── Hand raise / lower ── */
    case 'hand_raise': {
      if (!isMember) return;
      try {
        db.prepare('INSERT OR IGNORE INTO classroom_hands (session_id, user_id) VALUES (?,?)').run(sessionId, ws.userId);
      } catch { return; }
      broadcastToSession(sessionId, {
        type: 'classroom_hand_raised', sessionId,
        userId: ws.userId, userName: ws.userName,
      });
      break;
    }

    case 'hand_lower': {
      if (!isMember) return;
      const targetId = (isTeacher && msg.targetUserId) ? Number(msg.targetUserId) : ws.userId;
      try {
        db.prepare('DELETE FROM classroom_hands WHERE session_id=? AND user_id=?').run(sessionId, targetId);
      } catch { return; }
      broadcastToSession(sessionId, {
        type: 'classroom_hand_lowered', sessionId, userId: targetId,
      });
      break;
    }

    /* ── Page navigation (teacher) ── */
    case 'page_change': {
      if (!isTeacher) return;
      const pageNum = Number(msg.pageNum);
      if (!pageNum || pageNum < 1) return;
      db.prepare('UPDATE classroom_sessions SET current_page=? WHERE id=?').run(pageNum, sessionId);
      broadcastToSession(sessionId, {
        type: 'classroom_page_changed', sessionId, pageNum,
      }, true);
      break;
    }

    /* ── Page clear (teacher) ── */
    case 'page_clear': {
      if (!isTeacher) return;
      const pageNum = Number(msg.pageNum) || 1;
      db.prepare('DELETE FROM classroom_strokes WHERE session_id=? AND page_num=?').run(sessionId, pageNum);
      broadcastToSession(sessionId, {
        type: 'classroom_page_cleared', sessionId, pageNum,
      }, true);
      break;
    }

    /* ── Page template change (teacher) ── */
    case 'template_change': {
      if (!isTeacher) return;
      const pageNum  = Number(msg.pageNum) || 1;
      const template = (msg.template || 'blank').slice(0, 32);
      try {
        db.prepare(
          `INSERT INTO classroom_pages (session_id, page_num, template) VALUES (?,?,?)
           ON CONFLICT(session_id, page_num) DO UPDATE SET template=excluded.template`
        ).run(sessionId, pageNum, template);
      } catch { return; }
      broadcastToSession(sessionId, {
        type: 'classroom_template_changed', sessionId, pageNum, template,
      }, true);
      break;
    }

    /* ── Draw permission (teacher) ── */
    case 'allow_draw': {
      if (!isTeacher) return;
      const targetId = Number(msg.targetUserId);
      if (!targetId) return;
      try {
        db.prepare('INSERT OR IGNORE INTO classroom_draw_permissions (session_id, user_id) VALUES (?,?)').run(sessionId, targetId);
      } catch { return; }
      _invalidateDrawCache(sessionId, targetId);
      emitToUser(targetId, { type: 'classroom_draw_permitted', sessionId });
      break;
    }

    case 'revoke_draw': {
      if (!isTeacher) return;
      const targetId = Number(msg.targetUserId);
      if (!targetId) return;
      try {
        db.prepare('DELETE FROM classroom_draw_permissions WHERE session_id=? AND user_id=?').run(sessionId, targetId);
      } catch { return; }
      _invalidateDrawCache(sessionId, targetId);
      emitToUser(targetId, { type: 'classroom_draw_revoked', sessionId });
      break;
    }

    /* ── Mute peer (teacher) ── */
    case 'mute_peer': {
      if (!isTeacher) return;
      const targetId = Number(msg.targetUserId);
      if (!targetId) return;
      try {
        db.prepare('INSERT OR IGNORE INTO classroom_muted (session_id, user_id, muted_by) VALUES (?,?,?)').run(sessionId, targetId, ws.userId);
      } catch { return; }
      emitToUser(targetId, { type: 'classroom_muted', sessionId, by: ws.userId });
      break;
    }

    case 'unmute_peer': {
      if (!isTeacher) return;
      const targetId = Number(msg.targetUserId);
      if (!targetId) return;
      try {
        db.prepare('DELETE FROM classroom_muted WHERE session_id=? AND user_id=?').run(sessionId, targetId);
      } catch { return; }
      emitToUser(targetId, { type: 'classroom_unmuted', sessionId, by: ws.userId });
      break;
    }

    /* ── Screen share announce (teacher) ── */
    case 'screen_start': {
      if (!isTeacher) return;
      broadcastToSession(sessionId, { type: 'classroom_screen_started', sessionId });
      break;
    }

    case 'screen_stop': {
      if (!isTeacher) return;
      broadcastToSession(sessionId, { type: 'classroom_screen_stopped', sessionId });
      break;
    }
  }
}

/* ── WebSocket server ──────────────────────────────────────────────────── */
let _wss = null;
function attach(httpServer) {
  const wss = new WebSocketServer({ server: httpServer, path: '/ws' });
  _wss = wss;

  wss.on('connection', (ws) => {
    ws.userId   = null;
    ws.userName = '';
    ws.isAlive  = true;

    ws.on('pong', () => { ws.isAlive = true; });

    ws._msgCount       = 0;
    ws._msgWindowStart = Date.now();

    // Close unauthenticated connections after 5s
    ws._authTimer = setTimeout(() => {
      if (!ws.userId) { try { ws.close(4001, 'Auth timeout'); } catch {} }
    }, 5000);

    ws.on('message', raw => {
      // Rate-limit: max 120 messages per second per connection
      const now = Date.now();
      if (now - ws._msgWindowStart > 1000) { ws._msgCount = 0; ws._msgWindowStart = now; }
      if (++ws._msgCount > 120) return;

      let msg;
      try { msg = JSON.parse(raw); } catch { return; }

      // Pre-auth: only {type:'auth', token} accepted
      if (!ws.userId) {
        if (msg.type !== 'auth' || !msg.token) {
          try { ws.close(4001, 'Auth required'); } catch {}
          return;
        }
        try {
          const user = jwt.verify(msg.token, process.env.JWT_SECRET, { algorithms: ['HS256'] });
          ws.userId   = user.id;
          ws.userName = user.name || user.email || '';
          clearTimeout(ws._authTimer);
          ws._authTimer = null;
          try { ws.send(JSON.stringify({ type: 'auth_ok' })); } catch {}
        } catch {
          try { ws.close(4001, 'Invalid token'); } catch {}
        }
        return;
      }

      // Post-auth: normal handler
      try { _handleMessage(ws, msg); } catch (e) {
        // swallow — never crash the WS server on bad input
      }
    });

    ws.on('error', () => {});
    ws.on('close', () => {
      if (ws._authTimer) { clearTimeout(ws._authTimer); ws._authTimer = null; }
      if (ws.userId) _unregisterUser(ws);
    });
  });

  /* ── Ping/pong keepalive (30s) ── */
  const pingTimer = setInterval(() => {
    for (const ws of wss.clients) {
      if (!ws.isAlive) { ws.terminate(); continue; }
      ws.isAlive = false;
      try { ws.ping(); } catch {}
    }
  }, 30_000);
  wss.on('close', () => clearInterval(pingTimer));

  return wss;
}

function closeAll() {
  if (!_wss) return;
  for (const ws of _wss.clients) {
    try { ws.close(1001, 'Server shutting down'); } catch {}
  }
  try { _wss.close(); } catch {}
}

module.exports = {
  attach,
  broadcastToSession,
  emitToUser,
  invalidateSession: _invalidateSession,
  invalidateDrawCache: _invalidateDrawCache,
  closeAll,
};