server.js

import express from "express";
import http from "http";
import crypto from "crypto";
import { WebSocket, WebSocketServer } from "ws";

const app = express();
const server = http.createServer(app);
const maxPayload = Number(process.env.MAX_PAYLOAD_BYTES || 8 * 1024 * 1024);
const wss = new WebSocketServer({ server, path: "/ws", maxPayload });
const sessions = new Map();
const sockets = new Map();
const SESSION_TTL_MS = Number(process.env.SESSION_TTL_MS || 0);

app.get("/", (_req, res) => {
  res.type("text/plain").send("Consent Remote Support backend is running. WebSocket path: /ws");
});

app.get("/health", (_req, res) => {
  res.json({ ok: true, sessions: sessions.size, time: new Date().toISOString() });
});

wss.on("connection", (ws, req) => {
  ws.id = crypto.randomUUID();
  ws.role = "unknown";
  ws.isAlive = true;
  ws.remote = req.headers["x-forwarded-for"] || req.socket.remoteAddress;
  sockets.set(ws.id, ws);
  log("connect", ws.id, ws.remote);

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

  ws.on("message", (payload, isBinary) => {
    if (isBinary) {
      relayBinaryFrame(ws, payload);
      return;
    }
    let msg;
    try {
      msg = JSON.parse(payload.toString("utf8"));
    } catch {
      send(ws, { type: "error", message: "Invalid JSON" });
      return;
    }
    handleMessage(ws, msg);
  });

  ws.on("close", () => cleanupSocket(ws));
  ws.on("error", err => log("socket_error", ws.id, err.message));
});

function handleMessage(ws, msg) {
  const code = normalizeCode(msg.code);
  switch (msg.type) {
    case "register_client": {
      if (!code || !msg.token) return send(ws, { type: "error", message: "Missing client code or token" });
      const previous = sessions.get(code);
      if (previous?.client && previous.client !== ws) {
        send(previous.client, { type: "stop", message: "A new client registered this code" });
        previous.client.close(4001, "replaced by new client");
      }
      ws.role = "client";
      ws.code = code;
      ws.token = msg.token;
      sessions.set(code, {
        code,
        token: msg.token,
        client: ws,
        viewer: previous?.viewer || null,
        viewerId: previous?.viewerId || null,
        createdAt: Date.now(),
        lastSeen: Date.now()
      });
      send(ws, { type: "registered", code });
      if (previous?.viewer) pair(code, previous.viewer, previous.viewerId);
      break;
    }
    case "connect_viewer": {
      if (!code || !/^\d{8}$/.test(code)) return send(ws, { type: "error", message: "Enter a valid 8-digit code" });
      const session = sessions.get(code);
      if (!session?.client || session.client.readyState !== WebSocket.OPEN) {
        return send(ws, { type: "error", message: "Client is not online" });
      }
      if (session.viewer && session.viewer !== ws) {
        send(session.viewer, { type: "viewer_disconnected", message: "Another viewer connected" });
        session.viewer.close(4002, "replaced by new viewer");
      }
      pair(code, ws, msg.viewerId || `viewer-${crypto.randomUUID().slice(0, 8)}`);
      break;
    }
    case "control":
    case "camera_control":
    case "request_data":
    case "quality": {
      const session = sessions.get(code);
      if (session?.viewer === ws && session.client?.readyState === WebSocket.OPEN) {
        session.lastSeen = Date.now();
        send(session.client, msg);
      }
      break;
    }
    case "data":
    case "error":
    case "notification":
    case "camera_frame":
    case "frame_meta": {
      const session = sessions.get(ws.code);
      if (session?.client === ws && session.viewer?.readyState === WebSocket.OPEN) {
        session.lastSeen = Date.now();
        send(session.viewer, msg);
      }
      break;
    }
    case "client_disconnect":
    case "viewer_disconnect": {
      cleanupSocket(ws);
      ws.close(1000, "disconnect requested");
      break;
    }
    default:
      send(ws, { type: "error", message: `Unsupported message type: ${msg.type}` });
  }
}

function pair(code, viewer, viewerId) {
  const session = sessions.get(code);
  if (!session) return;
  viewer.role = "viewer";
  viewer.code = code;
  session.viewer = viewer;
  session.viewerId = viewerId;
  session.lastSeen = Date.now();
  send(viewer, { type: "paired", code, viewerId });
  send(session.client, { type: "viewer_connected", code, viewerId });
  log("paired", code, viewerId);
}

function relayBinaryFrame(ws, payload) {
  if (ws.role !== "client" || !ws.code) return;
  const session = sessions.get(ws.code);
  if (session?.client === ws && session.viewer?.readyState === WebSocket.OPEN) {
    session.lastSeen = Date.now();
    session.viewer.send(payload, { binary: true });
  }
}

function cleanupSocket(ws) {
  sockets.delete(ws.id);
  if (!ws.code) return;
  const session = sessions.get(ws.code);
  if (!session) return;
  if (session.client === ws) {
    send(session.viewer, { type: "client_disconnected" });
    session.viewer?.close(4003, "client disconnected");
    sessions.delete(ws.code);
  } else if (session.viewer === ws) {
    send(session.client, { type: "viewer_disconnected" });
    session.viewer = null;
    session.viewerId = null;
  }
  log("cleanup", ws.role, ws.code);
}

function send(ws, msg) {
  if (ws && ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify(msg));
}

function normalizeCode(code) {
  return typeof code === "string" ? code.replace(/\D/g, "").slice(0, 8) : null;
}

function log(...parts) {
  console.log(new Date().toISOString(), ...parts);
}

setInterval(() => {
  const now = Date.now();
  for (const ws of sockets.values()) {
    if (!ws.isAlive) {
      cleanupSocket(ws);
      ws.terminate();
      continue;
    }
    ws.isAlive = false;
    ws.ping();
  }
  if (SESSION_TTL_MS > 0) {
    for (const [code, session] of sessions.entries()) {
      if (now - session.lastSeen > SESSION_TTL_MS && !isOpen(session.client)) {
        session.viewer?.close(4004, "client offline timeout");
        sessions.delete(code);
      }
    }
  }
}, 30000);

const port = Number(process.env.PORT || 7860);
server.listen(port, "0.0.0.0", () => {
  log(`server listening on ${port}`);
});

function shutdown() {
  for (const ws of sockets.values()) {
    send(ws, { type: "server_shutdown", message: "Backend is restarting" });
    ws.close(1012, "server restart");
  }
  server.close(() => process.exit(0));
  setTimeout(() => process.exit(0), 3000).unref();
}

process.on("SIGTERM", shutdown);
process.on("SIGINT", shutdown);

function isOpen(ws) {
  return ws?.readyState === WebSocket.OPEN;
}