coder · code-asher · Sep 22, 2023 · Sep 18, 2023 · Sep 18, 2023 · Sep 18, 2023
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -32,7 +32,7 @@ jobs:
       - uses: gradle/wrapper-validation-action@v1.1.0
 
       # Run tests
-      - run: ./gradlew test
+      - run: ./gradlew test --info
 
       # Collect Tests Result of failed tests
       - if: ${{ failure() }}

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@
 
 ## Unreleased
 
+### Added
+- Add a setting for a command to run to get headers that will be set on all
+  requests to the Coder deployment.
+
 ## 2.6.0 - 2023-09-06
 
 ### Added

diff --git a/gradle.properties b/gradle.properties
@@ -3,7 +3,7 @@
 pluginGroup=com.coder.gateway
 pluginName=coder-gateway
 # SemVer format -> https://semver.org
-pluginVersion=2.6.0
+pluginVersion=2.7.0
 # See https://plugins.jetbrains.com/docs/intellij/build-number-ranges.html
 # for insight into build numbers and IntelliJ Platform versions.
 pluginSinceBuild=223.7571.70

diff --git a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
@@ -100,7 +100,7 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
             cli.login(client.token)
 
             indicator.text = "Configuring Coder CLI..."
-            cli.configSsh(workspaces.flatMap { it.toAgentModels() })
+            cli.configSsh(workspaces.flatMap { it.toAgentModels() }, settings.headerCommand)
 
             // TODO: Ask for these if missing.  Maybe we can reuse the second
             //  step of the wizard?  Could also be nice if we automatically used
@@ -150,7 +150,7 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
         if (token == null) { // User aborted.
             throw IllegalArgumentException("Unable to connect to $deploymentURL, $TOKEN is missing")
         }
-        val client = CoderRestClient(deploymentURL, token.first)
+        val client = CoderRestClient(deploymentURL, token.first, settings.headerCommand)
         return try {
             Pair(client, client.me().username)
         } catch (ex: AuthenticationResponseException) {

diff --git a/src/main/kotlin/com/coder/gateway/CoderSettingsConfigurable.kt b/src/main/kotlin/com/coder/gateway/CoderSettingsConfigurable.kt
@@ -66,6 +66,13 @@ class CoderSettingsConfigurable : BoundConfigurable("Coder") {
                         CoderGatewayBundle.message("gateway.connector.settings.enable-binary-directory-fallback.comment")
                     )
             }.layout(RowLayout.PARENT_GRID)
+            row(CoderGatewayBundle.message("gateway.connector.settings.header-command.title")) {
+                textField().resizableColumn().align(AlignX.FILL)
+                    .bindText(state::headerCommand)
+                    .comment(
+                        CoderGatewayBundle.message("gateway.connector.settings.header-command.comment")
+                    )
+            }.layout(RowLayout.PARENT_GRID)
         }
     }
 

diff --git a/src/main/kotlin/com/coder/gateway/sdk/CoderCLIManager.kt b/src/main/kotlin/com/coder/gateway/sdk/CoderCLIManager.kt
@@ -179,8 +179,9 @@ class CoderCLIManager @JvmOverloads constructor(
     /**
      * Configure SSH to use this binary.
      */
-    fun configSsh(workspaces: List<WorkspaceAgentModel>) {
-        writeSSHConfig(modifySSHConfig(readSSHConfig(), workspaces))
+    @JvmOverloads
+    fun configSsh(workspaces: List<WorkspaceAgentModel>, headerCommand: String? = null) {
+        writeSSHConfig(modifySSHConfig(readSSHConfig(), workspaces, headerCommand))
     }
 
     /**
@@ -199,11 +200,21 @@ class CoderCLIManager @JvmOverloads constructor(
      * 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? {
+    private fun modifySSHConfig(
+        contents: String?,
+        workspaces: List<WorkspaceAgentModel>,
+        headerCommand: String?,
+    ): String? {
         val host = getSafeHost(deploymentURL)
         val startBlock = "# --- START CODER JETBRAINS $host"
         val endBlock = "# --- END CODER JETBRAINS $host"
         val isRemoving = workspaces.isEmpty()
+        val proxyArgs = listOfNotNull(
+            escape(localBinaryPath.toString()),
+            "--global-config", escape(coderConfigPath.toString()),
+            if (!headerCommand.isNullOrBlank()) "--header-command" else null,
+            if (!headerCommand.isNullOrBlank()) escape(headerCommand) else null,
+           "ssh", "--stdio")
         val blockContent = workspaces.joinToString(
             System.lineSeparator(),
             startBlock + System.lineSeparator(),
@@ -212,7 +223,7 @@ class CoderCLIManager @JvmOverloads constructor(
                 """
                 Host ${getHostName(deploymentURL, it)}
                   HostName coder.${it.name}
-                  ProxyCommand "$localBinaryPath" --global-config "$coderConfigPath" ssh --stdio ${it.name}
+                  ProxyCommand ${proxyArgs.joinToString(" ")} ${it.name}
                   ConnectTimeout 0
                   StrictHostKeyChecking no
                   UserKnownHostsFile /dev/null
@@ -495,6 +506,24 @@ class CoderCLIManager @JvmOverloads constructor(
             // working binary and the binary directory does not.
             return if (cliMatches == null && dataCLIMatches != null) dataCLI else cli
         }
+
+        /**
+         * Escape a command argument to be used in the ProxyCommand of an SSH
+         * config.  Surround with double quotes if the argument contains
+         * whitespace and escape any existing double quotes.
+         *
+         * Throws if the argument is invalid.
+         */
+        @JvmStatic
+        fun escape(s: String): String {
+            if (s.contains("\n")) {
+                throw Exception("argument cannot contain newlines")
+            }
+            if (s.contains(" ") || s.contains("\t")) {
+                return "\"" + s.replace("\"", "\\\"") + "\""
+            }
+            return s.replace("\"", "\\\"")
+        }
     }
 }
 

diff --git a/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt b/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
@@ -20,6 +20,7 @@ import com.intellij.openapi.extensions.PluginId
 import com.intellij.openapi.util.SystemInfo
 import okhttp3.OkHttpClient
 import okhttp3.logging.HttpLoggingInterceptor
+import org.zeroturnaround.exec.ProcessExecutor
 import retrofit2.Retrofit
 import retrofit2.converter.gson.GsonConverterFactory
 import java.net.HttpURLConnection.HTTP_CREATED
@@ -41,16 +42,16 @@ class CoderRestClientService {
      *
      * @throws [AuthenticationResponseException] if authentication failed.
      */
-    fun initClientSession(url: URL, token: String): User {
-        client = CoderRestClient(url, token)
+    fun initClientSession(url: URL, token: String, headerCommand: String?): User {
+        client = CoderRestClient(url, token, headerCommand)
         me = client.me()
         buildVersion = client.buildInfo().version
         isReady = true
         return me
     }
 }
 
-class CoderRestClient(var url: URL, var token: String) {
+class CoderRestClient(var url: URL, var token: String, var headerCommand: String?) {
     private var httpClient: OkHttpClient
     private var retroRestClient: CoderV2RestFacade
 
@@ -61,6 +62,16 @@ class CoderRestClient(var url: URL, var token: String) {
         httpClient = OkHttpClient.Builder()
             .addInterceptor { it.proceed(it.request().newBuilder().addHeader("Coder-Session-Token", token).build()) }
             .addInterceptor { it.proceed(it.request().newBuilder().addHeader("User-Agent", "Coder Gateway/${pluginVersion.version} (${SystemInfo.getOsNameAndVersion()}; ${SystemInfo.OS_ARCH})").build()) }
+            .addInterceptor {
+                var request = it.request()
+                val headers = getHeaders(url, headerCommand)
+                if (headers.size > 0) {
+                    val builder = request.newBuilder()
+                    headers.forEach { h -> builder.addHeader(h.key, h.value) }
+                    request = builder.build()
+                }
+                it.proceed(request)
+            }
             // this should always be last if we want to see previous interceptors logged
             .addInterceptor(HttpLoggingInterceptor().apply { setLevel(HttpLoggingInterceptor.Level.BASIC) })
             .build()
@@ -141,4 +152,45 @@ class CoderRestClient(var url: URL, var token: String) {
 
         return buildResponse.body()!!
     }
+
+    companion object {
+        private val newlineRegex = "\r?\n".toRegex()
+        private val endingNewlineRegex = "\r?\n$".toRegex()
+
+        // TODO: This really only needs to be a private function, but
+        // unfortunately it is not possible to test the client because it fails
+        // on the plugin manager core call and I do not know how to fix it.  So,
+        // for now make this static and test it directly instead.
+        @JvmStatic
+        fun getHeaders(url: URL, headerCommand: String?): Map<String, String> {
+            if (headerCommand.isNullOrBlank()) {
+                return emptyMap()
+            }
+            val (shell, caller) = when (getOS()) {
+                OS.WINDOWS -> Pair("cmd.exe", "/c")
+                else -> Pair("sh", "-c")
+            }
+            return ProcessExecutor()
+                .command(shell, caller, headerCommand)
+                .environment("CODER_URL", url.toString())
+                .exitValues(0)
+                .readOutput(true)
+                .execute()
+                .outputUTF8()
+                .replaceFirst(endingNewlineRegex, "")
+                .split(newlineRegex)
+                .associate {
+                    // Header names cannot be blank or contain whitespace and
+                    // the Coder CLI requires that there be an equals sign (the
+                    // value can be blank though).  The second case is taken
+                    // care of by the destructure here, as it will throw if
+                    // there are not enough parts.
+                    val (name, value) = it.split("=", limit=2)
+                    if (name.contains(" ") || name == "") {
+                        throw Exception("\"$name\" is not a valid header name")
+                    }
+                    name to value
+                }
+        }
+    }
 }
diff --git a/src/main/kotlin/com/coder/gateway/services/CoderSettingsState.kt b/src/main/kotlin/com/coder/gateway/services/CoderSettingsState.kt
@@ -18,6 +18,7 @@ class CoderSettingsState : PersistentStateComponent<CoderSettingsState> {
     var dataDirectory: String = ""
     var enableDownloads: Boolean = true
     var enableBinaryDirectoryFallback: Boolean = false
+    var headerCommand: String = ""
     override fun getState(): CoderSettingsState {
         return this
     }

diff --git a/src/main/kotlin/com/coder/gateway/views/CoderGatewayRecentWorkspaceConnectionsView.kt b/src/main/kotlin/com/coder/gateway/views/CoderGatewayRecentWorkspaceConnectionsView.kt
@@ -13,6 +13,7 @@ import com.coder.gateway.sdk.toURL
 import com.coder.gateway.sdk.v2.models.WorkspaceStatus
 import com.coder.gateway.sdk.v2.models.toAgentModels
 import com.coder.gateway.services.CoderRecentWorkspaceConnectionsService
+import com.coder.gateway.services.CoderSettingsState
 import com.coder.gateway.toWorkspaceParams
 import com.intellij.icons.AllIcons
 import com.intellij.ide.BrowserUtil
@@ -72,6 +73,7 @@ data class DeploymentInfo(
 )
 
 class CoderGatewayRecentWorkspaceConnectionsView(private val setContentCallback: (Component) -> Unit) : GatewayRecentConnections, Disposable {
+    private val settings: CoderSettingsState = service()
     private val recentConnectionsService = service<CoderRecentWorkspaceConnectionsService>()
     private val cs = CoroutineScope(Dispatchers.Main)
 
@@ -259,7 +261,7 @@ class CoderGatewayRecentWorkspaceConnectionsView(private val setContentCallback:
                 deployments[dir] ?: try {
                     val url = Path.of(dir).resolve("url").readText()
                     val token = Path.of(dir).resolve("session").readText()
-                    DeploymentInfo(CoderRestClient(url.toURL(), token))
+                    DeploymentInfo(CoderRestClient(url.toURL(), token, settings.headerCommand))
                 } catch (e: Exception) {
                     logger.error("Unable to create client from $dir", e)
                     DeploymentInfo(error = "Error trying to read $dir: ${e.message}")

diff --git a/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt b/src/main/kotlin/com/coder/gateway/views/steps/CoderWorkspacesStepView.kt
@@ -527,7 +527,7 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
      */
     private fun authenticate(url: URL, token: String) {
         logger.info("Authenticating to $url...")
-        clientService.initClientSession(url, token)
+        clientService.initClientSession(url, token, settings.headerCommand)
 
         try {
             logger.info("Checking compatibility with Coder version ${clientService.buildVersion}...")
@@ -614,7 +614,7 @@ class CoderWorkspacesStepView(val setNextButtonEnabled: (Boolean) -> Unit) : Cod
         poller?.cancel()
 
         logger.info("Configuring Coder CLI...")
-        cli.configSsh(tableOfWorkspaces.items)
+        cli.configSsh(tableOfWorkspaces.items, settings.headerCommand)
 
         // The config directory can be used to pull the URL and token in
         // order to query this workspace's status in other flows, for

diff --git a/src/main/resources/messages/CoderGatewayBundle.properties b/src/main/resources/messages/CoderGatewayBundle.properties
@@ -87,3 +87,8 @@ gateway.connector.settings.enable-binary-directory-fallback.title=Fall back to d
 gateway.connector.settings.enable-binary-directory-fallback.comment=Checking this \
   box will allow the plugin to fall back to the data directory when the CLI \
   directory is not writable.
+gateway.connector.settings.header-command.title=Header command:
+gateway.connector.settings.header-command.comment=An external command that \
+  outputs additional HTTP headers added to all requests. The command must \
+  output each header as `key=value` on its own line. The following \
+  environment variables will be available to the process: CODER_URL.
diff --git a/src/test/fixtures/outputs/append-blank-newlines.conf b/src/test/fixtures/outputs/append-blank-newlines.conf
@@ -5,7 +5,7 @@
 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo-bar--test.coder.invalid
   HostName coder.foo-bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo-bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo-bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

diff --git a/src/test/fixtures/outputs/append-blank.conf b/src/test/fixtures/outputs/append-blank.conf
@@ -1,7 +1,7 @@
 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo-bar--test.coder.invalid
   HostName coder.foo-bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo-bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo-bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

diff --git a/src/test/fixtures/outputs/append-no-blocks.conf b/src/test/fixtures/outputs/append-no-blocks.conf
@@ -6,7 +6,7 @@ Host test2
 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo-bar--test.coder.invalid
   HostName coder.foo-bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo-bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo-bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

diff --git a/src/test/fixtures/outputs/append-no-newline.conf b/src/test/fixtures/outputs/append-no-newline.conf
@@ -5,7 +5,7 @@ Host test2
 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo-bar--test.coder.invalid
   HostName coder.foo-bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo-bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo-bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

diff --git a/src/test/fixtures/outputs/append-no-related-blocks.conf b/src/test/fixtures/outputs/append-no-related-blocks.conf
@@ -12,7 +12,7 @@ some jetbrains config
 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo-bar--test.coder.invalid
   HostName coder.foo-bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo-bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo-bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

diff --git a/src/test/fixtures/outputs/header-command-windows.conf b/src/test/fixtures/outputs/header-command-windows.conf
@@ -0,0 +1,10 @@
+# --- START CODER JETBRAINS test.coder.invalid
+Host coder-jetbrains--header--test.coder.invalid
+  HostName coder.header
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config --header-command "C:\Program Files\My Header Command\\"also has quotes\"\HeaderCommand.exe" ssh --stdio header
+  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/fixtures/outputs/header-command.conf b/src/test/fixtures/outputs/header-command.conf
@@ -0,0 +1,10 @@
+# --- START CODER JETBRAINS test.coder.invalid
+Host coder-jetbrains--header--test.coder.invalid
+  HostName coder.header
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config --header-command "my-header-command \"test\"" ssh --stdio header
+  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/fixtures/outputs/multiple-workspaces.conf b/src/test/fixtures/outputs/multiple-workspaces.conf
@@ -1,15 +1,15 @@
 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo--test.coder.invalid
   HostName coder.foo
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null
   LogLevel ERROR
   SetEnv CODER_SSH_SESSION_TYPE=JetBrains
 Host coder-jetbrains--bar--test.coder.invalid
   HostName coder.bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

diff --git a/src/test/fixtures/outputs/replace-end-no-newline.conf b/src/test/fixtures/outputs/replace-end-no-newline.conf
@@ -4,7 +4,7 @@ Host test2
   Port 443 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo-bar--test.coder.invalid
   HostName coder.foo-bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo-bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo-bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

diff --git a/src/test/fixtures/outputs/replace-end.conf b/src/test/fixtures/outputs/replace-end.conf
@@ -5,7 +5,7 @@ Host test2
 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo-bar--test.coder.invalid
   HostName coder.foo-bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo-bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo-bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null

diff --git a/src/test/fixtures/outputs/replace-middle-ignore-unrelated.conf b/src/test/fixtures/outputs/replace-middle-ignore-unrelated.conf
@@ -6,7 +6,7 @@ some coder config
 # --- START CODER JETBRAINS test.coder.invalid
 Host coder-jetbrains--foo-bar--test.coder.invalid
   HostName coder.foo-bar
-  ProxyCommand "/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64" --global-config "/tmp/coder-gateway/test.coder.invalid/config" ssh --stdio foo-bar
+  ProxyCommand /tmp/coder-gateway/test.coder.invalid/coder-linux-amd64 --global-config /tmp/coder-gateway/test.coder.invalid/config ssh --stdio foo-bar
   ConnectTimeout 0
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null