pr-preview/backend/src/services/ssh.ts

import { Client } from "ssh2";
import { createLogger } from "../lib/logger";

const log = createLogger("SSH");

export interface SshSession {
  exec(command: string, timeoutMs?: number): Promise<{ stdout: string; stderr: string; code: number }>;
  // Run a long-lived command (e.g. `tail -f`) and receive its stdout/stderr as
  // it arrives. Unlike exec(), this never buffers or times out — the channel
  // stays open until the returned handle's close() is called (or the connection
  // ends). Used to stream live app logs off a running instance.
  execStream(command: string, onData: (chunk: string) => void): { close(): void };
  close(): void;
  aborted: boolean;
  abort(): void;
}

export async function connectSsh(host: string, privateKey: string, maxWaitMs = 300_000): Promise<SshSession> {
  const start = Date.now();

  while (Date.now() - start < maxWaitMs) {
    try {
      const conn = await tryConnect(host, privateKey, 10000);
      let aborted = false;

      return {
        get aborted() { return aborted; },
        abort() {
          aborted = true;
          try { conn.end(); } catch {}
        },
        async exec(command: string, timeoutMs?: number) {
          if (aborted) throw new Error("SSH session aborted");
          return execOnConn(conn, command, timeoutMs);
        },
        execStream(command: string, onData: (chunk: string) => void) {
          let stream: any = null;
          let closed = false;
          conn.exec(command, (err, s) => {
            if (err) {
              if (!closed) onData(`[app-logs] stream error: ${err.message}\n`);
              return;
            }
            if (closed) { try { s.close(); } catch {} return; }
            stream = s;
            s.on("data", (d: Buffer) => onData(d.toString()));
            s.stderr.on("data", (d: Buffer) => onData(d.toString()));
          });
          return {
            close() {
              closed = true;
              try { stream?.close(); } catch {}
            },
          };
        },
        close() {
          try { conn.end(); } catch {}
        },
      };
    } catch (e: any) {
      if (Date.now() - start > maxWaitMs) throw e;
      log.debug({ host, error: e.message }, "SSH connect retry");
      await sleep(5000);
    }
  }
  throw new Error(`Could not SSH into ${host} within timeout`);
}

function tryConnect(host: string, privateKey: string, timeoutMs: number): Promise<Client> {
  return new Promise((resolve, reject) => {
    const conn = new Client();
    const timer = setTimeout(() => {
      conn.end();
      reject(new Error(`SSH connection to ${host} timed out`));
    }, timeoutMs);

    conn.on("ready", () => {
      clearTimeout(timer);
      resolve(conn);
    });
    conn.on("error", (e) => {
      clearTimeout(timer);
      reject(e);
    });
    conn.connect({
      host,
      port: 22,
      username: "ubuntu",
      privateKey,
      readyTimeout: timeoutMs,
      keepaliveInterval: 10000,
      keepaliveCountMax: 3,
      algorithms: {
        serverHostKey: ["ssh-rsa", "ecdsa-sha2-nistp256", "ecdsa-sha2-nistp384", "ecdsa-sha2-nistp521", "ssh-ed25519"],
      },
    });
  });
}

const EXEC_TIMEOUT_MS = 30_000;

function execOnConn(conn: Client, command: string, timeoutMs = EXEC_TIMEOUT_MS): Promise<{ stdout: string; stderr: string; code: number }> {
  return new Promise((resolve, reject) => {
    let stream: any = null;
    let settled = false;
    const settle = (fn: () => void) => {
      if (settled) return;
      settled = true;
      clearTimeout(timer);
      fn();
    };
    const timer = setTimeout(() => {
      settle(() => reject(new Error(`SSH exec timed out after ${Math.round(timeoutMs / 1000)}s: ${command.slice(0, 60)}`)));
      try { stream?.close(); } catch {}
    }, timeoutMs);

    conn.exec(command, (err, s) => {
      if (err) {
        return settle(() => reject(err));
      }
      stream = s;
      let stdout = "";
      let stderr = "";
      stream.on("data", (d: Buffer) => { stdout += d.toString(); });
      stream.stderr.on("data", (d: Buffer) => { stderr += d.toString(); });
      stream.on("close", (code: number) => {
        settle(() => resolve({ stdout, stderr, code: code ?? 0 }));
      });
    });
  });
}

// The `ubuntu` user accepts SSH almost immediately at boot — long before the
// cloud-init UserData bootstrap has finished disabling apt-daily and installing
// the small base dependency set PP needs before SSH-driven setup begins.
// Proceeding early can still cause apt-lock failures in the deploy steps. Block until the bootstrap writes its
// done marker; bail out fast (with the tail of its log) if it writes the failure
// marker instead. Uses short one-shot execs so a stale channel just retries.
export async function waitForBootstrap(session: SshSession, maxWaitMs = 600_000): Promise<void> {
  const start = Date.now();
  while (Date.now() - start < maxWaitMs) {
    if (session.aborted) throw new Error("SSH session aborted");
    try {
      const res = await session.exec(
        "if [ -f /var/lib/pp-bootstrap-failed ]; then echo failed; " +
        "elif [ -f /var/lib/pp-bootstrap-done ]; then echo done; else echo waiting; fi"
      );
      const state = res.stdout.trim();
      if (state === "done") return;
      if (state === "failed") {
        let tail = "";
        try {
          const logRes = await session.exec("tail -n 30 /var/log/pp-bootstrap.log 2>/dev/null || true");
          tail = logRes.stdout.trim();
        } catch {
          // best effort — the failure marker alone is enough to abort
        }
        throw new Error(`EC2 bootstrap failed${tail ? `:\n${tail}` : ""}`);
      }
    } catch (e: any) {
      // A bootstrap failure we detected above must propagate; only swallow
      // transient SSH exec errors (timeout, channel reset) and retry.
      if (e?.message?.startsWith("EC2 bootstrap failed")) throw e;
    }
    await sleep(10_000);
  }
  throw new Error("EC2 bootstrap did not complete within timeout");
}

function sleep(ms: number) {
  return new Promise(r => setTimeout(r, ms));
}