api/services/edge-broker/test/broker.test.mjs

import test from "node:test";
import assert from "node:assert/strict";
import WebSocket from "ws";

import { createBrokerServer } from "../server.mjs";

function waitForMessage(socket) {
  return new Promise((resolve) => {
    socket.once("message", (raw) => resolve(JSON.parse(raw.toString())));
  });
}

function collectMessages(socket) {
  const messages = [];
  socket.on("message", (raw) => {
    messages.push(JSON.parse(raw.toString()));
  });
  return messages;
}

function waitForClose(socket) {
  return new Promise((resolve) => {
    socket.once("close", resolve);
  });
}

function waitForCloseOrError(socket) {
  return new Promise((resolve) => {
    const onDone = () => {
      socket.off("error", onDone);
      socket.off("close", onDone);
      resolve();
    };
    socket.once("error", onDone);
    socket.once("close", onDone);
  });
}

async function waitFor(predicate, { timeoutMs = 1000, intervalMs = 10, description = "condition" } = {}) {
  const deadline = Date.now() + timeoutMs;

  while (Date.now() < deadline) {
    if (predicate()) {
      return;
    }

    await new Promise((resolve) => setTimeout(resolve, intervalMs));
  }

  throw new Error(`Timed out waiting for ${description}`);
}

test("broker dispatches commands to connected agents", async () => {
  const broker = createBrokerServer({ authMode: "stub", sharedSecret: "secret", commandTimeoutMs: 2000 });
  const address = await broker.listen(0);
  const port = address.port;
  const agent = new WebSocket(`ws://127.0.0.1:${port}/ws/agent?gatewayId=701&token=agent-token`);

  await new Promise((resolve) => agent.once("open", resolve));
  agent.on("message", (raw) => {
    const message = JSON.parse(raw.toString());
    if (message.type === "COMMAND") {
      agent.send(JSON.stringify({
        type: "COMMAND_RESULT",
        commandId: message.commandId,
        ok: true,
        payload: { online: true, on: true },
      }));
    }
  });

  const response = await fetch(`http://127.0.0.1:${port}/api/gateways/701/commands`, {
    method: "POST",
    headers: {
      "content-type": "application/json",
      "x-edge-broker-secret": "secret",
    },
    body: JSON.stringify({
      commandType: "GET_RELAY_STATUS",
      payload: { relayId: "M-7" },
    }),
  });
  const json = await response.json();

  assert.equal(response.status, 200);
  assert.equal(json.ok, true);
  assert.equal(json.payload.on, true);

  agent.terminate();
  await broker.close();
});

test("broker bridges browser shell sessions through the connected agent", async () => {
  const closedSessions = [];
  const broker = createBrokerServer({
    authMode: "stub",
    validateShellSession: async () => ({ id: "shell-1", gateway_id: "701", reason: "diagnostic" }),
    closeShellSession: async (_id, _token, transcript, reason) => {
      closedSessions.push({ transcript, reason });
    },
  });
  const address = await broker.listen(0);
  const port = address.port;

  const agent = new WebSocket(`ws://127.0.0.1:${port}/ws/agent?gatewayId=701&token=agent-token`);
  await new Promise((resolve) => agent.once("open", resolve));

  agent.on("message", (raw) => {
    const message = JSON.parse(raw.toString());
    if (message.type === "OPEN_ROOT_SHELL") {
      agent.send(JSON.stringify({ type: "SHELL_OPENED", sessionId: "shell-1" }));
      agent.send(JSON.stringify({ type: "SHELL_OUTPUT", sessionId: "shell-1", data: "root@pi:~# " }));
      agent.send(JSON.stringify({ type: "SHELL_EXIT", sessionId: "shell-1", code: 0 }));
    }
  });

  const browser = new WebSocket(`ws://127.0.0.1:${port}/ws/browser-shell?token=session-token`);
  const browserMessages = collectMessages(browser);
  await new Promise((resolve) => browser.once("open", resolve));
  await new Promise((resolve) => setTimeout(resolve, 100));

  assert.ok(browserMessages.some((message) => message.type === "opened"));
  assert.ok(browserMessages.some((message) => message.type === "output" && /root@pi/.test(message.data)));
  assert.ok(browserMessages.some((message) => message.type === "closed" && message.code === 0));
  assert.equal(closedSessions.length, 1);
  assert.equal(closedSessions[0].reason, "agent_exit");

  browser.terminate();
  agent.terminate();
  await broker.close();
});

test("broker forwards browser shell input, resize, and close events to the agent", async () => {
  const broker = createBrokerServer({
    authMode: "stub",
    validateShellSession: async () => ({ id: "shell-2", gateway_id: "701", reason: "diagnostic" }),
  });
  const address = await broker.listen(0);
  const port = address.port;

  const agent = new WebSocket(`ws://127.0.0.1:${port}/ws/agent?gatewayId=701&token=agent-token`);
  await new Promise((resolve) => agent.once("open", resolve));

  const agentMessages = collectMessages(agent);
  const browser = new WebSocket(`ws://127.0.0.1:${port}/ws/browser-shell?token=session-token`);
  await new Promise((resolve) => browser.once("open", resolve));
  await waitForMessage(agent);

  browser.send(JSON.stringify({ type: "input", data: "ls\r" }));
  browser.send(JSON.stringify({ type: "resize", cols: 140, rows: 44 }));
  browser.send(JSON.stringify({ type: "close" }));
  await new Promise((resolve) => setTimeout(resolve, 50));

  assert.ok(agentMessages.some((message) => message.type === "OPEN_ROOT_SHELL" && message.payload.sessionId === "shell-2"));
  assert.ok(
    agentMessages.some(
      (message) =>
        message.type === "SHELL_INPUT" &&
        message.payload.sessionId === "shell-2" &&
        message.payload.data === "ls\r"
    )
  );
  assert.ok(
    agentMessages.some(
      (message) =>
        message.type === "RESIZE_ROOT_SHELL" &&
        message.payload.sessionId === "shell-2" &&
        message.payload.cols === 140 &&
        message.payload.rows === 44
    )
  );
  assert.ok(agentMessages.some((message) => message.type === "CLOSE_ROOT_SHELL" && message.payload.sessionId === "shell-2"));

  browser.terminate();
  agent.terminate();
  await broker.close();
});

test("broker closes browser shell sessions immediately when no agent is connected", async () => {
  const closedSessions = [];
  const broker = createBrokerServer({
    authMode: "stub",
    validateShellSession: async () => ({ id: "shell-3", gateway_id: "701", reason: "diagnostic" }),
    closeShellSession: async (_id, _token, transcript, reason) => {
      closedSessions.push({ transcript, reason });
    },
  });
  const address = await broker.listen(0);
  const port = address.port;

  const browser = new WebSocket(`ws://127.0.0.1:${port}/ws/browser-shell?token=session-token`);
  await new Promise((resolve) => browser.once("open", resolve));
  await waitForClose(browser);
  await waitFor(() => closedSessions.length === 1, { description: "offline shell session close callback" });

  assert.deepEqual(closedSessions, [{ transcript: "", reason: "agent_offline" }]);

  await broker.close();
});

test("broker closes browser shell sessions when the agent disconnects before shell open", async () => {
  const closedSessions = [];
  const broker = createBrokerServer({
    authMode: "stub",
    validateShellSession: async () => ({ id: "shell-4", gateway_id: "701", reason: "diagnostic" }),
    closeShellSession: async (_id, _token, transcript, reason) => {
      closedSessions.push({ transcript, reason });
    },
  });
  const address = await broker.listen(0);
  const port = address.port;

  const agent = new WebSocket(`ws://127.0.0.1:${port}/ws/agent?gatewayId=701&token=agent-token`);
  await new Promise((resolve) => agent.once("open", resolve));

  const agentMessages = collectMessages(agent);
  const browser = new WebSocket(`ws://127.0.0.1:${port}/ws/browser-shell?token=session-token`);
  await new Promise((resolve) => browser.once("open", resolve));
  await waitForMessage(agent);

  agent.terminate();
  await waitForClose(browser);
  await waitFor(() => closedSessions.length === 1, { description: "disconnect shell session close callback" });

  assert.ok(agentMessages.some((message) => message.type === "OPEN_ROOT_SHELL" && message.payload.sessionId === "shell-4"));
  assert.deepEqual(closedSessions, [{ transcript: "", reason: "agent_disconnected" }]);

  await broker.close();
});

test("broker defaults to strict auth when no validators are configured", async () => {
  const broker = createBrokerServer();
  const address = await broker.listen(0);
  const port = address.port;

  const agent = new WebSocket(`ws://127.0.0.1:${port}/ws/agent?gatewayId=701&token=agent-token`);
  await waitForCloseOrError(agent);

  await broker.close();
});