pikasTech-unidesk/scripts/src/ssh.test.ts

import { createHash } from "node:crypto";
import { describe, expect, test } from "bun:test";
import {
  createPosixApplyPatchFileSystem,
  createSshStdoutForwarder,
  formatSshStdoutTruncationHint,
  parseSshInvocation,
  sshStdoutStreamMaxBytes,
  sshStdoutTruncationHint,
  windowsFsReadOnlyScript,
  windowsPowerShellScriptPrelude,
} from "./ssh";

function sha256Hex(text: string): string {
  return createHash("sha256").update(text, "utf8").digest("hex");
}

describe("ssh windows PowerShell safety prelude", () => {
  test("shadows ConvertTo-Json and strips filesystem ETS metadata", () => {
    const prelude = windowsPowerShellScriptPrelude("C:\\test");

    expect(prelude).toContain("function ConvertTo-UniDeskPlainJsonValue");
    expect(prelude).toContain("function ConvertTo-Json");
    expect(prelude).toContain("'PSPath','PSParentPath','PSChildName','PSDrive','PSProvider','ReadCount'");
    expect(prelude).toContain("Microsoft.PowerShell.Utility\\ConvertTo-Json");
    expect(prelude).toContain("Set-Location -LiteralPath 'C:\\test'");
  });
});

describe("ssh windows fs read-only operations", () => {
  test("routes common read-only commands through the Windows fs backend", () => {
    const invocation = parseSshInvocation("D601:win/c/test", ["cat", "hello.md"]);
    const rgInvocation = parseSshInvocation("D601:win/c/test", ["rg", "-i", "needle", "."]);

    expect(invocation.route.plane).toBe("win");
    expect(invocation.route.workspace).toBe("C:\\test");
    expect(invocation.parsed.invocationKind).toBe("helper");
    expect(invocation.parsed.requiresStdin).toBe(false);
    expect(invocation.parsed.remoteCommand).toContain("powershell.exe");
    expect(invocation.parsed.remoteCommand).toContain("-EncodedCommand");
    expect(rgInvocation.parsed.invocationKind).toBe("helper");
    expect(rgInvocation.parsed.requiresStdin).toBe(false);
  });

  test("routes git status, diff, and commit through Windows cmd convenience", () => {
    const statusInvocation = parseSshInvocation("D601:win/F/Work/ConStart", ["git", "status", "--short", "--branch"]);
    const diffInvocation = parseSshInvocation("D601:win/F/Work/ConStart", ["git", "diff", "--check"]);
    const commitInvocation = parseSshInvocation("D601:win/F/Work/ConStart", ["git", "commit", "-m", "fix: update docs"]);

    expect(statusInvocation.route.plane).toBe("win");
    expect(statusInvocation.route.workspace).toBe("F:\\Work\\ConStart");
    expect(statusInvocation.parsed.invocationKind).toBe("argv");
    expect(statusInvocation.parsed.requiresStdin).toBe(false);
    expect(statusInvocation.parsed.remoteCommand).toContain("powershell.exe");
    expect(diffInvocation.parsed.invocationKind).toBe("argv");
    expect(diffInvocation.parsed.requiresStdin).toBe(false);
    expect(commitInvocation.parsed.invocationKind).toBe("argv");
    expect(commitInvocation.parsed.requiresStdin).toBe(false);
  });

  test("builds bounded UTF-8 scripts for cat, ls, and rg", () => {
    const catScript = windowsFsReadOnlyScript("F:\\Work\\demo", "cat", ["--max-bytes", "4096", "中文.md"]);
    const lsScript = windowsFsReadOnlyScript("F:\\Work\\demo", "ls", ["-la", "--limit=5"]);
    const rgScript = windowsFsReadOnlyScript("F:\\Work\\demo", "rg", ["-i", "--max-count=5", "needle", "."]);

    expect(catScript).toContain("$operation = 'cat';");
    expect(catScript).toContain("F:\\Work\\demo");
    expect(catScript).toContain("binary file refused by windows fs read operation");
    expect(catScript).toContain("file is not valid UTF-8");
    expect(catScript).toContain("file too large for windows fs read operation");
    expect(lsScript).toContain("$operation = 'ls';");
    expect(lsScript).toContain("UNIDESK_WINDOWS_FS_TRUNCATED");
    expect(lsScript).toContain("TYPE BYTES UPDATED NAME");
    expect(rgScript).toContain("$operation = 'rg';");
    expect(rgScript).toContain("unsupported rg option on Windows route");
    expect(rgScript).toContain("UNIDESK_WINDOWS_FS_SKIPPED");
  });

  test("rejects unsupported Windows read operations instead of treating them as POSIX", () => {
    expect(() => parseSshInvocation("D601:win/c/test", ["sed", "-n", "1p", "hello.md"])).toThrow("unsupported ssh win operation: sed");
  });

  test("reports route-aware hints for repeated win operation and interactive git commit", () => {
    expect(() => parseSshInvocation("D601:win/F/Work/ConStart", ["win", "ps"])).toThrow("route D601:win/F/Work/ConStart already selects the Windows plane");
    expect(() => parseSshInvocation("D601:win/F/Work/ConStart", ["git", "commit"])).toThrow("ssh win git commit would open an editor");
  });
});

describe("ssh host apply-patch fs backend", () => {
  test("uses POSIX bulk fs operations for host routes", async () => {
    const invocation = parseSshInvocation("D601:/mnt/f/Work/ConStart", ["apply-patch"]);
    const text = "alpha\nold\nomega\n";
    const calls: Array<{ operation: string; args: string[]; input?: string }> = [];
    const fs = createPosixApplyPatchFileSystem(invocation, async (command, input) => {
      const operation = command[4] ?? "";
      const args = command.slice(5);
      calls.push({ operation, args, input });
      if (operation === "read-bulk-b64") {
        return {
          exitCode: 0,
          stderr: "",
          stdout: [
            "UNIDESK_APPLY_PATCH_V2_BULK_READ 1",
            [
              Buffer.from("docs/test.md", "utf8").toString("base64"),
              String(Buffer.byteLength(text, "utf8")),
              sha256Hex(text),
              Buffer.from(text, "utf8").toString("base64"),
            ].join(" "),
            "",
          ].join("\n"),
        };
      }
      if (operation === "apply-replacements-bulk-stdin") {
        expect(input).toContain(Buffer.from("docs/test.md", "utf8").toString("base64"));
        return { exitCode: 0, stdout: "", stderr: "" };
      }
      throw new Error(`unexpected operation ${operation}`);
    });

    const files = await fs.readFiles!(["docs/test.md"]);
    await fs.applyReplacementsBulk!(["docs/test.md"], new Map([[
      "docs/test.md",
      {
        path: "docs/test.md",
        originalBytes: Buffer.byteLength(text, "utf8"),
        originalSha256: sha256Hex(text),
        finalBytes: Buffer.byteLength("alpha\nnew\nomega\n", "utf8"),
        finalSha256: sha256Hex("alpha\nnew\nomega\n"),
        replacements: [[1, 1, ["new"]]],
      },
    ]]));

    expect(files.get("docs/test.md")).toBe(text);
    expect(calls.map((call) => call.operation)).toEqual(["read-bulk-b64", "apply-replacements-bulk-stdin"]);
    expect(calls[0]?.args).toEqual(["docs/test.md"]);
    expect(calls[1]?.args).toEqual(["1"]);
  });
});

describe("ssh stdout bounded streaming", () => {
  test("uses bounded defaults and clamps env override", () => {
    expect(sshStdoutStreamMaxBytes({} as NodeJS.ProcessEnv)).toBe(256 * 1024);
    expect(sshStdoutStreamMaxBytes({ UNIDESK_SSH_STDOUT_STREAM_MAX_BYTES: "8192" } as NodeJS.ProcessEnv)).toBe(8192);
    expect(sshStdoutStreamMaxBytes({ UNIDESK_SSH_STDOUT_STREAM_MAX_BYTES: "1" } as NodeJS.ProcessEnv)).toBe(4096);
    expect(sshStdoutStreamMaxBytes({ UNIDESK_SSH_STDOUT_STREAM_MAX_BYTES: String(64 * 1024 * 1024) } as NodeJS.ProcessEnv)).toBe(16 * 1024 * 1024);
  });

  test("formats truncation hint without echoing remote command", () => {
    const invocation = parseSshInvocation("D601:win", ["ps"]);
    expect(invocation.parsed.remoteCommand).not.toBeNull();
    const hint = sshStdoutTruncationHint({
      invocation,
      transport: "backend-core-broker",
      thresholdBytes: 4096,
      observedBytesAtTruncation: 5000,
      dumpPath: "/tmp/unidesk-cli-output/sample.stdout.bin",
    });
    const formatted = formatSshStdoutTruncationHint(hint);

    expect(formatted.startsWith("UNIDESK_SSH_STDOUT_TRUNCATED ")).toBe(true);
    const payload = JSON.parse(formatted.slice("UNIDESK_SSH_STDOUT_TRUNCATED ".length)) as Record<string, unknown>;
    expect(payload.code).toBe("ssh-stdout-truncated");
    expect(payload.providerId).toBe("D601");
    expect(payload.route).toBe("D601:win");
    expect(payload.thresholdBytes).toBe(4096);
    expect(formatted).not.toContain("Get-Content");
  });

  test("forwarder bounds stdout and emits one truncation hint", () => {
    const invocation = parseSshInvocation("D601:win", ["ps"]);
    const forwarded: Buffer[] = [];
    const stdout = {
      write(chunk: string | Buffer): boolean {
        forwarded.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
        return true;
      },
    } as NodeJS.WritableStream;
    const forwarder = createSshStdoutForwarder({
      invocation,
      transport: "frontend-websocket",
      maxBytes: 5,
      stdout,
    });

    expect(forwarder.write(Buffer.from("abc"))).toBeNull();
    const hint = forwarder.write(Buffer.from("defgh"));
    expect(hint).not.toBeNull();
    expect(forwarder.write(Buffer.from("ijkl"))).toBeNull();
    expect(Buffer.concat(forwarded).toString("utf8")).toBe("abcde");
    expect(hint).toContain("UNIDESK_SSH_STDOUT_TRUNCATED");
    expect(hint).toContain("\"transport\":\"frontend-websocket\"");
  });
});

describe("ssh removed shell aliases", () => {
  test("reports route-aware replacement examples for trans script", () => {
    const previousEntrypoint = process.env.UNIDESK_SSH_ENTRYPOINT;
    process.env.UNIDESK_SSH_ENTRYPOINT = "trans";
    try {
      parseSshInvocation("D601:/tmp", ["script", "--", "pwd"]);
      throw new Error("expected script alias to fail");
    } catch (error) {
      expect(error).toBeInstanceOf(Error);
      const payload = error as Error & {
        code?: string;
        entrypoint?: string;
        route?: string;
        operation?: string;
        replacementExamples?: {
          posixSh?: string;
          bash?: string;
        };
      };
      expect(payload.code).toBe("ssh-removed-shell-alias");
      expect(payload.entrypoint).toBe("trans");
      expect(payload.route).toBe("D601:/tmp");
      expect(payload.operation).toBe("script");
      expect(payload.replacementExamples?.posixSh).toBe("trans D601:/tmp sh -- 'pwd'");
      expect(payload.replacementExamples?.bash).toBe("trans D601:/tmp bash -- 'pwd'");
    } finally {
      if (previousEntrypoint === undefined) {
        delete process.env.UNIDESK_SSH_ENTRYPOINT;
      } else {
        process.env.UNIDESK_SSH_ENTRYPOINT = previousEntrypoint;
      }
    }
  });
});