kestrelos/server/plugins/cot.js

import { createServer as createTcpServer } from 'node:net'
import { createServer as createTlsServer } from 'node:tls'
import { readFileSync, existsSync } from 'node:fs'
import { updateFromCot } from '../utils/cotStore.js'
import { parseTakStreamFrame, parseTraditionalXmlFrame, parseCotPayload } from '../utils/cotParser.js'
import { validateCotAuth } from '../utils/cotAuth.js'
import { getCotSslPaths, getCotPort } from '../utils/cotSsl.js'
import { registerCleanup } from '../utils/shutdown.js'
import { COT_AUTH_TIMEOUT_MS } from '../utils/constants.js'
import { acquire } from '../utils/asyncLock.js'

const serverState = {
  tcpServer: null,
  tlsServer: null,
}
const relaySet = new Set()
const allSockets = new Set()
const socketBuffers = new WeakMap()
const socketAuthTimeout = new WeakMap()

function clearAuthTimeout(socket) {
  const t = socketAuthTimeout.get(socket)
  if (t) {
    clearTimeout(t)
    socketAuthTimeout.delete(socket)
  }
}

function removeFromRelay(socket) {
  relaySet.delete(socket)
  allSockets.delete(socket)
  clearAuthTimeout(socket)
  socketBuffers.delete(socket)
}

function broadcast(senderSocket, rawMessage) {
  for (const s of relaySet) {
    if (s !== senderSocket && !s.destroyed && s.writable) {
      try {
        s.write(rawMessage)
      }
      catch (err) {
        console.error('[cot] Broadcast write error:', err?.message)
      }
    }
  }
}

const createPreview = (payload) => {
  try {
    const str = payload.toString('utf8')
    if (str.startsWith('<')) {
      const s = str.length <= 120 ? str : str.slice(0, 120) + '...'
      // eslint-disable-next-line no-control-regex -- sanitize control chars for log preview
      return s.replace(/[\u0000-\u0008\v\f\u000E-\u001F]/g, '.')
    }
    return 'hex:' + payload.subarray(0, Math.min(40, payload.length)).toString('hex')
  }
  catch {
    return 'hex:' + payload.subarray(0, Math.min(40, payload.length)).toString('hex')
  }
}

async function processFrame(socket, rawMessage, payload, authenticated) {
  const requireAuth = socket._cotRequireAuth !== false
  const debug = socket._cotDebug === true
  const parsed = parseCotPayload(payload)
  if (debug) {
    const preview = createPreview(payload)
    console.log('[cot] payload length:', payload.length, 'parsed:', parsed ? parsed.type : null, 'preview:', preview)
  }
  if (!parsed) return

  if (parsed.type === 'auth') {
    if (authenticated) return
    console.log('[cot] auth attempt username=', parsed.username)
    // Use lock per socket to prevent concurrent auth attempts
    const socketKey = `cot-auth-${socket.remoteAddress || 'unknown'}-${socket.remotePort || 0}`
    await acquire(socketKey, async () => {
      // Re-check authentication state after acquiring lock
      if (socket._cotAuthenticated || socket.destroyed) return
      try {
        const valid = await validateCotAuth(parsed.username, parsed.password)
        console.log('[cot] auth result valid=', valid, 'for username=', parsed.username)
        if (!socket.writable || socket.destroyed) return
        if (valid) {
          clearAuthTimeout(socket)
          relaySet.add(socket)
          socket._cotAuthenticated = true
        }
        else {
          socket.destroy()
        }
      }
      catch (err) {
        console.log('[cot] auth validation error:', err?.message)
        if (!socket.destroyed) socket.destroy()
      }
    }).catch((err) => {
      console.log('[cot] auth lock error:', err?.message)
      if (!socket.destroyed) socket.destroy()
    })
    return
  }

  if (parsed.type === 'cot') {
    if (requireAuth && !authenticated) {
      socket.destroy()
      return
    }
    updateFromCot(parsed).catch((err) => {
      console.error('[cot] Error updating from CoT:', err?.message)
    })
    if (authenticated) broadcast(socket, rawMessage)
  }
}

const parseFrame = (buf) => {
  const takResult = parseTakStreamFrame(buf)
  if (takResult) return { result: takResult, frameType: 'tak' }
  if (buf[0] === 0x3C) {
    const xmlResult = parseTraditionalXmlFrame(buf)
    if (xmlResult) return { result: xmlResult, frameType: 'traditional' }
  }
  return { result: null, frameType: null }
}

const processBufferedData = async (socket, buf, authenticated) => {
  if (buf.length === 0) return buf
  const { result, frameType } = parseFrame(buf)
  if (result && socket._cotDebug) {
    console.log('[cot] frame parsed as', frameType, 'bytesConsumed=', result.bytesConsumed)
  }
  if (!result) return buf
  const { payload, bytesConsumed } = result
  const rawMessage = buf.subarray(0, bytesConsumed)
  await processFrame(socket, rawMessage, payload, authenticated)
  if (socket.destroyed) return null
  const remainingBuf = buf.subarray(bytesConsumed)
  socketBuffers.set(socket, remainingBuf)
  return processBufferedData(socket, remainingBuf, authenticated)
}

async function onData(socket, data) {
  const existingBuf = socketBuffers.get(socket)
  const buf = Buffer.concat([existingBuf || Buffer.alloc(0), data])
  socketBuffers.set(socket, buf)
  const authenticated = Boolean(socket._cotAuthenticated)

  if (socket._cotDebug && buf.length > 0 && !socket._cotFirstChunkLogged) {
    socket._cotFirstChunkLogged = true
    const hex = buf.subarray(0, Math.min(80, buf.length)).toString('hex')
    console.log('[cot] first chunk len=', buf.length, 'first bytes (hex):', hex, 'starts with 0xBF:', buf[0] === 0xBF, 'starts with <:', buf[0] === 0x3C)
  }
  await processBufferedData(socket, buf, authenticated)
}

function setupSocket(socket, tls = false) {
  const remote = socket.remoteAddress || 'unknown'
  console.log('[cot] client connected', tls ? '(TLS)' : '(TCP)', 'from', remote)
  allSockets.add(socket)
  const config = useRuntimeConfig()
  socket._cotDebug = Boolean(config.cotDebug)
  socket._cotRequireAuth = config.cotRequireAuth !== false
  if (socket._cotRequireAuth) {
    const timeout = setTimeout(() => {
      if (!socket._cotAuthenticated && !socket.destroyed) {
        console.log('[cot] auth timeout, closing connection from', remote)
        socket.destroy()
      }
    }, COT_AUTH_TIMEOUT_MS)
    socketAuthTimeout.set(socket, timeout)
  }
  else {
    socket._cotAuthenticated = true
    relaySet.add(socket)
  }

  socket.on('data', data => onData(socket, data))
  socket.on('error', (err) => {
    console.error('[cot] Socket error:', err?.message)
  })
  socket.on('close', () => {
    console.log('[cot] client disconnected', socket._cotAuthenticated ? '(was authenticated)' : '', 'from', remote)
    removeFromRelay(socket)
  })
}

function startCotServers() {
  const config = useRuntimeConfig()
  const { certPath, keyPath } = getCotSslPaths(config) || {}
  const hasTls = certPath && keyPath && existsSync(certPath) && existsSync(keyPath)
  const port = getCotPort()

  try {
    if (hasTls) {
      const tlsOpts = {
        cert: readFileSync(certPath),
        key: readFileSync(keyPath),
        rejectUnauthorized: false,
      }
      serverState.tlsServer = createTlsServer(tlsOpts, socket => setupSocket(socket, true))
      serverState.tlsServer.on('error', err => console.error('[cot] TLS server error:', err?.message))
      serverState.tlsServer.listen(port, '0.0.0.0', () => {
        console.log('[cot] CoT server listening on 0.0.0.0:' + port + ' (TLS) - use this port in ATAK/iTAK and enable SSL')
      })
    }
    else {
      serverState.tcpServer = createTcpServer(socket => setupSocket(socket, false))
      serverState.tcpServer.on('error', err => console.error('[cot] TCP server error:', err?.message))
      serverState.tcpServer.listen(port, '0.0.0.0', () => {
        console.log('[cot] CoT server listening on 0.0.0.0:' + port + ' (plain TCP) - use this port in ATAK/iTAK with SSL disabled')
      })
    }
  }
  catch (err) {
    console.error('[cot] Failed to start CoT server:', err?.message)
    if (err?.code === 'EADDRINUSE') {
      console.error('[cot] Port', port, 'is already in use. Stop the other process or set COT_PORT to a different port.')
    }
  }
}

export default defineNitroPlugin((nitroApp) => {
  nitroApp.hooks.hook('ready', startCotServers)
  // Start immediately so CoT is up before first request in dev; ready may fire late in some setups.
  setImmediate(startCotServers)

  const cleanupServers = () => {
    if (serverState.tcpServer) {
      serverState.tcpServer.close()
      serverState.tcpServer = null
    }
    if (serverState.tlsServer) {
      serverState.tlsServer.close()
      serverState.tlsServer = null
    }
  }

  const cleanupSockets = () => {
    for (const s of allSockets) {
      try {
        s.destroy()
      }
      catch {
        /* ignore */
      }
    }
    allSockets.clear()
    relaySet.clear()
  }

  registerCleanup(async () => {
    cleanupSockets()
    cleanupServers()
  })

  nitroApp.hooks.hook('close', async () => {
    cleanupSockets()
    cleanupServers()
  })
})