feat: add provider ssh bridge

scripts/src/ssh.ts

@@ -0,0 +1,186 @@
+import { spawn } from "node:child_process";
+import { type UniDeskConfig, repoRoot } from "./config";
+
+interface ParsedSshArgs {
+  remoteCommand: string | null;
+}
+
+const sshOptionsWithValue = new Set([
+  "-B", "-b", "-c", "-D", "-E", "-e", "-F", "-I", "-i", "-J", "-L", "-l", "-m", "-O", "-o", "-p", "-Q", "-R", "-S", "-W", "-w",
+]);
+
+function parseSshArgs(args: string[]): ParsedSshArgs {
+  const remote: string[] = [];
+  let remoteStarted = false;
+  for (let index = 0; index < args.length; index += 1) {
+    const arg = args[index] ?? "";
+    if (remoteStarted) {
+      remote.push(arg);
+      continue;
+    }
+    if (arg === "--") {
+      remoteStarted = true;
+      continue;
+    }
+    if (arg.startsWith("-") && arg !== "-") {
+      if (sshOptionsWithValue.has(arg) && index + 1 < args.length) index += 1;
+      continue;
+    }
+    remoteStarted = true;
+    remote.push(arg);
+  }
+  return { remoteCommand: remote.length === 0 ? null : remote.join(" ") };
+}
+
+function brokerSource(): string {
+  return String.raw`
+const open = JSON.parse(process.argv[2] || process.argv[1] || "{}");
+const token = process.env.PROVIDER_TOKEN || "";
+const url = "ws://127.0.0.1:8080/ws/ssh?token=" + encodeURIComponent(token);
+const ws = new WebSocket(url);
+let exitCode = 255;
+let canSend = false;
+let opened = false;
+const pending = [];
+const openTimer = setTimeout(() => {
+  if (opened) return;
+  process.stderr.write("unidesk ssh bridge timed out waiting for provider session\n");
+  try { ws.close(); } catch {}
+  process.exit(255);
+}, Number(open.openTimeoutMs || 15000));
+
+function send(value) {
+  const text = JSON.stringify(value);
+  if (!canSend || ws.readyState !== WebSocket.OPEN) {
+    pending.push(text);
+    return;
+  }
+  ws.send(text);
+}
+
+function flush() {
+  while (pending.length > 0 && ws.readyState === WebSocket.OPEN) {
+    ws.send(pending.shift());
+  }
+}
+
+function decodeData(data) {
+  return typeof data === "string" ? data : Buffer.from(data).toString("utf8");
+}
+
+ws.addEventListener("open", () => {
+  canSend = true;
+  send({
+    type: "ssh.open",
+    providerId: open.providerId,
+    command: open.command || undefined,
+    cwd: open.cwd || undefined,
+    cols: open.cols || 100,
+    rows: open.rows || 30,
+  });
+  flush();
+});
+
+ws.addEventListener("message", (event) => {
+  const message = JSON.parse(decodeData(event.data));
+  if (message.type === "ssh.data") {
+    opened = true;
+    const chunk = Buffer.from(message.data || "", "base64");
+    if (message.stream === "stderr") process.stderr.write(chunk);
+    else process.stdout.write(chunk);
+    return;
+  }
+  if (message.type === "ssh.opened") {
+    opened = true;
+    clearTimeout(openTimer);
+    return;
+  }
+  if (message.type === "ssh.dispatched") return;
+  if (message.type === "ssh.error") {
+    clearTimeout(openTimer);
+    process.stderr.write(String(message.message || "ssh bridge error") + "\n");
+    exitCode = 255;
+    ws.close();
+    return;
+  }
+  if (message.type === "ssh.exit") {
+    clearTimeout(openTimer);
+    exitCode = Number.isInteger(message.exitCode) ? message.exitCode : 255;
+    ws.close();
+  }
+});
+
+ws.addEventListener("close", () => {
+  process.exit(exitCode);
+});
+
+ws.addEventListener("error", () => {
+  process.stderr.write("unidesk ssh bridge websocket error\n");
+  process.exit(255);
+});
+
+process.stdin.on("data", (chunk) => {
+  send({ type: "ssh.input", data: Buffer.from(chunk).toString("base64"), encoding: "base64" });
+  flush();
+});
+process.stdin.on("end", () => {
+  send({ type: "ssh.eof" });
+  flush();
+});
+`;
+}
+
+function terminalSize(): { cols: number; rows: number } {
+  return {
+    cols: Number(process.stdout.columns) > 0 ? Number(process.stdout.columns) : 100,
+    rows: Number(process.stdout.rows) > 0 ? Number(process.stdout.rows) : 30,
+  };
+}
+
+export async function runSsh(config: UniDeskConfig, providerId: string, args: string[]): Promise<number> {
+  if (!providerId) throw new Error("ssh requires provider id, for example: bun scripts/cli.ts ssh D518");
+  const parsed = parseSshArgs(args);
+  const size = terminalSize();
+  const payload = {
+    providerId,
+    command: parsed.remoteCommand,
+    cols: size.cols,
+    rows: size.rows,
+  };
+  const child = spawn("docker", [
+    "exec",
+    "-i",
+    "unidesk-backend-core",
+    "bun",
+    "-e",
+    brokerSource(),
+    "--",
+    JSON.stringify(payload),
+  ], {
+    cwd: repoRoot,
+    stdio: ["pipe", "pipe", "pipe"],
+  });
+
+  const rawMode = parsed.remoteCommand === null && process.stdin.isTTY;
+  if (rawMode) process.stdin.setRawMode(true);
+  process.stdin.resume();
+  process.stdin.pipe(child.stdin);
+  child.stdout.pipe(process.stdout);
+  child.stderr.pipe(process.stderr);
+
+  return await new Promise<number>((resolve) => {
+    const restore = (): void => {
+      process.stdin.unpipe(child.stdin);
+      if (rawMode) process.stdin.setRawMode(false);
+    };
+    child.on("error", (error) => {
+      restore();
+      process.stderr.write(`unidesk ssh failed to start broker: ${error.message}\n`);
+      resolve(255);
+    });
+    child.on("close", (code) => {
+      restore();
+      resolve(code ?? 255);
+    });
+  });
+}