Support multiple deployments

code-asher · code-asher · commit 9eaeb0f7404f · 2023-04-14T11:09:32.000-08:00
This also lets us set a custom environment variable to track JetBrains
sessions.
diff --git a/src/main/kotlin/com/coder/gateway/models/CoderWorkspacesWizardModel.kt b/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,
 )
diff --git a/src/main/kotlin/com/coder/gateway/sdk/CoderCLIManager.kt b/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,25 @@ import javax.xml.bind.annotation.adapters.HexBinaryAdapter
 /**
  * Manage the CLI for a single deployment.
  */
-class CoderCLIManager @JvmOverloads constructor(private val deployment: URL, destinationDir: Path = getDataDir()) {
+class CoderCLIManager @JvmOverloads constructor(private val deploymentURL: URL, destinationDir: Path = getDataDir()) {
     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
+        val host = getSafeHost(deploymentURL)
+        val subdir = if (deploymentURL.port > 0) "${host}-${deploymentURL.port}" else host
         localBinaryPath = destinationDir.resolve(subdir).resolve(binaryName)
+        coderConfigPath = destinationDir.resolve(subdir).resolve("config")
     }
 
     /**
@@ -146,17 +149,93 @@ class CoderCLIManager @JvmOverloads constructor(private val deployment: URL, des
      * Use the provided token to authenticate the CLI.
      */
     fun login(token: String): String {
-        return exec("login", deployment.toString(), "--token", token)
+        logger.info("Storing CLI credentials in $coderConfigPath")
+        return exec(
+            "login",
+            deploymentURL.toString(),
+            "--token",
+            token,
+            "--global-config",
+            coderConfigPath.toAbsolutePath().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>,
+        sshConfigPath: Path = Path.of(System.getProperty("user.home")).resolve(".ssh/config"),
+    ) {
+        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.toAbsolutePath()}" --global-config "${coderConfigPath.toAbsolutePath()}" ssh --stdio ${it.name}
+                  ConnectTimeout 0
+                  StrictHostKeyChecking no
+                  UserKnownHostsFile /dev/null
+                  LogLevel ERROR
+                  SetEnv CODER_SSH_SESSION_TYPE=JetBrains
+                """.trimIndent().replace("\n", System.lineSeparator())
+            })
+        Files.createDirectories(sshConfigPath.parent)
+        try {
+            val contents = sshConfigPath.toFile().readText()
+            val start = "(\\s*)$startBlock".toRegex().find(contents)
+            val end = "$endBlock(\\s*)".toRegex().find(contents)
+            if (start == null && end == null && isRemoving) {
+                logger.info("Leaving $sshConfigPath alone since there are no workspaces and no config to remove")
+            } else if (start == null && end == null) {
+                logger.info("Appending config to $sshConfigPath")
+                sshConfigPath.toFile().writeText(
+                    if (contents.isEmpty()) blockContent else listOf(
+                        contents,
+                        blockContent
+                    ).joinToString(System.lineSeparator())
+                )
+            } else if (start == null) {
+                throw SSHConfigFormatException("End block exists but no start block")
+            } else if (end == null) {
+                throw SSHConfigFormatException("Start block exists but no end block")
+            } else if (start.range.first > end.range.first) {
+                throw SSHConfigFormatException("Start block found after end block")
+            } else if (isRemoving) {
+                logger.info("Removing config from $sshConfigPath")
+                sshConfigPath.toFile().writeText(
+                    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("")
+                )
+            } else {
+                logger.info("Replacing config in $sshConfigPath")
+                sshConfigPath.toFile().writeText(
+                    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("")
+                )
+            }
+        } catch (e: FileNotFoundException) {
+            logger.info("Writing config to $sshConfigPath")
+            sshConfigPath.toFile().writeText(blockContent)
+        }
     }
 
     /**
@@ -241,6 +320,15 @@ class CoderCLIManager @JvmOverloads constructor(private val deployment: URL, des
                 }
             }
         }
+
+        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 +343,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)
diff --git a/src/main/kotlin/com/coder/gateway/views/steps/CoderLocateRemoteProjectStepView.kt b/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}")
             )
diff --git a/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt b/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt
@@ -440,6 +440,9 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
         poller?.cancel()
         listTableModelOfWorkspaces.items = emptyList()
 
+        val deploymentURL = localWizardModel.coderURL.toURL()
+        val token = localWizardModel.token
+
         // Authenticate and load in a background process with progress.
         // TODO: Make this cancelable.
         LifetimeDefinition().launchUnderBackgroundProgress(
@@ -449,35 +452,32 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
         ) {
             try {
                 this.indicator.text = "Authenticating client..."
-                authenticate(localWizardModel.coderURL.toURL(), localWizardModel.token)
+                authenticate(deploymentURL, token)
                 // Remember these in order to default to them for future attempts.
-                appPropertiesService.setValue(CODER_URL_KEY, localWizardModel.coderURL)
-                appPropertiesService.setValue(SESSION_TOKEN, localWizardModel.token)
+                appPropertiesService.setValue(CODER_URL_KEY, deploymentURL.toString())
+                appPropertiesService.setValue(SESSION_TOKEN, token)
 
                 this.indicator.text = "Retrieving workspaces..."
                 loadWorkspaces()
 
                 this.indicator.text = "Downloading Coder CLI..."
-                val cliManager = CoderCLIManager(localWizardModel.coderURL.toURL())
+                val cliManager = CoderCLIManager(deploymentURL)
                 cliManager.downloadCLI()
 
                 this.indicator.text = "Authenticating Coder CLI..."
-                cliManager.login(localWizardModel.token)
-
-                this.indicator.text = "Configuring SSH..."
-                cliManager.configSsh()
+                cliManager.login(token)
 
                 updateWorkspaceActions()
                 triggerWorkspacePolling(false)
             } catch (e: AuthenticationResponseException) {
-                logger.error("Token was rejected by ${localWizardModel.coderURL}; has your token expired?", e)
-                askTokenAndConnect(openBrowser) // Try again but no more opening browser windows.
+                logger.error("Token was rejected by $deploymentURL; has your token expired?", e)
+                askTokenAndConnect(openBrowser)
             } catch (e: SocketTimeoutException) {
-                logger.error("Unable to connect to ${localWizardModel.coderURL}; is it up?", e)
+                logger.error("Unable to connect to $deploymentURL; is it up?", e)
             } catch (e: ResponseException) {
                 logger.error("Failed to download Coder CLI", e)
             } catch (e: Exception) {
-                logger.error("Failed to configure connection to ${localWizardModel.coderURL}", e)
+                logger.error("Failed to configure connection to $deploymentURL", e)
             }
         }
     }
@@ -705,6 +705,11 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
         if (workspace != null) {
             wizardModel.selectedWorkspace = workspace
             poller?.cancel()
+
+            logger.info("Configuring Coder CLI...")
+            val cliManager = CoderCLIManager(wizardModel.coderURL.toURL())
+            cliManager.configSsh(listTableModelOfWorkspaces.items)
+
             logger.info("Opening IDE and Project Location window for ${workspace.name}")
             return true
         }
diff --git a/src/test/groovy/CoderCLIManagerTest.groovy b/src/test/groovy/CoderCLIManagerTest.groovy

Original file line number	Diff line number	Diff line change
`@@ -4,5 +4,5 @@ data class CoderWorkspacesWizardModel(`
`4`	`4`	`var coderURL: String = "https://coder.example.com",`
`5`	`5`	`var token: String = "",`
`6`	`6`	`var selectedWorkspace: WorkspaceAgentModel? = null,`
`7`		`- var useExistingToken: Boolean = false`
	`7`	`+ var useExistingToken: Boolean = false,`
`8`	`8`	`)`