SSH2ConnectionPool.ts

/**
 * SSH2 Connection Pool
 *
 * Manages persistent ssh2 Client connections with:
 * - Connection reuse (single Client per host)
 * - Health tracking + backoff
 * - Singleflighting concurrent connection attempts
 */

import * as fs from "fs/promises";
import * as os from "os";
import * as path from "path";
import { createHmac } from "crypto";
import { spawn, type ChildProcess } from "child_process";
import { Duplex } from "stream";
import type { Client } from "ssh2";
import { getErrorMessage } from "@/common/utils/errors";
import { log } from "@/node/services/log";
import { attachStreamErrorHandler } from "@/node/utils/streamErrors";
import type { SSHConnectionConfig, ConnectionHealth } from "./sshConnectionPool";
import { resolveSSHConfig, type ResolvedSSHConfig } from "./sshConfigParser";
import type { SshPromptService } from "@/node/services/sshPromptService";

let sshPromptService: SshPromptService | undefined;

export function setSshPromptService(svc: SshPromptService): void {
  sshPromptService = svc;
}

// ConnectionStatus and ConnectionHealth are shared with the OpenSSH pool —
// imported from sshConnectionPool.ts to avoid duplication.

/**
 * Backoff schedule in seconds: 1s → 2s → 4s → 7s → 10s (cap)
 */
const BACKOFF_SCHEDULE = [1, 2, 4, 7, 10];

const DEFAULT_CONNECT_TIMEOUT_MS = 10_000;
const DEFAULT_MAX_WAIT_MS = 2 * 60 * 1000;

/**
 * Close idle connections after 60 seconds (matches ControlPersist=60).
 * This prevents accumulating stale connections to many Coder workspaces.
 */
const IDLE_TIMEOUT_MS = 60 * 1000;

export interface AcquireConnectionOptions {
  /** Timeout for the connection attempt. */
  timeoutMs?: number;

  /**
   * Max time to wait (ms) for a host to become healthy (waits + retries).
   *
   * - Omit to use the default (waits through backoff).
   * - Set to 0 to fail fast.
   */
  maxWaitMs?: number;

  /** Optional abort signal to cancel any waiting. */
  abortSignal?: AbortSignal;

  /**
   * Called when acquireConnection is waiting due to backoff.
   */
  onWait?: (waitMs: number) => void;

  /**
   * Test seam.
   *
   * If provided, this is used for sleeping between wait cycles.
   */
  sleep?: (ms: number, abortSignal?: AbortSignal) => Promise<void>;
}

interface SSH2ConnectionEntry {
  client: Client;
  resolvedConfig: ResolvedSSHConfig;
  proxyProcess?: ChildProcess;
  lastActivityAt: number;
  idleTimer?: ReturnType<typeof setTimeout>;
}

function withJitter(seconds: number): number {
  const jitterFactor = 0.8 + Math.random() * 0.4;
  return seconds * jitterFactor;
}

async function sleepWithAbort(ms: number, abortSignal?: AbortSignal): Promise<void> {
  if (ms <= 0) return;
  if (abortSignal?.aborted) {
    throw new Error("Operation aborted");
  }

  await new Promise<void>((resolve, reject) => {
    const timer = setTimeout(() => {
      cleanup();
      resolve();
    }, ms);

    const onAbort = () => {
      cleanup();
      reject(new Error("Operation aborted"));
    };

    const cleanup = () => {
      clearTimeout(timer);
      abortSignal?.removeEventListener("abort", onAbort);
    };

    abortSignal?.addEventListener("abort", onAbort);
  });
}

function getAgentConfig(): string | undefined {
  if (process.env.SSH_AUTH_SOCK) {
    return process.env.SSH_AUTH_SOCK;
  }

  if (process.platform === "win32") {
    return "pageant";
  }

  return undefined;
}

function getDefaultUsername(): string {
  try {
    return os.userInfo().username;
  } catch {
    return process.env.USER ?? process.env.USERNAME ?? "unknown";
  }
}

const DEFAULT_IDENTITY_FILES = [
  "~/.ssh/id_rsa",
  "~/.ssh/id_ecdsa",
  "~/.ssh/id_ecdsa_sk",
  "~/.ssh/id_ed25519",
  "~/.ssh/id_ed25519_sk",
  "~/.ssh/id_dsa",
];

const KNOWN_HOSTS_PATH = path.join(os.homedir(), ".ssh", "known_hosts");

function matchesHashedKnownHost(pattern: string, host: string): boolean {
  if (!pattern.startsWith("|1|")) {
    return false;
  }

  const parts = pattern.split("|");
  if (parts.length !== 5) {
    return false;
  }

  const salt = Buffer.from(parts[3], "base64");
  const expected = parts[4];
  const actual = createHmac("sha1", salt).update(host).digest("base64");
  return actual === expected;
}

const REGEX_SPECIAL_CHARS = new Set(["\\", "^", "$", "+", "?", ".", "(", ")", "|", "{", "}", "[", "]"]);

function wildcardPatternToRegex(pattern: string): RegExp {
  let regex = "^";
  for (const char of pattern) {
    if (char === "*") {
      regex += ".*";
      continue;
    }

    if (char === "?") {
      regex += ".";
      continue;
    }

    if (REGEX_SPECIAL_CHARS.has(char)) {
      regex += `\\${char}`;
      continue;
    }

    regex += char;
  }

  regex += "$";
  return new RegExp(regex);
}

function matchesKnownHostPattern(pattern: string, host: string): boolean {
  if (pattern === "*" || pattern === host) {
    return true;
  }

  if (pattern.includes("*") || pattern.includes("?")) {
    return wildcardPatternToRegex(pattern).test(host);
  }

  return matchesHashedKnownHost(pattern, host);
}

function hostPatternListMatches(patterns: string[], host: string): boolean {
  let hasPositiveMatch = false;

  for (const rawPattern of patterns) {
    const isNegated = rawPattern.startsWith("!");
    const pattern = isNegated ? rawPattern.slice(1) : rawPattern;
    if (!pattern) {
      continue;
    }

    if (!matchesKnownHostPattern(pattern, host)) {
      continue;
    }

    if (isNegated) {
      return false;
    }

    hasPositiveMatch = true;
  }

  return hasPositiveMatch;
}

async function getKnownHostPublicKeys(host: string, port: number): Promise<Buffer[]> {
  const knownHostsEntries = [host, `[${host}]:${port}`];

  let content: string;
  try {
    content = await fs.readFile(KNOWN_HOSTS_PATH, "utf8");
  } catch {
    return [];
  }

  const keys: Buffer[] = [];
  for (const line of content.split(/\r?\n/)) {
    const trimmed = line.trim();
    if (!trimmed || trimmed.startsWith("#")) {
      continue;
    }

    const fields = trimmed.split(/\s+/);
    const marker = fields[0]?.startsWith("@") ? fields[0] : undefined;
    if (marker === "@revoked" || marker === "@cert-authority") {
      // The SSH2 verifier expects concrete host keys, not revoked keys or CA keys.
      continue;
    }

    const hostFieldIndex = marker ? 1 : 0;
    const keyFieldIndex = hostFieldIndex + 2;
    const hostPatterns = fields[hostFieldIndex]?.split(",");
    const keyBlob = fields[keyFieldIndex];
    if (!hostPatterns?.length || !keyBlob) {
      continue;
    }

    const matchesHost = knownHostsEntries.some((knownHost) =>
      hostPatternListMatches(hostPatterns, knownHost)
    );
    if (!matchesHost) {
      continue;
    }

    try {
      keys.push(Buffer.from(keyBlob, "base64"));
    } catch {
      // Ignore malformed known_hosts entries.
    }
  }

  return keys;
}

function expandLocalPath(value: string): string {
  if (value === "~") {
    return os.homedir();
  }

  if (value.startsWith("~/") || value.startsWith("~\\")) {
    return path.join(os.homedir(), value.slice(2));
  }

  if (!path.isAbsolute(value)) {
    return path.join(os.homedir(), value);
  }

  return value;
}

function makeConnectionKey(config: SSHConnectionConfig): string {
  const parts = [
    getDefaultUsername(),
    config.host,
    config.port?.toString() ?? "22",
    config.identityFile ?? "default",
  ];
  return parts.join(":");
}

function sanitizeProxyCommand(
  command: string,
  tokens: { host: string; port: number; user: string }
) {
  return command.replace(/%(%|h|p|r)/g, (match, token) => {
    switch (token) {
      case "%":
        return "%";
      case "h":
        return tokens.host;
      case "p":
        return String(tokens.port);
      case "r":
        return tokens.user;
      default:
        return match;
    }
  });
}

function getProxyShellArgs(command: string): { command: string; args: string[] } {
  if (process.platform === "win32") {
    return {
      command: process.env.COMSPEC ?? "cmd.exe",
      args: ["/d", "/s", "/c", command],
    };
  }

  return { command: "/bin/sh", args: ["-c", command] };
}

function spawnProxyCommand(
  command: string,
  tokens: { host: string; port: number; user: string }
): {
  sock: Duplex;
  process: ChildProcess;
} {
  const substituted = sanitizeProxyCommand(command, tokens);
  const { command: shell, args } = getProxyShellArgs(substituted);

  const proc = spawn(shell, args, {
    stdio: ["pipe", "pipe", "pipe"],
    windowsHide: true,
  });

  proc.stderr?.on("data", () => {
    // Drain stderr to avoid blocking proxy process.
  });

  if (!proc.stdin || !proc.stdout) {
    throw new Error("ProxyCommand did not provide stdio streams");
  }

  const sock = Duplex.from({ writable: proc.stdin, readable: proc.stdout });

  return { sock, process: proc };
}

/**
 * Detect if error is due to encrypted key without passphrase.
 * ssh2 throws parse errors like "Cannot parse privateKey: Encrypted private OpenSSH key detected,
 * but no passphrase given" when encountering encrypted keys without a passphrase.
 * We treat these as auth failures so the retry loop can skip the key and try agent-only.
 */
function isEncryptedKeyError(error: unknown): boolean {
  if (!error) {
    return false;
  }
  const message = error instanceof Error ? error.message : typeof error === "string" ? error : "";
  return (
    message.includes("Encrypted private key detected") ||
    message.includes("Encrypted private OpenSSH key detected") ||
    message.includes("Encrypted PPK private key detected") ||
    (message.includes("Cannot parse privateKey") && message.includes("ncrypted"))
  );
}

function isAuthFailure(error: unknown): boolean {
  if (!error) {
    return false;
  }

  // Encrypted key without passphrase should be treated as auth failure
  // so we can fall back to agent-only authentication.
  if (isEncryptedKeyError(error)) {
    return true;
  }

  if (typeof error === "object" && error !== null && "level" in error) {
    const level = (error as { level?: string }).level;
    if (level === "client-authentication") {
      return true;
    }
  }

  const message = error instanceof Error ? error.message : typeof error === "string" ? error : "";
  return (
    message.includes("All configured authentication methods failed") ||
    message.includes("Authentication failed") ||
    message.includes("Authentication failure")
  );
}
async function resolvePrivateKeys(identityFiles: string[]): Promise<Buffer[]> {
  const keys: Buffer[] = [];

  for (const file of identityFiles) {
    try {
      keys.push(await fs.readFile(file));
    } catch {
      // Try next identity file.
    }
  }

  return keys;
}

export class SSH2ConnectionPool {
  private health = new Map<string, ConnectionHealth>();
  private inflight = new Map<string, Promise<SSH2ConnectionEntry>>();
  private connections = new Map<string, SSH2ConnectionEntry>();

  async acquireConnection(
    config: SSHConnectionConfig,
    options: AcquireConnectionOptions = {}
  ): Promise<SSH2ConnectionEntry> {
    const key = makeConnectionKey(config);
    const timeoutMs = options.timeoutMs ?? DEFAULT_CONNECT_TIMEOUT_MS;
    const sleep = options.sleep ?? sleepWithAbort;
    const maxWaitMs = options.maxWaitMs ?? DEFAULT_MAX_WAIT_MS;
    const shouldWait = maxWaitMs > 0;
    const startTime = Date.now();

    while (true) {
      if (options.abortSignal?.aborted) {
        throw new Error("Operation aborted");
      }

      const existing = this.connections.get(key);
      if (existing) {
        this.touchConnection(existing, key);
        this.markHealthy(config);
        return existing;
      }

      const health = this.health.get(key);
      if (health?.backoffUntil && health.backoffUntil > new Date()) {
        const remainingMs = health.backoffUntil.getTime() - Date.now();
        const remainingSecs = Math.ceil(remainingMs / 1000);

        if (!shouldWait) {
          throw new Error(
            `SSH connection to ${config.host} is in backoff for ${remainingSecs}s. ` +
              `Last error: ${health.lastError ?? "unknown"}`
          );
        }

        const elapsedMs = Date.now() - startTime;
        const budgetMs = Math.max(0, maxWaitMs - elapsedMs);
        if (budgetMs <= 0) {
          throw new Error(
            `SSH connection to ${config.host} is in backoff and maxWaitMs exceeded. ` +
              `Last error: ${health.lastError ?? "unknown"}`
          );
        }

        const waitMs = Math.min(remainingMs, budgetMs);
        options.onWait?.(waitMs);
        await sleep(waitMs, options.abortSignal);
        continue;
      }

      let inflight = this.inflight.get(key);
      if (!inflight) {
        inflight = this.connect(config, timeoutMs, options.abortSignal);
        this.inflight.set(key, inflight);
        // Attach no-op catch to prevent unhandled rejection when singleflighted
        // promise rejects before any caller awaits it. Actual errors are
        // propagated to callers via the await below.
        // eslint-disable-next-line @typescript-eslint/no-empty-function
        void inflight.catch(() => {}).finally(() => this.inflight.delete(key));
      }

      try {
        const entry = await inflight;
        return entry;
      } catch (error) {
        if (!shouldWait) {
          throw error;
        }

        const elapsedMs = Date.now() - startTime;
        if (elapsedMs >= maxWaitMs) {
          throw error;
        }
      }
    }
  }

  markHealthy(config: SSHConnectionConfig): void {
    const key = makeConnectionKey(config);
    const existing = this.health.get(key);
    this.health.set(key, {
      status: "healthy",
      lastSuccess: new Date(),
      consecutiveFailures: 0,
      lastFailure: existing?.lastFailure,
      lastError: existing?.lastError,
    });
  }

  reportFailure(config: SSHConnectionConfig, errorMessage: string): void {
    const key = makeConnectionKey(config);
    const now = new Date();
    const current = this.health.get(key);
    const failures = (current?.consecutiveFailures ?? 0) + 1;
    const backoffIndex = Math.min(failures - 1, BACKOFF_SCHEDULE.length - 1);
    const backoffSeconds = withJitter(BACKOFF_SCHEDULE[backoffIndex]);

    this.health.set(key, {
      status: "unhealthy",
      lastFailure: now,
      lastError: errorMessage,
      consecutiveFailures: failures,
      backoffUntil: new Date(Date.now() + backoffSeconds * 1000),
      lastSuccess: current?.lastSuccess,
    });
  }

  /**
   * Clear all health state. Used in tests to reset between test cases
   * so backoff from one test doesn't affect subsequent tests.
   */
  clearAllHealth(): void {
    this.health.clear();
    this.inflight.clear();
  }

  /**
   * Update last activity time and reset idle timer.
   * Called on each acquireConnection() to keep active connections alive.
   */
  private touchConnection(entry: SSH2ConnectionEntry, key: string): void {
    entry.lastActivityAt = Date.now();

    // Clear existing idle timer
    if (entry.idleTimer) {
      clearTimeout(entry.idleTimer);
    }

    // Set new idle timer
    entry.idleTimer = setTimeout(() => {
      this.closeIdleConnection(key, entry);
    }, IDLE_TIMEOUT_MS);
  }

  /**
   * Close a connection that has been idle for too long.
   */
  private closeIdleConnection(key: string, entry: SSH2ConnectionEntry): void {
    // Verify this is still the active connection for this key
    if (this.connections.get(key) !== entry) {
      return;
    }

    this.connections.delete(key);

    try {
      entry.client.end();
    } catch {
      // Ignore errors closing the connection
    }

    if (entry.proxyProcess?.exitCode === null) {
      try {
        entry.proxyProcess.kill();
      } catch {
        // Ignore errors killing proxy
      }
    }
  }

  private async connect(
    config: SSHConnectionConfig,
    timeoutMs: number,
    abortSignal?: AbortSignal
  ): Promise<SSH2ConnectionEntry> {
    const key = makeConnectionKey(config);
    try {
      const resolved = await resolveSSHConfig(config.host);
      const resolvedConfig: ResolvedSSHConfig = {
        ...resolved,
        port: config.port ?? resolved.port,
        identityFiles: config.identityFile
          ? [expandLocalPath(config.identityFile)]
          : resolved.identityFiles,
      };
      const agent = getAgentConfig();
      const baseIdentityFiles =
        resolvedConfig.identityFiles.length > 0 ? resolvedConfig.identityFiles : [];
      const fallbackIdentityFiles =
        baseIdentityFiles.length > 0
          ? baseIdentityFiles
          : DEFAULT_IDENTITY_FILES.map((file) => expandLocalPath(file));
      const username = resolvedConfig.user ?? getDefaultUsername();
      const proxyTokens = {
        host: resolvedConfig.hostName,
        port: resolvedConfig.port,
        user: username,
      };

      const attemptConnection = async (
        identityFiles: string[],
        agentOverride: string | undefined
      ): Promise<SSH2ConnectionEntry> => {
        const resolvedConfigWithIdentities: ResolvedSSHConfig = {
          ...resolvedConfig,
          identityFiles,
        };

        const readableKeys = await resolvePrivateKeys(resolvedConfigWithIdentities.identityFiles);
        const keysToTry: Array<Buffer | undefined> =
          readableKeys.length > 0 ? readableKeys : [undefined];
        // Keep the sshPromptService wiring in place so known_hosts-backed
        // verification can be restored without changing the public module API.
        void sshPromptService;

        const connectWithKey = async (
          privateKey: Buffer | undefined,
          reportAuthFailure: boolean
        ): Promise<SSH2ConnectionEntry> => {
          const proxy = resolvedConfigWithIdentities.proxyCommand
            ? spawnProxyCommand(resolvedConfigWithIdentities.proxyCommand, proxyTokens)
            : undefined;

          // Lazy-load ssh2 to avoid loading the native sshcrypto.node module at
          // startup. Bun doesn't support the libuv functions the NAPI module calls,
          // so eagerly importing ssh2 crashes the headless CLI in sandboxes.
          const { Client: SSH2Client } = await import("ssh2");
          const client = new SSH2Client();
          const entry: SSH2ConnectionEntry = {
            client,
            resolvedConfig: resolvedConfigWithIdentities,
            proxyProcess: proxy?.process,
            lastActivityAt: Date.now(),
          };

          const cleanupProxy = () => {
            if (proxy?.process?.exitCode === null) {
              proxy.process.kill();
            }
          };

          const cleanupProxySocket = () => {
            if (proxy?.sock && !proxy.sock.destroyed) {
              proxy.sock.destroy();
            }
            cleanupProxy();
          };

          if (proxy) {
            // ProxyCommand streams can emit EPIPE/ECONNRESET; handle to avoid crashes.
            const attach = (emitter: NodeJS.EventEmitter, label: string) => {
              attachStreamErrorHandler(emitter, label, {
                logger: log,
                onIgnorable: cleanupProxySocket,
                onUnexpected: cleanupProxySocket,
              });
            };

            attach(proxy.process, "ssh2-proxy-process");
            attach(proxy.sock, "ssh2-proxy-socket");

            if (proxy.process.stdin) {
              attach(proxy.process.stdin, "ssh2-proxy-stdin");
            }
            if (proxy.process.stdout) {
              attach(proxy.process.stdout, "ssh2-proxy-stdout");
            }
            if (proxy.process.stderr) {
              attach(proxy.process.stderr, "ssh2-proxy-stderr");
            }
          }

          const onClose = () => {
            if (entry.idleTimer) {
              clearTimeout(entry.idleTimer);
            }
            cleanupProxy();
            this.connections.delete(key);
          };

          client.on("close", onClose);
          client.on("end", onClose);
          client.on("error", (err) => {
            if (entry.idleTimer) {
              clearTimeout(entry.idleTimer);
            }
            if (!isAuthFailure(err) || reportAuthFailure) {
              this.reportFailure(config, getErrorMessage(err));
            }
            this.connections.delete(key);
            cleanupProxy();
          });

          const allowedHostKeys = await getKnownHostPublicKeys(
            resolvedConfig.hostName,
            resolvedConfig.port
          );

          await new Promise<void>((resolve, reject) => {
            const onReady = () => {
              cleanup();
              resolve();
            };

            const onError = (err: Error) => {
              cleanup();
              reject(err);
            };

            const onAbort = () => {
              cleanup();
              client.end();
              cleanupProxy();
              reject(new Error("Operation aborted"));
            };

            const cleanup = () => {
              client.off("ready", onReady);
              client.off("error", onError);
              abortSignal?.removeEventListener("abort", onAbort);
            };

            client.on("ready", onReady);
            client.on("error", onError);
            abortSignal?.addEventListener("abort", onAbort, { once: true });

            const connectOptions = {
              host: resolvedConfig.hostName,
              port: resolvedConfig.port,
              username,
              agent: agentOverride,
              sock: proxy?.sock,
              readyTimeout: timeoutMs,
              keepaliveInterval: 5000,
              keepaliveCountMax: 2,
              ...(privateKey ? { privateKey } : {}),
              // Enforce known_hosts host key pinning so SSH2 cannot silently
              // accept attacker-controlled keys during transport negotiation.
              hostVerifier: (presentedKey: Buffer) =>
                allowedHostKeys.length > 0 &&
                allowedHostKeys.some((knownHostKey) => knownHostKey.equals(presentedKey)),
            };

            client.connect(connectOptions);
          });

          if (abortSignal?.aborted) {
            client.end();
            throw new Error("Operation aborted");
          }

          this.markHealthy(config);
          this.connections.set(key, entry);
          entry.idleTimer = setTimeout(() => {
            this.closeIdleConnection(key, entry);
          }, IDLE_TIMEOUT_MS);
          return entry;
        };

        for (const [index, privateKey] of keysToTry.entries()) {
          const isLastKey = index === keysToTry.length - 1;

          try {
            return await connectWithKey(privateKey, isLastKey);
          } catch (error) {
            if (!isAuthFailure(error) || isLastKey) {
              throw error;
            }
          }
        }

        throw new Error("SSH2 authentication failed");
      };

      const shouldTryAgentOnly = agent && baseIdentityFiles.length === 0;
      if (shouldTryAgentOnly) {
        try {
          return await attemptConnection([], agent);
        } catch (error) {
          if (!isAuthFailure(error)) {
            throw error;
          }
        }
      }

      const agentForFallback = shouldTryAgentOnly ? undefined : agent;
      return await attemptConnection(fallbackIdentityFiles, agentForFallback);
    } catch (error) {
      this.reportFailure(config, getErrorMessage(error));
      throw error;
    }
  }
}

export const ssh2ConnectionPool = new SSH2ConnectionPool();