ForgeFlow/tests/ssh-connection-pool.test.mjs

import test from "node:test";
import assert from "node:assert/strict";
import { EventEmitter } from "node:events";
import { createRequire } from "node:module";

const require = createRequire(import.meta.url);

const ssh2Path = require.resolve("ssh2");
const realSsh2 = require("ssh2");

// SshService resolves ssh2 lazily and after an await, so the replacement has to
// stay in place until the whole operation settles.
async function withFakeSsh2(Client, operation) {
  require.cache[ssh2Path] = { id: ssh2Path, filename: ssh2Path, loaded: true, exports: { ...realSsh2, Client } };
  try {
    return await operation();
  } finally {
    require.cache[ssh2Path] = { id: ssh2Path, filename: ssh2Path, loaded: true, exports: realSsh2 };
  }
}

const { SshService } = require("../src/main/ssh-service.cjs");

function makeStore(overrides = {}) {
  const server = {
    id: "unraid",
    host: "tower",
    port: 22,
    username: "root",
    authType: "password",
    basePath: "/mnt/user/appdata",
    hostFingerprint: "SHA256:trusted",
    ...overrides,
  };
  return {
    server,
    getServer: () => server,
    getServerCredentials: () => ({ password: "secret", passphrase: "" }),
  };
}

// A client that reports what the pool does to it: how often it connected, how
// many channels it opened, and whether it was closed.
function fakeClientFactory({ execBehaviour = () => ({ ok: true }) } = {}) {
  const state = { connects: 0, execs: 0, ends: 0, instances: [] };
  class FakeClient extends EventEmitter {
    constructor() {
      super();
      this.ended = false;
      state.instances.push(this);
    }
    connect(options) {
      state.connects += 1;
      options.hostVerifier(Buffer.from("host key"));
      setImmediate(() => this.emit("ready"));
    }
    exec(command, callback) {
      state.execs += 1;
      const outcome = execBehaviour(state.execs, this);
      if (outcome.channelError) {
        setImmediate(() => callback(outcome.channelError));
        return;
      }
      const stream = new EventEmitter();
      stream.stderr = new EventEmitter();
      // A channel that closes without an exit status reports null, which is how
      // a connection lost mid-command surfaces. That is not the same as 0.
      const closeCode = Object.hasOwn(outcome, "exitCode") ? outcome.exitCode : 0;
      setImmediate(() => {
        stream.emit("data", Buffer.from(outcome.stdout ?? "ok"));
        stream.emit("close", closeCode, null);
      });
      callback(null, stream);
    }
    end() {
      if (this.ended) return;
      this.ended = true;
      state.ends += 1;
      setImmediate(() => this.emit("close"));
    }
  }
  return { FakeClient, state };
}

function service(store, options = {}) {
  return new SshService({ store, diagnostics: null, ...options });
}

const run = withFakeSsh2;

test("a sequence of commands to one server shares a single connection", async () => {
  const { FakeClient, state } = fakeClientFactory();
  const store = makeStore();
  const ssh = service(store);

  for (let index = 0; index < 5; index += 1) {
    await run(FakeClient, () => ssh.exec("unraid", `echo ${index}`));
  }

  assert.equal(state.execs, 5);
  assert.equal(state.connects, 1, "five commands, one handshake");
  ssh.closeAll();
});

test("concurrent commands share the connection and it survives until the last one finishes", async () => {
  const { FakeClient, state } = fakeClientFactory();
  const ssh = service(makeStore());

  await run(FakeClient, () => Promise.all([
    ssh.exec("unraid", "one"),
    ssh.exec("unraid", "two"),
    ssh.exec("unraid", "three"),
  ]));

  assert.equal(state.connects, 1);
  assert.equal(state.execs, 3);
  assert.equal(state.ends, 0, "the shared connection is not closed while it is idle in the pool");
  ssh.closeAll();
  assert.equal(state.ends, 1);
});

test("a connection that died while pooled is replaced and the command runs once", async () => {
  const { FakeClient, state } = fakeClientFactory({
    execBehaviour: (call, client) => (call === 2 && !client.reopened
      ? { channelError: Object.assign(new Error("channel open failure"), { code: "ERR_CHANNEL" }) }
      : { ok: true }),
  });
  const ssh = service(makeStore());

  await run(FakeClient, () => ssh.exec("unraid", "first"));
  const result = await run(FakeClient, () => ssh.exec("unraid", "second"));

  assert.equal(result.stdout, "ok");
  assert.equal(state.connects, 2, "the stale connection is replaced");
  assert.equal(state.execs, 3, "the failed attempt never reached the server, so it is retried once");
  ssh.closeAll();
});

test("a command that reached the server is never retried, not even on a reused connection", async () => {
  let deployAttempts = 0;
  const { FakeClient, state } = fakeClientFactory({
    execBehaviour: (call) => {
      if (call === 1) return { ok: true };
      deployAttempts += 1;
      return { exitCode: 1, stdout: "docker compose failed" };
    },
  });
  const ssh = service(makeStore());

  // The first command establishes the pooled connection, so the deployment below
  // runs on a reused one - the case where a retry would be tempting.
  await run(FakeClient, () => ssh.exec("unraid", "true"));
  await assert.rejects(() => run(FakeClient, () => ssh.exec("unraid", "docker compose up -d")), (error) => {
    assert.equal(error.code, "SSH_COMMAND_FAILED");
    return true;
  });

  assert.equal(deployAttempts, 1, "a deployment command is never repeated by the pool");
  assert.equal(state.connects, 1);
  ssh.closeAll();
});

test("a connection lost while a command was running is not retried either", async () => {
  let attempts = 0;
  const { FakeClient, state } = fakeClientFactory({
    execBehaviour: (call) => {
      if (call === 1) return { ok: true };
      attempts += 1;
      // The stream opened, so the server may already be acting on this command.
      return { exitCode: null, stdout: "" };
    },
  });
  const ssh = service(makeStore());

  await run(FakeClient, () => ssh.exec("unraid", "true"));
  await assert.rejects(() => run(FakeClient, () => ssh.exec("unraid", "docker compose up -d")), (error) => {
    assert.equal(error.code, "SSH_COMMAND_FAILED");
    return true;
  });

  assert.equal(attempts, 1);
  assert.equal(state.connects, 1);
  ssh.closeAll();
});

test("a first connection that cannot be established is reported without a retry", async () => {
  class RefusingClient extends EventEmitter {
    connect() {
      setImmediate(() => this.emit("error", Object.assign(new Error("ECONNREFUSED"), { code: "ECONNREFUSED" })));
    }
    end() {}
  }
  const ssh = service(makeStore());

  await assert.rejects(() => run(RefusingClient, () => ssh.exec("unraid", "true")), /SSH connection failed/);
  assert.equal(ssh.sessions.size, 0, "a failed connection is not pooled");
});

test("a trust-on-first-use connection is never pooled or reused", async () => {
  const { FakeClient, state } = fakeClientFactory();
  const ssh = service(makeStore({ hostFingerprint: "" }));

  await run(FakeClient, () => ssh.test("unraid", { trustOnFirstUse: true }));
  await run(FakeClient, () => ssh.test("unraid", { trustOnFirstUse: true }));

  assert.equal(state.connects, 2, "an unverified connection is opened fresh every time");
  assert.equal(ssh.sessions.size, 0);
  assert.equal(state.ends, 2, "and closed immediately afterwards");
});

test("changing the server identity or credentials invalidates the pooled connection", async () => {
  const { FakeClient, state } = fakeClientFactory();
  const store = makeStore();
  const ssh = service(store);

  await run(FakeClient, () => ssh.exec("unraid", "before"));
  assert.equal(state.connects, 1);

  store.server.hostFingerprint = "SHA256:rotated";
  await run(FakeClient, () => ssh.exec("unraid", "after"));
  assert.equal(state.connects, 2, "the previous connection is not reused across an identity change");
  ssh.closeAll();
});

test("an idle connection is closed after its lifetime and reopened on demand", async () => {
  const { FakeClient, state } = fakeClientFactory();
  const ssh = service(makeStore(), { idleConnectionMs: 40 });

  await run(FakeClient, () => ssh.exec("unraid", "one"));
  assert.equal(state.ends, 0);

  await new Promise((resolve) => setTimeout(resolve, 120));
  assert.equal(state.ends, 1, "the idle connection is released");
  assert.equal(ssh.sessions.size, 0);

  await run(FakeClient, () => ssh.exec("unraid", "two"));
  assert.equal(state.connects, 2);
  ssh.closeAll();
});

test("an error on an idle pooled connection is absorbed instead of terminating the process", async () => {
  const { FakeClient, state } = fakeClientFactory();
  const ssh = service(makeStore());

  await run(FakeClient, () => ssh.exec("unraid", "one"));
  const pooled = state.instances.at(-1);

  pooled.emit("error", new Error("read ECONNRESET"));
  await new Promise((resolve) => setTimeout(resolve, 20));

  assert.equal(ssh.sessions.size, 0, "the dead connection leaves the pool");
  await run(FakeClient, () => ssh.exec("unraid", "two"));
  assert.equal(state.connects, 2);
  ssh.closeAll();
});