CubeAdmin/lib/socket/server.ts

import type { Server, Socket } from "socket.io"
import os from "node:os"
import type { auth as AuthType } from "@/lib/auth/index"
import { mcProcessManager } from "@/lib/minecraft/process"
import { rconClient } from "@/lib/minecraft/rcon"

// Shell metacharacters - same set as rcon.ts for defence-in-depth
const SHELL_METACHAR_RE = /[;&|`$<>\\(){}\[\]!#~]/

// ---- Auth middleware ---------------------------------------------------------

/**
 * Returns a Socket.io middleware that validates the Better Auth session token
 * from the `better-auth.session_token` cookie (or the `auth` handshake header).
 */
function makeAuthMiddleware(auth: typeof AuthType) {
  return async (socket: Socket, next: (err?: Error) => void) => {
    try {
      // Prefer the cookie sent during the upgrade handshake
      const rawCookie: string =
        (socket.handshake.headers.cookie as string | undefined) ?? ""

      const token = parseCookieToken(rawCookie) ?? socket.handshake.auth?.token

      if (!token || typeof token !== "string") {
        return next(new Error("Authentication required"))
      }

      // Use Better Auth's built-in session verification
      const session = await auth.api.getSession({
        headers: new Headers({
          cookie: `better-auth.session_token=${token}`,
        }),
      })

      if (!session?.user) {
        return next(new Error("Invalid or expired session"))
      }

      // Attach user info to the socket for later use
      ;(socket.data as Record<string, unknown>).user = session.user
      next()
    } catch (err) {
      next(
        new Error(
          `Auth error: ${err instanceof Error ? err.message : "unknown"}`,
        ),
      )
    }
  }
}

/** Extract the value of `better-auth.session_token` from a Cookie header string */
function parseCookieToken(cookieHeader: string): string | null {
  for (const part of cookieHeader.split(";")) {
    const [rawKey, ...rest] = part.trim().split("=")
    if (rawKey?.trim() === "better-auth.session_token") {
      return decodeURIComponent(rest.join("=").trim())
    }
  }
  return null
}

// ---- /console namespace -----------------------------------------------------

function setupConsoleNamespace(io: Server): void {
  const consoleNsp = io.of("/console")

  consoleNsp.on("connection", (socket: Socket) => {
    console.log(`[Socket /console] Client connected: ${socket.id}`)

    // Send buffered output so the client gets historical lines immediately
    const history = mcProcessManager.getOutput()
    socket.emit("history", history)

    // Stream live output
    const unsubscribe = mcProcessManager.onOutput((line: string) => {
      socket.emit("output", line)
    })

    // Forward process lifecycle events to this client
    const onStarted = (data: unknown) => socket.emit("server:started", data)
    const onStopped = (data: unknown) => socket.emit("server:stopped", data)
    const onCrash = (data: unknown) => socket.emit("server:crash", data)

    mcProcessManager.on("started", onStarted)
    mcProcessManager.on("stopped", onStopped)
    mcProcessManager.on("crash", onCrash)

    // Handle commands sent by the client
    socket.on("command", async (rawCommand: unknown) => {
      if (typeof rawCommand !== "string" || !rawCommand.trim()) {
        socket.emit("error", { message: "Command must be a non-empty string" })
        return
      }

      const command = rawCommand.trim()

      if (SHELL_METACHAR_RE.test(command)) {
        socket.emit("error", { message: "Command contains disallowed characters" })
        return
      }

      if (command.length > 1024) {
        socket.emit("error", { message: "Command too long" })
        return
      }

      try {
        let response: string
        if (rconClient.isConnected()) {
          response = await rconClient.sendCommand(command)
        } else {
          // Fallback: write directly to stdin
          mcProcessManager.writeStdin(command)
          response = "(sent via stdin)"
        }
        socket.emit("command:response", { command, response })
      } catch (err) {
        socket.emit("error", {
          message: `Command failed: ${err instanceof Error ? err.message : String(err)}`,
        })
      }
    })

    socket.on("disconnect", () => {
      console.log(`[Socket /console] Client disconnected: ${socket.id}`)
      unsubscribe()
      mcProcessManager.off("started", onStarted)
      mcProcessManager.off("stopped", onStopped)
      mcProcessManager.off("crash", onCrash)
    })
  })
}

// ---- /monitoring namespace --------------------------------------------------

interface MonitoringStats {
  cpu: number
  ram: { usedMB: number; totalMB: number }
  uptime: number
  server: {
    running: boolean
    pid?: number
    uptimeSecs?: number
  }
}

/** Compute CPU usage as a percentage across all cores (averaged over a 100 ms window) */
function getCpuPercent(): Promise<number> {
  return new Promise((resolve) => {
    const cpus1 = os.cpus()

    setTimeout(() => {
      const cpus2 = os.cpus()
      let idleDiff = 0
      let totalDiff = 0

      for (let i = 0; i < cpus1.length; i++) {
        const t1 = cpus1[i]!.times
        const t2 = cpus2[i]!.times

        const idle = t2.idle - t1.idle
        const total =
          t2.user + t2.nice + t2.sys + t2.idle + t2.irq -
          (t1.user + t1.nice + t1.sys + t1.idle + t1.irq)

        idleDiff += idle
        totalDiff += total
      }

      const percent = totalDiff === 0 ? 0 : (1 - idleDiff / totalDiff) * 100
      resolve(Math.round(percent * 10) / 10)
    }, 100)
  })
}

function setupMonitoringNamespace(io: Server): void {
  const monitoringNsp = io.of("/monitoring")

  // Interval for the polling loop; only active when at least one client is connected
  let interval: ReturnType<typeof setInterval> | null = null
  let clientCount = 0

  const startPolling = () => {
    if (interval) return
    interval = setInterval(async () => {
      try {
        const [cpuPercent] = await Promise.all([getCpuPercent()])

        const totalMem = os.totalmem()
        const freeMem = os.freemem()
        const usedMem = totalMem - freeMem

        const status = mcProcessManager.getStatus()

        const stats: MonitoringStats = {
          cpu: cpuPercent,
          ram: {
            usedMB: Math.round(usedMem / 1024 / 1024),
            totalMB: Math.round(totalMem / 1024 / 1024),
          },
          uptime: os.uptime(),
          server: {
            running: status.running,
            pid: status.pid,
            uptimeSecs: status.uptime,
          },
        }

        monitoringNsp.emit("stats", stats)
      } catch (err) {
        console.error("[Socket /monitoring] Stats collection error:", err)
      }
    }, 2000)
  }

  const stopPolling = () => {
    if (interval) {
      clearInterval(interval)
      interval = null
    }
  }

  monitoringNsp.on("connection", (socket: Socket) => {
    console.log(`[Socket /monitoring] Client connected: ${socket.id}`)
    clientCount++
    startPolling()

    socket.on("disconnect", () => {
      console.log(`[Socket /monitoring] Client disconnected: ${socket.id}`)
      clientCount--
      if (clientCount <= 0) {
        clientCount = 0
        stopPolling()
      }
    })
  })
}

// ---- Public setup function --------------------------------------------------

/**
 * Attach all Socket.io namespaces to the given `io` server.
 * Must be called after the HTTP server has been created.
 */
export function setupSocketServer(io: Server, auth: typeof AuthType): void {
  const authMiddleware = makeAuthMiddleware(auth)

  // Apply auth middleware to both namespaces
  io.of("/console").use(authMiddleware)
  io.of("/monitoring").use(authMiddleware)

  setupConsoleNamespace(io)
  setupMonitoringNamespace(io)

  console.log("[Socket.io] Namespaces /console and /monitoring ready")
}