feat: Add extension settings for customizing ssh config (#74)

Emyrk · web-flow · commit 3e9cfec684dc · 2023-03-31T08:47:01.000-05:00
diff --git a/package.json b/package.json
@@ -29,6 +29,22 @@
   ],
   "main": "./dist/extension.js",
   "contributes": {
+    "configuration": {
+      "title": "Coder",
+      "properties": {
+        "coder.sshConfig": {
+          "markdownDescription": "These values will be included in the ssh config file. Eg: `'ConnectTimeout=10'` will set the timeout to 10 seconds. Any values included here will override anything provided by default or by the deployment. To unset a value that is written by default, set the value to the empty string, Eg: `'ConnectTimeout='` will unset it.",
+          "type": "array",
+          "items": {
+            "title": "SSH Config Value",
+            "type": "string",
+            "pattern": "^[a-zA-Z0-9-]+[=\\s].*$"
+          },
+          "scope": "machine",
+          "default": []
+        }
+      }
+    },
     "viewsContainers": {
       "activitybar": [
         {
@@ -140,4 +156,4 @@
     "ws": "^8.11.0",
     "yaml": "^1.10.0"
   }
-}
+}
diff --git a/src/remote.ts b/src/remote.ts
@@ -8,7 +8,7 @@ import {
   startWorkspace,
   getDeploymentSSHConfig,
 } from "coder/site/src/api/api"
-import { ProvisionerJobLog, SSHConfigResponse, Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated"
+import { ProvisionerJobLog, Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated"
 import EventSource from "eventsource"
 import find from "find-process"
 import * as fs from "fs/promises"
@@ -19,7 +19,7 @@ import prettyBytes from "pretty-bytes"
 import * as semver from "semver"
 import * as vscode from "vscode"
 import * as ws from "ws"
-import { SSHConfig, defaultSSHConfigResponse } from "./sshConfig"
+import { SSHConfig, defaultSSHConfigResponse, mergeSSHConfigValues } from "./sshConfig"
 import { Storage } from "./storage"
 
 export class Remote {
@@ -441,9 +441,10 @@ export class Remote {
   // updateSSHConfig updates the SSH configuration with a wildcard that handles
   // all Coder entries.
   private async updateSSHConfig() {
-    let deploymentConfig: SSHConfigResponse = defaultSSHConfigResponse
+    let deploymentSSHConfig = defaultSSHConfigResponse
     try {
-      deploymentConfig = await getDeploymentSSHConfig()
+      const deploymentConfig = await getDeploymentSSHConfig()
+      deploymentSSHConfig = deploymentConfig.ssh_config_options
     } catch (error) {
       if (!axios.isAxiosError(error)) {
         throw error
@@ -452,7 +453,6 @@ export class Remote {
         case 404: {
           // Deployment does not support overriding ssh config yet. Likely an
           // older version, just use the default.
-          deploymentConfig = defaultSSHConfigResponse
           break
         }
         case 401: {
@@ -464,6 +464,27 @@ export class Remote {
       }
     }
 
+    // deploymentConfig is now set from the remote coderd deployment.
+    // Now override with the user's config.
+    const userConfigSSH = vscode.workspace.getConfiguration("coder").get<string[]>("sshConfig") || []
+    // Parse the user's config into a Record<string, string>.
+    const userConfig = userConfigSSH.reduce((acc, line) => {
+      let i = line.indexOf("=")
+      if (i === -1) {
+        i = line.indexOf(" ")
+        if (i === -1) {
+          // This line is malformed. The setting is incorrect, and does not match
+          // the pattern regex in the settings schema.
+          return acc
+        }
+      }
+      const key = line.slice(0, i)
+      const value = line.slice(i + 1)
+      acc[key] = value
+      return acc
+    }, {} as Record<string, string>)
+    const sshConfigOverrides = mergeSSHConfigValues(deploymentSSHConfig, userConfig)
+
     let sshConfigFile = vscode.workspace.getConfiguration().get<string>("remote.SSH.configFile")
     if (!sshConfigFile) {
       sshConfigFile = path.join(os.homedir(), ".ssh", "config")
@@ -504,7 +525,7 @@ export class Remote {
       SetEnv: "CODER_SSH_SESSION_TYPE=vscode",
     }
 
-    await sshConfig.update(sshValues, deploymentConfig)
+    await sshConfig.update(sshValues, sshConfigOverrides)
   }
 
   // showNetworkUpdates finds the SSH process ID that is being used by this
diff --git a/src/sshConfig.test.ts b/src/sshConfig.test.ts
@@ -30,11 +30,11 @@ it("creates a new file and adds the config", async () => {
 
   const expectedOutput = `# --- START CODER VSCODE ---
 Host coder-vscode--*
-  ProxyCommand some-command-here
   ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null
-  LogLevel ERROR
 # --- END CODER VSCODE ---`
 
   expect(mockFileSystem.readFile).toBeCalledWith(sshFilePath, expect.anything())
@@ -43,12 +43,12 @@ Host coder-vscode--*
 
 it("adds a new coder config in an existent SSH configuration", async () => {
   const existentSSHConfig = `Host coder.something
-  HostName coder.something
   ConnectTimeout=0
-  StrictHostKeyChecking=no
-  UserKnownHostsFile=/dev/null
   LogLevel ERROR
-  ProxyCommand command`
+  HostName coder.something
+  ProxyCommand command
+  StrictHostKeyChecking=no
+  UserKnownHostsFile=/dev/null`
   mockFileSystem.readFile.mockResolvedValueOnce(existentSSHConfig)
 
   const sshConfig = new SSHConfig(sshFilePath, mockFileSystem)
@@ -66,11 +66,11 @@ it("adds a new coder config in an existent SSH configuration", async () => {
 
 # --- START CODER VSCODE ---
 Host coder-vscode--*
-  ProxyCommand some-command-here
   ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null
-  LogLevel ERROR
 # --- END CODER VSCODE ---`
 
   expect(mockFileSystem.writeFile).toBeCalledWith(sshFilePath, expectedOutput, {
@@ -90,11 +90,11 @@ it("updates an existent coder config", async () => {
 
 # --- START CODER VSCODE ---
 Host coder-vscode--*
-  ProxyCommand some-command-here
   ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null
-  LogLevel ERROR
 # --- END CODER VSCODE ---`
   mockFileSystem.readFile.mockResolvedValueOnce(existentSSHConfig)
 
@@ -119,11 +119,11 @@ Host coder-vscode--*
 
 # --- START CODER VSCODE ---
 Host coder--updated--vscode--*
-  ProxyCommand some-command-here
   ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null
-  LogLevel ERROR
 # --- END CODER VSCODE ---`
 
   expect(mockFileSystem.writeFile).toBeCalledWith(sshFilePath, expectedOutput, {
@@ -134,12 +134,12 @@ Host coder--updated--vscode--*
 
 it("removes old coder SSH config and adds the new one", async () => {
   const existentSSHConfig = `Host coder-vscode--*
-  HostName coder.something
   ConnectTimeout=0
-  StrictHostKeyChecking=no
-  UserKnownHostsFile=/dev/null
+  HostName coder.something
   LogLevel ERROR
-  ProxyCommand command`
+  ProxyCommand command
+  StrictHostKeyChecking=no
+  UserKnownHostsFile=/dev/null`
   mockFileSystem.readFile.mockResolvedValueOnce(existentSSHConfig)
 
   const sshConfig = new SSHConfig(sshFilePath, mockFileSystem)
@@ -155,11 +155,11 @@ it("removes old coder SSH config and adds the new one", async () => {
 
   const expectedOutput = `# --- START CODER VSCODE ---
 Host coder-vscode--*
-  ProxyCommand some-command-here
   ConnectTimeout 0
+  LogLevel ERROR
+  ProxyCommand some-command-here
   StrictHostKeyChecking no
   UserKnownHostsFile /dev/null
-  LogLevel ERROR
 # --- END CODER VSCODE ---`
 
   expect(mockFileSystem.writeFile).toBeCalledWith(sshFilePath, expectedOutput, {
@@ -182,29 +182,26 @@ it("override values", async () => {
       LogLevel: "ERROR",
     },
     {
-      ssh_config_options: {
-        loglevel: "DEBUG", // This tests case insensitive
-        ConnectTimeout: "500",
-        ExtraKey: "ExtraValue",
-        Foo: "bar",
-        Buzz: "baz",
-        // Remove this key
-        StrictHostKeyChecking: "",
-        ExtraRemove: "",
-      },
-      hostname_prefix: "",
+      loglevel: "DEBUG", // This tests case insensitive
+      ConnectTimeout: "500",
+      ExtraKey: "ExtraValue",
+      Foo: "bar",
+      Buzz: "baz",
+      // Remove this key
+      StrictHostKeyChecking: "",
+      ExtraRemove: "",
     },
   )
 
   const expectedOutput = `# --- START CODER VSCODE ---
 Host coder-vscode--*
-  ProxyCommand some-command-here
-  ConnectTimeout 500
-  UserKnownHostsFile /dev/null
-  LogLevel DEBUG
   Buzz baz
+  ConnectTimeout 500
   ExtraKey ExtraValue
   Foo bar
+  ProxyCommand some-command-here
+  UserKnownHostsFile /dev/null
+  loglevel DEBUG
 # --- END CODER VSCODE ---`
 
   expect(mockFileSystem.readFile).toBeCalledWith(sshFilePath, expect.anything())
diff --git a/src/sshConfig.ts b/src/sshConfig.ts
@@ -31,10 +31,54 @@ const defaultFileSystem: FileSystem = {
   writeFile,
 }
 
-export const defaultSSHConfigResponse: SSHConfigResponse = {
-  ssh_config_options: {},
-  // The prefix is not used by the vscode-extension
-  hostname_prefix: "coder.",
+export const defaultSSHConfigResponse: Record<string, string> = {}
+
+// mergeSSHConfigValues will take a given ssh config and merge it with the overrides
+// provided. The merge handles key case insensitivity, so casing in the "key" does
+// not matter.
+export function mergeSSHConfigValues(
+  config: Record<string, string>,
+  overrides: Record<string, string>,
+): Record<string, string> {
+  const merged: Record<string, string> = {}
+
+  // We need to do a case insensitive match for the overrides as ssh config keys are case insensitive.
+  // To get the correct key:value, use:
+  //   key = caseInsensitiveOverrides[key.toLowerCase()]
+  //   value = overrides[key]
+  const caseInsensitiveOverrides: Record<string, string> = {}
+  Object.keys(overrides).forEach((key) => {
+    caseInsensitiveOverrides[key.toLowerCase()] = key
+  })
+
+  Object.keys(config).forEach((key) => {
+    const lower = key.toLowerCase()
+    // If the key is in overrides, use the override value.
+    if (caseInsensitiveOverrides[lower]) {
+      const correctCaseKey = caseInsensitiveOverrides[lower]
+      const value = overrides[correctCaseKey]
+      delete caseInsensitiveOverrides[lower]
+
+      // If the value is empty, do not add the key. It is being removed.
+      if (value === "") {
+        return
+      }
+      merged[correctCaseKey] = value
+      return
+    }
+    // If no override, take the original value.
+    if (config[key] !== "") {
+      merged[key] = config[key]
+    }
+  })
+
+  // Add remaining overrides.
+  Object.keys(caseInsensitiveOverrides).forEach((lower) => {
+    const correctCaseKey = caseInsensitiveOverrides[lower]
+    merged[correctCaseKey] = overrides[correctCaseKey]
+  })
+
+  return merged
 }
 
 export class SSHConfig {
@@ -58,15 +102,15 @@ export class SSHConfig {
     }
   }
 
-  async update(values: SSHValues, overrides: SSHConfigResponse = defaultSSHConfigResponse) {
+  async update(values: SSHValues, overrides: Record<string, string> = defaultSSHConfigResponse) {
     // We should remove this in March 2023 because there is not going to have
     // old configs
     this.cleanUpOldConfig()
     const block = this.getBlock()
     if (block) {
       this.eraseBlock(block)
     }
-    this.appendBlock(values, overrides.ssh_config_options)
+    this.appendBlock(values, overrides)
     await this.save()
   }
 
@@ -122,43 +166,16 @@ export class SSHConfig {
    */
   private appendBlock({ Host, ...otherValues }: SSHValues, overrides: Record<string, string>) {
     const lines = [this.startBlockComment, `Host ${Host}`]
-    // We need to do a case insensitive match for the overrides as ssh config keys are case insensitive.
-    // To get the correct key:value, use:
-    //   key = caseInsensitiveOverrides[key.toLowerCase()]
-    //   value = overrides[key]
-    const caseInsensitiveOverrides: Record<string, string> = {}
-    Object.keys(overrides).forEach((key) => {
-      caseInsensitiveOverrides[key.toLowerCase()] = key
-    })
 
-    const keys = Object.keys(otherValues) as Array<keyof typeof otherValues>
+    // configValues is the merged values of the defaults and the overrides.
+    const configValues = mergeSSHConfigValues(otherValues, overrides)
+
+    // keys is the sorted keys of the merged values.
+    const keys = (Object.keys(configValues) as Array<keyof typeof configValues>).sort()
     keys.forEach((key) => {
-      const lower = key.toLowerCase()
-      if (caseInsensitiveOverrides[lower]) {
-        const correctCaseKey = caseInsensitiveOverrides[lower]
-        const value = overrides[correctCaseKey]
-        // Remove the key from the overrides so we don't write it again.
-        delete caseInsensitiveOverrides[lower]
-        if (value === "") {
-          // If the value is empty, don't write it. Prevent writing the default
-          // value as well.
-          return
-        }
-        // If the key is in overrides, use the override value.
-        // Doing it this way maintains the default order of the keys.
-        lines.push(this.withIndentation(`${key} ${value}`))
-        return
-      }
-      lines.push(this.withIndentation(`${key} ${otherValues[key]}`))
-    })
-    // Write remaining overrides that have not been written yet. Sort to maintain deterministic order.
-    const remainingKeys = (Object.keys(caseInsensitiveOverrides) as Array<keyof typeof caseInsensitiveOverrides>).sort()
-    remainingKeys.forEach((key) => {
-      const correctKey = caseInsensitiveOverrides[key]
-      const value = overrides[correctKey]
-      // Only write the value if it is not empty.
+      const value = configValues[key]
       if (value !== "") {
-        lines.push(this.withIndentation(`${correctKey} ${value}`))
+        lines.push(this.withIndentation(`${key} ${value}`))
       }
     })
 
