pikasTech-unidesk/scripts/src/ssh.ts

import { spawn } from "node:child_process";
import { type UniDeskConfig, repoRoot } from "./config";

export 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",
]);

export 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);
    });
  });
}