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

electron/c2s-auth.cjs

@@ -0,0 +1,78 @@
+const { randomUUID } = require("node:crypto");
+
+function createC2SAuthBridge(ipcMain, getWindow) {
+  const pending = new Map();
+  ipcMain.handle("answer-c2s-auth", (event, id, answer) => {
+    const item = pending.get(id);
+    if (!item || item.sender !== event.sender) return false;
+    if (answer !== null && (typeof answer !== "string" || answer.length > 8192))
+      return false;
+    item.finish(answer);
+    return true;
+  });
+  return function attach(ws, tunnelName) {
+    const active = new Set();
+    const close = () => {
+      for (const id of active) pending.get(id)?.finish(null);
+      ws.off("message", onMessage);
+      ws.off("close", close);
+      ws.off("error", close);
+    };
+    const onMessage = (raw, binary) => {
+      if (binary) return;
+      let message;
+      try {
+        message = JSON.parse(raw.toString());
+      } catch {
+        return;
+      }
+      if (message.type !== "auth-prompt") return;
+      const window = getWindow();
+      if (
+        !window ||
+        window.isDestroyed() ||
+        active.size ||
+        typeof message.requestId !== "string" ||
+        !["totp", "input", "browser"].includes(message.request?.kind)
+      ) {
+        ws.close();
+        return;
+      }
+      const sender = window.webContents;
+      const id = randomUUID();
+      const finish = (answer) => {
+        if (!pending.delete(id)) return;
+        active.delete(id);
+        clearTimeout(timer);
+        sender.off("destroyed", cancelled);
+        sender.off("did-start-loading", cancelled);
+        if (!sender.isDestroyed())
+          sender.send("c2s-auth-prompt", { id, closed: true });
+        if (ws.readyState === 1)
+          ws.send(
+            JSON.stringify({
+              type: "auth-response",
+              requestId: message.requestId,
+              answer,
+            }),
+          );
+        if (answer === null) ws.close();
+      };
+      const cancelled = () => finish(null);
+      const timer = setTimeout(cancelled, 60_000);
+      active.add(id);
+      pending.set(id, { sender, finish });
+      sender.once("destroyed", cancelled);
+      sender.once("did-start-loading", cancelled);
+      sender.send("c2s-auth-prompt", {
+        id,
+        tunnelName,
+        request: message.request,
+      });
+    };
+    ws.on("message", onMessage);
+    ws.once("close", close);
+    ws.once("error", close);
+  };
+}
+module.exports = { createC2SAuthBridge };

electron/main.cjs

@@ -2020,6 +2020,11 @@ async function getC2SRelayHeaders(relayUrl) {
   return headers;
 }
 
+const attachC2SAuth = require("./c2s-auth.cjs").createC2SAuthBridge(
+  ipcMain,
+  () => mainWindow,
+);
+
 function getC2STunnelName(tunnel, index = 0) {
   if (tunnel.name) return tunnel.name;
   const localAddress = getC2SLocalAddress(tunnel);
@@ -2136,6 +2141,7 @@ async function openC2SRelay(
     relayUrl,
     getWebSocketOptions(relayUrl, { headers }),
   );
+  attachC2SAuth(ws, tunnel.name || getC2STunnelName(tunnel));
   const pendingChunks = [];
   let ready = false;
   let closed = false;
@@ -2225,22 +2231,33 @@ async function openC2SRelay(
   });
 }
 
-async function testC2SRelay(tunnel, targetHost, targetPort, authToken) {
+async function testC2SRelay(
+  tunnel,
+  targetHost,
+  targetPort,
+  authToken,
+  runtime,
+) {
   const tunnelConfig = normalizeC2STunnelAddresses(tunnel);
   const relayUrl = getC2SRelayUrl();
   const headers = await getC2SRelayHeaders(relayUrl, authToken);
+  if (runtime && c2sTunnelRuntimes.get(tunnel.name) !== runtime) {
+    return { success: false, error: "Tunnel stopped" };
+  }
   const ws = new WebSocket(
     relayUrl,
     getWebSocketOptions(relayUrl, { headers }),
   );
+  attachC2SAuth(ws, tunnel.name || getC2STunnelName(tunnel));
 
+  if (runtime) runtime.authSocket = ws;
   return new Promise((resolve) => {
     let settled = false;
     const settle = (result) => {
       if (settled) return;
       settled = true;
       try {
-        ws.close();
+        if (!runtime || !result.success) ws.close();
       } catch {
         // expected during shutdown
       }
@@ -2249,12 +2266,13 @@ async function testC2SRelay(tunnel, targetHost, targetPort, authToken) {
 
     const timer = setTimeout(() => {
       settle({ success: false, error: "Tunnel test timed out" });
-    }, 15000);
+    }, 180_000);
 
     ws.on("open", () => {
       ws.send(
         JSON.stringify({
           type: "test",
+          keepAlive: !!runtime,
           tunnelConfig,
           targetHost,
           targetPort,
@@ -2283,6 +2301,9 @@ async function testC2SRelay(tunnel, targetHost, targetPort, authToken) {
       settle(createC2SFailure(error.message, "Tunnel test failed"));
     });
     ws.on("close", () => {
+      if (runtime && c2sTunnelRuntimes.get(tunnel.name) === runtime) {
+        setC2STunnelError(tunnel.name, "Endpoint SSH connection closed");
+      }
       clearTimeout(timer);
       settle({ success: false, error: "Tunnel test connection closed" });
     });
@@ -2535,6 +2556,7 @@ async function startC2SRemoteTunnel(tunnel, index = 0, authToken) {
     relayUrl,
     getWebSocketOptions(relayUrl, { headers }),
   );
+  attachC2SAuth(ws, tunnel.name || getC2STunnelName(tunnel));
   const sockets = new Map();
   let closed = false;
 
@@ -2723,7 +2745,10 @@ async function startC2STunnel(tunnel, index = 0, authToken) {
   const resolvedTunnel = await resolveC2SRemoteSourceHost(tunnel);
   const mode = resolvedTunnel.mode || resolvedTunnel.tunnelType || "local";
   const tunnelName = getC2STunnelName(resolvedTunnel, index);
-  const tunnelConfig = normalizeC2STunnelAddresses(resolvedTunnel);
+  const tunnelConfig = {
+    ...normalizeC2STunnelAddresses(resolvedTunnel),
+    sessionId: crypto.randomUUID(),
+  };
   const bindHost = getC2SLocalAddress(tunnelConfig);
   const sourcePort = Number(resolvedTunnel.sourcePort);
   logToFile(`[c2s] starting tunnel ${tunnelName}`, {
@@ -2774,6 +2799,10 @@ async function startC2STunnel(tunnel, index = 0, authToken) {
 
   const sockets = new Set();
   const server = net.createServer((socket) => {
+    if (!c2sTunnelRuntimes.get(tunnelName)?.status.connected) {
+      socket.destroy();
+      return;
+    }
     sockets.add(socket);
     socket.on("close", () => sockets.delete(socket));
     if (mode === "dynamic") {
@@ -2816,33 +2845,38 @@ async function startC2STunnel(tunnel, index = 0, authToken) {
         reason: "Verifying endpoint SSH connection",
       });
 
-      const verifyTunnel =
-        mode === "dynamic"
-          ? testC2SRelay(
-              { ...tunnel, name: `${tunnelName}::verify`, mode },
-              undefined,
-              undefined,
-            )
-          : testC2SRelay(
-              { ...tunnel, name: `${tunnelName}::verify`, mode },
-              tunnel.targetHost || "127.0.0.1",
-              Number(tunnel.endpointPort),
-            );
+      const runtime = c2sTunnelRuntimes.get(tunnelName);
+      const verifyTunnel = testC2SRelay(
+        { ...tunnelConfig, name: tunnelName, mode },
+        mode === "dynamic" ? undefined : getC2SRemoteAddress(tunnelConfig),
+        mode === "dynamic" ? undefined : Number(tunnelConfig.endpointPort),
+        authToken,
+        runtime,
+      );
 
-      verifyTunnel.then((result) => {
-        if (!c2sTunnelRuntimes.has(tunnelName)) return;
-        if (result.success) {
-          setC2STunnelStatus(tunnelName, {
-            connected: true,
-            status: "CONNECTED",
-          });
-        } else {
-          setC2STunnelError(
-            tunnelName,
-            result.error || "Endpoint SSH connection failed",
-          );
-        }
-      });
+      verifyTunnel
+        .then((result) => {
+          if (c2sTunnelRuntimes.get(tunnelName) !== runtime) return;
+          if (result.success) {
+            setC2STunnelStatus(tunnelName, {
+              connected: true,
+              status: "CONNECTED",
+            });
+          } else {
+            setC2STunnelError(
+              tunnelName,
+              result.error || "Endpoint SSH connection failed",
+            );
+          }
+        })
+        .catch((error) => {
+          if (c2sTunnelRuntimes.get(tunnelName) === runtime) {
+            setC2STunnelError(
+              tunnelName,
+              error.message || "Endpoint SSH connection failed",
+            );
+          }
+        });
 
       resolve({ success: true, tunnelName });
     });
@@ -2859,6 +2893,7 @@ async function stopC2STunnel(tunnelName) {
     connected: false,
     status: "DISCONNECTING",
   });
+  runtime.authSocket?.close();
 
   return new Promise((resolve) => {
     if (typeof runtime.close === "function") {
@@ -2883,6 +2918,7 @@ async function stopC2STunnel(tunnelName) {
 function stopAllC2STunnels() {
   for (const [tunnelName, runtime] of c2sTunnelRuntimes.entries()) {
     try {
+      runtime.authSocket?.close();
       if (typeof runtime.close === "function") {
         runtime.close();
       } else {

electron/preload.js

@@ -32,6 +32,13 @@ contextBridge.exposeInMainWorld("electronAPI", {
 
   getSetting: (key) => ipcRenderer.invoke("get-setting", key),
   setSetting: (key, value) => ipcRenderer.invoke("set-setting", key, value),
+  answerC2SAuth: (id, answer) =>
+    ipcRenderer.invoke("answer-c2s-auth", id, answer),
+  onC2SAuthPrompt: (callback) => {
+    const listener = (_event, prompt) => callback(prompt);
+    ipcRenderer.on("c2s-auth-prompt", listener);
+    return () => ipcRenderer.removeListener("c2s-auth-prompt", listener);
+  },
   getC2STunnelConfig: () => ipcRenderer.invoke("get-c2s-tunnel-config"),
   saveC2STunnelConfig: (config) =>
     ipcRenderer.invoke("save-c2s-tunnel-config", config),