OJ2/apps/api/src/routes/judge-server.ts

import { createHash, timingSafeEqual } from 'node:crypto'

import { eq } from 'drizzle-orm'
import { Hono, type Context } from 'hono'
import { z } from 'zod'

import { config } from '../config'
import { db, schema } from '../db'
import { failure, readJson } from '../http'

const heartbeatSchema = z.object({
  hostname: z.string().min(1).max(128),
  judger_version: z.string().min(1).max(32),
  cpu_core: z.number().int().positive(),
  memory: z.number().min(0).max(100),
  cpu: z.number().min(0).max(100),
  action: z.literal('heartbeat'),
  service_url: z.string().min(1).max(256),
})

export const judgeServerRoutes = new Hono()

function tokenMatches(value: string | undefined) {
  if (!value) return false
  const expected = createHash('sha256').update(config.judgeServerToken).digest('hex')
  const actualBuffer = Buffer.from(value)
  const expectedBuffer = Buffer.from(expected)
  return actualBuffer.length === expectedBuffer.length && timingSafeEqual(actualBuffer, expectedBuffer)
}

async function heartbeat(c: Context) {
  if (!tokenMatches(c.req.header('X-Judge-Server-Token'))) {
    return failure(c, 403, 'invalid-judge-token', 'Invalid token')
  }

  const parsed = await readJson(c, heartbeatSchema)
  if (!parsed.success) {
    return failure(c, 400, 'invalid-heartbeat', 'Invalid heartbeat payload')
  }

  const now = new Date().toISOString()
  const [existing] = await db
    .select({ id: schema.judgeServer.id })
    .from(schema.judgeServer)
    .where(eq(schema.judgeServer.hostname, parsed.data.hostname))
    .limit(1)

  const common = {
    ip: c.req.header('x-forwarded-for')?.split(',')[0]?.trim() || c.req.header('x-real-ip') || null,
    judgerVersion: parsed.data.judger_version,
    cpuCore: parsed.data.cpu_core,
    memoryUsage: parsed.data.memory,
    cpuUsage: parsed.data.cpu,
    lastHeartbeat: now,
    serviceUrl: parsed.data.service_url,
  }

  if (existing) {
    await db.update(schema.judgeServer).set(common).where(eq(schema.judgeServer.id, existing.id))
  } else {
    await db.insert(schema.judgeServer).values({
      ...common,
      hostname: parsed.data.hostname,
      createTime: now,
      taskNumber: 0,
      isDisabled: false,
    })
  }

  // JudgeServer 1.6.1 healthcheck still reads the legacy { error, data } envelope.
  return c.json({ error: null, data: null })
}

judgeServerRoutes.post('/judge-server/heartbeat', heartbeat)
judgeServerRoutes.post('/judge-server/heartbeat/', heartbeat)