Support multiple deployments

src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt

@@ -21,6 +21,8 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
 
     override suspend fun connect(parameters: Map<String, String>, requestor: ConnectionRequestor): GatewayConnectionHandle? {
         val clientLifetime = LifetimeDefinition()
+        // TODO: If this fails determine if it is an auth error and if so prompt
+        // for a new token, configure the CLI, then try again.
         clientLifetime.launchUnderBackgroundProgress(CoderGatewayBundle.message("gateway.connector.coder.connection.provider.title"), canBeCancelled = true, isIndeterminate = true, project = null) {
             val context = SshMultistagePanelContext(parameters.toHostDeployInputs())
             logger.info("Deploying and starting IDE with $context")
@@ -43,4 +45,4 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
     companion object {
         val logger = Logger.getInstance(CoderGatewayConnectionProvider::class.java.simpleName)
     }
-}
+}

src/main/kotlin/com/coder/gateway/models/CoderWorkspacesWizardModel.kt

@@ -4,5 +4,5 @@ data class CoderWorkspacesWizardModel(
     var coderURL: String = "https://coder.example.com",
     var token: String = "",
     var selectedWorkspace: WorkspaceAgentModel? = null,
-    var useExistingToken: Boolean = false
+    var useExistingToken: Boolean = false,
 )

src/main/kotlin/com/coder/gateway/sdk/CoderCLIManager.kt

@@ -1,5 +1,6 @@
 package com.coder.gateway.sdk
 
+import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.views.steps.CoderWorkspacesStepView
 import com.intellij.openapi.diagnostic.Logger
 import org.zeroturnaround.exec.ProcessExecutor
@@ -23,23 +24,29 @@ import javax.xml.bind.annotation.adapters.HexBinaryAdapter
 /**
  * Manage the CLI for a single deployment.
  */
-class CoderCLIManager @JvmOverloads constructor(deployment: URL, destinationDir: Path = getDataDir()) {
+class CoderCLIManager @JvmOverloads constructor(
+    private val deploymentURL: URL,
+    destinationDir: Path = getDataDir(),
+    private val sshConfigPath: Path = Path.of(System.getProperty("user.home")).resolve(".ssh/config"),
+) {
     private var remoteBinaryUrl: URL
     var localBinaryPath: Path
+    private var coderConfigPath: Path
 
     init {
         val binaryName = getCoderCLIForOS(getOS(), getArch())
         remoteBinaryUrl = URL(
-            deployment.protocol,
-            deployment.host,
-            deployment.port,
+            deploymentURL.protocol,
+            deploymentURL.host,
+            deploymentURL.port,
             "/bin/$binaryName"
         )
         // Convert IDN to ASCII in case the file system cannot support the
         // necessary character set.
-        val host = IDN.toASCII(deployment.host, IDN.ALLOW_UNASSIGNED)
-        val subdir = if (deployment.port > 0) "${host}-${deployment.port}" else host
-        localBinaryPath = destinationDir.resolve(subdir).resolve(binaryName)
+        val host = getSafeHost(deploymentURL)
+        val subdir = if (deploymentURL.port > 0) "${host}-${deploymentURL.port}" else host
+        localBinaryPath = destinationDir.resolve(subdir).resolve(binaryName).toAbsolutePath()
+        coderConfigPath = destinationDir.resolve(subdir).resolve("config").toAbsolutePath()
     }
 
     /**
@@ -81,7 +88,7 @@ class CoderCLIManager @JvmOverloads constructor(deployment: URL, destinationDir:
         val etag = getBinaryETag()
         val conn = remoteBinaryUrl.openConnection() as HttpURLConnection
         if (etag != null) {
-            logger.info("Found existing binary at ${localBinaryPath.toAbsolutePath()}; calculated hash as $etag")
+            logger.info("Found existing binary at $localBinaryPath; calculated hash as $etag")
             conn.setRequestProperty("If-None-Match", "\"$etag\"")
         }
         conn.setRequestProperty("Accept-Encoding", "gzip")
@@ -91,7 +98,7 @@ class CoderCLIManager @JvmOverloads constructor(deployment: URL, destinationDir:
             logger.info("GET ${conn.responseCode} $remoteBinaryUrl")
             when (conn.responseCode) {
                 HttpURLConnection.HTTP_OK -> {
-                    logger.info("Downloading binary to ${localBinaryPath.toAbsolutePath()}")
+                    logger.info("Downloading binary to $localBinaryPath")
                     Files.createDirectories(localBinaryPath.parent)
                     conn.inputStream.use {
                         Files.copy(
@@ -110,7 +117,7 @@ class CoderCLIManager @JvmOverloads constructor(deployment: URL, destinationDir:
                 }
 
                 HttpURLConnection.HTTP_NOT_MODIFIED -> {
-                    logger.info("Using cached binary at ${localBinaryPath.toAbsolutePath()}")
+                    logger.info("Using cached binary at $localBinaryPath")
                     return false
                 }
             }
@@ -137,26 +144,133 @@ class CoderCLIManager @JvmOverloads constructor(deployment: URL, destinationDir:
         } catch (e: FileNotFoundException) {
             null
         } catch (e: Exception) {
-            logger.warn("Unable to calculate hash for ${localBinaryPath.toAbsolutePath()}", e)
+            logger.warn("Unable to calculate hash for $localBinaryPath", e)
             null
         }
     }
 
     /**
-     * Use the provided credentials to authenticate the CLI.
+     * Use the provided token to authenticate the CLI.
      */
-    fun login(url: String, token: String): String {
-        return exec("login", url, "--token", token)
+    fun login(token: String): String {
+        logger.info("Storing CLI credentials in $coderConfigPath")
+        return exec(
+            "login",
+            deploymentURL.toString(),
+            "--token",
+            token,
+            "--global-config",
+            coderConfigPath.toString(),
+        )
     }
 
     /**
      * Configure SSH to use this binary.
-     *
-     * TODO: Support multiple deployments; currently they will clobber each
-     *       other.
      */
-    fun configSsh(): String {
-        return exec("config-ssh", "--yes", "--use-previous-options")
+    fun configSsh(workspaces: List<WorkspaceAgentModel>) {
+        writeSSHConfig(modifySSHConfig(readSSHConfig(), workspaces))
+    }
+
+    /**
+     * Return the contents of the SSH config or null if it does not exist.
+     */
+    private fun readSSHConfig(): String? {
+        return try {
+            sshConfigPath.toFile().readText()
+        } catch (e: FileNotFoundException) {
+            null
+        }
+    }
+
+    /**
+     * Given an existing SSH config modify it to add or remove the config for
+     * this deployment and return the modified config or null if it does not
+     * need to be modified.
+     */
+    private fun modifySSHConfig(contents: String?, workspaces: List<WorkspaceAgentModel>): String? {
+        val host = getSafeHost(deploymentURL)
+        val startBlock = "# --- START CODER JETBRAINS $host"
+        val endBlock = "# --- END CODER JETBRAINS $host"
+        val isRemoving = workspaces.isEmpty()
+        val blockContent = workspaces.joinToString(
+            System.lineSeparator(),
+            startBlock + System.lineSeparator(),
+            System.lineSeparator() + endBlock,
+            transform = {
+                """
+                Host ${getHostName(deploymentURL, it)}
+                  HostName coder.${it.name}
+                  ProxyCommand "$localBinaryPath" --global-config "$coderConfigPath" ssh --stdio ${it.name}
+                  ConnectTimeout 0
+                  StrictHostKeyChecking no
+                  UserKnownHostsFile /dev/null
+                  LogLevel ERROR
+                  SetEnv CODER_SSH_SESSION_TYPE=JetBrains
+                """.trimIndent().replace("\n", System.lineSeparator())
+            })
+
+        if (contents == null) {
+            logger.info("No existing SSH config to modify")
+            return blockContent + System.lineSeparator()
+        }
+
+        val start = "(\\s*)$startBlock".toRegex().find(contents)
+        val end = "$endBlock(\\s*)".toRegex().find(contents)
+
+        if (start == null && end == null && isRemoving) {
+            logger.info("No workspaces and no existing config blocks to remove")
+            return null
+        }
+
+        if (start == null && end == null) {
+            logger.info("Appending config block")
+            val toAppend = if (contents.isEmpty()) blockContent else listOf(
+                contents,
+                blockContent
+            ).joinToString(System.lineSeparator())
+            return toAppend + System.lineSeparator()
+        }
+
+        if (start == null) {
+            throw SSHConfigFormatException("End block exists but no start block")
+        }
+        if (end == null) {
+            throw SSHConfigFormatException("Start block exists but no end block")
+        }
+        if (start.range.first > end.range.first) {
+            throw SSHConfigFormatException("Start block found after end block")
+        }
+
+        if (isRemoving) {
+            logger.info("No workspaces; removing config block")
+            return listOf(
+                contents.substring(0, start.range.first),
+                // Need to keep the trailing newline(s) if we are not at the
+                // front of the file otherwise the before and after lines would
+                // get joined.
+                if (start.range.first > 0) end.groupValues[1] else "",
+                contents.substring(end.range.last + 1)
+            ).joinToString("")
+        }
+
+        logger.info("Replacing existing config block")
+        return listOf(
+            contents.substring(0, start.range.first),
+            start.groupValues[1], // Leading newline(s).
+            blockContent,
+            end.groupValues[1], // Trailing newline(s).
+            contents.substring(end.range.last + 1)
+        ).joinToString("")
+    }
+
+    /**
+     * Write the provided SSH config or do nothing if null.
+     */
+    private fun writeSSHConfig(contents: String?) {
+        if (contents != null) {
+            Files.createDirectories(sshConfigPath.parent)
+            sshConfigPath.toFile().writeText(contents)
+        }
     }
 
     /**
@@ -241,6 +355,15 @@ class CoderCLIManager @JvmOverloads constructor(deployment: URL, destinationDir:
                 }
             }
         }
+
+        private fun getSafeHost(url: URL): String {
+            return IDN.toASCII(url.host, IDN.ALLOW_UNASSIGNED)
+        }
+
+        @JvmStatic
+        fun getHostName(url: URL, ws: WorkspaceAgentModel): String {
+            return "coder-jetbrains--${ws.name}--${getSafeHost(url)}"
+        }
     }
 }
 
@@ -255,3 +378,5 @@ class Environment(private val env: Map<String, String> = emptyMap()) {
 }
 
 class ResponseException(message: String, val code: Int) : Exception(message)
+
+class SSHConfigFormatException(message: String) : Exception(message)

src/main/kotlin/com/coder/gateway/views/steps/CoderLocateRemoteProjectStepView.kt

@@ -5,8 +5,10 @@ import com.coder.gateway.icons.CoderIcons
 import com.coder.gateway.models.CoderWorkspacesWizardModel
 import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.sdk.Arch
+import com.coder.gateway.sdk.CoderCLIManager
 import com.coder.gateway.sdk.CoderRestClientService
 import com.coder.gateway.sdk.OS
+import com.coder.gateway.sdk.toURL
 import com.coder.gateway.sdk.withPath
 import com.coder.gateway.toWorkspaceParams
 import com.coder.gateway.views.LazyBrowserLink
@@ -30,7 +32,12 @@ import com.intellij.ui.AnimatedIcon
 import com.intellij.ui.ColoredListCellRenderer
 import com.intellij.ui.DocumentAdapter
 import com.intellij.ui.components.JBTextField
-import com.intellij.ui.dsl.builder.*
+import com.intellij.ui.dsl.builder.AlignX
+import com.intellij.ui.dsl.builder.BottomGap
+import com.intellij.ui.dsl.builder.RightGap
+import com.intellij.ui.dsl.builder.RowLayout
+import com.intellij.ui.dsl.builder.TopGap
+import com.intellij.ui.dsl.builder.panel
 import com.intellij.util.ui.JBFont
 import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.UIUtil
@@ -61,7 +68,7 @@ import kotlinx.coroutines.withContext
 import java.awt.Component
 import java.awt.FlowLayout
 import java.time.Duration
-import java.util.*
+import java.util.Locale
 import javax.swing.ComboBoxModel
 import javax.swing.DefaultComboBoxModel
 import javax.swing.JLabel
@@ -151,9 +158,11 @@ class CoderLocateRemoteProjectStepView(private val setNextButtonEnabled: (Boolea
     override fun onInit(wizardModel: CoderWorkspacesWizardModel) {
         cbIDE.renderer = IDECellRenderer()
         ideComboBoxModel.removeAllElements()
+        val deploymentURL = wizardModel.coderURL.toURL()
         val selectedWorkspace = wizardModel.selectedWorkspace
         if (selectedWorkspace == null) {
-            logger.warn("No workspace was selected. Please go back to the previous step and select a Coder Workspace")
+            // TODO: Should be impossible, tweak the types/flow to enforce this.
+            logger.warn("No workspace was selected. Please go back to the previous step and select a workspace")
             return
         }
 
@@ -163,7 +172,9 @@ class CoderLocateRemoteProjectStepView(private val setNextButtonEnabled: (Boolea
 
         ideResolvingJob = cs.launch {
             try {
-                val executor = withTimeout(Duration.ofSeconds(60)) { createRemoteExecutor(selectedWorkspace) }
+                val executor = withTimeout(Duration.ofSeconds(60)) {
+                    createRemoteExecutor(CoderCLIManager.getHostName(deploymentURL, selectedWorkspace))
+                }
                 retrieveIDES(executor, selectedWorkspace)
                 if (ComponentValidator.getInstance(tfProject).isEmpty) {
                     installRemotePathValidator(executor)
@@ -235,10 +246,10 @@ class CoderLocateRemoteProjectStepView(private val setNextButtonEnabled: (Boolea
         })
     }
 
-    private suspend fun createRemoteExecutor(selectedWorkspace: WorkspaceAgentModel): HighLevelHostAccessor {
+    private suspend fun createRemoteExecutor(host: String): HighLevelHostAccessor {
         return HighLevelHostAccessor.create(
             RemoteCredentialsHolder().apply {
-                setHost("coder.${selectedWorkspace.name}")
+                setHost(host)
                 userName = "coder"
                 port = 22
                 authType = AuthType.OPEN_SSH
@@ -310,11 +321,18 @@ class CoderLocateRemoteProjectStepView(private val setNextButtonEnabled: (Boolea
     override fun onNext(wizardModel: CoderWorkspacesWizardModel): Boolean {
         val selectedIDE = cbIDE.selectedItem ?: return false
         logger.info("Going to launch the IDE")
+        val deploymentURL = wizardModel.coderURL.toURL()
+        val selectedWorkspace = wizardModel.selectedWorkspace
+        if (selectedWorkspace == null) {
+            // TODO: Should be impossible, tweak the types/flow to enforce this.
+            logger.warn("No workspace was selected. Please go back to the previous step and select a workspace")
+            return false
+        }
         cs.launch {
             GatewayUI.getInstance().connect(
                 selectedIDE
                     .toWorkspaceParams()
-                    .withWorkspaceHostname("coder.${wizardModel.selectedWorkspace?.name}")
+                    .withWorkspaceHostname(CoderCLIManager.getHostName(deploymentURL, selectedWorkspace))
                     .withProjectPath(tfProject.text)
                     .withWebTerminalLink("${terminalLink.url}")
             )