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

plugins/tunnels/src/backend/c2s-auth.ts

@@ -0,0 +1,53 @@
+import { randomUUID } from "node:crypto";
+import type { WebSocket } from "ws";
+import type { PluginSshPromptChannel } from "@termix/plugin-sdk/backend";
+
+/** Answers belong to this authenticated socket and one outstanding challenge. */
+export function createC2SPrompt(ws: WebSocket): PluginSshPromptChannel {
+  return {
+    ask(request) {
+      if (ws.readyState !== 1)
+        return Promise.reject(new Error("Tunnel closed"));
+      const requestId = randomUUID();
+      return new Promise((resolve, reject) => {
+        const finish = (answer: string | null, error?: Error) => {
+          clearTimeout(timer);
+          ws.off("message", onMessage);
+          ws.off("close", onClose);
+          ws.off("error", onClose);
+          if (error) reject(error);
+          else resolve(answer);
+        };
+        const onClose = () =>
+          finish(null, new Error("Tunnel authentication cancelled"));
+        const onMessage = (raw: Buffer, binary: boolean) => {
+          if (binary) return;
+          let value;
+          try {
+            value = JSON.parse(raw.toString());
+          } catch {
+            return;
+          }
+          if (value.type !== "auth-response" || value.requestId !== requestId)
+            return;
+          if (value.answer === null) {
+            onClose();
+            ws.close();
+          } else if (
+            typeof value.answer === "string" &&
+            value.answer.length <= 8192
+          )
+            finish(value.answer);
+        };
+        const timer = setTimeout(() => {
+          finish(null, new Error("Tunnel authentication timed out"));
+          ws.close();
+        }, 60_000);
+        ws.on("message", onMessage);
+        ws.once("close", onClose);
+        ws.once("error", onClose);
+        ws.send(JSON.stringify({ type: "auth-prompt", requestId, request }));
+      });
+    },
+  };
+}

plugins/tunnels/src/backend/c2s-connections.ts

@@ -0,0 +1,94 @@
+import { randomUUID } from "node:crypto";
+import type { Client } from "ssh2";
+import type { WebSocket } from "ws";
+import type {
+  PluginContext,
+  PluginSshConnection,
+} from "@termix/plugin-sdk/backend";
+import { createC2SPrompt } from "./c2s-auth.js";
+
+/** A desktop runtime holds one authenticated SSH connection across SOCKS streams. */
+export function createC2SConnections(ctx: PluginContext) {
+  type Entry = {
+    sockets: Set<WebSocket>;
+    source: Promise<PluginSshConnection<Client>>;
+  };
+  const entries = new Map<string, Entry>();
+  ctx.disposables.add(() => {
+    for (const entry of entries.values())
+      for (const ws of entry.sockets) ws.close();
+    entries.clear();
+  });
+  return async (ws: WebSocket, hostId: number, sessionId?: string) => {
+    if (ws.readyState !== 1) throw new Error("Tunnel closed");
+    const key = JSON.stringify([
+      ctx.currentActor(),
+      hostId,
+      typeof sessionId === "string" && sessionId.length <= 128
+        ? sessionId
+        : randomUUID(),
+    ]);
+    let entry = entries.get(key);
+    if (!entry) {
+      const source = ctx.ssh.connect<Client>(hostId, {
+        purpose: "tunnel",
+        profile: "forward",
+        timeoutMs: 60_000,
+        prompt: createC2SPrompt(ws),
+      });
+      entry = { sockets: new Set(), source };
+      entries.set(key, entry);
+      const own = entry;
+      void source.then(
+        (connection) => {
+          const close = () => {
+            if (entries.get(key) === own) entries.delete(key);
+            for (const socket of own.sockets) socket.close();
+          };
+          connection.client.once("close", close);
+          connection.client.once("error", close);
+        },
+        () => {
+          if (entries.get(key) === own) entries.delete(key);
+          for (const socket of own.sockets) {
+            if (socket.readyState === 1)
+              socket.send(
+                JSON.stringify({
+                  type: "error",
+                  error: "Endpoint SSH authentication failed",
+                }),
+              );
+            socket.close();
+          }
+        },
+      );
+    }
+    const own = entry;
+    own.sockets.add(ws);
+    let released = false;
+    const dispose = () => {
+      if (released) return;
+      released = true;
+      ws.off("close", dispose);
+      ws.off("error", dispose);
+      own.sockets.delete(ws);
+      if (!own.sockets.size) {
+        if (entries.get(key) === own) entries.delete(key);
+        void own.source.then(
+          (source) => source.dispose(),
+          () => {},
+        );
+      }
+    };
+    ws.once("close", dispose);
+    ws.once("error", dispose);
+    try {
+      const source = await own.source;
+      if (released || ws.readyState !== 1) throw new Error("Tunnel closed");
+      return { ...source, dispose };
+    } catch (error) {
+      dispose();
+      throw error;
+    }
+  };
+}

plugins/tunnels/src/backend/c2s-relay.ts

@@ -1,10 +1,8 @@
 import type { Duplex } from "node:stream";
-import type { Client, ClientChannel } from "ssh2";
+import type { ClientChannel } from "ssh2";
 import type { WebSocket } from "ws";
-import type {
-  PluginContext,
-  PluginSshConnection,
-} from "@termix/plugin-sdk/backend";
+import type { PluginContext } from "@termix/plugin-sdk/backend";
+import { createC2SConnections } from "./c2s-connections.js";
 import type { TunnelMode } from "./types.js";
 import {
   bindForwardIn,
@@ -16,6 +14,7 @@ import { getTunnelMode } from "./utils.js";
 /** What the desktop app sends in the first message on the relay socket. */
 export interface C2SRelayTunnel {
   name?: string;
+  sessionId?: string;
   mode?: TunnelMode;
   tunnelType?: "local" | "remote";
   localAddress?: string;
@@ -30,6 +29,7 @@ export interface C2SRelayTunnel {
 
 export interface C2SOpenMessage {
   type: "open" | "test";
+  keepAlive?: boolean;
   tunnelConfig?: C2SRelayTunnel;
   targetHost?: string;
   targetPort?: number;
@@ -151,6 +151,7 @@ function address(value: unknown, fallback: string): string {
 
 interface ResolvedRelay {
   name: string;
+  sessionId?: string;
   hostId: number;
   mode: TunnelMode;
   remoteAddress: string;
@@ -165,6 +166,7 @@ interface ResolvedRelay {
  */
 export function createC2SRelay(ctx: PluginContext) {
   let streamCounter = 0;
+  const connectSource = createC2SConnections(ctx);
 
   async function findHostIdBySyncId(syncId: string): Promise<number | null> {
     /* eslint-disable @typescript-eslint/no-explicit-any */
@@ -188,6 +190,7 @@ export function createC2SRelay(ctx: PluginContext) {
     return {
       name: tunnel.name || `c2s:${hostId}`,
       hostId,
+      sessionId: tunnel.sessionId,
       mode: getTunnelMode({
         mode: tunnel.mode,
         tunnelType: tunnel.tunnelType,
@@ -201,14 +204,6 @@ export function createC2SRelay(ctx: PluginContext) {
     };
   }
 
-  function connectSource(hostId: number): Promise<PluginSshConnection<Client>> {
-    return ctx.ssh.connect<Client>(hostId, {
-      purpose: "tunnel",
-      profile: "forward",
-      timeoutMs: 60_000,
-    });
-  }
-
   async function openRemote(ws: WebSocket, relay: ResolvedRelay) {
     const bindHost = relay.remoteAddress;
     const bindPort = relay.sourcePort;
@@ -216,7 +211,7 @@ export function createC2SRelay(ctx: PluginContext) {
       throw new Error("Invalid remote port");
     }
 
-    const source = await connectSource(relay.hostId);
+    const source = await connectSource(ws, relay.hostId, relay.sessionId);
     const sourceClient = source.client;
     let actualPort: number;
     try {
@@ -226,6 +221,11 @@ export function createC2SRelay(ctx: PluginContext) {
       throw error;
     }
 
+    if (ws.readyState !== 1) {
+      unbindForwardIn(sourceClient, bindHost, actualPort, ctx.log);
+      source.dispose();
+      return;
+    }
     const streams = new Map<string, ClientChannel>();
     let closed = false;
 
@@ -340,7 +340,7 @@ export function createC2SRelay(ctx: PluginContext) {
       throw new Error("Invalid client tunnel target");
     }
 
-    const source = await connectSource(relay.hostId);
+    const source = await connectSource(ws, relay.hostId, relay.sessionId);
     let outbound: ClientChannel;
     try {
       outbound = await forwardOut(source.client, targetHost, targetPort);
@@ -349,6 +349,11 @@ export function createC2SRelay(ctx: PluginContext) {
       throw error;
     }
 
+    if (ws.readyState !== 1) {
+      outbound.destroy();
+      source.dispose();
+      return;
+    }
     const close = () => {
       try {
         outbound.destroy();
@@ -384,7 +389,7 @@ export function createC2SRelay(ctx: PluginContext) {
 
   async function test(ws: WebSocket, message: C2SOpenMessage): Promise<void> {
     const relay = await resolve(message.tunnelConfig ?? {});
-    const source = await connectSource(relay.hostId);
+    const source = await connectSource(ws, relay.hostId, relay.sessionId);
     try {
       if (relay.mode === "remote") {
         const bindPort = relay.sourcePort;
@@ -410,7 +415,7 @@ export function createC2SRelay(ctx: PluginContext) {
       }
       sendMessage(ws, { type: "ready" });
     } finally {
-      source.dispose();
+      if (!message.keepAlive) source.dispose();
     }
   }
 

plugins/tunnels/src/frontend/TunnelAuthPrompts.tsx

@@ -0,0 +1,135 @@
+import { useEffect, useRef, useState } from "react";
+import type { PluginSshPromptRequest } from "@termix/plugin-sdk/backend";
+import { useTranslation } from "@termix/plugin-sdk/frontend";
+import {
+  Button,
+  Dialog,
+  DialogContent,
+  DialogHeader,
+  DialogTitle,
+  Input,
+} from "@termix/plugin-sdk/ui";
+
+type Challenge = {
+  id: string;
+  tunnelName?: string;
+  request: PluginSshPromptRequest;
+};
+
+export function TunnelAuthPrompts() {
+  const [queue, setQueue] = useState<Challenge[]>([]);
+  const pending = useRef(new Set<string>());
+  useEffect(() => {
+    const active = pending.current;
+    const api = window.electronAPI;
+    if (!api?.onC2SAuthPrompt) return;
+    const dispose = api.onC2SAuthPrompt((event) => {
+      if (event.closed) {
+        active.delete(event.id);
+        setQueue((items) => items.filter((item) => item.id !== event.id));
+      } else if (event.request && !active.has(event.id)) {
+        active.add(event.id);
+        setQueue((items) => [...items, { ...event, request: event.request! }]);
+      }
+    });
+    return () => {
+      dispose();
+      for (const id of active) void api.answerC2SAuth(id, null);
+      active.clear();
+    };
+  }, []);
+  const current = queue[0];
+  return current ? (
+    <AuthDialog
+      key={current.id}
+      challenge={current}
+      onAnswer={(answer) => {
+        pending.current.delete(current.id);
+        setQueue((items) => items.filter((item) => item.id !== current.id));
+        void window.electronAPI.answerC2SAuth(current.id, answer);
+      }}
+    />
+  ) : null;
+}
+
+function AuthDialog({
+  challenge,
+  onAnswer,
+}: {
+  challenge: Challenge;
+  onAnswer: (answer: string | null) => void;
+}) {
+  const { t } = useTranslation();
+  const [answer, setAnswer] = useState("");
+  const request = challenge.request;
+  const browser = request.kind === "browser";
+  const url =
+    browser && request.url && /^https?:\/\//i.test(request.url)
+      ? request.url
+      : null;
+  return (
+    <Dialog
+      open
+      onOpenChange={(open) => {
+        if (!open) onAnswer(null);
+      }}
+    >
+      <DialogContent>
+        <DialogHeader>
+          <DialogTitle>{t("tunnels.authTitle")}</DialogTitle>
+        </DialogHeader>
+        <p className="text-sm break-all">{challenge.tunnelName}</p>
+        <form
+          onSubmit={(event) => {
+            event.preventDefault();
+            onAnswer(answer);
+          }}
+          className="space-y-3"
+        >
+          <p className="text-sm whitespace-pre-wrap">
+            {browser ? request.instructions : request.prompt}
+          </p>
+          {request.kind === "totp" && request.retry && (
+            <p role="alert">{t("tunnels.authRetry")}</p>
+          )}
+          {browser ? (
+            <>
+              <p>{request.code}</p>
+              {url && (
+                <a
+                  href={url}
+                  target="_blank"
+                  rel="noreferrer"
+                  className="break-all underline"
+                >
+                  {url}
+                </a>
+              )}
+            </>
+          ) : (
+            <Input
+              autoFocus
+              aria-label={t("tunnels.authAnswer")}
+              autoComplete="off"
+              type={
+                request.kind === "input" && request.echo ? "text" : "password"
+              }
+              value={answer}
+              onChange={(event) => setAnswer(event.target.value)}
+            />
+          )}
+          <div className="flex justify-end gap-2">
+            <Button
+              type="button"
+              variant="outline"
+              onClick={() => onAnswer(null)}
+            >
+              {t("common.cancel")}
+            </Button>
+            <Button type="submit">{t("tunnels.authContinue")}</Button>
+          </div>
+        </form>
+      </DialogContent>
+    </Dialog>
+  );
+}

plugins/tunnels/src/frontend/index.tsx

@@ -1,6 +1,7 @@
 import type { ComponentType } from "react";
 import { Network, Plug } from "lucide-react";
 import type { TabProps, TermixApp } from "@termix/plugin-sdk/frontend";
+import { TunnelAuthPrompts } from "./TunnelAuthPrompts";
 import { TunnelTab } from "./TunnelTab";
 import { HostTunnelsSection } from "./HostTunnelsSection";
 import {
@@ -27,6 +28,12 @@ export function activate(app: TermixApp): void {
   if (typeof window !== "undefined" && window.electronAPI?.setC2SRelayPath) {
     void window.electronAPI.setC2SRelayPath(C2S_RELAY_PATH);
   }
+  app.registerSlotContribution("shell.overlay", {
+    actionId: "tunnels.authPrompts",
+    titleKey: "tunnels.authTitle",
+    kind: "component",
+    component: TunnelAuthPrompts as ComponentType<Record<string, unknown>>,
+  });
   // apiFor("remote") hands back app.api itself outside the desktop app.
   const remote = app.apiFor("remote");
   setTunnelsApi(app.api, {