Read url and token from cli config

Instead of pulling the url and token from the currently logged-in
deployment, we use the url and token as they exist on the disk organized
by the label (~/.config/.../coder.coder-remote/dev.coder.com for
example).

The label we extract from the host (coder-vscode.dev.coder.com-- for
example).  Because the new format is a little different, the parsing has
also changed.  It has been moved to utils for easier testing.

This way, we pull the right url and token for the host even if the user
is logged into a different deployment currently.

If there is no label because it is the older format, we read the old
deployment-unaware path instead (~/.config/.../coder.coder-remote).

Since this now uses the same url and token the cli itself will use, this
has the extra advantage of preventing a potential desync between what
exists on disk and what exists in VS Code's memory.

Author: code-asher

src/commands.ts

@@ -7,7 +7,7 @@ import { extractAgents } from "./api-helper"
 import { CertificateError } from "./error"
 import { Remote } from "./remote"
 import { Storage } from "./storage"
-import { toSafeHost } from "./util"
+import { AuthorityPrefix, toSafeHost } from "./util"
 import { OpenableTreeItem } from "./workspacesProvider"
 
 export class Commands {
@@ -475,7 +475,7 @@ async function openWorkspace(
 ) {
   // A workspace can have multiple agents, but that's handled
   // when opening a workspace unless explicitly specified.
-  let remoteAuthority = `ssh-remote+${Remote.Prefix}.${toSafeHost(baseUrl)}--${workspaceOwner}--${workspaceName}`
+  let remoteAuthority = `ssh-remote+${AuthorityPrefix}.${toSafeHost(baseUrl)}--${workspaceOwner}--${workspaceName}`
   if (workspaceAgent) {
     remoteAuthority += `--${workspaceAgent}`
   }

src/remote.ts

@@ -17,15 +17,11 @@ import { getHeaderCommand } from "./headers"
 import { SSHConfig, SSHValues, mergeSSHConfigValues } from "./sshConfig"
 import { computeSSHProperties, sshSupportsSetEnv } from "./sshSupport"
 import { Storage } from "./storage"
-import { toSafeHost } from "./util"
+import { parseRemoteAuthority } from "./util"
 import { supportsCoderAgentLogDirFlag } from "./version"
 import { WorkspaceAction } from "./workspaceAction"
 
 export class Remote {
-  // Prefix is a magic string that is prepended to SSH hosts to indicate that
-  // they should be handled by this extension.
-  public static readonly Prefix = "coder-vscode"
-
   public constructor(
     private readonly vscodeProposed: typeof vscode,
     private readonly storage: Storage,
@@ -34,34 +30,19 @@ export class Remote {
   ) {}
 
   public async setup(remoteAuthority: string): Promise<vscode.Disposable | undefined> {
-    const authorityParts = remoteAuthority.split("+")
-    // If the URI passed doesn't have the proper prefix ignore it. We don't need
-    // to do anything special, because this isn't trying to open a Coder
-    // workspace.
-    if (!authorityParts[1].startsWith(Remote.Prefix)) {
+    const parts = parseRemoteAuthority(remoteAuthority)
+    if (!parts) {
+      // Not a Coder host.
       return
     }
-    const sshAuthority = authorityParts[1].substring(Remote.Prefix.length)
-
-    // Authorities are in one of two formats:
-    // coder-vscode--<username>--<workspace>--<agent> (old style)
-    // coder-vscode.<label>--<username>--<workspace>--<agent>
-    // The agent can be omitted; the user will be prompted for it instead.
-    const parts = sshAuthority.split("--")
-    if (parts.length !== 2 && parts.length !== 3) {
-      throw new Error(`Invalid Coder SSH authority. Must be: <username>--<workspace>--<agent?>`)
-    }
-    const workspaceName = `${parts[0]}/${parts[1]}`
 
-    // It is possible to connect to any previously connected workspace, which
-    // might not belong to the deployment the plugin is currently logged into.
-    // For that reason, create a separate REST client instead of using the
-    // global one generally used by the plugin.  For now this is not actually
-    // useful because we are using the the current URL and token anyway, but in
-    // a future PR we will store these per deployment and grab the right one
-    // based on the host name of the workspace to which we are connecting.
-    const baseUrlRaw = this.storage.getUrl()
-    if (!baseUrlRaw) {
+    const workspaceName = `${parts.username}/${parts.workspace}`
+
+    // Get the URL and token belonging to this host.
+    const { url: baseUrlRaw, token } = await this.storage.readCliConfig(parts.label)
+
+    // It could be that the cli config was deleted.  If so, ask for the url.
+    if (!baseUrlRaw || !token) {
       const result = await this.vscodeProposed.window.showInformationMessage(
         "You are not logged in...",
         {
@@ -76,15 +57,16 @@ export class Remote {
         await this.closeRemote()
       } else {
         // Log in then try again.
-        await vscode.commands.executeCommand("coder.login")
+        await vscode.commands.executeCommand("coder.login", baseUrlRaw)
         await this.setup(remoteAuthority)
       }
       return
     }
 
-    const baseUrl = new URL(baseUrlRaw)
-    const safeHost = toSafeHost(baseUrlRaw) // Deployment label.
-    const token = await this.storage.getSessionToken()
+    // It is possible to connect to any previously connected workspace, which
+    // might not belong to the deployment the plugin is currently logged into.
+    // For that reason, create a separate REST client instead of using the
+    // global one generally used by the plugin.
     const restClient = await makeCoderSdk(baseUrlRaw, token, this.storage)
     // Store for use in commands.
     this.commands.workspaceRestClient = restClient
@@ -116,7 +98,7 @@ export class Remote {
     // Next is to find the workspace from the URI scheme provided.
     let workspace: Workspace
     try {
-      workspace = await restClient.getWorkspaceByOwnerAndName(parts[0], parts[1])
+      workspace = await restClient.getWorkspaceByOwnerAndName(parts.username, parts.workspace)
       this.commands.workspace = workspace
     } catch (error) {
       if (!isAxiosError(error)) {
@@ -235,24 +217,30 @@ export class Remote {
         path += `&after=${logs[logs.length - 1].id}`
       }
       await new Promise<void>((resolve, reject) => {
-        const proto = baseUrl.protocol === "https:" ? "wss:" : "ws:"
-        const socket = new ws.WebSocket(new URL(`${proto}//${baseUrl.host}${path}`), {
-          headers: {
-            "Coder-Session-Token": token,
-          },
-        })
-        socket.binaryType = "nodebuffer"
-        socket.on("message", (data) => {
-          const buf = data as Buffer
-          const log = JSON.parse(buf.toString()) as ProvisionerJobLog
-          writeEmitter.fire(log.output + "\r\n")
-        })
-        socket.on("error", (err) => {
-          reject(err)
-        })
-        socket.on("close", () => {
-          resolve()
-        })
+        try {
+          const baseUrl = new URL(baseUrlRaw)
+          const proto = baseUrl.protocol === "https:" ? "wss:" : "ws:"
+          const socket = new ws.WebSocket(new URL(`${proto}//${baseUrl.host}${path}`), {
+            headers: {
+              "Coder-Session-Token": token,
+            },
+          })
+          socket.binaryType = "nodebuffer"
+          socket.on("message", (data) => {
+            const buf = data as Buffer
+            const log = JSON.parse(buf.toString()) as ProvisionerJobLog
+            writeEmitter.fire(log.output + "\r\n")
+          })
+          socket.on("error", (err) => {
+            reject(err)
+          })
+          socket.on("close", () => {
+            resolve()
+          })
+        } catch (error) {
+          // If this errors, it is probably a malformed URL.
+          reject(error)
+        }
       })
       writeEmitter.fire("Build complete")
       workspace = await restClient.getWorkspace(workspace.id)
@@ -268,7 +256,7 @@ export class Remote {
           `This workspace is stopped!`,
           {
             modal: true,
-            detail: `Click below to start and open ${parts[0]}/${parts[1]}.`,
+            detail: `Click below to start and open ${workspaceName}.`,
             useCustom: true,
           },
           "Start Workspace",
@@ -286,28 +274,26 @@ export class Remote {
       return acc.concat(resource.agents || [])
     }, [] as WorkspaceAgent[])
 
-    let agent: WorkspaceAgent | undefined
-
-    if (parts.length === 2) {
+    // With no agent specified, pick the first one.  Otherwise choose the
+    // matching agent.
+    let agent: WorkspaceAgent
+    if (!parts.agent) {
       if (agents.length === 1) {
         agent = agents[0]
+      } else {
+        // TODO: Show the agent selector here instead.
+        throw new Error("Invalid Coder SSH authority. An agent must be specified when there are multiple.")
       }
-
-      // If there are multiple agents, we should select one here! TODO: Support
-      // multiple agents!
-    }
-
-    if (!agent) {
-      const matchingAgents = agents.filter((agent) => agent.name === parts[2])
+    } else {
+      const matchingAgents = agents.filter((agent) => agent.name === parts.agent)
       if (matchingAgents.length !== 1) {
-        // TODO: Show the agent selector here instead!
-        throw new Error(`Invalid Coder SSH authority. Agent not found!`)
+        // TODO: Show the agent selector here instead.
+        throw new Error("Invalid Coder SSH authority. Agent not found.")
       }
       agent = matchingAgents[0]
     }
 
     // Do some janky setting manipulation.
-    const hostname = authorityParts[1]
     const remotePlatforms = this.vscodeProposed.workspace
       .getConfiguration()
       .get<Record<string, string>>("remote.SSH.remotePlatform", {})
@@ -330,8 +316,8 @@ export class Remote {
     // Add the remote platform for this host to bypass a step where VS Code asks
     // the user for the platform.
     let mungedPlatforms = false
-    if (!remotePlatforms[hostname] || remotePlatforms[hostname] !== agent.operating_system) {
-      remotePlatforms[hostname] = agent.operating_system
+    if (!remotePlatforms[parts.host] || remotePlatforms[parts.host] !== agent.operating_system) {
+      remotePlatforms[parts.host] = agent.operating_system
       settingsContent = jsonc.applyEdits(
         settingsContent,
         jsonc.modify(settingsContent, ["remote.SSH.remotePlatform"], remotePlatforms, {}),
@@ -512,7 +498,7 @@ export class Remote {
     // If we didn't write to the SSH config file, connecting would fail with
     // "Host not found".
     try {
-      await this.updateSSHConfig(restClient, safeHost, authorityParts[1], hasCoderLogs)
+      await this.updateSSHConfig(restClient, parts.label, parts.host, hasCoderLogs)
     } catch (error) {
       this.storage.writeToCoderOutputChannel(`Failed to configure SSH: ${error}`)
       throw error

src/storage.ts

@@ -118,7 +118,7 @@ export class Storage {
   public async fetchBinary(restClient: Api, label: string | string): Promise<string> {
     const baseUrl = restClient.getAxiosInstance().defaults.baseURL
     this.output.appendLine(`Using deployment URL: ${baseUrl}`)
-    this.output.appendLine(`Using deployment label: ${label}`)
+    this.output.appendLine(`Using deployment label: ${label || "n/a"}`)
 
     // Settings can be undefined when set to their defaults (true in this case),
     // so explicitly check against false.
@@ -459,6 +459,26 @@ export class Storage {
     }
   }
 
+  /**
+   * Read the CLI config for a deployment with the provided label.
+   *
+   * IF a config file does not exist, return an empty string.
+   *
+   * If the label is empty, read the old deployment-unaware config.
+   */
+  public async readCliConfig(label: string): Promise<{ url: string; token: string }> {
+    const urlPath = this.getURLPath(label)
+    const tokenPath = this.getSessionTokenPath(label)
+    const [url, token] = await Promise.allSettled([
+      fs.readFile(urlPath, "utf8"),
+      fs.readFile(tokenPath, "utf8"),
+    ])
+    return {
+      url: url.status === "fulfilled" ? url.value : "",
+      token: token.status === "fulfilled" ? token.value : "",
+    }
+  }
+
   /**
    * Run the header command and return the generated headers.
    */

src/util.test.ts

@@ -1,5 +1,62 @@
 import { it, expect } from "vitest"
-import { toSafeHost } from "./util"
+import { parseRemoteAuthority, toSafeHost } from "./util"
+
+it("ignore unrelated authorities", async () => {
+  const tests = [
+    "vscode://ssh-remote+some-unrelated-host.com",
+    "vscode://ssh-remote+coder-vscode",
+    "vscode://ssh-remote+coder-vscode-test",
+    "vscode://ssh-remote+coder-vscode-test--foo--bar",
+    "vscode://ssh-remote+coder-vscode-foo--bar",
+    "vscode://ssh-remote+coder--foo--bar",
+  ]
+  for (const test of tests) {
+    expect(parseRemoteAuthority(test)).toBe(null)
+  }
+})
+
+it("should error on invalid authorities", async () => {
+  const tests = [
+    "vscode://ssh-remote+coder-vscode--foo",
+    "vscode://ssh-remote+coder-vscode--",
+    "vscode://ssh-remote+coder-vscode--foo--",
+    "vscode://ssh-remote+coder-vscode--foo--bar--",
+  ]
+  for (const test of tests) {
+    expect(() => parseRemoteAuthority(test)).toThrow("Invalid")
+  }
+})
+
+it("should parse authority", async () => {
+  expect(parseRemoteAuthority("vscode://ssh-remote+coder-vscode--foo--bar")).toStrictEqual({
+    agent: "",
+    host: "coder-vscode--foo--bar",
+    label: "",
+    username: "foo",
+    workspace: "bar",
+  })
+  expect(parseRemoteAuthority("vscode://ssh-remote+coder-vscode--foo--bar--baz")).toStrictEqual({
+    agent: "baz",
+    host: "coder-vscode--foo--bar--baz",
+    label: "",
+    username: "foo",
+    workspace: "bar",
+  })
+  expect(parseRemoteAuthority("vscode://ssh-remote+coder-vscode.dev.coder.com--foo--bar")).toStrictEqual({
+    agent: "",
+    host: "coder-vscode.dev.coder.com--foo--bar",
+    label: "dev.coder.com",
+    username: "foo",
+    workspace: "bar",
+  })
+  expect(parseRemoteAuthority("vscode://ssh-remote+coder-vscode.dev.coder.com--foo--bar--baz")).toStrictEqual({
+    agent: "baz",
+    host: "coder-vscode.dev.coder.com--foo--bar--baz",
+    label: "dev.coder.com",
+    username: "foo",
+    workspace: "bar",
+  })
+})
 
 it("escapes url host", async () => {
   expect(toSafeHost("https://foobar:8080")).toBe("foobar")

src/util.ts

@@ -1,5 +1,54 @@
 import url from "url"
 
+export interface AuthorityParts {
+  agent: string | undefined
+  host: string
+  label: string
+  username: string
+  workspace: string
+}
+
+// Prefix is a magic string that is prepended to SSH hosts to indicate that
+// they should be handled by this extension.
+export const AuthorityPrefix = "coder-vscode"
+
+/**
+ * Given an authority, parse into the expected parts.
+ *
+ * If this is not a Coder host, return null.
+ *
+ * Throw an error if the host is invalid.
+ */
+export function parseRemoteAuthority(authority: string): AuthorityParts | null {
+  // The authority looks like: vscode://ssh-remote+<ssh host name>
+  const authorityParts = authority.split("+")
+
+  // We create SSH host names in one of two formats:
+  // coder-vscode--<username>--<workspace>--<agent?> (old style)
+  // coder-vscode.<label>--<username>--<workspace>--<agent>
+  // The agent can be omitted; the user will be prompted for it instead.
+  // Anything else is unrelated to Coder and can be ignored.
+  const parts = authorityParts[1].split("--")
+  if (parts.length <= 1 || (parts[0] !== AuthorityPrefix && !parts[0].startsWith(`${AuthorityPrefix}.`))) {
+    return null
+  }
+
+  // It has the proper prefix, so this is probably a Coder host name.
+  // Validate the SSH host name.  Including the prefix, we expect at least
+  // three parts, or four if including the agent.
+  if ((parts.length !== 3 && parts.length !== 4) || parts.some((p) => !p)) {
+    throw new Error(`Invalid Coder SSH authority. Must be: <username>--<workspace>--<agent?>`)
+  }
+
+  return {
+    agent: parts[3] ?? "",
+    host: authorityParts[1],
+    label: parts[0].replace(/^coder-vscode\.?/, ""),
+    username: parts[1],
+    workspace: parts[2],
+  }
+}
+
 /**
  * Given a URL, return the host in a format that is safe to write.
  */