Scope CLI and CLI config by label

Every time you log in or out we would add or remove the cli and its config in a single directory, meaning logging into another deployment would wipe out the binary and config for the previous one. Now, similar to the ssh config blocks, the cli and its config files are scoped by their label, which is derived from the deployment URL, so they can coexist.
coder · code-asher · Jun 4, 2024 · Jun 3, 2024 · Jun 4, 2024 · Jun 3, 2024
commit 996611404adc8ab347bcb0f1d75700182cbd56f3
diff --git a/src/commands.ts b/src/commands.ts
@@ -154,10 +154,13 @@ export class Commands {
     this.restClient.setHost(url)
     this.restClient.setSessionToken(token)
 
-    // Store these to be used in later sessions and in the cli.
+    // Store these to be used in later sessions.
     await this.storage.setURL(url)
     await this.storage.setSessionToken(token)
 
+    // Store on disk to be used by the cli.
+    await this.storage.configureCli(toSafeHost(url), url, token)
+
     await vscode.commands.executeCommand("setContext", "coder.authenticated", true)
     if (user.roles.find((role) => role.name === "owner")) {
       await vscode.commands.executeCommand("setContext", "coder.isOwner", true)
@@ -198,6 +201,12 @@ export class Commands {
    * Log out from the currently logged-in deployment.
    */
   public async logout(): Promise<void> {
+    const url = this.storage.getUrl()
+    if (!url) {
+      // Sanity check; command should not be available if no url.
+      throw new Error("You are not logged in")
+    }
+
     // Clear from the REST client.  An empty url will indicate to other parts of
     // the code that we are logged out.
     this.restClient.setHost("")
@@ -207,6 +216,9 @@ export class Commands {
     await this.storage.setURL(undefined)
     await this.storage.setSessionToken(undefined)
 
+    // Clear from disk.
+    await this.storage.configureCli(toSafeHost(url), undefined, undefined)
+
     await vscode.commands.executeCommand("setContext", "coder.authenticated", false)
     vscode.window.showInformationMessage("You've been logged out of Coder!", "Login").then((action) => {
       if (action === "Login") {

diff --git a/src/extension.ts b/src/extension.ts
@@ -8,6 +8,7 @@ import { Commands } from "./commands"
 import { CertificateError, getErrorDetail } from "./error"
 import { Remote } from "./remote"
 import { Storage } from "./storage"
+import { toSafeHost } from "./util"
 import { WorkspaceQuery, WorkspaceProvider } from "./workspacesProvider"
 
 export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
@@ -108,6 +109,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
         // hit enter and move on.
         const url = await commands.maybeAskUrl(params.get("url"), storage.getUrl())
         if (url) {
+          restClient.setHost(url)
           await storage.setURL(url)
         } else {
           throw new Error("url must be provided or specified as a query parameter")
@@ -117,9 +119,13 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
         // prompted to sign in again, so we do not need to ensure it is set.
         const token = params.get("token")
         if (token) {
+          restClient.setSessionToken(token)
           await storage.setSessionToken(token)
         }
 
+        // Store on disk to be used by the cli.
+        await storage.configureCli(toSafeHost(url), url, token)
+
         vscode.commands.executeCommand("coder.open", owner, workspace, agent, folder, openRecent)
       }
     },

diff --git a/src/remote.ts b/src/remote.ts
@@ -610,15 +610,15 @@ export class Remote {
 
     let binaryPath: string | undefined
     if (this.mode === vscode.ExtensionMode.Production) {
-      binaryPath = await this.storage.fetchBinary(restClient)
+      binaryPath = await this.storage.fetchBinary(restClient, label)
     } else {
       try {
         // In development, try to use `/tmp/coder` as the binary path.
         // This is useful for debugging with a custom bin!
         binaryPath = path.join(os.tmpdir(), "coder")
         await fs.stat(binaryPath)
       } catch (ex) {
-        binaryPath = await this.storage.fetchBinary(restClient)
+        binaryPath = await this.storage.fetchBinary(restClient, label)
       }
     }
 
@@ -647,8 +647,8 @@ export class Remote {
       Host: label ? `${Remote.Prefix}.${label}--*` : `${RemotePrefix}--*`,
       ProxyCommand: `${escape(binaryPath)}${headerArg} vscodessh --network-info-dir ${escape(
         this.storage.getNetworkInfoPath(),
-      )}${logArg} --session-token-file ${escape(this.storage.getSessionTokenPath())} --url-file ${escape(
-        this.storage.getURLPath(),
+      )}${logArg} --session-token-file ${escape(this.storage.getSessionTokenPath(label))} --url-file ${escape(
+        this.storage.getURLPath(label),
       )} %h`,
       ConnectTimeout: "0",
       StrictHostKeyChecking: "no",

diff --git a/src/storage.ts b/src/storage.ts
@@ -1,7 +1,6 @@
 import { Api } from "coder/site/src/api/api"
 import { createWriteStream } from "fs"
 import fs from "fs/promises"
-import { ensureDir } from "fs-extra"
 import { IncomingMessage } from "http"
 import path from "path"
 import prettyBytes from "pretty-bytes"
@@ -24,14 +23,13 @@ export class Storage {
 
   /**
    * Add the URL to the list of recently accessed URLs in global storage, then
-   * set it as the last used URL and update it on disk for the cli.
+   * set it as the last used URL.
    *
-   * If the URL is falsey, then remove it as the currently accessed URL and do
-   * not touch the history.
+   * If the URL is falsey, then remove it as the last used URL and do not touch
+   * the history.
    */
   public async setURL(url?: string): Promise<void> {
     await this.memento.update("url", url)
-    this.updateUrl(url)
     if (url) {
       const history = this.withUrlHistory(url)
       await this.memento.update("urlHistory", history)
@@ -64,15 +62,13 @@ export class Storage {
   }
 
   /**
-   * Set or unset the last used token and update it on disk for the cli.
+   * Set or unset the last used token.
    */
   public async setSessionToken(sessionToken?: string): Promise<void> {
     if (!sessionToken) {
       await this.secrets.delete("sessionToken")
-      this.updateSessionToken(undefined)
     } else {
       await this.secrets.store("sessionToken", sessionToken)
-      this.updateSessionToken(sessionToken)
     }
   }
 
@@ -109,16 +105,20 @@ export class Storage {
   }
 
   /**
-   * Download and return the path to a working binary using the provided client.
+   * Download and return the path to a working binary for the deployment with
+   * the provided label using the provided client.  If the label is empty, use
+   * the old deployment-unaware path instead.
+   *
    * If there is already a working binary and it matches the server version,
    * return that, skipping the download.  If it does not match but downloads are
    * disabled, return whatever we have and log a warning.  Otherwise throw if
    * unable to download a working binary, whether because of network issues or
    * downloads being disabled.
    */
-  public async fetchBinary(restClient: Api): Promise<string> {
+  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}`)
 
     // Settings can be undefined when set to their defaults (true in this case),
     // so explicitly check against false.
@@ -133,7 +133,7 @@ export class Storage {
     // Check if there is an existing binary and whether it looks valid.  If it
     // is valid and matches the server, or if it does not match the server but
     // downloads are disabled, we can return early.
-    const binPath = path.join(this.getBinaryCachePath(), cli.name())
+    const binPath = path.join(this.getBinaryCachePath(label), cli.name())
     this.output.appendLine(`Using binary path: ${binPath}`)
     const stat = await cli.stat(binPath)
     if (stat === undefined) {
@@ -351,13 +351,21 @@ export class Storage {
     }
   }
 
-  // getBinaryCachePath returns the path where binaries are cached.
-  // The caller must ensure it exists before use.
-  public getBinaryCachePath(): string {
+  /**
+   * Return the directory for a deployment with the provided label to where its
+   * binary is cached.
+   *
+   * If the label is empty, read the old deployment-unaware config instead.
+   *
+   * The caller must ensure this directory exists before use.
+   */
+  public getBinaryCachePath(label: string): string {
     const configPath = vscode.workspace.getConfiguration().get("coder.binaryDestination")
     return configPath && String(configPath).trim().length > 0
       ? path.resolve(String(configPath))
-      : path.join(this.globalStorageUri.fsPath, "bin")
+      : label
+        ? path.join(this.globalStorageUri.fsPath, label, "bin")
+        : path.join(this.globalStorageUri.fsPath, "bin")
   }
 
   // getNetworkInfoPath returns the path where network information
@@ -377,17 +385,31 @@ export class Storage {
   }
 
   /**
-   * Return the path to the session token file for the cli.
+   * Return the directory for the deployment with the provided label to where
+   * its session token is stored.
+   *
+   * If the label is empty, read the old deployment-unaware config instead.
+   *
+   * The caller must ensure this directory exists before use.
    */
-  public getSessionTokenPath(): string {
-    return path.join(this.globalStorageUri.fsPath, "session_token")
+  public getSessionTokenPath(label: string): string {
+    return label
+      ? path.join(this.globalStorageUri.fsPath, label, "session_token")
+      : path.join(this.globalStorageUri.fsPath, "session_token")
   }
 
   /**
-   * Return the path to the URL file for the cli.
+   * Return the directory for the deployment with the provided label to where
+   * its url is stored.
+   *
+   * If the label is empty, read the old deployment-unaware config instead.
+   *
+   * The caller must ensure this directory exists before use.
    */
-  public getURLPath(): string {
-    return path.join(this.globalStorageUri.fsPath, "url")
+  public getURLPath(label: string): string {
+    return label
+      ? path.join(this.globalStorageUri.fsPath, label, "url")
+      : path.join(this.globalStorageUri.fsPath, "url")
   }
 
   public writeToCoderOutputChannel(message: string) {
@@ -399,28 +421,41 @@ export class Storage {
   }
 
   /**
-   * Update or remove the URL on disk which can be used by the CLI via
-   * --url-file.
+   * Configure the CLI for the deployment with the provided label.
    */
-  private async updateUrl(url: string | undefined): Promise<void> {
+  public async configureCli(label: string, url: string | undefined, token: string | undefined | null) {
+    await Promise.all([this.updateUrlForCli(label, url), this.updateTokenForCli(label, token)])
+  }
+
+  /**
+   * Update or remove the URL for the deployment with the provided label on disk
+   * which can be used by the CLI via --url-file.
+   *
+   * If the label is empty, read the old deployment-unaware config instead.
+   */
+  private async updateUrlForCli(label: string, url: string | undefined): Promise<void> {
+    const urlPath = this.getURLPath(label)
     if (url) {
-      await ensureDir(this.globalStorageUri.fsPath)
-      await fs.writeFile(this.getURLPath(), url)
+      await fs.mkdir(path.dirname(urlPath), { recursive: true })
+      await fs.writeFile(urlPath, url)
     } else {
-      await fs.rm(this.getURLPath(), { force: true })
+      await fs.rm(urlPath, { force: true })
     }
   }
 
   /**
-   * Update or remove the session token on disk which can be used by the CLI
-   * via --session-token-file.
+   * Update or remove the session token for a deployment with the provided label
+   * on disk which can be used by the CLI via --session-token-file.
+   *
+   * If the label is empty, read the old deployment-unaware config instead.
    */
-  private async updateSessionToken(token: string | undefined) {
+  private async updateTokenForCli(label: string, token: string | undefined | null) {
+    const tokenPath = this.getSessionTokenPath(label)
     if (token) {
-      await ensureDir(this.globalStorageUri.fsPath)
-      await fs.writeFile(this.getSessionTokenPath(), token)
+      await fs.mkdir(path.dirname(tokenPath), { recursive: true })
+      await fs.writeFile(tokenPath, token)
     } else {
-      await fs.rm(this.getSessionTokenPath(), { force: true })
+      await fs.rm(tokenPath, { force: true })
     }
   }