my-pi/pi-ssh/src/ssh2-transport.ts

/// <reference path="./ssh2-shim.d.ts" />
import { createHash, randomBytes } from "node:crypto";
import { readFileSync } from "node:fs";
import { posix as posixPath } from "node:path";
import { Client, type ClientChannel, type ConnectConfig, type SFTPWrapper } from "ssh2";
import { expandUserPath, type SshHostConfig } from "./config.ts";

export interface RemoteExecOptions {
  onData: (data: Buffer) => void;
  signal?: AbortSignal;
  timeout?: number;
}

export interface RemoteTransport {
  connect(): Promise<void>;
  dispose(): Promise<void>;
  exec(command: string, cwd: string, options: RemoteExecOptions): Promise<{ exitCode: number | null }>;
  capture(command: string, cwd?: string, timeout?: number): Promise<{ exitCode: number | null; output: Buffer }>;
  readFile(remotePath: string): Promise<Buffer>;
  ensureReadable(remotePath: string): Promise<void>;
  ensureReadableWritable(remotePath: string): Promise<void>;
  detectImageMimeType(remotePath: string): Promise<string | null>;
  mkdir(remoteDir: string): Promise<void>;
  writeFile(remotePath: string, content: Buffer): Promise<void>;
}

const DEFAULT_TIMEOUT_SECONDS = 300;

class CommandQueue {
  private tail: Promise<void> = Promise.resolve();

  enqueue<T>(task: () => Promise<T>): Promise<T> {
    const run = this.tail.then(task, task);
    this.tail = run.then(() => undefined, () => undefined);
    return run;
  }
}

function shellQuote(value: string): string {
  return `'${value.replace(/'/g, `'"'"'`)}'`;
}

function keyAlgorithm(key: Buffer): string {
  if (key.length < 4) return "unknown";
  const length = key.readUInt32BE(0);
  if (length < 1 || length > key.length - 4) return "unknown";
  return key.subarray(4, 4 + length).toString("ascii");
}

export function fingerprintHostKey(key: Buffer): { algorithm: string; fingerprint: string } {
  return {
    algorithm: keyAlgorithm(key),
    fingerprint: `SHA256:${createHash("sha256").update(key).digest("base64").replace(/=+$/, "")}`,
  };
}

function buildConnectConfig(host: SshHostConfig): ConnectConfig {
  const config: ConnectConfig = {
    host: host.hostName,
    port: host.port,
    username: host.user,
    readyTimeout: 20_000,
    keepaliveInterval: 15_000,
    keepaliveCountMax: 3,
    hostVerifier: (key) => fingerprintHostKey(key).fingerprint === host.hostKey.fingerprint,
  };
  if (host.auth.type === "password") {
    config.password = host.auth.password;
    config.tryKeyboard = host.auth.method !== "password";
  } else {
    config.privateKey = readFileSync(expandUserPath(host.auth.identityFile));
    if (host.auth.passphrase) config.passphrase = host.auth.passphrase;
  }
  return config;
}

function errorMessage(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}

export class Ssh2Transport implements RemoteTransport {
  private readonly client: Client;
  private readonly queue = new CommandQueue();
  private connected = false;
  private disposed = false;
  private disconnectError: Error | null = null;
  private sftpClient: SFTPWrapper | null = null;
  private readonly host: SshHostConfig;

  constructor(host: SshHostConfig, client: Client = new Client()) {
    this.host = host;
    this.client = client;
  }

  async connect(): Promise<void> {
    if (this.connected) return;
    if (this.disposed) throw new Error("SSH2 transport is disposed");
    await new Promise<void>((resolve, reject) => {
      let settled = false;
      const succeed = () => {
        if (settled) return;
        settled = true;
        this.connected = true;
        resolve();
      };
      const fail = (error: unknown) => {
        const normalized = error instanceof Error ? error : new Error(String(error));
        this.disconnectError = normalized;
        if (settled) return;
        settled = true;
        reject(normalized);
      };
      this.client.once("ready", succeed);
      this.client.on("error", fail);
      this.client.on("close", () => {
        this.connected = false;
        this.sftpClient = null;
        if (!this.disposed && !this.disconnectError) this.disconnectError = new Error("SSH2 connection closed unexpectedly");
      });
      if (this.host.auth.type === "password" && this.host.auth.method !== "password") {
        this.client.on("keyboard-interactive", (_name, _instructions, _language, prompts, finish) => {
          finish(prompts.map(() => this.host.auth.type === "password" ? this.host.auth.password : ""));
        });
      }
      try {
        this.client.connect(buildConnectConfig(this.host));
      } catch (error) {
        fail(error);
      }
    });
  }

  async dispose(): Promise<void> {
    this.disposed = true;
    this.connected = false;
    this.sftpClient = null;
    this.client.end();
  }

  private assertConnected(): void {
    if (!this.connected) throw this.disconnectError ?? new Error("SSH2 connection is not active");
  }

  exec(command: string, cwd: string, options: RemoteExecOptions): Promise<{ exitCode: number | null }> {
    return this.queue.enqueue(() => this.execUnqueued(command, cwd, options));
  }

  private async execUnqueued(command: string, cwd: string, options: RemoteExecOptions): Promise<{ exitCode: number | null }> {
    this.assertConnected();
    const remoteCommand = `cd -- ${shellQuote(cwd)} && bash -lc ${shellQuote(command)} </dev/null`;
    return new Promise((resolve, reject) => {
      let channel: ClientChannel | undefined;
      let settled = false;
      let timeoutHandle: NodeJS.Timeout | undefined;
      const effectiveTimeout = options.timeout ?? DEFAULT_TIMEOUT_SECONDS;
      const cleanup = () => {
        if (timeoutHandle) clearTimeout(timeoutHandle);
        if (options.signal) options.signal.removeEventListener("abort", onAbort);
      };
      const fail = (error: Error) => {
        if (settled) return;
        settled = true;
        cleanup();
        reject(error);
      };
      const onAbort = () => {
        try {
          channel?.signal("KILL");
          channel?.close();
        } catch {
          // channel may already be closed
        }
        fail(new Error("SSH command aborted"));
      };

      if (options.signal?.aborted) {
        fail(new Error("SSH command aborted"));
        return;
      }
      options.signal?.addEventListener("abort", onAbort, { once: true });
      if (effectiveTimeout > 0) {
        timeoutHandle = setTimeout(() => {
          try {
            channel?.signal("KILL");
            channel?.close();
          } catch {
            // channel may already be closed
          }
          fail(new Error(`SSH command timed out after ${effectiveTimeout}s`));
        }, effectiveTimeout * 1000);
      }

      this.client.exec(remoteCommand, (error, stream) => {
        if (error) {
          fail(error);
          return;
        }
        channel = stream;
        stream.on("data", (data: Buffer | string) => options.onData(Buffer.isBuffer(data) ? data : Buffer.from(data)));
        stream.stderr.on("data", (data: Buffer | string) => options.onData(Buffer.isBuffer(data) ? data : Buffer.from(data)));
        stream.once("error", fail);
        stream.once("close", (code: number | undefined) => {
          if (settled) return;
          settled = true;
          cleanup();
          resolve({ exitCode: typeof code === "number" ? code : null });
        });
      });
    });
  }

  async capture(command: string, cwd = ".", timeout = 30): Promise<{ exitCode: number | null; output: Buffer }> {
    const chunks: Buffer[] = [];
    const result = await this.exec(command, cwd, { timeout, onData: (data) => chunks.push(data) });
    return { exitCode: result.exitCode, output: Buffer.concat(chunks) };
  }

  private async sftp(): Promise<SFTPWrapper> {
    this.assertConnected();
    if (this.sftpClient) return this.sftpClient;
    this.sftpClient = await new Promise<SFTPWrapper>((resolve, reject) => {
      this.client.sftp((error, sftp) => error ? reject(error) : resolve(sftp));
    });
    return this.sftpClient;
  }

  async readFile(remotePath: string): Promise<Buffer> {
    return this.queue.enqueue(async () => {
      const sftp = await this.sftp();
      return new Promise<Buffer>((resolve, reject) => {
        sftp.readFile(remotePath, (error, data) => error ? reject(error) : resolve(data));
      });
    });
  }

  private async ensureOpen(remotePath: string, flags: "r" | "r+"): Promise<void> {
    return this.queue.enqueue(async () => {
      const sftp = await this.sftp();
      await new Promise<void>((resolve, reject) => {
        sftp.open(remotePath, flags, (error, handle) => {
          if (error) {
            reject(error);
            return;
          }
          sftp.close(handle, (closeError) => closeError ? reject(closeError) : resolve());
        });
      });
    });
  }

  ensureReadable(remotePath: string): Promise<void> {
    return this.ensureOpen(remotePath, "r");
  }

  ensureReadableWritable(remotePath: string): Promise<void> {
    return this.ensureOpen(remotePath, "r+");
  }

  async detectImageMimeType(remotePath: string): Promise<string | null> {
    const content = await this.readFile(remotePath);
    if (content.subarray(0, 3).equals(Buffer.from([0xff, 0xd8, 0xff]))) return "image/jpeg";
    if (content.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))) return "image/png";
    const prefix = content.subarray(0, 6).toString("ascii");
    if (prefix === "GIF87a" || prefix === "GIF89a") return "image/gif";
    if (content.subarray(0, 4).toString("ascii") === "RIFF" && content.subarray(8, 12).toString("ascii") === "WEBP") return "image/webp";
    return null;
  }

  async mkdir(remoteDir: string): Promise<void> {
    return this.queue.enqueue(() => this.mkdirUnqueued(remoteDir));
  }

  private async mkdirUnqueued(remoteDir: string): Promise<void> {
    const sftp = await this.sftp();
    const normalized = posixPath.normalize(remoteDir);
    const segments = normalized.split("/").filter(Boolean);
    let current = normalized.startsWith("/") ? "/" : "";
    for (const segment of segments) {
      current = current === "/" ? `/${segment}` : current ? `${current}/${segment}` : segment;
      await new Promise<void>((resolve, reject) => {
        sftp.mkdir(current, (error) => {
          if (!error) {
            resolve();
            return;
          }
          sftp.stat(current, (statError) => statError ? reject(error) : resolve());
        });
      });
    }
  }

  async writeFile(remotePath: string, content: Buffer): Promise<void> {
    return this.queue.enqueue(async () => {
      const sftp = await this.sftp();
      await this.mkdirUnqueued(posixPath.dirname(remotePath));
      const temporary = `${remotePath}.pi-ssh-${randomBytes(8).toString("hex")}.tmp`;
      const write = (path: string, data: Buffer) => new Promise<void>((resolve, reject) => {
        sftp.writeFile(path, data, (error) => error ? reject(error) : resolve());
      });
      const standardRename = (from: string, to: string) => new Promise<void>((resolve, reject) => {
        sftp.rename(from, to, (error) => error ? reject(error) : resolve());
      });
      const atomicRename = typeof sftp.ext_openssh_rename === "function"
        ? (from: string, to: string) => new Promise<void>((resolve, reject) => {
            sftp.ext_openssh_rename?.(from, to, (error) => error ? reject(error) : resolve());
          })
        : undefined;
      try {
        await write(temporary, content);
        if (atomicRename) {
          await atomicRename(temporary, remotePath);
        } else {
          try {
            await standardRename(temporary, remotePath);
          } catch {
            // SFTP v3 rename commonly refuses to replace an existing target.
            // Fall back to a direct write without deleting the existing file first.
            await write(remotePath, content);
            await new Promise<void>((resolve) => sftp.unlink(temporary, () => resolve()));
          }
        }
      } catch (error) {
        await new Promise<void>((resolve) => sftp.unlink(temporary, () => resolve()));
        throw new Error(`remote write failed: ${errorMessage(error)}`);
      }
    });
  }
}

export async function probeHostKey(
  hostName: string,
  port: number,
  user: string,
  client: Client = new Client(),
  timeoutMs = 15_000,
): Promise<{ algorithm: string; fingerprint: string }> {
  return new Promise((resolve, reject) => {
    let observed: { algorithm: string; fingerprint: string } | undefined;
    const timer = setTimeout(() => {
      client.destroy();
      reject(new Error("timed out while obtaining SSH host key"));
    }, timeoutMs);
    client.on("error", (error) => {
      clearTimeout(timer);
      client.destroy();
      observed ? resolve(observed) : reject(error);
    });
    client.on("close", () => {
      clearTimeout(timer);
      if (observed) resolve(observed);
    });
    client.connect({
      host: hostName,
      port,
      username: user,
      readyTimeout: Math.min(timeoutMs, 12_000),
      hostVerifier: (key) => {
        observed = fingerprintHostKey(key);
        return false;
      },
    });
  });
}