fix(tunnels): support TOTP for client tunnels and jump hosts (#1508)

src/backend/hosts/connect/connect-host.ts

@@ -120,6 +120,7 @@ export async function connectHost(
   } else {
     ({ jumpClient } = await openSshTransport(host, config, {
       log: options.log,
+      prompt: options.prompt,
       ...options.transport,
     }));
   }

src/backend/hosts/connect/transport.ts

@@ -19,6 +19,7 @@ import type {
   MutableConnectConfig,
   SshAuthLog,
   SshConnectHost,
+  SshPromptChannel,
 } from "./types.js";
 
 function getHostSocks5Config(host: SshConnectHost): SOCKS5Config | null {
@@ -48,6 +49,7 @@ class SshTransportError extends Error {
 }
 
 export interface OpenTransportOptions {
+  prompt?: SshPromptChannel;
   /** Knock before connecting when the host has a sequence. Default true. */
   portKnock?: boolean;
   /** Resolve DNS up front for a direct connection. Default true. */
@@ -111,7 +113,11 @@ export async function openSshTransport(
 
   const jumpUserId = host.userId || "";
   if (host.jumpHosts && host.jumpHosts.length > 0 && jumpUserId) {
-    const jumpClient = await createJumpHostChain(host.jumpHosts, jumpUserId);
+    const jumpClient = await createJumpHostChain(
+      host.jumpHosts,
+      jumpUserId,
+      options.prompt,
+    );
     if (!jumpClient) {
       throw new SshTransportError(
         "Failed to connect through jump hosts",

src/backend/hosts/jump-host-chain.ts

@@ -4,8 +4,11 @@ import { createSocks5Connection } from "../utils/socks5-helper.js";
 import { getErrorMessage } from "../utils/error-message.js";
 import { getJumpHostSocks5Config } from "./jump-host-proxy.js";
 import { buildConnectConfig } from "./connect/build-connect-config.js";
-import { createAutoKeyboardInteractiveHandler } from "./connect/keyboard-interactive.js";
-import type { SshConnectHost } from "./connect/types.js";
+import {
+  createAutoKeyboardInteractiveHandler,
+  createPromptKeyboardInteractiveHandler,
+} from "./connect/keyboard-interactive.js";
+import type { SshConnectHost, SshPromptChannel } from "./connect/types.js";
 import { resolveHostById } from "./host-resolver.js";
 
 type JumpHostConfig = {
@@ -61,6 +64,7 @@ class JumpHostChainError extends Error {
 export async function createJumpHostChain(
   jumpHosts: Array<{ hostId: number }>,
   userId: string,
+  prompt?: SshPromptChannel,
 ): Promise<SSHClient | null> {
   if (!jumpHosts || jumpHosts.length === 0) {
     return null;
@@ -68,6 +72,12 @@ export async function createJumpHostChain(
 
   let currentClient: SSHClient | null = null;
   const clients: SSHClient[] = [];
+  let closed = false;
+  const closeChain = () => {
+    if (closed) return;
+    closed = true;
+    for (const client of clients) client.end();
+  };
 
   try {
     const jumpHostConfigs: Array<Awaited<ReturnType<typeof resolveJumpHost>>> =
@@ -87,7 +97,7 @@ export async function createJumpHostChain(
           hopIndex: i,
           totalHops,
         });
-        clients.forEach((c) => c.end());
+        closeChain();
         throw new JumpHostChainError(
           `Jump host ${i + 1} of ${totalHops} was not found`,
           i,
@@ -108,10 +118,12 @@ export async function createJumpHostChain(
     }
 
     for (let i = 0; i < jumpHostConfigs.length; i++) {
+      if (closed) throw new Error("Jump host chain closed");
       const jumpHostConfig = jumpHostConfigs[i]!;
 
       const jumpClient = new SSHClient();
       clients.push(jumpClient);
+      jumpClient.once("close", closeChain);
 
       let lastError: Error | null = null;
 
@@ -128,6 +140,11 @@ export async function createJumpHostChain(
           jumpClient.end();
         }, readyTimeoutMs + 5000);
 
+        jumpClient.once("close", () => {
+          clearTimeout(timeout);
+          lastError = new Error("Jump host connection closed");
+          resolve(false);
+        });
         jumpClient.on("ready", () => {
           clearTimeout(timeout);
           resolve(true);
@@ -161,8 +178,20 @@ export async function createJumpHostChain(
         // every auth type, not just password, key and agent.
         const built = await buildConnectConfig(
           jumpHostConfig as unknown as SshConnectHost,
-          { userId, purpose: "jump-host", profile: "jump", client: jumpClient },
+          {
+            userId,
+            purpose: "jump-host",
+            profile: "jump",
+            client: jumpClient,
+            interactive: !!prompt,
+            hostKeySocket: prompt?.hostKeySocket ?? null,
+          },
         );
+        if (closed) {
+          clearTimeout(timeout);
+          resolve(false);
+          return;
+        }
         if (built.outcome.status !== "ready") {
           clearTimeout(timeout);
           lastError = new Error(
@@ -176,9 +205,28 @@ export async function createJumpHostChain(
 
         jumpClient.on(
           "keyboard-interactive",
-          createAutoKeyboardInteractiveHandler(
-            jumpHostConfig as unknown as SshConnectHost,
-          ),
+          prompt
+            ? createPromptKeyboardInteractiveHandler(
+                jumpHostConfig as unknown as SshConnectHost,
+                {
+                  ...prompt,
+                  ask: (request) =>
+                    prompt.ask(
+                      request.kind === "browser"
+                        ? {
+                            ...request,
+                            instructions: `Jump host ${i + 1}/${totalHops} (${jumpHostConfig.ip}): ${request.instructions}`,
+                          }
+                        : {
+                            ...request,
+                            prompt: `Jump host ${i + 1}/${totalHops} (${jumpHostConfig.ip}): ${request.prompt}`,
+                          },
+                    ),
+                },
+              )
+            : createAutoKeyboardInteractiveHandler(
+                jumpHostConfig as unknown as SshConnectHost,
+              ),
         );
 
         if (currentClient) {
@@ -207,7 +255,7 @@ export async function createJumpHostChain(
       });
 
       if (!connected) {
-        clients.forEach((c) => c.end());
+        closeChain();
         throw new JumpHostChainError(
           getErrorMessage(
             lastError,
@@ -227,7 +275,7 @@ export async function createJumpHostChain(
     fileLogger.error("Failed to create jump host chain", error, {
       operation: "jump_host_chain",
     });
-    clients.forEach((c) => c.end());
+    closeChain();
     return null;
   }
 }

src/backend/tests/electron/c2s-auth.test.ts

@@ -0,0 +1,86 @@
+import { createRequire } from "node:module";
+import { EventEmitter } from "node:events";
+import { afterEach, describe, expect, it, vi } from "vitest";
+const require = createRequire(import.meta.url);
+const { createC2SAuthBridge } = require("../../../../electron/c2s-auth.cjs");
+
+function setup() {
+  const sender = Object.assign(new EventEmitter(), {
+    send: vi.fn(),
+    isDestroyed: () => false,
+  });
+  const ws = Object.assign(new EventEmitter(), {
+    readyState: 1,
+    send: vi.fn(),
+    close: vi.fn(() => {
+      ws.readyState = 3;
+      ws.emit("close");
+    }),
+  });
+  let answer!: (
+    event: { sender: unknown },
+    id: string,
+    value: unknown,
+  ) => boolean;
+  const attach = createC2SAuthBridge(
+    {
+      handle: (_name: string, fn: typeof answer) => {
+        answer = fn;
+      },
+    },
+    () => ({ isDestroyed: () => false, webContents: sender }),
+  );
+  attach(ws, "dynamic");
+  const prompt = () =>
+    ws.emit(
+      "message",
+      Buffer.from(
+        JSON.stringify({
+          type: "auth-prompt",
+          requestId: "server-challenge",
+          request: { kind: "totp", prompt: "Code:", retry: false },
+        }),
+      ),
+      false,
+    );
+  return { sender, ws, answer, prompt };
+}
+afterEach(() => vi.useRealTimers());
+describe("desktop tunnel prompt bridge", () => {
+  it("binds answers to the renderer and challenge, and rejects replay", () => {
+    const { sender, ws, answer, prompt } = setup();
+    prompt();
+    const id = sender.send.mock.calls[0][1].id;
+    expect(answer({ sender: {} }, id, "123456")).toBe(false);
+    expect(answer({ sender }, "unknown", "123456")).toBe(false);
+    expect(answer({ sender }, id, { secret: "bad" })).toBe(false);
+    expect(ws.send).not.toHaveBeenCalled();
+    expect(answer({ sender }, id, "123456")).toBe(true);
+    expect(JSON.parse(ws.send.mock.calls[0][0])).toEqual({
+      type: "auth-response",
+      requestId: "server-challenge",
+      answer: "123456",
+    });
+    expect(answer({ sender }, id, "again")).toBe(false);
+    expect(sender.listenerCount("destroyed")).toBe(0);
+    ws.close();
+  });
+  it.each(["close", "reload", "timeout"])(
+    "clears a pending dialog on %s",
+    (reason) => {
+      vi.useFakeTimers();
+      const { sender, ws, answer, prompt } = setup();
+      prompt();
+      const id = sender.send.mock.calls[0][1].id;
+      if (reason === "close") ws.close();
+      else if (reason === "reload") sender.emit("did-start-loading");
+      else vi.advanceTimersByTime(60_000);
+      expect(sender.send).toHaveBeenLastCalledWith("c2s-auth-prompt", {
+        id,
+        closed: true,
+      });
+      expect(answer({ sender }, id, "late")).toBe(false);
+      expect(ws.close).toHaveBeenCalled();
+    },
+  );
+});

src/backend/tests/electron/c2s-runtime.test.ts

@@ -0,0 +1,143 @@
+import { EventEmitter } from "node:events";
+import { readFileSync } from "node:fs";
+import vm from "node:vm";
+import { describe, expect, it, vi } from "vitest";
+
+const main = readFileSync(
+  new URL("../../../../electron/main.cjs", import.meta.url),
+  "utf8",
+);
+const section = (from: string, to: string) =>
+  main.slice(main.indexOf(from), main.indexOf(to, main.indexOf(from)));
+function setup() {
+  const sockets: Socket[] = [];
+  class Socket extends EventEmitter {
+    readyState = 1;
+    send = vi.fn();
+    close = vi.fn(() => {
+      if (this.readyState === 3) return;
+      this.readyState = 3;
+      this.emit("close");
+    });
+    constructor() {
+      super();
+      sockets.push(this);
+      queueMicrotask(() => this.emit("open"));
+    }
+    ready() {
+      this.emit(
+        "message",
+        Buffer.from(JSON.stringify({ type: "ready" })),
+        false,
+      );
+    }
+  }
+  const runtimes = new Map();
+  let incoming!: (socket: EventEmitter) => void;
+  const server = Object.assign(new EventEmitter(), {
+    listen: (_options: unknown, ready: () => void) => ready(),
+    close: vi.fn((done?: () => void) => done?.()),
+  });
+  const headers = vi.fn(async () => ({}));
+  const dynamic = vi.fn();
+  const context = vm.createContext({
+    setTimeout,
+    clearTimeout,
+    WebSocket: Socket,
+    crypto: { randomUUID: () => "runtime-token" },
+    net: {
+      createServer: (callback: typeof incoming) => {
+        incoming = callback;
+        return server;
+      },
+    },
+    c2sTunnelRuntimes: runtimes,
+    normalizeC2STunnelAddresses: (t: unknown) => t,
+    resolveC2SRemoteSourceHost: async (t: unknown) => t,
+    getC2STunnelName: (t: { name: string }) => t.name,
+    getC2SLocalAddress: () => "127.0.0.1",
+    getC2SRemoteAddress: () => "127.0.0.1",
+    getC2SRelayUrl: () => "wss://server/relay",
+    getC2SRelayHeaders: headers,
+    getWebSocketOptions: (_url: string, options: unknown) => options,
+    attachC2SAuth: vi.fn(),
+    logToFile: vi.fn(),
+    emitC2STunnelStatuses: vi.fn(),
+    checkLocalPortAvailable: async () => ({ available: true }),
+    handleC2SDynamicConnection: dynamic,
+    setC2STunnelStatus: (name: string, status: unknown) => {
+      runtimes.get(name).status = status;
+    },
+    setC2STunnelError: (name: string, error: string) => {
+      runtimes.get(name).status = { connected: false, reason: error };
+    },
+    createC2SFailure: (error: string) => ({ success: false, error }),
+  });
+  vm.runInContext(
+    section("async function testC2SRelay(", "async function testC2STunnel(") +
+      section("async function startC2STunnel(", "function stopAllC2STunnels("),
+    context,
+  );
+  return {
+    context,
+    runtimes,
+    sockets,
+    server,
+    headers,
+    dynamic,
+    incoming: () => incoming,
+  };
+}
+
+describe("client tunnel runtime authentication", () => {
+  it("holds preflight open, shares its runtime identity and token, and releases it when stopped", async () => {
+    const { context, runtimes, sockets, server, headers, dynamic, incoming } =
+      setup();
+    const tunnel = {
+      name: "socks",
+      mode: "dynamic",
+      sourceHostId: 7,
+      sourcePort: 1080,
+    };
+    await context.startC2STunnel(tunnel, 0, "local-auth");
+    await vi.waitFor(() => expect(sockets).toHaveLength(1));
+    const premature = Object.assign(new EventEmitter(), { destroy: vi.fn() });
+    incoming()(premature);
+    expect(premature.destroy).toHaveBeenCalledOnce();
+    expect(headers).toHaveBeenCalledWith("wss://server/relay", "local-auth");
+    expect(JSON.parse(sockets[0].send.mock.calls[0][0])).toMatchObject({
+      keepAlive: true,
+      tunnelConfig: { sessionId: "runtime-token" },
+    });
+    sockets[0].ready();
+    await vi.waitFor(() =>
+      expect(runtimes.get("socks").status.connected).toBe(true),
+    );
+    expect(sockets[0].close).not.toHaveBeenCalled();
+    const client = Object.assign(new EventEmitter(), { destroy: vi.fn() });
+    incoming()(client);
+    expect(dynamic).toHaveBeenCalledWith(
+      expect.objectContaining({ sessionId: "runtime-token" }),
+      client,
+      "local-auth",
+    );
+    await context.stopC2STunnel("socks");
+    expect(sockets[0].close).toHaveBeenCalledOnce();
+    expect(client.destroy).toHaveBeenCalledOnce();
+    expect(server.close).toHaveBeenCalledOnce();
+    expect(runtimes.size).toBe(0);
+  });
+
+  it("does not retain a socket for a standalone test or a stopped runtime", async () => {
+    const { context, sockets } = setup();
+    const test = context.testC2SRelay({ name: "test" });
+    await vi.waitFor(() => expect(sockets).toHaveLength(1));
+    sockets[0].ready();
+    expect(await test).toEqual({ success: true });
+    expect(sockets[0].close).toHaveBeenCalledOnce();
+    expect(
+      await context.testC2SRelay({ name: "stopped" }, null, null, null, {}),
+    ).toEqual({ success: false, error: "Tunnel stopped" });
+    expect(sockets).toHaveLength(1);
+  });
+});

src/backend/tests/hosts/connect/pipeline.test.ts

@@ -518,6 +518,89 @@ describe("connectHost", () => {
     dispose();
   });
 
+  it("answers TOTP on key and password jump hops and closes the entire chain", async () => {
+    const { createJumpHostChain } = await vi.importActual<
+      typeof import("../../../hosts/jump-host-chain.js")
+    >("../../../hosts/jump-host-chain.js");
+    mocks.resolveHostById
+      .mockResolvedValueOnce(
+        host({ id: 2, authType: "key", key: plainKey, password: null }),
+      )
+      .mockResolvedValueOnce(host({ id: 3, password: "hop-password" }));
+    FakeSshClient.nextBehaviour = ["hang", "hang"];
+    const prompt = {
+      ask: vi
+        .fn()
+        .mockResolvedValueOnce("111111")
+        .mockResolvedValueOnce("222222"),
+    };
+    const pending = createJumpHostChain(
+      [{ hostId: 2 }, { hostId: 3 }],
+      "user-1",
+      prompt,
+    );
+    await vi.waitFor(() =>
+      expect(FakeSshClient.instances[0]?.connectConfig).toBeTruthy(),
+    );
+    const first = FakeSshClient.instances[0];
+    expect(first.connectConfig?.privateKey).toBeTruthy();
+    const firstAnswer = vi.fn();
+    first.emit(
+      "keyboard-interactive",
+      "",
+      "",
+      "",
+      [{ prompt: "Verification code:", echo: false }],
+      firstAnswer,
+    );
+    await vi.waitFor(() =>
+      expect(firstAnswer).toHaveBeenCalledWith(["111111"]),
+    );
+    first.emit("ready");
+    await vi.waitFor(() =>
+      expect(FakeSshClient.instances[1]?.connectConfig).toBeTruthy(),
+    );
+    const second = FakeSshClient.instances[1];
+    const secondAnswer = vi.fn();
+    second.emit(
+      "keyboard-interactive",
+      "",
+      "",
+      "",
+      [
+        { prompt: "Password:", echo: false },
+        { prompt: "Verification code:", echo: false },
+      ],
+      secondAnswer,
+    );
+    await vi.waitFor(() =>
+      expect(secondAnswer).toHaveBeenCalledWith(["hop-password", "222222"]),
+    );
+    second.emit("ready");
+    expect(await pending).toBe(second);
+    expect(prompt.ask.mock.calls[0][0].prompt).toContain("Jump host 1/2");
+    expect(prompt.ask.mock.calls[1][0].prompt).toContain("Jump host 2/2");
+    second.emit("close");
+    expect(first.ended).toBe(true);
+    expect(second.ended).toBe(true);
+  });
+
+  it("passes the interactive prompt channel through to jump hosts", async () => {
+    const jumpClient = new FakeSshClient();
+    mocks.createJumpHostChain.mockResolvedValueOnce(jumpClient);
+    const prompt = { ask: vi.fn().mockResolvedValue("123456") };
+    const connection = await connectHost(host({ jumpHosts: [{ hostId: 2 }] }), {
+      userId: "user-1",
+      prompt,
+    });
+    expect(mocks.createJumpHostChain).toHaveBeenCalledWith(
+      [{ hostId: 2 }],
+      "owner-1",
+      prompt,
+    );
+    connection.dispose();
+  });
+
   it("goes through the jump chain and forwards to the target", async () => {
     const jumpClient = new FakeSshClient();
     mocks.createJumpHostChain.mockResolvedValueOnce(jumpClient);
@@ -528,6 +611,7 @@ describe("connectHost", () => {
     expect(mocks.createJumpHostChain).toHaveBeenCalledWith(
       [{ hostId: 2 }, { hostId: 3 }],
       "owner-1",
+      undefined,
     );
     expect(jumpClient.forwardOut).toHaveBeenCalledWith(
       "127.0.0.1",