Add customizable header command

.eslintrc.json

@@ -35,7 +35,13 @@
     }],
     "import/no-unresolved": ["error", {
       "ignore": ["vscode"]
-    }]
+    }],
+    "@typescript-eslint/no-unused-vars": [
+      "error",
+      {
+        "varsIgnorePattern": "^_"
+      }
+    ]
   },
   "ignorePatterns": [
     "out",

package.json

@@ -57,6 +57,11 @@
           "markdownDescription": "The full path of the directory into which the Coder CLI will be downloaded. Defaults to the extension's global storage directory.",
           "type": "string",
           "default": ""
+        },
+        "coder.headerCommand": {
+          "markdownDescription": "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`.",
+          "type": "string",
+          "default": ""
         }
       }
     },

src/commands.ts

@@ -70,8 +70,10 @@ export class Commands {
                   severity: vscode.InputBoxValidationSeverity.Error,
                 }
               }
+              // This could be something like the header command erroring or an
+              // invalid session token.
               return {
-                message: "Invalid session token! (" + message + ")",
+                message: "Failed to authenticate: " + message,
                 severity: vscode.InputBoxValidationSeverity.Error,
               }
             })

src/extension.ts

@@ -46,7 +46,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
   axios.interceptors.response.use(
     (r) => r,
     async (err) => {
-      throw await CertificateError.maybeWrap(err, err.config.baseURL, storage)
+      throw await CertificateError.maybeWrap(err, axios.getUri(err.config), storage)
     },
   )
 
@@ -59,27 +59,42 @@ 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()
 
+  // Add headers from the header command.
+  axios.interceptors.request.use(async (config) => {
+    Object.entries(await storage.getHeaders(config.baseURL || axios.getUri(config))).forEach(([key, value]) => {
+      config.headers[key] = value
+    })
+    return config
+  })
+
   const myWorkspacesProvider = new WorkspaceProvider(WorkspaceQuery.Mine, storage)
   const allWorkspacesProvider = new WorkspaceProvider(WorkspaceQuery.All, storage)
 
   vscode.window.registerTreeDataProvider("myWorkspaces", myWorkspacesProvider)
   vscode.window.registerTreeDataProvider("allWorkspaces", allWorkspacesProvider)
 
-  getAuthenticatedUser()
-    .then(async (user) => {
-      if (user) {
-        vscode.commands.executeCommand("setContext", "coder.authenticated", true)
-        if (user.roles.find((role) => role.name === "owner")) {
-          await vscode.commands.executeCommand("setContext", "coder.isOwner", true)
+  const url = storage.getURL()
+  if (url) {
+    getAuthenticatedUser()
+      .then(async (user) => {
+        if (user) {
+          vscode.commands.executeCommand("setContext", "coder.authenticated", true)
+          if (user.roles.find((role) => role.name === "owner")) {
+            await vscode.commands.executeCommand("setContext", "coder.isOwner", true)
+          }
         }
-      }
-    })
-    .catch(() => {
-      // Not authenticated!
-    })
-    .finally(() => {
-      vscode.commands.executeCommand("setContext", "coder.loaded", true)
-    })
+      })
+      .catch((error) => {
+        // This should be a failure to make the request, like the header command
+        // errored.
+        vscodeProposed.window.showErrorMessage("Failed to check user authentication: " + error.message)
+      })
+      .finally(() => {
+        vscode.commands.executeCommand("setContext", "coder.loaded", true)
+      })
+  } else {
+    vscode.commands.executeCommand("setContext", "coder.loaded", true)
+  }
 
   vscode.window.registerUriHandler({
     handleUri: async (uri) => {

src/headers.test.ts

@@ -0,0 +1,57 @@
+import * as os from "os"
+import { it, expect } from "vitest"
+import { getHeaders } from "./headers"
+
+const logger = {
+  writeToCoderOutputChannel() {
+    // no-op
+  },
+}
+
+it("should return no headers", async () => {
+  await expect(getHeaders(undefined, undefined, logger)).resolves.toStrictEqual({})
+  await expect(getHeaders("localhost", undefined, logger)).resolves.toStrictEqual({})
+  await expect(getHeaders(undefined, "command", logger)).resolves.toStrictEqual({})
+  await expect(getHeaders("localhost", "", logger)).resolves.toStrictEqual({})
+  await expect(getHeaders("", "command", logger)).resolves.toStrictEqual({})
+  await expect(getHeaders("localhost", "  ", logger)).resolves.toStrictEqual({})
+  await expect(getHeaders("  ", "command", logger)).resolves.toStrictEqual({})
+})
+
+it("should return headers", async () => {
+  await expect(getHeaders("localhost", "printf foo=bar'\n'baz=qux", logger)).resolves.toStrictEqual({
+    foo: "bar",
+    baz: "qux",
+  })
+  await expect(getHeaders("localhost", "printf foo=bar'\r\n'baz=qux", logger)).resolves.toStrictEqual({
+    foo: "bar",
+    baz: "qux",
+  })
+  await expect(getHeaders("localhost", "printf foo=bar'\r\n'", logger)).resolves.toStrictEqual({ foo: "bar" })
+  await expect(getHeaders("localhost", "printf foo=bar", logger)).resolves.toStrictEqual({ foo: "bar" })
+  await expect(getHeaders("localhost", "printf foo=bar=", logger)).resolves.toStrictEqual({ foo: "bar=" })
+  await expect(getHeaders("localhost", "printf foo=bar=baz", logger)).resolves.toStrictEqual({ foo: "bar=baz" })
+  await expect(getHeaders("localhost", "printf foo=", logger)).resolves.toStrictEqual({ foo: "" })
+})
+
+it("should error on malformed or empty lines", async () => {
+  await expect(getHeaders("localhost", "printf foo=bar'\r\n\r\n'", logger)).rejects.toMatch(/Malformed/)
+  await expect(getHeaders("localhost", "printf '\r\n'foo=bar", logger)).rejects.toMatch(/Malformed/)
+  await expect(getHeaders("localhost", "printf =foo", logger)).rejects.toMatch(/Malformed/)
+  await expect(getHeaders("localhost", "printf foo", logger)).rejects.toMatch(/Malformed/)
+  await expect(getHeaders("localhost", "printf '  =foo'", logger)).rejects.toMatch(/Malformed/)
+  await expect(getHeaders("localhost", "printf 'foo  =bar'", logger)).rejects.toMatch(/Malformed/)
+  await expect(getHeaders("localhost", "printf 'foo  foo=bar'", logger)).rejects.toMatch(/Malformed/)
+  await expect(getHeaders("localhost", "printf ''", logger)).rejects.toMatch(/Malformed/)
+})
+
+it("should have access to environment variables", async () => {
+  const coderUrl = "dev.coder.com"
+  await expect(
+    getHeaders(coderUrl, os.platform() === "win32" ? "printf url=%CODER_URL" : "printf url=$CODER_URL", logger),
+  ).resolves.toStrictEqual({ url: coderUrl })
+})
+
+it("should error on non-zero exit", async () => {
+  await expect(getHeaders("localhost", "exit 10", logger)).rejects.toMatch(/exited unexpectedly with code 10/)
+})

src/headers.ts

@@ -0,0 +1,64 @@
+import * as cp from "child_process"
+import * as util from "util"
+
+export interface Logger {
+  writeToCoderOutputChannel(message: string): void
+}
+
+interface ExecException {
+  code?: number
+  stderr?: string
+  stdout?: string
+}
+
+function isExecException(err: unknown): err is ExecException {
+  return typeof (err as ExecException).code !== "undefined"
+}
+
+// TODO: getHeaders might make more sense to directly implement on Storage
+// but it is difficult to test Storage right now since we use vitest instead of
+// the standard extension testing framework which would give us access to vscode
+// APIs.  We should revert the testing framework then consider moving this.
+
+// getHeaders executes the header command and parses the headers from stdout.
+// Both stdout and stderr are logged on error but stderr is otherwise ignored.
+// Throws an error if the process exits with non-zero or the JSON is invalid.
+// Returns undefined if there is no header command set.  No effort is made to
+// validate the JSON other than making sure it can be parsed.
+export async function getHeaders(
+  url: string | undefined,
+  command: string | undefined,
+  logger: Logger,
+): Promise<Record<string, string>> {
+  const headers: Record<string, string> = {}
+  if (typeof url === "string" && url.trim().length > 0 && typeof command === "string" && command.trim().length > 0) {
+    let result: { stdout: string; stderr: string }
+    try {
+      result = await util.promisify(cp.exec)(command, {
+        env: {
+          ...process.env,
+          CODER_URL: url,
+        },
+      })
+    } catch (error) {
+      if (isExecException(error)) {
+        logger.writeToCoderOutputChannel(`Header command exited unexpectedly with code ${error.code}`)
+        logger.writeToCoderOutputChannel(`stdout: ${error.stdout}`)
+        logger.writeToCoderOutputChannel(`stderr: ${error.stderr}`)
+        throw new Error(`Header command exited unexpectedly with code ${error.code}`)
+      }
+      throw new Error(`Header command exited unexpectedly: ${error}`)
+    }
+    const lines = result.stdout.replace(/\r?\n$/, "").split(/\r?\n/)
+    for (let i = 0; i < lines.length; ++i) {
+      const [key, value] = lines[i].split(/=(.*)/)
+      // 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).
+      if (key.length === 0 || key.indexOf(" ") !== -1 || typeof value === "undefined") {
+        throw new Error(`Malformed line from header command: [${lines[i]}] (out: ${result.stdout})`)
+      }
+      headers[key] = value
+    }
+  }
+  return headers
+}

src/remote.ts

@@ -508,9 +508,17 @@ export class Remote {
     }
 
     const escape = (str: string): string => `"${str.replace(/"/g, '\\"')}"`
+
+    // Add headers from the header command.
+    let headerArg = ""
+    const headerCommand = vscode.workspace.getConfiguration().get("coder.headerCommand")
+    if (typeof headerCommand === "string" && headerCommand.trim().length > 0) {
+      headerArg = ` --header-command ${escape(headerCommand)}`
+    }
+
     const sshValues: SSHValues = {
       Host: `${Remote.Prefix}*`,
-      ProxyCommand: `${escape(binaryPath)} vscodessh --network-info-dir ${escape(
+      ProxyCommand: `${escape(binaryPath)}${headerArg} vscodessh --network-info-dir ${escape(
         this.storage.getNetworkInfoPath(),
       )} --session-token-file ${escape(this.storage.getSessionTokenPath())} --url-file ${escape(
         this.storage.getURLPath(),

src/sshSupport.test.ts

@@ -29,12 +29,14 @@ it("computes the config for a host", () => {
 Host coder-vscode--*
   StrictHostKeyChecking no
   Another=true
+  ProxyCommand=/tmp/coder --header="X-FOO=bar" coder.dev
 # --- END CODER VSCODE ---
 `,
   )
 
   expect(properties).toEqual({
     Another: "true",
     StrictHostKeyChecking: "yes",
+    ProxyCommand: '/tmp/coder --header="X-FOO=bar" coder.dev',
   })
 })

src/sshSupport.ts

@@ -52,7 +52,11 @@ export function computeSSHProperties(host: string, config: string): Record<strin
     if (line === "") {
       return
     }
-    const [key, ...valueParts] = line.split(/\s+|=/)
+    // The capture group here will include the captured portion in the array
+    // which we need to join them back up with their original values.  The first
+    // separate is ignored since it splits the key and value but is not part of
+    // the value itself.
+    const [key, _, ...valueParts] = line.split(/(\s+|=)/)
     if (key.startsWith("#")) {
       // Ignore comments!
       return
@@ -62,15 +66,15 @@ export function computeSSHProperties(host: string, config: string): Record<strin
         configs.push(currentConfig)
       }
       currentConfig = {
-        Host: valueParts.join(" "),
+        Host: valueParts.join(""),
         properties: {},
       }
       return
     }
     if (!currentConfig) {
       return
     }
-    currentConfig.properties[key] = valueParts.join(" ")
+    currentConfig.properties[key] = valueParts.join("")
   })
   if (currentConfig) {
     configs.push(currentConfig)

src/storage.ts

@@ -11,6 +11,7 @@ import os from "os"
 import path from "path"
 import prettyBytes from "pretty-bytes"
 import * as vscode from "vscode"
+import { getHeaders } from "./headers"
 
 export class Storage {
   public workspace?: Workspace
@@ -391,6 +392,10 @@ export class Storage {
       await fs.rm(this.getSessionTokenPath(), { force: true })
     }
   }
+
+  public async getHeaders(url = this.getURL()): Promise<Record<string, string>> {
+    return getHeaders(url, vscode.workspace.getConfiguration().get("coder.headerCommand"), this)
+  }
 }
 
 // goos returns the Go format for the current platform.