stream-control/packages/server/src/agentGateway.ts

import { randomUUID } from 'node:crypto';
import type { IncomingMessage } from 'node:http';
import type { RawData, WebSocket } from 'ws';
import type { AgentToServer, ServerToAgent } from '@stream-control/shared';
import { PROTOCOL_VERSION, detectPlatform, isAgentEvent, safeJsonParse } from '@stream-control/shared';
import { config } from './config.ts';
import { extractBearer, generateToken, hashToken, safeEqual } from './auth.ts';
import { agentsRepo, type AgentRecord } from './db.ts';
import { hub } from './hub.ts';

const HELLO_TIMEOUT_MS = 10_000;

export type AgentAuth =
  | { mode: 'known'; record: AgentRecord }
  | { mode: 'enroll' };

/**
 * Authentifie une tentative de connexion agent à partir du jeton porté par la
 * requête d'upgrade. Renvoie null si le jeton est inconnu.
 */
export function authenticateAgent(req: IncomingMessage): AgentAuth | null {
  const url = new URL(req.url ?? '/', 'http://localhost');
  const token =
    extractBearer(req.headers.authorization) ??
    url.searchParams.get('token') ??
    (typeof req.headers['x-agent-token'] === 'string' ? req.headers['x-agent-token'] : null);

  if (!token) return null;

  const record = agentsRepo.findByTokenHash(hashToken(token));
  if (record) return { mode: 'known', record };

  if (config.enrollmentToken && safeEqual(token, config.enrollmentToken)) {
    return { mode: 'enroll' };
  }
  return null;
}

function send(socket: WebSocket, message: ServerToAgent): void {
  if (socket.readyState === socket.OPEN) socket.send(JSON.stringify(message));
}

export function handleAgentConnection(
  socket: WebSocket,
  auth: AgentAuth,
  remoteAddress: string,
): void {
  let agentId: string | null = null;

  const helloTimer = setTimeout(() => {
    if (!agentId) socket.close(4002, 'Message hello absent');
  }, HELLO_TIMEOUT_MS);

  socket.on('message', (raw: RawData) => {
    const message = safeJsonParse<AgentToServer>(raw.toString());
    if (!message || typeof message.type !== 'string') {
      hub.log(agentId, 'warn', 'Message agent illisible, ignoré');
      return;
    }

    // Tant que l'agent ne s'est pas présenté, seul `hello` est accepté.
    if (!agentId && message.type !== 'hello') return;

    switch (message.type) {
      case 'hello': {
        if (agentId) return; // hello dupliqué
        clearTimeout(helloTimer);

        if (message.protocol !== PROTOCOL_VERSION) {
          hub.log(
            null,
            'warn',
            `Agent « ${message.name} » en protocole v${message.protocol}, serveur en v${PROTOCOL_VERSION}`,
          );
        }

        const platform = detectPlatform(message.platform);
        let record: AgentRecord;
        let issuedToken: string | undefined;

        if (auth.mode === 'enroll') {
          issuedToken = generateToken();
          record = agentsRepo.create({
            id: message.agentId?.trim() || randomUUID(),
            name: message.name || message.hostname || 'agent',
            tokenHash: hashToken(issuedToken),
            hostname: message.hostname,
            platform,
            agentVersion: message.agentVersion,
          });
          hub.log(
            record.id,
            'info',
            `Nouvel agent enrôlé depuis ${remoteAddress}`,
            Date.now(),
            'agent.enrolled',
          );
        } else {
          record = auth.record;
          agentsRepo.updateIdentity(record.id, {
            hostname: message.hostname,
            platform,
            agentVersion: message.agentVersion,
          });
          record = agentsRepo.get(record.id) ?? record;
        }

        agentId = record.id;
        hub.attachAgent(record.id, socket);
        hub.log(
          record.id,
          'info',
          `Agent connecté (${platform}, ${remoteAddress})`,
          Date.now(),
          'agent.connected',
        );

        send(socket, {
          type: 'welcome',
          agentId: record.id,
          token: issuedToken,
          obs: record.obs,
          statusIntervalMs: config.statusIntervalMs,
          autoConnectObs: record.autoConnectObs,
          watch: record.watch,
          browser: record.browser,
          recording: record.recording,
        });
        break;
      }

      case 'status': {
        if (!agentId) return;
        hub.updateStatus(agentId, { ...message.status, updatedAt: Date.now() });
        agentsRepo.touch(agentId);
        break;
      }

      case 'result': {
        if (!agentId) return;
        hub.markSeen(agentId);
        hub.resolveCommand(agentId, message.requestId, message.ok, message.data, message.error);
        break;
      }

      case 'log': {
        hub.log(
          agentId,
          message.level,
          message.message,
          message.ts || Date.now(),
          isAgentEvent(message.event) ? message.event : null,
        );
        break;
      }

      case 'pong': {
        if (agentId) hub.markSeen(agentId);
        break;
      }
    }
  });

  socket.on('close', () => {
    clearTimeout(helloTimer);
    if (agentId) {
      hub.detachAgent(agentId, socket);
      hub.log(agentId, 'info', 'Agent déconnecté', Date.now(), 'agent.disconnected');
    }
  });

  socket.on('error', (err: Error) => {
    hub.log(agentId, 'error', `Erreur socket agent : ${err.message}`);
  });
}