edubox/server/lib/websocket.ts

import { WebSocketServer, WebSocket } from "ws";
import { prisma } from "./prisma";

interface NodeMessage {
  action: string;
  nodeId?: string;
  code?: string;
  instanceId?: string;
  type?: string;
  port?: number;
  composeConfig?: string;
  studentName?: string;
  error?: string;
  tailscaleIp?: string;
}

const nodes = new Map<string, WebSocket>();

export function initWebSocketServer(wss: WebSocketServer) {
  wss.on("connection", (ws: WebSocket) => {
    let nodeId: string | null = null;
    console.log("[WS] New connection");

    ws.on("message", async (raw) => {
      try {
        const msg: NodeMessage = JSON.parse(raw.toString());
        console.log("[WS] Received:", msg.action, "from", msg.nodeId || nodeId);

        if (msg.action === "register" && msg.nodeId) {
          nodeId = msg.nodeId;
          nodes.set(nodeId, ws);
          await prisma.node.upsert({
            where: { id: nodeId },
            update: { status: "online", lastSeen: new Date() },
            create: { id: nodeId, status: "online", lastSeen: new Date() },
          });
          ws.send(JSON.stringify({ action: "registered" }));
          return;
        }

        if (msg.action === "activate" && msg.code && msg.nodeId) {
          nodeId = msg.nodeId;
          const student = await prisma.student.findUnique({
            where: { activationCode: msg.code },
          });
          if (!student) {
            console.log("[WS] Invalid code:", msg.code);
            ws.send(JSON.stringify({ action: "activation_failed", error: "Invalid code" }));
            return;
          }
          await prisma.node.upsert({
            where: { id: nodeId },
            update: { studentId: student.id, status: "online", lastSeen: new Date() },
            create: { id: nodeId, studentId: student.id, status: "online", lastSeen: new Date() },
          });
          console.log("[WS] Activated:", student.firstName, student.lastName, "on", nodeId);
          ws.send(JSON.stringify({
            action: "activated",
            studentId: student.id,
            studentName: `${student.firstName} ${student.lastName}`,
            headscaleUrl: process.env.HEADSCALE_URL,
            headscaleAuthKey: process.env.HEADSCALE_AUTH_KEY,
          }));
          return;
        }

        if (msg.action === "heartbeat" && nodeId) {
          await prisma.node.upsert({
            where: { id: nodeId },
            update: { lastSeen: new Date() },
            create: { id: nodeId, status: "online", lastSeen: new Date() },
          });
          return;
        }

        if (msg.action === "tailscale_ip" && nodeId && msg.tailscaleIp) {
          await prisma.node.update({
            where: { id: nodeId },
            data: { tailscaleIp: msg.tailscaleIp },
          });
          console.log("[WS] Tailscale IP updated for", nodeId, ":", msg.tailscaleIp);
          return;
        }

        if (msg.action === "instance_started" && msg.instanceId) {
          await prisma.instance.update({
            where: { id: msg.instanceId },
            data: { status: "running" },
          });
          return;
        }

        if (msg.action === "instance_error" && msg.instanceId) {
          await prisma.instance.update({
            where: { id: msg.instanceId },
            data: { status: "error" },
          });
          return;
        }
      } catch (err) {
        console.error("WS error:", err);
      }
    });

    ws.on("close", async () => {
      console.log("[WS] Connection closed for", nodeId);
      if (nodeId) {
        nodes.delete(nodeId);
        await prisma.node.upsert({
          where: { id: nodeId },
          update: { status: "offline" },
          create: { id: nodeId, status: "offline" },
        });
      }
    });
  });
}

export function sendToNode(nodeId: string, message: object) {
  const ws = nodes.get(nodeId);
  if (ws && ws.readyState === WebSocket.OPEN) {
    ws.send(JSON.stringify(message));
    return true;
  }
  return false;
}