feat: Add feature to allow config-ssh value overrides (#69)

Emyrk · web-flow · commit 3dfa5885f522 · 2023-03-28T09:38:56.000-05:00
diff --git a/src/commands.ts b/src/commands.ts
@@ -1,5 +1,5 @@
 import axios from "axios"
-import { getUser, getWorkspaces, updateWorkspaceVersion } from "coder/site/src/api/api"
+import { getAuthenticatedUser, getWorkspaces, updateWorkspaceVersion } from "coder/site/src/api/api"
 import { Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated"
 import * as vscode from "vscode"
 import { Remote } from "./remote"
@@ -73,21 +73,28 @@ export class Commands {
 
     await this.storage.setURL(url)
     await this.storage.setSessionToken(token)
-    const user = await getUser()
-    await vscode.commands.executeCommand("setContext", "coder.authenticated", true)
-    vscode.window
-      .showInformationMessage(
-        `Welcome to Coder, ${user.username}!`,
-        {
-          detail: "You can now use the Coder extension to manage your Coder instance.",
-        },
-        "Open Workspace",
-      )
-      .then((action) => {
-        if (action === "Open Workspace") {
-          vscode.commands.executeCommand("coder.open")
-        }
-      })
+    try {
+      const user = await getAuthenticatedUser()
+      if (!user) {
+        throw new Error("Failed to get authenticated user")
+      }
+      await vscode.commands.executeCommand("setContext", "coder.authenticated", true)
+      vscode.window
+        .showInformationMessage(
+          `Welcome to Coder, ${user.username}!`,
+          {
+            detail: "You can now use the Coder extension to manage your Coder instance.",
+          },
+          "Open Workspace",
+        )
+        .then((action) => {
+          if (action === "Open Workspace") {
+            vscode.commands.executeCommand("coder.open")
+          }
+        })
+    } catch (error) {
+      vscode.window.showErrorMessage("Failed to authenticate with Coder: " + error)
+    }
   }
 
   public async logout(): Promise<void> {
diff --git a/src/extension.ts b/src/extension.ts
@@ -1,6 +1,6 @@
 "use strict"
 
-import { getUser } from "coder/site/src/api/api"
+import { getAuthenticatedUser } from "coder/site/src/api/api"
 import * as module from "module"
 import * as vscode from "vscode"
 import { Commands } from "./commands"
@@ -12,7 +12,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
   const storage = new Storage(output, ctx.globalState, ctx.secrets, ctx.globalStorageUri, ctx.logUri)
   await storage.init()
 
-  getUser()
+  getAuthenticatedUser()
     .then(() => {
       vscode.commands.executeCommand("setContext", "coder.authenticated", true)
     })
diff --git a/src/remote.ts b/src/remote.ts
@@ -6,8 +6,9 @@ import {
   getWorkspaceBuildLogs,
   getWorkspaceByOwnerAndName,
   startWorkspace,
+  getDeploymentSSHConfig,
 } from "coder/site/src/api/api"
-import { ProvisionerJobLog, Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated"
+import { ProvisionerJobLog, SSHConfigResponse, Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated"
 import EventSource from "eventsource"
 import find from "find-process"
 import * as fs from "fs/promises"
@@ -18,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 } from "./sshConfig"
+import { SSHConfig, defaultSSHConfigResponse } from "./sshConfig"
 import { Storage } from "./storage"
 
 export class Remote {
@@ -440,6 +441,29 @@ export class Remote {
   // updateSSHConfig updates the SSH configuration with a wildcard that handles
   // all Coder entries.
   private async updateSSHConfig() {
+    let deploymentConfig: SSHConfigResponse = defaultSSHConfigResponse
+    try {
+      deploymentConfig = await getDeploymentSSHConfig()
+    } catch (error) {
+      if (!axios.isAxiosError(error)) {
+        throw error
+      }
+      switch (error.response?.status) {
+        case 404: {
+          // Deployment does not support overriding ssh config yet. Likely an
+          // older version, just use the default.
+          deploymentConfig = defaultSSHConfigResponse
+          break
+        }
+        case 401: {
+          await this.vscodeProposed.window.showErrorMessage("Your session expired...")
+          throw error
+        }
+        default:
+          throw error
+      }
+    }
+
     let sshConfigFile = vscode.workspace.getConfiguration().get<string>("remote.SSH.configFile")
     if (!sshConfigFile) {
       sshConfigFile = path.join(os.homedir(), ".ssh", "config")
@@ -480,7 +504,7 @@ export class Remote {
       SetEnv: "CODER_SSH_SESSION_TYPE=vscode",
     }
 
-    await sshConfig.update(sshValues)
+    await sshConfig.update(sshValues, deploymentConfig)
   }
 
   // showNetworkUpdates finds the SSH process ID that is being used by this
diff --git a/src/sshConfig.test.ts b/src/sshConfig.test.ts
@@ -167,3 +167,46 @@ Host coder-vscode--*
     mode: 384,
   })
 })
+
+it("override values", async () => {
+  mockFileSystem.readFile.mockRejectedValueOnce("No file found")
+  const sshConfig = new SSHConfig(sshFilePath, mockFileSystem)
+  await sshConfig.load()
+  await sshConfig.update(
+    {
+      Host: "coder-vscode--*",
+      ProxyCommand: "some-command-here",
+      ConnectTimeout: "0",
+      StrictHostKeyChecking: "no",
+      UserKnownHostsFile: "/dev/null",
+      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: "",
+    },
+  )
+
+  const expectedOutput = `# --- START CODER VSCODE ---
+Host coder-vscode--*
+  ProxyCommand some-command-here
+  ConnectTimeout 500
+  UserKnownHostsFile /dev/null
+  LogLevel DEBUG
+  Buzz baz
+  ExtraKey ExtraValue
+  Foo bar
+# --- END CODER VSCODE ---`
+
+  expect(mockFileSystem.readFile).toBeCalledWith(sshFilePath, expect.anything())
+  expect(mockFileSystem.writeFile).toBeCalledWith(sshFilePath, expectedOutput, expect.anything())
+})
diff --git a/src/sshConfig.ts b/src/sshConfig.ts
@@ -1,3 +1,4 @@
+import { SSHConfigResponse } from "coder/site/src/api/typesGenerated"
 import { writeFile, readFile } from "fs/promises"
 import { ensureDir } from "fs-extra"
 import path from "path"
@@ -30,6 +31,12 @@ const defaultFileSystem: FileSystem = {
   writeFile,
 }
 
+export const defaultSSHConfigResponse: SSHConfigResponse = {
+  ssh_config_options: {},
+  // The prefix is not used by the vscode-extension
+  hostname_prefix: "coder.",
+}
+
 export class SSHConfig {
   private filePath: string
   private fileSystem: FileSystem
@@ -51,15 +58,15 @@ export class SSHConfig {
     }
   }
 
-  async update(values: SSHValues) {
+  async update(values: SSHValues, overrides: SSHConfigResponse = 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)
+    this.appendBlock(values, overrides.ssh_config_options)
     await this.save()
   }
 
@@ -102,12 +109,59 @@ export class SSHConfig {
     this.raw = this.getRaw().replace(block.raw, "")
   }
 
-  private appendBlock({ Host, ...otherValues }: SSHValues) {
+  /**
+   *
+   * appendBlock builds the ssh config block. The order of the keys is determinstic based on the input.
+   * Expected values are always in a consistent order followed by any additional overrides in sorted order.
+   *
+   * @param param0 - SSHValues are the expected SSH values for using ssh with coder.
+   * @param overrides - Overrides typically come from the deployment api and are used to override the default values.
+   *                    The overrides are given as key:value pairs where the key is the ssh config file key.
+   *                    If the key matches an expected value, the expected value is overridden. If it does not
+   *                    match an expected value, it is appended to the end of the block.
+   */
+  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>
     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.
+      if (value !== "") {
+        lines.push(this.withIndentation(`${correctKey} ${value}`))
+      }
+    })
+
     lines.push(this.endBlockComment)
     const raw = this.getRaw()
 
diff --git a/yarn.lock b/yarn.lock
@@ -1418,7 +1418,7 @@ co@3.1.0:
 
 "coder@https://github.com/coder/coder":
   version "0.0.0"
-  resolved "https://github.com/coder/coder#7a1731b6205d9c68f6308ee362ff2d62124b6950"
+  resolved "https://github.com/coder/coder#a6fa8cac582f2fc54eca0191bd54fd43d6d67ac2"
 
 collapse-white-space@^1.0.2:
   version "1.0.6"
