Add header command setting

This will be called before requests and added to the SSH config.

Author: code-asher

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

@@ -59,6 +59,17 @@ 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) => {
+    return {
+      ...config,
+      headers: {
+        ...(await storage.getHeaders()),
+        ...creds.headers,
+      },
+    }
+  })
+
   const myWorkspacesProvider = new WorkspaceProvider(WorkspaceQuery.Mine, storage)
   const allWorkspacesProvider = new WorkspaceProvider(WorkspaceQuery.All, storage)
 
@@ -74,8 +85,10 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
         }
       }
     })
-    .catch(() => {
-      // Not authenticated!
+    .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)

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,16 @@ export class Remote {
     }
 
     const escape = (str: string): string => `"${str.replace(/"/g, '\\"')}"`
+
+    // Add headers from the header command.
+    let headerCommand = vscode.workspace.getConfiguration().get("coder.headerCommand")
+    if (headerCommand) {
+      headerCommand = ` --header-command ${headerCommand}`
+    }
+
     const sshValues: SSHValues = {
       Host: `${Remote.Prefix}*`,
-      ProxyCommand: `${escape(binaryPath)} vscodessh --network-info-dir ${escape(
+      ProxyCommand: `${escape(binaryPath)}${headerCommand} vscodessh --network-info-dir ${escape(
         this.storage.getNetworkInfoPath(),
       )} --session-token-file ${escape(this.storage.getSessionTokenPath())} --url-file ${escape(
         this.storage.getURLPath(),

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> | undefined> {
+    return getHeaders(url, vscode.workspace.getConfiguration().get("coder.headerCommand"), this)
+  }
 }
 
 // goos returns the Go format for the current platform.