feat: add provider ssh bridge

2026-05-05 01:24:32 +00:00
parent 2d6829ed44
commit f6d0bd1e3b
18 changed files with 1014 additions and 34 deletions
@@ -8,9 +8,19 @@ import {
  type ApiNodeDockerStatus,
  type ApiNodeSystemStatus,
  type ApiTask,
+  type CoreHostSshCloseMessage,
+  type CoreHostSshEofMessage,
+  type CoreHostSshInputMessage,
+  type CoreHostSshOpenMessage,
+  type CoreHostSshResizeMessage,
  type CoreDispatchMessage,
  type JsonValue,
+  type ProviderHostSshDataMessage,
+  type ProviderHostSshErrorMessage,
+  type ProviderHostSshExitMessage,
+  type ProviderHostSshOpenedMessage,
  type ProviderLabels,
+  isProviderDispatchCommand,
  type ProviderToCoreMessage,
  isProviderToCoreMessage,
 } from "../../shared/src/index";
@@ -27,6 +37,8 @@ interface RuntimeConfig {

 interface WsData {
  providerId?: string;
+  channel?: "provider" | "ssh";
+  sessionId?: string;
 }

 type ProviderSocket = ServerWebSocket<WsData>;
@@ -35,6 +47,7 @@ type SqlClient = ReturnType<typeof postgres>;

 const recentLogs: unknown[] = [];
 const activeProviders = new Map<string, ProviderSocket>();
+const activeSshClients = new Map<string, ProviderSocket>();
 const serviceStartedAt = new Date();
 const config = readConfig();
 const logger = createLogger("backend-core", config.logFile);
@@ -281,11 +294,30 @@ async function upsertNodeOnline(providerId: string, name: string, labels: Provid
 async function touchHeartbeat(providerId: string, labels: ProviderLabels): Promise<void> {
  await sql`
    UPDATE unidesk_nodes
-    SET labels = ${sql.json(labels)}, status = 'online', last_heartbeat = now(), updated_at = now()
+    SET labels = unidesk_nodes.labels || ${sql.json(labels)}, status = 'online', last_heartbeat = now(), updated_at = now()
    WHERE provider_id = ${providerId}
  `;
 }

+async function providerCapabilities(providerId: string): Promise<string[]> {
+  if (!dbReady) return [];
+  const rows = await sql<Array<{ labels: unknown }>>`
+    SELECT labels
+    FROM unidesk_nodes
+    WHERE provider_id = ${providerId}
+    LIMIT 1
+  `;
+  const labels = rows[0]?.labels;
+  if (typeof labels !== "object" || labels === null || Array.isArray(labels)) return [];
+  const capabilities = (labels as Record<string, unknown>).unideskCapabilities;
+  return Array.isArray(capabilities) ? capabilities.filter((item): item is string => typeof item === "string") : [];
+}
+
+async function providerSupports(providerId: string, capability: string): Promise<boolean> {
+  const capabilities = await providerCapabilities(providerId);
+  return capabilities.includes(capability);
+}
+
 async function upsertDockerStatus(providerId: string, status: JsonValue, collectedAt: string): Promise<void> {
  await sql`
    INSERT INTO unidesk_node_docker_status (provider_id, status, collected_at, updated_at)
@@ -415,17 +447,75 @@ function parseMessage(raw: string | Buffer): ProviderToCoreMessage {
  return parsed;
 }

+function safeSessionId(): string {
+  return `ssh_${Date.now()}_${Math.random().toString(16).slice(2)}`;
+}
+
+function wsSendJson(ws: ProviderSocket, value: unknown): void {
+  if (ws.readyState === WebSocket.OPEN) ws.send(JSON.stringify(value));
+}
+
+function sshClientFor(sessionId: string): ProviderSocket | null {
+  return activeSshClients.get(sessionId) ?? null;
+}
+
+function providerForSsh(providerId: string): ProviderSocket | null {
+  return activeProviders.get(providerId) ?? null;
+}
+
+function closeSshClient(sessionId: string, code = 1000, reason = "ssh session closed"): void {
+  const client = sshClientFor(sessionId);
+  activeSshClients.delete(sessionId);
+  if (client !== null && client.readyState === WebSocket.OPEN) client.close(code, reason);
+}
+
+function forwardSshProviderMessage(
+  message: ProviderHostSshOpenedMessage | ProviderHostSshDataMessage | ProviderHostSshExitMessage | ProviderHostSshErrorMessage,
+): void {
+  const client = sshClientFor(message.sessionId);
+  if (client === null) {
+    logger("warn", "ssh_client_missing", { providerId: message.providerId, sessionId: message.sessionId, type: message.type });
+    return;
+  }
+  if (message.type === "host_ssh_opened") {
+    wsSendJson(client, { type: "ssh.opened", providerId: message.providerId, sessionId: message.sessionId });
+    return;
+  }
+  if (message.type === "host_ssh_data") {
+    wsSendJson(client, { type: "ssh.data", stream: message.stream, data: message.data, encoding: message.encoding });
+    return;
+  }
+  if (message.type === "host_ssh_exit") {
+    wsSendJson(client, { type: "ssh.exit", exitCode: message.exitCode, signal: message.signal });
+    setTimeout(() => closeSshClient(message.sessionId), 50);
+    return;
+  }
+  wsSendJson(client, { type: "ssh.error", message: message.message });
+  setTimeout(() => closeSshClient(message.sessionId, 1011, "ssh session error"), 50);
+}
+
 async function handleProviderMessage(ws: ProviderSocket, raw: string | Buffer): Promise<void> {
  const message = parseMessage(raw);
  ws.data.providerId = message.providerId;
  activeProviders.set(message.providerId, ws);

+  if (
+    message.type === "host_ssh_opened" ||
+    message.type === "host_ssh_data" ||
+    message.type === "host_ssh_exit" ||
+    message.type === "host_ssh_error"
+  ) {
+    forwardSshProviderMessage(message);
+    return;
+  }
+
  if (message.type === "register") {
-    await upsertNodeOnline(message.providerId, message.name, message.labels);
+    const labels = { ...message.labels, unideskCapabilities: message.capabilities };
+    await upsertNodeOnline(message.providerId, message.name, labels);
    await recordEvent("provider_registered", message.providerId, {
      providerId: message.providerId,
      name: message.name,
-      labels: message.labels,
+      labels,
      capabilities: message.capabilities,
    });
    ws.send(JSON.stringify({ type: "ack", requestId: "register", ok: true, message: "registered" }));
@@ -648,11 +738,17 @@ async function getOverview(): Promise<JsonValue> {
 async function dispatchTask(req: Request): Promise<Response> {
  const body = (await req.json()) as { providerId?: unknown; command?: unknown; payload?: unknown };
  const providerId = typeof body.providerId === "string" ? body.providerId : "";
-  const command = body.command === "docker.ps" || body.command === "provider.upgrade" || body.command === "echo" ? body.command : "echo";
+  const command = isProviderDispatchCommand(body.command) ? body.command : null;
  const payload = typeof body.payload === "object" && body.payload !== null ? (body.payload as Record<string, JsonValue>) : {};
  if (!providerId) {
    return jsonResponse({ ok: false, error: "providerId is required" }, 400);
  }
+  if (command === null) {
+    return jsonResponse({ ok: false, error: "command must be one of docker.ps, provider.upgrade, host.ssh, echo" }, 400);
+  }
+  if (command === "host.ssh" && !(await providerSupports(providerId, "host.ssh"))) {
+    return jsonResponse({ ok: false, error: `provider does not declare host.ssh capability: ${providerId}` }, 409);
+  }
  const taskId = `task_${Date.now()}_${Math.random().toString(16).slice(2)}`;
  await sql`
    INSERT INTO unidesk_tasks (id, provider_id, command, status, payload, result)
@@ -674,11 +770,113 @@ async function dispatchTask(req: Request): Promise<Response> {
  return jsonResponse({ ok: true, taskId, status: "dispatched", providerOnline: true });
 }

-async function route(req: Request): Promise<Response> {
+function numberFromUnknown(value: unknown, fallback: number, min: number, max: number): number {
+  const parsed = typeof value === "number" ? value : typeof value === "string" ? Number(value) : fallback;
+  if (!Number.isFinite(parsed)) return fallback;
+  return Math.max(min, Math.min(max, Math.floor(parsed)));
+}
+
+async function handleSshClientMessage(ws: ProviderSocket, raw: string | Buffer): Promise<void> {
+  const text = typeof raw === "string" ? raw : raw.toString("utf8");
+  const message = JSON.parse(text) as { type?: unknown; providerId?: unknown; cwd?: unknown; command?: unknown; data?: unknown; encoding?: unknown; cols?: unknown; rows?: unknown };
+  if (message.type === "ssh.open") {
+    const providerId = typeof message.providerId === "string" ? message.providerId : "";
+    if (providerId.length === 0) {
+      wsSendJson(ws, { type: "ssh.error", message: "providerId is required" });
+      return;
+    }
+    const provider = providerForSsh(providerId);
+    if (provider === null) {
+      wsSendJson(ws, { type: "ssh.error", message: `provider is not online: ${providerId}` });
+      return;
+    }
+    if (!(await providerSupports(providerId, "host.ssh"))) {
+      wsSendJson(ws, { type: "ssh.error", message: `provider does not declare host.ssh capability: ${providerId}` });
+      return;
+    }
+    const sessionId = safeSessionId();
+    ws.data.channel = "ssh";
+    ws.data.providerId = providerId;
+    ws.data.sessionId = sessionId;
+    activeSshClients.set(sessionId, ws);
+    const openMessage: CoreHostSshOpenMessage = {
+      type: "host_ssh_open",
+      sessionId,
+      cols: numberFromUnknown(message.cols, 100, 20, 300),
+      rows: numberFromUnknown(message.rows, 30, 8, 120),
+    };
+    if (typeof message.cwd === "string" && message.cwd.length > 0) openMessage.cwd = message.cwd;
+    if (typeof message.command === "string" && message.command.length > 0) openMessage.command = message.command;
+    provider.send(JSON.stringify(openMessage));
+    wsSendJson(ws, { type: "ssh.dispatched", providerId, sessionId });
+    logger("info", "ssh_session_dispatched", { providerId, sessionId, hasCommand: typeof openMessage.command === "string" });
+    return;
+  }
+
+  const sessionId = ws.data.sessionId;
+  const providerId = ws.data.providerId;
+  if (sessionId === undefined || providerId === undefined) {
+    if (message.type === "ssh.eof" || message.type === "ssh.close") return;
+    wsSendJson(ws, { type: "ssh.error", message: "ssh session is not open" });
+    return;
+  }
+  const provider = providerForSsh(providerId);
+  if (provider === null) {
+    wsSendJson(ws, { type: "ssh.error", message: `provider went offline: ${providerId}` });
+    return;
+  }
+
+  if (message.type === "ssh.input") {
+    if (typeof message.data !== "string" || message.encoding !== "base64") {
+      wsSendJson(ws, { type: "ssh.error", message: "ssh.input requires base64 data" });
+      return;
+    }
+    const inputMessage: CoreHostSshInputMessage = { type: "host_ssh_input", sessionId, data: message.data, encoding: "base64" };
+    provider.send(JSON.stringify(inputMessage));
+    return;
+  }
+
+  if (message.type === "ssh.resize") {
+    const resizeMessage: CoreHostSshResizeMessage = {
+      type: "host_ssh_resize",
+      sessionId,
+      cols: numberFromUnknown(message.cols, 100, 20, 300),
+      rows: numberFromUnknown(message.rows, 30, 8, 120),
+    };
+    provider.send(JSON.stringify(resizeMessage));
+    return;
+  }
+
+  if (message.type === "ssh.eof") {
+    const eofMessage: CoreHostSshEofMessage = { type: "host_ssh_eof", sessionId };
+    provider.send(JSON.stringify(eofMessage));
+    return;
+  }
+
+  if (message.type === "ssh.close") {
+    const closeMessage: CoreHostSshCloseMessage = { type: "host_ssh_close", sessionId };
+    provider.send(JSON.stringify(closeMessage));
+    closeSshClient(sessionId);
+    return;
+  }
+
+  wsSendJson(ws, { type: "ssh.error", message: `unsupported ssh client message: ${String(message.type)}` });
+}
+
+async function sshRoute(req: Request, server: Server<WsData>): Promise<Response | undefined> {
+  const url = new URL(req.url);
+  const token = url.searchParams.get("token") ?? req.headers.get("x-provider-token");
+  if (token !== config.providerToken) return jsonResponse({ ok: false, error: "invalid ssh bridge token" }, 401);
+  const upgraded = server.upgrade(req, { data: { channel: "ssh" } satisfies WsData });
+  return upgraded ? undefined : jsonResponse({ ok: false, error: "websocket upgrade failed" }, 400);
+}
+
+async function route(req: Request, server: Server<WsData>): Promise<Response | undefined> {
  const url = new URL(req.url);
  if (req.method === "OPTIONS") return jsonResponse({ ok: true });

  try {
+    if (url.pathname === "/ws/ssh") return sshRoute(req, server);
    if (url.pathname === "/" || url.pathname === "/health") {
      return jsonResponse({ ok: true, service: "unidesk-core", dbReady, startedAt: serviceStartedAt.toISOString() });
    }
@@ -714,7 +912,7 @@ async function providerRoute(req: Request, server: Server<WsData>): Promise<Resp
    await recordEvent("provider_auth_failed", "unknown", { remote: req.headers.get("x-forwarded-for") ?? null });
    return jsonResponse({ ok: false, error: "invalid provider token" }, 401);
  }
-  const upgraded = server.upgrade(req, { data: {} satisfies WsData });
+  const upgraded = server.upgrade(req, { data: { channel: "provider" } satisfies WsData });
  return upgraded ? undefined : jsonResponse({ ok: false, error: "websocket upgrade failed" }, 400);
 }

@@ -733,6 +931,36 @@ const apiServer = Bun.serve<WsData>({
  port: config.port,
  hostname: "0.0.0.0",
  fetch: route,
+  websocket: {
+    open(ws) {
+      if (ws.data.channel === "ssh") logger("info", "ssh_client_open");
+    },
+    message(ws, raw) {
+      if (ws.data.channel !== "ssh") {
+        ws.close(1008, "unsupported websocket channel");
+        return;
+      }
+      handleSshClientMessage(ws, raw).catch((error) => {
+        logger("error", "ssh_client_message_failed", { providerId: ws.data.providerId ?? null, sessionId: ws.data.sessionId ?? null, error: errorToJson(error) });
+        wsSendJson(ws, { type: "ssh.error", message: error instanceof Error ? error.message : String(error) });
+      });
+    },
+    close(ws) {
+      if (ws.data.channel !== "ssh") return;
+      const { sessionId, providerId } = ws.data;
+      logger("warn", "ssh_client_close", { providerId: providerId ?? null, sessionId: sessionId ?? null });
+      if (sessionId !== undefined) {
+        activeSshClients.delete(sessionId);
+        if (providerId !== undefined) {
+          const provider = providerForSsh(providerId);
+          if (provider !== null) {
+            const closeMessage: CoreHostSshCloseMessage = { type: "host_ssh_close", sessionId };
+            provider.send(JSON.stringify(closeMessage));
+          }
+        }
+      }
+    },
+  },
 });

 const providerServer = Bun.serve<WsData>({