openlattice-ssh 0.0.3

package/src/ssh-provider.ts

@@ -0,0 +1,559 @@
+import { readFileSync } from "fs";
+import { Client as SSHClient } from "ssh2";
+import type { ClientChannel, SFTPWrapper } from "ssh2";
+import type {
+  ComputeProvider,
+  ComputeSpec,
+  ExecOpts,
+  ExecResult,
+  ExtensionMap,
+  FileEntry,
+  FileExtension,
+  HealthStatus,
+  NetworkExtension,
+  ProviderCapabilities,
+  ProviderNode,
+  ProviderNodeStatus,
+} from "openlattice";
+import type { SSHProviderConfig, SSHHostConfig } from "./config";
+
+/** Internal state for a provisioned SSH node. */
+interface NodeState {
+  host: SSHHostConfig;
+  workdir: string;
+  pid?: number;
+  startedAt: Date;
+  /** Original command for restart support. */
+  command?: string;
+  /** Original env prefix for restart support. */
+  envPrefix?: string;
+}
+
+let nextSessionId = 1;
+
+export class SSHProvider implements ComputeProvider {
+  readonly name = "ssh";
+  readonly capabilities: ProviderCapabilities;
+  private readonly config: SSHProviderConfig;
+  private readonly connections = new Map<string, SSHClient>();
+  private readonly nodes = new Map<string, NodeState>();
+  private roundRobinIdx = 0;
+
+  constructor(config: SSHProviderConfig) {
+    if (!config.hosts || config.hosts.length === 0) {
+      throw new Error("[ssh] at least one host must be configured");
+    }
+    // Load private key from file path if defaultKeyPath is set and no key content
+    if (config.defaultKeyPath && !config.defaultPrivateKey) {
+      config.defaultPrivateKey = readFileSync(config.defaultKeyPath);
+    }
+    this.config = config;
+    const hasGpu = config.hosts.some((h) => h.gpuAvailable);
+    this.capabilities = {
+      restart: true,
+      pause: false, // SIGSTOP unreliable via SSH
+      snapshot: false,
+      gpu: hasGpu,
+      logs: false,
+      tailscale: true,
+      coldStartMs: 1000,
+      maxConcurrent: 0,
+      architectures: ["x86_64", "arm64"],
+      persistentStorage: true,
+    };
+  }
+
+  // ── Required methods ────────────────────────────────────────────
+
+  async provision(spec: ComputeSpec): Promise<ProviderNode> {
+    const host = this.selectHost(spec);
+    const conn = await this.getConnection(host);
+    const sessionId = `ssh-${nextSessionId++}`;
+    const hostKey = hostKeyOf(host);
+    const externalId = `${hostKey}/${sessionId}`;
+    const workdir = `/tmp/openlattice/${sessionId}`;
+
+    // Create working directory and validate connectivity
+    await this.sshExec(conn, `mkdir -p ${workdir}`);
+
+    // Store node state
+    const state: NodeState = {
+      host,
+      workdir,
+      startedAt: new Date(),
+    };
+    this.nodes.set(externalId, state);
+
+    // Ensure Tailscale is installed and running if requested (authKey implies tailscale)
+    if (spec.network?.tailscale || spec.network?.tailscaleAuthKey) {
+      await this.ensureTailscale(conn);
+      if (spec.network?.tailscaleAuthKey) {
+        await this.tailscaleUp(conn, spec.network.tailscaleAuthKey);
+      }
+    }
+
+    // Run initial command in background if specified
+    if (spec.runtime.command && spec.runtime.command.length > 0) {
+      const cmd = spec.runtime.command.join(" ");
+      const envPrefix = buildEnvPrefix(spec.runtime.env);
+      state.command = cmd;
+      state.envPrefix = envPrefix;
+      const result = await this.sshExec(
+        conn,
+        `cd ${workdir} && ${envPrefix}nohup ${cmd} > ${workdir}/.stdout 2> ${workdir}/.stderr & echo $!`
+      );
+      const pid = parseInt(result.stdout.trim(), 10);
+      if (!isNaN(pid)) {
+        state.pid = pid;
+      }
+    }
+
+    return {
+      externalId,
+      endpoints: [
+        {
+          type: "ssh",
+          host: host.host,
+          port: host.port ?? 22,
+        },
+      ],
+      metadata: { hostKey, sessionId, workdir },
+    };
+  }
+
+  async exec(
+    externalId: string,
+    command: string[],
+    opts?: ExecOpts
+  ): Promise<ExecResult> {
+    const state = this.getNodeState(externalId);
+    const conn = await this.getConnection(state.host);
+    const cmdStr = command.join(" ");
+    const cwd = opts?.cwd ?? state.workdir;
+    const envPrefix = buildEnvPrefix(opts?.env);
+    const fullCmd = `cd ${cwd} && ${envPrefix}${cmdStr}`;
+
+    return this.sshExec(conn, fullCmd, opts);
+  }
+
+  async destroy(externalId: string): Promise<void> {
+    const state = this.nodes.get(externalId);
+    if (!state) return; // idempotent
+
+    try {
+      const conn = await this.getConnection(state.host);
+
+      // Kill managed process if any
+      if (state.pid) {
+        await this.sshExec(conn, `kill ${state.pid} 2>/dev/null || true`);
+      }
+
+      // Remove working directory
+      await this.sshExec(conn, `rm -rf ${state.workdir}`);
+    } catch {
+      // Best-effort cleanup
+    }
+
+    this.nodes.delete(externalId);
+  }
+
+  async inspect(externalId: string): Promise<ProviderNodeStatus> {
+    const state = this.nodes.get(externalId);
+    if (!state) {
+      return { status: "terminated" };
+    }
+
+    try {
+      const conn = await this.getConnection(state.host);
+
+      if (state.pid) {
+        // Check if process is still running
+        const result = await this.sshExec(
+          conn,
+          `kill -0 ${state.pid} 2>/dev/null && echo running || echo stopped`
+        );
+        const status = result.stdout.trim() === "running" ? "running" : "stopped";
+        return {
+          status,
+          startedAt: state.startedAt,
+        };
+      }
+
+      // No PID tracked — just check SSH connectivity
+      await this.sshExec(conn, "echo ok");
+      return {
+        status: "running",
+        startedAt: state.startedAt,
+      };
+    } catch {
+      return { status: "unknown" };
+    }
+  }
+
+  // ── Optional: stop / start ──────────────────────────────────────
+
+  async stop(externalId: string): Promise<void> {
+    const state = this.getNodeState(externalId);
+    if (!state.pid) return;
+
+    const conn = await this.getConnection(state.host);
+    await this.sshExec(conn, `kill ${state.pid} 2>/dev/null || true`);
+    state.pid = undefined;
+  }
+
+  async start(externalId: string): Promise<void> {
+    const state = this.getNodeState(externalId);
+    const conn = await this.getConnection(state.host);
+
+    // Re-run the original command if one was stored
+    if (state.command) {
+      const envPrefix = state.envPrefix ?? "";
+      const result = await this.sshExec(
+        conn,
+        `cd ${state.workdir} && ${envPrefix}nohup ${state.command} > ${state.workdir}/.stdout 2> ${state.workdir}/.stderr & echo $!`
+      );
+      const pid = parseInt(result.stdout.trim(), 10);
+      if (!isNaN(pid)) {
+        state.pid = pid;
+      }
+    } else {
+      // No command to restart, just verify connectivity
+      await this.sshExec(conn, "echo ok");
+    }
+  }
+
+  // ── Optional: healthCheck ───────────────────────────────────────
+
+  async healthCheck(): Promise<HealthStatus> {
+    const start = Date.now();
+    const results: Array<{ host: string; ok: boolean; error?: string }> = [];
+
+    for (const host of this.config.hosts) {
+      try {
+        const conn = await this.getConnection(host);
+        await this.sshExec(conn, "echo ok");
+        results.push({ host: host.host, ok: true });
+      } catch (err: unknown) {
+        results.push({
+          host: host.host,
+          ok: false,
+          error: err instanceof Error ? err.message : String(err),
+        });
+      }
+    }
+
+    const allHealthy = results.every((r) => r.ok);
+    const unhealthy = results.filter((r) => !r.ok);
+
+    return {
+      healthy: allHealthy,
+      latencyMs: Date.now() - start,
+      message: allHealthy
+        ? undefined
+        : `Unhealthy hosts: ${unhealthy.map((r) => `${r.host} (${r.error})`).join(", ")}`,
+    };
+  }
+
+  // ── Optional: extensions ────────────────────────────────────────
+
+  getExtension<K extends keyof ExtensionMap>(
+    externalId: string,
+    extension: K
+  ): ExtensionMap[K] | undefined {
+    if (extension === "files") {
+      return this.createFileExtension(externalId) as ExtensionMap[K];
+    }
+    if (extension === "network") {
+      return this.createNetworkExtension(externalId) as ExtensionMap[K];
+    }
+    return undefined;
+  }
+
+  /** Close all SSH connections. Call when shutting down the provider. */
+  async close(): Promise<void> {
+    for (const conn of this.connections.values()) {
+      conn.end();
+    }
+    this.connections.clear();
+    this.nodes.clear();
+  }
+
+  // ── Private helpers ─────────────────────────────────────────────
+
+  private createNetworkExtension(externalId: string): NetworkExtension {
+    const state = this.getNodeState(externalId);
+    const host = state.host.host;
+    return {
+      async getUrl(port: number): Promise<string> {
+        return `http://${host}:${port}`;
+      },
+    };
+  }
+
+  private selectHost(spec: ComputeSpec): SSHHostConfig {
+    let candidates = [...this.config.hosts];
+
+    // Filter by GPU if requested
+    if (spec.gpu && spec.gpu.count > 0) {
+      candidates = candidates.filter((h) => h.gpuAvailable);
+      if (candidates.length === 0) {
+        throw new Error("[ssh] no hosts with GPU available");
+      }
+    }
+
+    // Filter by labels if specified
+    if (spec.labels) {
+      candidates = candidates.filter((h) => {
+        if (!h.labels) return false;
+        return Object.entries(spec.labels!).every(
+          ([k, v]) => h.labels![k] === v
+        );
+      });
+      if (candidates.length === 0) {
+        throw new Error("[ssh] no hosts matching labels");
+      }
+    }
+
+    if (candidates.length === 0) {
+      throw new Error("[ssh] no hosts available");
+    }
+
+    // Round-robin selection
+    const host = candidates[this.roundRobinIdx % candidates.length];
+    this.roundRobinIdx++;
+    return host;
+  }
+
+  private getNodeState(externalId: string): NodeState {
+    const state = this.nodes.get(externalId);
+    if (!state) {
+      throw new Error(`[ssh] node not found: ${externalId}`);
+    }
+    return state;
+  }
+
+  private async getConnection(host: SSHHostConfig): Promise<SSHClient> {
+    const key = hostKeyOf(host);
+    const existing = this.connections.get(key);
+    if (existing) return existing;
+
+    const conn = new SSHClient();
+
+    return new Promise<SSHClient>((resolve, reject) => {
+      conn.on("ready", () => {
+        this.connections.set(key, conn);
+        resolve(conn);
+      });
+      conn.on("error", (err) => {
+        this.connections.delete(key);
+        reject(new Error(`[ssh] connection error (${key}): ${err.message}`));
+      });
+      conn.on("close", () => {
+        this.connections.delete(key);
+      });
+
+      conn.connect({
+        host: host.host,
+        port: host.port ?? 22,
+        username: host.username ?? this.config.defaultUser,
+        privateKey: host.privateKey ?? this.config.defaultPrivateKey,
+        password: host.password,
+        readyTimeout: this.config.connectTimeoutMs ?? 20_000,
+        keepaliveInterval: this.config.keepaliveIntervalMs ?? 10_000,
+        keepaliveCountMax: 3,
+      });
+    });
+  }
+
+  private async ensureTailscale(conn: SSHClient): Promise<void> {
+    // Check if tailscale is installed
+    const check = await this.sshExec(conn, "which tailscale 2>/dev/null");
+    if (check.exitCode !== 0) {
+      // Install Tailscale via official install script
+      const install = await this.sshExec(
+        conn,
+        "curl -fsSL https://tailscale.com/install.sh | sh",
+        { timeoutMs: 120_000 }
+      );
+      if (install.exitCode !== 0) {
+        throw new Error(
+          `[ssh] failed to install tailscale: ${install.stderr.trim()}`
+        );
+      }
+    }
+
+    // Ensure tailscaled is running
+    const daemonCheck = await this.sshExec(
+      conn,
+      "pgrep tailscaled >/dev/null 2>&1 || sudo tailscaled --state=/var/lib/tailscale/tailscaled.state &",
+      { timeoutMs: 10_000 }
+    );
+    if (daemonCheck.exitCode !== 0) {
+      // Try systemd as fallback
+      await this.sshExec(conn, "sudo systemctl start tailscaled 2>/dev/null || true", {
+        timeoutMs: 10_000,
+      });
+    }
+  }
+
+  private async tailscaleUp(conn: SSHClient, authKey: string): Promise<void> {
+    const result = await this.sshExec(
+      conn,
+      `sudo tailscale up --authkey=${authKey}`,
+      { timeoutMs: 30_000 }
+    );
+    if (result.exitCode !== 0) {
+      throw new Error(`[ssh] tailscale up failed: ${result.stderr.trim()}`);
+    }
+  }
+
+  private sshExec(
+    conn: SSHClient,
+    command: string,
+    opts?: ExecOpts
+  ): Promise<ExecResult> {
+    return new Promise<ExecResult>((resolve, reject) => {
+      let settled = false;
+      let timer: ReturnType<typeof setTimeout> | undefined;
+
+      const settle = (fn: () => void) => {
+        if (settled) return;
+        settled = true;
+        if (timer) clearTimeout(timer);
+        fn();
+      };
+
+      conn.exec(command, (err: Error | undefined, stream: ClientChannel) => {
+        if (err) {
+          return settle(() => reject(new Error(`[ssh] exec failed: ${err.message}`)));
+        }
+
+        let stdout = "";
+        let stderr = "";
+
+        // Enforce timeout
+        if (opts?.timeoutMs && opts.timeoutMs > 0) {
+          timer = setTimeout(() => {
+            stream.close();
+            settle(() =>
+              resolve({
+                exitCode: 124, // Conventional timeout exit code
+                stdout,
+                stderr: stderr + `\n[ssh] command timed out after ${opts.timeoutMs}ms`,
+              })
+            );
+          }, opts.timeoutMs);
+        }
+
+        stream.on("data", (data: Buffer) => {
+          const str = data.toString();
+          stdout += str;
+          opts?.onStdout?.(str);
+        });
+
+        stream.stderr.on("data", (data: Buffer) => {
+          const str = data.toString();
+          stderr += str;
+          opts?.onStderr?.(str);
+        });
+
+        stream.on("close", (code: number | null) => {
+          settle(() => resolve({ exitCode: code ?? 1, stdout, stderr }));
+        });
+      });
+    });
+  }
+
+  private getSftp(conn: SSHClient): Promise<SFTPWrapper> {
+    return new Promise((resolve, reject) => {
+      conn.sftp((err, sftp) => {
+        if (err) return reject(new Error(`[ssh] SFTP failed: ${err.message}`));
+        resolve(sftp);
+      });
+    });
+  }
+
+  private createFileExtension(externalId: string): FileExtension {
+    const provider = this;
+    return {
+      async read(path: string): Promise<string | Buffer> {
+        const state = provider.getNodeState(externalId);
+        const conn = await provider.getConnection(state.host);
+        const sftp = await provider.getSftp(conn);
+        return new Promise((resolve, reject) => {
+          sftp.readFile(path, (err, data) => {
+            if (err) return reject(err);
+            resolve(data.toString());
+          });
+        });
+      },
+      async write(path: string, content: string | Buffer): Promise<void> {
+        const state = provider.getNodeState(externalId);
+        const conn = await provider.getConnection(state.host);
+        const sftp = await provider.getSftp(conn);
+        return new Promise((resolve, reject) => {
+          sftp.writeFile(
+            path,
+            typeof content === "string" ? content : content,
+            (err) => {
+              if (err) return reject(err);
+              resolve();
+            }
+          );
+        });
+      },
+      async list(dirPath: string): Promise<FileEntry[]> {
+        const state = provider.getNodeState(externalId);
+        const conn = await provider.getConnection(state.host);
+        const sftp = await provider.getSftp(conn);
+        return new Promise((resolve, reject) => {
+          sftp.readdir(dirPath, (err, list) => {
+            if (err) return reject(err);
+            resolve(
+              list.map((entry) => ({
+                name: entry.filename,
+                path: `${dirPath}/${entry.filename}`,
+                type: (entry.longname.startsWith("d")
+                  ? "directory"
+                  : "file") as "file" | "directory",
+                size: entry.attrs.size,
+              }))
+            );
+          });
+        });
+      },
+      async remove(path: string): Promise<void> {
+        const state = provider.getNodeState(externalId);
+        const conn = await provider.getConnection(state.host);
+        const sftp = await provider.getSftp(conn);
+        return new Promise((resolve, reject) => {
+          sftp.unlink(path, (err) => {
+            if (err) return reject(err);
+            resolve();
+          });
+        });
+      },
+      async mkdir(dirPath: string): Promise<void> {
+        const state = provider.getNodeState(externalId);
+        const conn = await provider.getConnection(state.host);
+        // Use exec with mkdir -p for recursive directory creation
+        await provider.sshExec(conn, `mkdir -p '${dirPath.replace(/'/g, "'\\''")}'`);
+      },
+    };
+  }
+}
+
+// ── Utility functions ───────────────────────────────────────────────
+
+function hostKeyOf(host: SSHHostConfig): string {
+  return `${host.host}:${host.port ?? 22}`;
+}
+
+function buildEnvPrefix(env?: Record<string, string>): string {
+  if (!env) return "";
+  return (
+    Object.entries(env)
+      .map(([k, v]) => `${k}='${v.replace(/'/g, "'\\''")}'`)
+      .join(" ") + " "
+  );
+}

package/tests/conformance.test.ts

@@ -0,0 +1,28 @@
+import { describe, it, expect, beforeEach, afterEach } from "vitest";
+import { runProviderConformanceTests } from "openlattice/testing";
+import { SSHProvider } from "../src/ssh-provider";
+
+const HAS_SSH = process.env.TEST_SSH === "1";
+
+describe.skipIf(!HAS_SSH)("SSHProvider conformance", () => {
+  runProviderConformanceTests(
+    {
+      createProvider: () =>
+        new SSHProvider({
+          hosts: [
+            {
+              host: process.env.SSH_HOST ?? "localhost",
+              port: parseInt(process.env.SSH_PORT ?? "22", 10),
+              username: process.env.SSH_USER ?? "root",
+              privateKey: process.env.SSH_KEY,
+            },
+          ],
+        }),
+      createSpec: () => ({
+        runtime: { image: "any" },
+      }),
+      timeoutMs: 30_000,
+    },
+    { describe, it, expect, beforeEach, afterEach }
+  );
+});