feat: add provider ssh bridge

src/components/provider-gateway/src/index.ts

@@ -2,6 +2,11 @@ import { appendFileSync, existsSync, mkdirSync, readFileSync } from "node:fs";
 import { dirname } from "node:path";
 import {
   type CoreDispatchMessage,
+  type CoreHostSshCloseMessage,
+  type CoreHostSshEofMessage,
+  type CoreHostSshInputMessage,
+  type CoreHostSshOpenMessage,
+  type CoreHostSshResizeMessage,
   type DockerContainerSummary,
   type DockerImageSummary,
   type DockerNetworkSummary,
@@ -33,6 +38,12 @@ interface RuntimeConfig {
   upgradeComposeProject: string;
   upgradeService: string;
   upgradeRunnerImage: string;
+  hostSshHost: string | null;
+  hostSshPort: number | null;
+  hostSshUser: string | null;
+  hostSshKey: string | null;
+  hostRemoteCwd: string | null;
+  hostLoginShell: string | null;
   logFile: string;
 }
 
@@ -49,6 +60,18 @@ let previousCpuSample: { idle: number; total: number } | null = null;
 let reconnectAttempt = 0;
 let stopping = false;
 
+interface HostSshSession {
+  proc: ReturnType<typeof Bun.spawn>;
+  openedAt: number;
+}
+
+interface HostSshStdin {
+  write(chunk: Uint8Array): unknown;
+  end(): unknown;
+}
+
+const hostSshSessions = new Map<string, HostSshSession>();
+
 function requiredEnv(name: string): string {
   const value = process.env[name];
   if (value === undefined || value.length === 0) {
@@ -66,6 +89,22 @@ function readNumberEnv(name: string): number {
   return parsed;
 }
 
+function readOptionalStringEnv(name: string): string | null {
+  const value = process.env[name];
+  if (value === undefined || value.length === 0) return null;
+  return value;
+}
+
+function readOptionalNumberEnv(name: string): number | null {
+  const raw = readOptionalStringEnv(name);
+  if (raw === null) return null;
+  const parsed = Number(raw);
+  if (!Number.isFinite(parsed) || parsed <= 0) {
+    throw new Error(`Environment variable ${name} must be a positive number, got ${raw}`);
+  }
+  return parsed;
+}
+
 function readBooleanEnv(name: string): boolean {
   const raw = requiredEnv(name);
   if (raw === "true") return true;
@@ -93,6 +132,12 @@ function readConfig(): RuntimeConfig {
     upgradeComposeProject: requiredEnv("PROVIDER_UPGRADE_COMPOSE_PROJECT"),
     upgradeService: requiredEnv("PROVIDER_UPGRADE_SERVICE"),
     upgradeRunnerImage: requiredEnv("PROVIDER_UPGRADE_RUNNER_IMAGE"),
+    hostSshHost: readOptionalStringEnv("HOST_SSH_HOST"),
+    hostSshPort: readOptionalNumberEnv("HOST_SSH_PORT"),
+    hostSshUser: readOptionalStringEnv("HOST_SSH_USER"),
+    hostSshKey: readOptionalStringEnv("HOST_SSH_KEY"),
+    hostRemoteCwd: readOptionalStringEnv("HOST_REMOTE_CWD"),
+    hostLoginShell: readOptionalStringEnv("HOST_LOGIN_SHELL"),
     logFile: requiredEnv("LOG_FILE"),
   };
 }
@@ -121,9 +166,13 @@ function withToken(rawUrl: string, token: string): string {
 }
 
 function currentLabels(): ProviderLabels {
+  const hostSshConfigured = isHostSshConfigured();
   return {
     ...config.labels,
     dockerSocketPresent: existsSync(config.dockerSocketPath),
+    hostSshConfigured,
+    hostSshKeyPresent: config.hostSshKey !== null && existsSync(config.hostSshKey),
+    hostSshTarget: hostSshConfigured ? `${config.hostSshUser}@${config.hostSshHost}:${config.hostSshPort}` : "not-configured",
     runtime: "bun",
     gatewayUptimeSeconds: Math.floor((Date.now() - startedAt.getTime()) / 1000),
   };
@@ -135,13 +184,15 @@ function sendJson(value: unknown): void {
 }
 
 function sendRegister(): void {
+  const capabilities = ["heartbeat", "system.status", "docker.status", "docker.ps", "provider.upgrade", "echo"];
+  if (isHostSshConfigured()) capabilities.push("host.ssh");
   sendJson({
     type: "register",
     providerId: config.providerId,
     name: config.providerName,
     labels: currentLabels(),
     startedAt: startedAt.toISOString(),
-    capabilities: ["heartbeat", "system.status", "docker.status", "docker.ps", "provider.upgrade", "echo"],
+    capabilities,
   });
 }
 
@@ -209,19 +260,23 @@ async function sendTaskStatus(taskId: string, status: ProviderTaskStatusMessage[
   });
 }
 
-async function runProcessCommand(command: string, args: string[], timeoutMs = 6000): Promise<{ ok: boolean; stdout: string; stderr: string; exitCode: number }> {
+async function runProcessCommand(command: string, args: string[], timeoutMs = 6000): Promise<{ ok: boolean; stdout: string; stderr: string; exitCode: number; timedOut: boolean }> {
   const proc = Bun.spawn([command, ...args], {
     stdout: "pipe",
     stderr: "pipe",
   });
-  const timeout = setTimeout(() => proc.kill("SIGKILL"), timeoutMs);
+  let timedOut = false;
+  const timeout = setTimeout(() => {
+    timedOut = true;
+    proc.kill("SIGKILL");
+  }, timeoutMs);
   const [stdout, stderr, exitCode] = await Promise.all([
     new Response(proc.stdout).text(),
     new Response(proc.stderr).text(),
     proc.exited,
   ]);
   clearTimeout(timeout);
-  return { ok: exitCode === 0, stdout, stderr, exitCode };
+  return { ok: exitCode === 0 && !timedOut, stdout, stderr, exitCode, timedOut };
 }
 
 async function runDockerCommand(args: string[], timeoutMs = 6000): Promise<{ ok: boolean; stdout: string; stderr: string; exitCode: number }> {
@@ -508,6 +563,274 @@ function shellQuote(value: string): string {
   return `'${value.replace(/'/g, `'\\''`)}'`;
 }
 
+function isHostSshConfigured(): boolean {
+  return config.hostSshHost !== null && config.hostSshPort !== null && config.hostSshUser !== null && config.hostSshKey !== null;
+}
+
+function missingHostSshFields(): string[] {
+  const missing: string[] = [];
+  if (config.hostSshHost === null) missing.push("HOST_SSH_HOST");
+  if (config.hostSshPort === null) missing.push("HOST_SSH_PORT");
+  if (config.hostSshUser === null) missing.push("HOST_SSH_USER");
+  if (config.hostSshKey === null) missing.push("HOST_SSH_KEY");
+  return missing;
+}
+
+function payloadString(payload: Record<string, JsonValue>, key: string): string | null {
+  const value = payload[key];
+  return typeof value === "string" && value.length > 0 ? value : null;
+}
+
+function payloadTimeoutMs(payload: Record<string, JsonValue>): number {
+  const raw = payload.timeoutMs;
+  const value = typeof raw === "number" ? raw : typeof raw === "string" ? Number(raw) : 8000;
+  if (!Number.isFinite(value) || value <= 0) return 8000;
+  return Math.max(1000, Math.min(15_000, Math.floor(value)));
+}
+
+function truncateText(value: string, maxLength = 4000): string {
+  return value.length > maxLength ? `${value.slice(0, maxLength)}...<truncated:${value.length}>` : value;
+}
+
+function safeKnownHostsName(providerId: string): string {
+  return providerId.replace(/[^a-zA-Z0-9_.-]/g, "-").slice(0, 80) || "provider";
+}
+
+function defaultHostSshProbeCommand(): string {
+  return "printf 'UNIDESK_SSH_TEST user=%s host=%s bridge=%s cwd=%s\\n' \"$(whoami)\" \"$(hostname)\" \"${UNIDESK_BRIDGE:-}\" \"$(pwd)\"";
+}
+
+async function runHostSsh(payload: Record<string, JsonValue>): Promise<JsonValue> {
+  if (!isHostSshConfigured()) {
+    throw new Error(`host SSH bridge is not configured; missing ${missingHostSshFields().join(", ")}`);
+  }
+  const host = config.hostSshHost ?? "";
+  const port = config.hostSshPort ?? 22;
+  const user = config.hostSshUser ?? "";
+  const key = config.hostSshKey ?? "";
+  if (!existsSync(key)) {
+    throw new Error(`host SSH key is not mounted at ${key}`);
+  }
+
+  const mode = payload.mode === "exec" ? "exec" : "probe";
+  const timeoutMs = payloadTimeoutMs(payload);
+  const requestedCommand = payloadString(payload, "command");
+  const command = mode === "exec" && requestedCommand !== null ? requestedCommand : defaultHostSshProbeCommand();
+  if (command.length > 4000) {
+    throw new Error(`host SSH command is too long: ${command.length} bytes`);
+  }
+  const cwd = payloadString(payload, "cwd") ?? config.hostRemoteCwd;
+  const scriptParts = [
+    "set -e",
+    cwd === null ? null : `cd ${shellQuote(cwd)}`,
+    "export UNIDESK_BRIDGE=host.ssh",
+    `export UNIDESK_PROVIDER_ID=${shellQuote(config.providerId)}`,
+    config.hostLoginShell === null
+      ? command
+      : `${shellQuote(config.hostLoginShell)} -lc ${shellQuote(command)}`,
+  ].filter((part): part is string => part !== null);
+  const remoteScript = scriptParts.join("; ");
+  const result = await runProcessCommand("ssh", [
+    "-T",
+    "-i",
+    key,
+    "-p",
+    String(port),
+    "-o",
+    "BatchMode=yes",
+    "-o",
+    "StrictHostKeyChecking=accept-new",
+    "-o",
+    `UserKnownHostsFile=/tmp/unidesk-host-known-hosts-${safeKnownHostsName(config.providerId)}`,
+    "-o",
+    "ServerAliveInterval=20",
+    "-o",
+    "ServerAliveCountMax=3",
+    `${user}@${host}`,
+    remoteScript,
+  ], timeoutMs);
+  const stdout = truncateText(result.stdout);
+  const stderr = truncateText(result.stderr);
+  const probeLine = stdout.split("\n").find((line) => line.includes("UNIDESK_SSH_TEST")) ?? "";
+  return {
+    ok: result.ok,
+    mode,
+    host,
+    port,
+    user,
+    cwd: cwd ?? "",
+    timeoutMs,
+    timedOut: result.timedOut,
+    exitCode: result.exitCode,
+    command: truncateText(command, 1000),
+    stdout,
+    stderr,
+    probeLine,
+    sshKeyPresent: true,
+  };
+}
+
+function hostSshRemoteScript(command: string | null, cwd: string | null, cols?: number, rows?: number): string {
+  const fallbackCwd = config.hostRemoteCwd ?? `/home/${config.hostSshUser ?? "root"}`;
+  const requestedCwd = cwd ?? fallbackCwd;
+  const loginShell = config.hostLoginShell ?? "/bin/bash";
+  const resize = Number.isFinite(cols) && Number.isFinite(rows)
+    ? `stty rows ${Math.max(8, Math.min(120, Math.floor(rows ?? 30)))} cols ${Math.max(20, Math.min(300, Math.floor(cols ?? 100)))} 2>/dev/null || true`
+    : "true";
+  const enterCwd = `cd ${shellQuote(requestedCwd)} 2>/dev/null || cd ${shellQuote(fallbackCwd)} 2>/dev/null || cd`;
+  const exports = `export UNIDESK_BRIDGE=host.ssh; export UNIDESK_PROVIDER_ID=${shellQuote(config.providerId)}`;
+  const execPart = command === null || command.length === 0
+    ? `exec ${shellQuote(loginShell)} -l`
+    : `${shellQuote(loginShell)} -lc ${shellQuote(command)}`;
+  return `${enterCwd}; ${exports}; ${resize}; ${execPart}`;
+}
+
+function hostSshArgs(remoteScript: string): string[] {
+  if (!isHostSshConfigured()) {
+    throw new Error(`host SSH bridge is not configured; missing ${missingHostSshFields().join(", ")}`);
+  }
+  const key = config.hostSshKey ?? "";
+  if (!existsSync(key)) {
+    throw new Error(`host SSH key is not mounted at ${key}`);
+  }
+  return [
+    "-tt",
+    "-i",
+    key,
+    "-p",
+    String(config.hostSshPort ?? 22),
+    "-o",
+    "BatchMode=yes",
+    "-o",
+    "StrictHostKeyChecking=accept-new",
+    "-o",
+    `UserKnownHostsFile=/tmp/unidesk-host-known-hosts-${safeKnownHostsName(config.providerId)}`,
+    "-o",
+    "ServerAliveInterval=20",
+    "-o",
+    "ServerAliveCountMax=3",
+    `${config.hostSshUser}@${config.hostSshHost}`,
+    remoteScript,
+  ];
+}
+
+async function pumpHostSshOutput(sessionId: string, streamName: "stdout" | "stderr", stream: ReadableStream<Uint8Array> | null): Promise<void> {
+  if (stream === null) return;
+  const reader = stream.getReader();
+  while (true) {
+    const chunk = await reader.read();
+    if (chunk.done) return;
+    sendJson({
+      type: "host_ssh_data",
+      providerId: config.providerId,
+      sessionId,
+      stream: streamName,
+      data: Buffer.from(chunk.value).toString("base64"),
+      encoding: "base64",
+      at: new Date().toISOString(),
+    });
+  }
+}
+
+function sendHostSshError(sessionId: string, error: unknown): void {
+  const message = error instanceof Error ? error.message : String(error);
+  logger("error", "host_ssh_session_error", { sessionId, error: message });
+  sendJson({
+    type: "host_ssh_error",
+    providerId: config.providerId,
+    sessionId,
+    message,
+    at: new Date().toISOString(),
+  });
+}
+
+function startHostSshSession(message: CoreHostSshOpenMessage): void {
+  if (hostSshSessions.has(message.sessionId)) {
+    sendHostSshError(message.sessionId, `host SSH session already exists: ${message.sessionId}`);
+    return;
+  }
+  try {
+    const remoteScript = hostSshRemoteScript(message.command ?? null, message.cwd ?? null, message.cols, message.rows);
+    const proc = Bun.spawn(["ssh", ...hostSshArgs(remoteScript)], {
+      stdin: "pipe",
+      stdout: "pipe",
+      stderr: "pipe",
+    });
+    hostSshSessions.set(message.sessionId, { proc, openedAt: Date.now() });
+    sendJson({
+      type: "host_ssh_opened",
+      providerId: config.providerId,
+      sessionId: message.sessionId,
+      at: new Date().toISOString(),
+    });
+    const stdoutDone = pumpHostSshOutput(message.sessionId, "stdout", proc.stdout);
+    const stderrDone = pumpHostSshOutput(message.sessionId, "stderr", proc.stderr);
+    Promise.all([stdoutDone, stderrDone, proc.exited])
+      .then(([, , exitCode]) => {
+        hostSshSessions.delete(message.sessionId);
+        sendJson({
+          type: "host_ssh_exit",
+          providerId: config.providerId,
+          sessionId: message.sessionId,
+          exitCode,
+          signal: null,
+          at: new Date().toISOString(),
+        });
+      })
+      .catch((error) => {
+        hostSshSessions.delete(message.sessionId);
+        sendHostSshError(message.sessionId, error);
+      });
+    logger("info", "host_ssh_session_started", { sessionId: message.sessionId, hasCommand: typeof message.command === "string", cwd: message.cwd ?? null });
+  } catch (error) {
+    sendHostSshError(message.sessionId, error);
+  }
+}
+
+function writeHostSshInput(message: CoreHostSshInputMessage): void {
+  const session = hostSshSessions.get(message.sessionId);
+  if (session === undefined) {
+    sendHostSshError(message.sessionId, `unknown host SSH session: ${message.sessionId}`);
+    return;
+  }
+  try {
+    const stdin = session.proc.stdin as HostSshStdin | undefined;
+    if (stdin === undefined) throw new Error("host SSH stdin is not available");
+    stdin.write(Buffer.from(message.data, "base64"));
+  } catch (error) {
+    sendHostSshError(message.sessionId, error);
+  }
+}
+
+function resizeHostSshSession(message: CoreHostSshResizeMessage): void {
+  logger("debug", "host_ssh_resize_requested", { sessionId: message.sessionId, cols: message.cols, rows: message.rows });
+}
+
+function eofHostSshSession(message: CoreHostSshEofMessage): void {
+  const session = hostSshSessions.get(message.sessionId);
+  if (session === undefined) return;
+  try {
+    const stdin = session.proc.stdin as HostSshStdin | undefined;
+    if (stdin !== undefined) stdin.end();
+  } catch (error) {
+    sendHostSshError(message.sessionId, error);
+  }
+}
+
+function closeHostSshSession(message: CoreHostSshCloseMessage): void {
+  const session = hostSshSessions.get(message.sessionId);
+  if (session === undefined) return;
+  hostSshSessions.delete(message.sessionId);
+  session.proc.kill("SIGTERM");
+  setTimeout(() => {
+    try {
+      session.proc.kill("SIGKILL");
+    } catch {
+      /* process may already be gone */
+    }
+  }, 1000);
+}
+
 function safeDockerName(value: string): string {
   return value.replace(/[^a-zA-Z0-9_.-]/g, "-").slice(0, 80);
 }
@@ -603,6 +926,15 @@ async function handleDispatch(message: CoreDispatchMessage): Promise<void> {
       await sendTaskStatus(message.taskId, "succeeded", "provider upgrade command completed", result);
       return;
     }
+    if (message.command === "host.ssh") {
+      const result = await runHostSsh(message.payload);
+      if ((result as { ok?: unknown }).ok !== true) {
+        await sendTaskStatus(message.taskId, "failed", "host SSH command failed", result);
+        return;
+      }
+      await sendTaskStatus(message.taskId, "succeeded", "host SSH command completed", result);
+      return;
+    }
     await sendTaskStatus(message.taskId, "succeeded", "echo completed", { echo: message.payload });
   } catch (error) {
     const text = error instanceof Error ? `${error.name}: ${error.message}` : String(error);
@@ -618,6 +950,26 @@ function handleMessage(raw: MessageEvent<string>): void {
       handleDispatch(parsed as CoreDispatchMessage).catch((error) => logger("error", "dispatch_handler_failed", { error: String(error) }));
       return;
     }
+    if (parsed.type === "host_ssh_open") {
+      startHostSshSession(parsed as CoreHostSshOpenMessage);
+      return;
+    }
+    if (parsed.type === "host_ssh_input") {
+      writeHostSshInput(parsed as CoreHostSshInputMessage);
+      return;
+    }
+    if (parsed.type === "host_ssh_resize") {
+      resizeHostSshSession(parsed as CoreHostSshResizeMessage);
+      return;
+    }
+    if (parsed.type === "host_ssh_eof") {
+      eofHostSshSession(parsed as CoreHostSshEofMessage);
+      return;
+    }
+    if (parsed.type === "host_ssh_close") {
+      closeHostSshSession(parsed as CoreHostSshCloseMessage);
+      return;
+    }
     logger("debug", "core_message", parsed as JsonValue);
   } catch (error) {
     logger("error", "core_message_parse_failed", { error: String(error) });
@@ -676,6 +1028,11 @@ process.on("SIGTERM", () => {
   if (heartbeatTimer !== null) clearInterval(heartbeatTimer);
   if (systemStatusTimer !== null) clearInterval(systemStatusTimer);
   if (dockerStatusTimer !== null) clearInterval(dockerStatusTimer);
+  for (const [sessionId, session] of hostSshSessions) {
+    logger("warn", "host_ssh_session_terminated_by_shutdown", { sessionId });
+    session.proc.kill("SIGTERM");
+  }
+  hostSshSessions.clear();
   socket?.close(1000, "provider shutdown");
   process.exit(0);
 });