EduBox Dev
327 lines
11 KiB
TypeScript
327 lines
11 KiB
TypeScript
import { WebSocketServer, WebSocket } from "ws";
|
|
import { randomBytes } from "crypto";
|
|
import type { IncomingMessage } from "http";
|
|
import { prisma } from "./prisma";
|
|
import { createEphemeralPreAuthKey, getHeadscaleUserId } from "./headscale";
|
|
|
|
interface NodeMessage {
|
|
action: string;
|
|
nodeId?: string;
|
|
code?: string;
|
|
instanceId?: string;
|
|
type?: string;
|
|
port?: number;
|
|
composeConfig?: string;
|
|
studentName?: string;
|
|
error?: string;
|
|
tailscaleIp?: string;
|
|
token?: string;
|
|
}
|
|
|
|
const nodes = new Map<string, WebSocket>();
|
|
|
|
interface AttemptWindow {
|
|
count: number;
|
|
firstAttempt: number;
|
|
}
|
|
|
|
const activationAttemptsByCode = new Map<string, AttemptWindow>();
|
|
const activationAttemptsByNode = new Map<string, AttemptWindow>();
|
|
const MAX_ACTIVATION_ATTEMPTS = 5;
|
|
const ACTIVATION_WINDOW_MS = 15 * 60 * 1000;
|
|
|
|
const HEADSCALE_USER = "studioe5";
|
|
const HEADSCALE_AGENT_TAG = "tag:student-agent";
|
|
const HEADSCALE_KEY_EXPIRATION_MINUTES = 15;
|
|
|
|
let headscaleUserIdCache: string | null = null;
|
|
|
|
function recordActivationAttempt(map: Map<string, AttemptWindow>, key: string): boolean {
|
|
const now = Date.now();
|
|
const win = map.get(key);
|
|
if (!win || now - win.firstAttempt > ACTIVATION_WINDOW_MS) {
|
|
map.set(key, { count: 1, firstAttempt: now });
|
|
return true;
|
|
}
|
|
win.count++;
|
|
return win.count <= MAX_ACTIVATION_ATTEMPTS;
|
|
}
|
|
|
|
function clearActivationAttempts(code: string, nodeId: string) {
|
|
activationAttemptsByCode.delete(code);
|
|
activationAttemptsByNode.delete(nodeId);
|
|
}
|
|
|
|
function generateNodeToken(): string {
|
|
return randomBytes(32).toString("hex");
|
|
}
|
|
|
|
function getBearerToken(req: IncomingMessage): string | null {
|
|
const auth = req.headers.authorization || "";
|
|
const match = auth.match(/^Bearer\s+(\S+)$/i);
|
|
return match ? match[1] : null;
|
|
}
|
|
|
|
function close(ws: WebSocket, code: number, reason: string) {
|
|
try {
|
|
ws.close(code, reason);
|
|
} catch {
|
|
// ignore
|
|
}
|
|
}
|
|
|
|
export function initWebSocketServer(wss: WebSocketServer) {
|
|
wss.on("connection", (ws: WebSocket, req: IncomingMessage) => {
|
|
let nodeId: string | null = null;
|
|
let authenticated = false;
|
|
const token = getBearerToken(req);
|
|
|
|
console.log("[WS] New connection", token ? "(token provided)" : "(no token)");
|
|
|
|
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) {
|
|
const id = msg.nodeId;
|
|
const existing = await prisma.node.findUnique({ where: { id } });
|
|
|
|
if (token) {
|
|
// Token supplied: it must match the stored token for this node.
|
|
if (!existing || existing.token !== token) {
|
|
console.log("[WS] Invalid token for node", id);
|
|
close(ws, 1008, "invalid token");
|
|
return;
|
|
}
|
|
authenticated = true;
|
|
} else if (existing && existing.token) {
|
|
// Existing node has a token but none was supplied.
|
|
console.log("[WS] Missing token for node", id);
|
|
close(ws, 1008, "missing token");
|
|
return;
|
|
} else if (existing) {
|
|
// Migration path: existing node without a token gets one on first register.
|
|
const newToken = generateNodeToken();
|
|
await prisma.node.update({
|
|
where: { id },
|
|
data: { token: newToken, status: "online", lastSeen: new Date() },
|
|
});
|
|
ws.send(JSON.stringify({ action: "set_token", token: newToken }));
|
|
authenticated = true;
|
|
}
|
|
// If the node does not exist yet, we stay unauthenticated until activation.
|
|
|
|
nodeId = id;
|
|
if (authenticated) {
|
|
const existing = nodes.get(id);
|
|
if (existing && existing !== ws && existing.readyState === WebSocket.OPEN) {
|
|
console.log("[WS] Superseding previous connection for", id);
|
|
existing.close(1008, "superseded");
|
|
}
|
|
nodes.set(id, ws);
|
|
await prisma.node.upsert({
|
|
where: { id },
|
|
update: { status: "online", lastSeen: new Date() },
|
|
create: { id, status: "online", lastSeen: new Date() },
|
|
});
|
|
}
|
|
ws.send(JSON.stringify({ action: "registered" }));
|
|
return;
|
|
}
|
|
|
|
if (msg.action === "activate" && msg.code && msg.nodeId) {
|
|
const id = msg.nodeId;
|
|
nodeId = id;
|
|
|
|
if (!recordActivationAttempt(activationAttemptsByCode, msg.code) ||
|
|
!recordActivationAttempt(activationAttemptsByNode, id)) {
|
|
console.log("[WS] Too many activation attempts for code/node", msg.code, id);
|
|
ws.send(JSON.stringify({ action: "activation_failed", error: "Too many attempts" }));
|
|
return;
|
|
}
|
|
|
|
const existing = await prisma.node.findUnique({ where: { id } });
|
|
if (existing && existing.token && (!authenticated || nodeId !== id)) {
|
|
console.log("[WS] Node already activated and not authenticated:", id);
|
|
ws.send(JSON.stringify({ action: "activation_failed", error: "Node already activated" }));
|
|
return;
|
|
}
|
|
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;
|
|
}
|
|
if (!student.activationCodeExpiresAt || student.activationCodeExpiresAt < new Date()) {
|
|
console.log("[WS] Expired code:", msg.code);
|
|
ws.send(JSON.stringify({ action: "activation_failed", error: "Code expired" }));
|
|
return;
|
|
}
|
|
|
|
const newToken = generateNodeToken();
|
|
await prisma.node.upsert({
|
|
where: { id },
|
|
update: {
|
|
studentId: student.id,
|
|
status: "online",
|
|
lastSeen: new Date(),
|
|
token: newToken,
|
|
},
|
|
create: {
|
|
id,
|
|
studentId: student.id,
|
|
status: "online",
|
|
lastSeen: new Date(),
|
|
token: newToken,
|
|
},
|
|
});
|
|
|
|
// Invalidate the activation code so it cannot be reused.
|
|
await prisma.student.update({
|
|
where: { id: student.id },
|
|
data: { activationCode: null, activationCodeExpiresAt: null },
|
|
});
|
|
clearActivationAttempts(msg.code, id);
|
|
|
|
authenticated = true;
|
|
const previous = nodes.get(id);
|
|
if (previous && previous !== ws && previous.readyState === WebSocket.OPEN) {
|
|
console.log("[WS] Superseding previous connection for", id);
|
|
previous.close(1008, "superseded");
|
|
}
|
|
nodes.set(id, ws);
|
|
const headscaleUrl = process.env.HEADSCALE_URL;
|
|
const headscaleApiKey = process.env.HEADSCALE_API_KEY;
|
|
const reusableAuthKey = process.env.HEADSCALE_AUTH_KEY;
|
|
|
|
if (!headscaleUrl) {
|
|
console.log("[WS] HEADSCALE_URL missing");
|
|
ws.send(JSON.stringify({ action: "activation_failed", error: "Server misconfiguration" }));
|
|
return;
|
|
}
|
|
|
|
let headscaleAuthKey: string;
|
|
try {
|
|
if (headscaleApiKey) {
|
|
if (!headscaleUserIdCache) {
|
|
headscaleUserIdCache = await getHeadscaleUserId(headscaleUrl, headscaleApiKey, HEADSCALE_USER);
|
|
}
|
|
headscaleAuthKey = await createEphemeralPreAuthKey(headscaleUrl, headscaleApiKey, headscaleUserIdCache, {
|
|
expirationMinutes: HEADSCALE_KEY_EXPIRATION_MINUTES,
|
|
aclTags: [HEADSCALE_AGENT_TAG],
|
|
});
|
|
console.log("[WS] Generated ephemeral Headscale key for", id);
|
|
} else if (reusableAuthKey) {
|
|
console.log("[WS] HEADSCALE_API_KEY not set, falling back to reusable HEADSCALE_AUTH_KEY");
|
|
headscaleAuthKey = reusableAuthKey;
|
|
} else {
|
|
console.log("[WS] No Headscale key available");
|
|
ws.send(JSON.stringify({ action: "activation_failed", error: "Server misconfiguration" }));
|
|
return;
|
|
}
|
|
} catch (err) {
|
|
console.error("[WS] Failed to create ephemeral Headscale key:", err);
|
|
ws.send(JSON.stringify({ action: "activation_failed", error: "Failed to create VPN key" }));
|
|
return;
|
|
}
|
|
|
|
console.log("[WS] Activated:", student.firstName, student.lastName, "on", id);
|
|
ws.send(JSON.stringify({
|
|
action: "activated",
|
|
studentId: student.id,
|
|
studentName: `${student.firstName} ${student.lastName}`,
|
|
headscaleUrl,
|
|
headscaleAuthKey,
|
|
token: newToken,
|
|
}));
|
|
return;
|
|
}
|
|
|
|
if (!authenticated || !nodeId) {
|
|
console.log("[WS] Unauthenticated message", msg.action, "ignored");
|
|
return;
|
|
}
|
|
|
|
if (msg.action === "heartbeat") {
|
|
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" && 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) {
|
|
const { count } = await prisma.instance.updateMany({
|
|
where: { id: msg.instanceId },
|
|
data: { status: "running" },
|
|
});
|
|
if (count) console.log("[WS] Instance started:", msg.instanceId);
|
|
return;
|
|
}
|
|
|
|
if (msg.action === "instance_stopped" && msg.instanceId) {
|
|
const { count } = await prisma.instance.updateMany({
|
|
where: { id: msg.instanceId },
|
|
data: { status: "stopped" },
|
|
});
|
|
if (count) console.log("[WS] Instance stopped:", msg.instanceId);
|
|
return;
|
|
}
|
|
|
|
if (msg.action === "instance_deleted" && msg.instanceId) {
|
|
const { count } = await prisma.instance.deleteMany({
|
|
where: { id: msg.instanceId },
|
|
});
|
|
if (count) console.log("[WS] Instance deleted:", msg.instanceId);
|
|
return;
|
|
}
|
|
|
|
if (msg.action === "instance_error" && msg.instanceId) {
|
|
const { count } = await prisma.instance.updateMany({
|
|
where: { id: msg.instanceId },
|
|
data: { status: "error" },
|
|
});
|
|
if (count) console.log("[WS] Instance error:", msg.instanceId);
|
|
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;
|
|
}
|