Use wildcard SSH config Host entries

This simplifies the written SSH config and avoids the need to make an API request for every workspace the filter returns. This can remove minutes from the "Configuring Coder CLI..." step when the user has access to many workspaces (for example, an admin who wants the option of connecting to anyone's workspace on a large deployment). Depends on coder/coder#16088
coder · Jan 14, 2025 · ebf86fa · ebf86fa
1 parent faddd24
commit ebf86fa
Show file tree

Hide file tree

Showing 5 changed files with 124 additions and 44 deletions.
diff --git a/src/main/kotlin/com/coder/gateway/cli/CoderCLIManager.kt b/src/main/kotlin/com/coder/gateway/cli/CoderCLIManager.kt
@@ -115,6 +115,7 @@ fun ensureCLI(
 data class Features(
     val disableAutostart: Boolean = false,
     val reportWorkspaceUsage: Boolean = false,
+    val wildcardSSH: Boolean = false,
 )
 
 /**
@@ -285,37 +286,57 @@ class CoderCLIManager(
             } else {
                 ""
             }
+        val sshOpts = """
+            ConnectTimeout 0
+            StrictHostKeyChecking no
+            UserKnownHostsFile /dev/null
+            LogLevel ERROR
+            SetEnv CODER_SSH_SESSION_TYPE=JetBrains
+        """.trimIndent()
         val blockContent =
+            if (feats.wildcardSSH) {
+                startBlock + System.lineSeparator() +
+                    """
+                    Host ${getHostPrefix()}--*
+                      ProxyCommand ${proxyArgs.joinToString(" ")} --ssh-host-prefix ${getHostPrefix()}-- %h
+                    """.trimIndent()
+                        .plus("\n" + sshOpts.prependIndent("  "))
+                        .plus(extraConfig)
+                        .plus("\n\n")
+                        .plus(
+                            """
+                            Host ${getHostPrefix()}-bg--*
+                              ProxyCommand ${backgroundProxyArgs.joinToString(" ")} --ssh-host-prefix ${getHostPrefix()}-bg-- %h
+                            """.trimIndent()
+                            .plus("\n" + sshOpts.prependIndent("  "))
+                            .plus(extraConfig),
+                        ).replace("\n", System.lineSeparator()) +
+                    System.lineSeparator() + endBlock
+
+        } else {
             workspaceNames.joinToString(
                 System.lineSeparator(),
                 startBlock + System.lineSeparator(),
                 System.lineSeparator() + endBlock,
                 transform = {
                     """
-                    Host ${getHostName(deploymentURL, it.first, currentUser, it.second)}
+                    Host ${getHostName(it.first, currentUser, it.second)}
                       ProxyCommand ${proxyArgs.joinToString(" ")} ${getWorkspaceParts(it.first, it.second)}
-                      ConnectTimeout 0
-                      StrictHostKeyChecking no
-                      UserKnownHostsFile /dev/null
-                      LogLevel ERROR
-                      SetEnv CODER_SSH_SESSION_TYPE=JetBrains
                     """.trimIndent()
+                        .plus("\n" + sshOpts.prependIndent("  "))
                         .plus(extraConfig)
                         .plus("\n")
                         .plus(
                             """
-                            Host ${getBackgroundHostName(deploymentURL, it.first, currentUser, it.second)}
+                            Host ${getBackgroundHostName(it.first, currentUser, it.second)}
                               ProxyCommand ${backgroundProxyArgs.joinToString(" ")} ${getWorkspaceParts(it.first, it.second)}
-                              ConnectTimeout 0
-                              StrictHostKeyChecking no
-                              UserKnownHostsFile /dev/null
-                              LogLevel ERROR
-                              SetEnv CODER_SSH_SESSION_TYPE=JetBrains
                             """.trimIndent()
+                                .plus("\n" + sshOpts.prependIndent("  "))
                                 .plus(extraConfig),
                         ).replace("\n", System.lineSeparator())
                 },
             )
+        }
 
         if (contents == null) {
             logger.info("No existing SSH config to modify")
@@ -489,40 +510,53 @@ class CoderCLIManager(
                 Features(
                     disableAutostart = version >= SemVer(2, 5, 0),
                     reportWorkspaceUsage = version >= SemVer(2, 13, 0),
+                    wildcardSSH = version >= SemVer(2, 19, 0),
                 )
             }
         }
 
+    /*
+     * This function returns the ssh-host-prefix used for Host entries.
+     */
+    fun getHostPrefix(): String =
+        "coder-jetbrains-${deploymentURL.safeHost()}"
+
+    /**
+     * This function returns the ssh host name generated for connecting to the workspace.
+     */
+    fun getHostName(
+        workspace: Workspace,
+        currentUser: User,
+        agent: WorkspaceAgent,
+    ): String =
+        if (features.wildcardSSH) {
+           "${getHostPrefix()}--${workspace.ownerName}--${workspace.name}.${agent.name}"
+        } else {
+           // For a user's own workspace, we use the old syntax without a username for backwards compatibility,
+           // since the user might have recent connections that still use the old syntax.
+           if (currentUser.username == workspace.ownerName) {
+               "coder-jetbrains--${workspace.name}.${agent.name}--${deploymentURL.safeHost()}"
+           } else {
+               "coder-jetbrains--${workspace.ownerName}--${workspace.name}.${agent.name}--${deploymentURL.safeHost()}"
+        }
+    }
+
+    fun getBackgroundHostName(
+        workspace: Workspace,
+        currentUser: User,
+        agent: WorkspaceAgent,
+    ): String =
+        if (features.wildcardSSH) {
+            "${getHostPrefix()}-bg--${workspace.ownerName}--${workspace.name}.${agent.name}"
+        } else {
+            getHostName(workspace, currentUser, agent) + "--bg"
+        }
+
     companion object {
         val logger = Logger.getInstance(CoderCLIManager::class.java.simpleName)
 
         private val tokenRegex = "--token [^ ]+".toRegex()
 
-        /**
-         * This function returns the ssh host name generated for connecting to the workspace.
-         */
-        @JvmStatic
-        fun getHostName(
-            url: URL,
-            workspace: Workspace,
-            currentUser: User,
-            agent: WorkspaceAgent,
-        ): String =
-            // For a user's own workspace, we use the old syntax without a username for backwards compatibility,
-            // since the user might have recent connections that still use the old syntax.
-            if (currentUser.username == workspace.ownerName) {
-                "coder-jetbrains--${workspace.name}.${agent.name}--${url.safeHost()}"
-            } else {
-                "coder-jetbrains--${workspace.ownerName}--${workspace.name}.${agent.name}--${url.safeHost()}"
-            }
-
-        fun getBackgroundHostName(
-            url: URL,
-            workspace: Workspace,
-            currentUser: User,
-            agent: WorkspaceAgent,
-        ): String = getHostName(url, workspace, currentUser, agent) + "--bg"
-
         /**
          * This function returns the identifier for the workspace to pass to the
          * coder ssh proxy command.
@@ -536,6 +570,18 @@ class CoderCLIManager(
         @JvmStatic
         fun getBackgroundHostName(
             hostname: String,
-        ): String = hostname + "--bg"
+        ): String {
+            val parts = hostname.split("--").toMutableList()
+            if (parts.size < 2) {
+                throw SSHConfigFormatException("Invalid hostname: $hostname")
+            }
+            // non-wildcard case
+            if (parts[0] == "coder-jetbrains") {
+                return hostname + "--bg"
+            }
+            // wildcard case
+            parts[0] += "-bg"
+            return parts.joinToString("--")
+        }
     }
 }
diff --git a/src/main/kotlin/com/coder/gateway/util/LinkHandler.kt b/src/main/kotlin/com/coder/gateway/util/LinkHandler.kt
@@ -111,7 +111,11 @@ open class LinkHandler(
         }
 
         indicator?.invoke("Configuring Coder CLI...")
-        cli.configSsh(workspacesAndAgents = client.withAgents(workspaces), currentUser = client.me)
+        if (cli.features.wildcardSSH) {
+            cli.configSsh(workspacesAndAgents = emptySet(), currentUser = client.me)
+        } else {
+            cli.configSsh(workspacesAndAgents = client.withAgents(workspaces), currentUser = client.me)
+        }
 
         val openDialog =
             parameters.ideProductCode().isNullOrBlank() ||
@@ -127,7 +131,7 @@ open class LinkHandler(
             verifyDownloadLink(parameters)
             WorkspaceProjectIDE.fromInputs(
                 name = CoderCLIManager.getWorkspaceParts(workspace, agent),
-                hostname = CoderCLIManager.getHostName(deploymentURL.toURL(), workspace, client.me, agent),
+                hostname = CoderCLIManager(deploymentURL.toURL(), settings).getHostName(workspace, client.me, agent),
                 projectPath = parameters.folder(),
                 ideProductCode = parameters.ideProductCode(),
                 ideBuildNumber = parameters.ideBuildNumber(),

diff --git a/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspaceProjectIDEStepView.kt b/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspaceProjectIDEStepView.kt
@@ -199,7 +199,11 @@ class CoderWorkspaceProjectIDEStepView(
                     logger.info("Configuring Coder CLI...")
                     cbIDE.renderer = IDECellRenderer("Configuring Coder CLI...")
                     withContext(Dispatchers.IO) {
-                        data.cliManager.configSsh(data.client.withAgents(data.workspaces), data.client.me)
+                        if (data.cliManager.features.wildcardSSH) {
+                            data.cliManager.configSsh(emptySet(), data.client.me)
+                        } else {
+                            data.cliManager.configSsh(data.client.withAgents(data.workspaces), data.client.me)
+                        }
                     }
 
                     val ides =
@@ -214,7 +218,7 @@ class CoderWorkspaceProjectIDEStepView(
                                     } else {
                                         IDECellRenderer(CoderGatewayBundle.message("gateway.connector.view.coder.connect-ssh"))
                                     }
-                                val executor = createRemoteExecutor(CoderCLIManager.getBackgroundHostName(data.client.url, data.workspace, data.client.me, data.agent))
+                                val executor = createRemoteExecutor(CoderCLIManager(data.client.url).getBackgroundHostName(data.workspace, data.client.me, data.agent))
 
                                 if (ComponentValidator.getInstance(tfProject).isEmpty) {
                                     logger.info("Installing remote path validator...")
@@ -404,7 +408,7 @@ class CoderWorkspaceProjectIDEStepView(
     override fun data(): WorkspaceProjectIDE = withoutNull(cbIDE.selectedItem, state) { selectedIDE, state ->
         selectedIDE.withWorkspaceProject(
             name = CoderCLIManager.getWorkspaceParts(state.workspace, state.agent),
-            hostname = CoderCLIManager.getHostName(state.client.url, state.workspace, state.client.me, state.agent),
+            hostname = CoderCLIManager(state.client.url).getHostName(state.workspace, state.client.me, state.agent),
             projectPath = tfProject.text,
             deploymentURL = state.client.url,
         )

diff --git a/src/test/fixtures/outputs/wildcard.conf b/src/test/fixtures/outputs/wildcard.conf
@@ -0,0 +1,17 @@
+# --- START CODER JETBRAINS test.coder.invalid
+Host coder-jetbrains-test.coder.invalid--*
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --ssh-host-prefix coder-jetbrains-test.coder.invalid-- %h
+  ConnectTimeout 0
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+  LogLevel ERROR
+  SetEnv CODER_SSH_SESSION_TYPE=JetBrains
+
+Host coder-jetbrains-test.coder.invalid-bg--*
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config --url https://test.coder.invalid ssh --stdio --ssh-host-prefix coder-jetbrains-test.coder.invalid-bg-- %h
+  ConnectTimeout 0
+  StrictHostKeyChecking no
+  UserKnownHostsFile /dev/null
+  LogLevel ERROR
+  SetEnv CODER_SSH_SESSION_TYPE=JetBrains
+# --- END CODER JETBRAINS test.coder.invalid
diff --git a/src/test/kotlin/com/coder/gateway/cli/CoderCLIManagerTest.kt b/src/test/kotlin/com/coder/gateway/cli/CoderCLIManagerTest.kt
@@ -419,6 +419,15 @@ internal class CoderCLIManagerTest {
                     output = "multiple-agents",
                     remove = "blank",
                 ),
+                SSHTest(
+                    listOf(workspace),
+                    input = null,
+                    output = "wildcard",
+                    remove = "blank",
+                    features = Features(
+                        wildcardSSH = true,
+                    ),
+                ),
             )
 
         val newlineRe = "\r?\n".toRegex()
@@ -804,7 +813,7 @@ internal class CoderCLIManagerTest {
             listOf(
                 Pair("2.5.0", Features(true)),
                 Pair("2.13.0", Features(true, true)),
-                Pair("4.9.0", Features(true, true)),
+                Pair("4.9.0", Features(true, true, true)),
                 Pair("2.4.9", Features(false)),
                 Pair("1.0.1", Features(false)),
             )