fix: specify --header-command when running coder start (#526)

aslilac · web-flow · commit d70a5868e9a8 · 2025-06-05T11:16:41.000-06:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Use `--header-command` properly when starting a workspace.
+
 ## [v1.9.1](https://github.com/coder/vscode-coder/releases/tag/v1.9.1) 2025-05-27
 
 ### Fixed
diff --git a/src/api.ts b/src/api.ts
@@ -9,6 +9,7 @@ import * as vscode from "vscode"
 import * as ws from "ws"
 import { errToStr } from "./api-helper"
 import { CertificateError } from "./error"
+import { getHeaderArgs } from "./headers"
 import { getProxyForUrl } from "./proxy"
 import { Storage } from "./storage"
 import { expandPath } from "./util"
@@ -168,6 +169,7 @@ export async function startWorkspaceIfStoppedOrFailed(
     const startArgs = [
       "--global-config",
       globalConfigDir,
+      ...getHeaderArgs(vscode.workspace.getConfiguration()),
       "start",
       "--yes",
       workspace.owner_name + "/" + workspace.name,
diff --git a/src/headers.ts b/src/headers.ts
@@ -1,7 +1,8 @@
 import * as cp from "child_process"
+import * as os from "os"
 import * as util from "util"
-
-import { WorkspaceConfiguration } from "vscode"
+import type { WorkspaceConfiguration } from "vscode"
+import { escapeCommandArg } from "./util"
 
 export interface Logger {
   writeToCoderOutputChannel(message: string): void
@@ -25,6 +26,23 @@ export function getHeaderCommand(config: WorkspaceConfiguration): string | undef
   return cmd
 }
 
+export function getHeaderArgs(config: WorkspaceConfiguration): string[] {
+  // Escape a command line to be executed by the Coder binary, so ssh doesn't substitute variables.
+  const escapeSubcommand: (str: string) => string =
+    os.platform() === "win32"
+      ? // On Windows variables are %VAR%, and we need to use double quotes.
+        (str) => escapeCommandArg(str).replace(/%/g, "%%")
+      : // On *nix we can use single quotes to escape $VARS.
+        // Note single quotes cannot be escaped inside single quotes.
+        (str) => `'${str.replace(/'/g, "'\\''")}'`
+
+  const command = getHeaderCommand(config)
+  if (!command) {
+    return []
+  }
+  return ["--header-command", escapeSubcommand(command)]
+}
+
 // 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
diff --git a/src/remote.ts b/src/remote.ts
@@ -14,12 +14,12 @@ import { extractAgents } from "./api-helper"
 import * as cli from "./cliManager"
 import { Commands } from "./commands"
 import { featureSetForVersion, FeatureSet } from "./featureSet"
-import { getHeaderCommand } from "./headers"
+import { getHeaderArgs } from "./headers"
 import { Inbox } from "./inbox"
 import { SSHConfig, SSHValues, mergeSSHConfigValues } from "./sshConfig"
 import { computeSSHProperties, sshSupportsSetEnv } from "./sshSupport"
 import { Storage } from "./storage"
-import { AuthorityPrefix, expandPath, findPort, parseRemoteAuthority } from "./util"
+import { AuthorityPrefix, escapeCommandArg, expandPath, findPort, parseRemoteAuthority } from "./util"
 import { WorkspaceMonitor } from "./workspaceMonitor"
 
 export interface RemoteDetails extends vscode.Disposable {
@@ -611,32 +611,18 @@ export class Remote {
     const sshConfig = new SSHConfig(sshConfigFile)
     await sshConfig.load()
 
-    const escape = (str: string): string => `"${str.replace(/"/g, '\\"')}"`
-    // Escape a command line to be executed by the Coder binary, so ssh doesn't substitute variables.
-    const escapeSubcommand: (str: string) => string =
-      os.platform() === "win32"
-        ? // On Windows variables are %VAR%, and we need to use double quotes.
-          (str) => escape(str).replace(/%/g, "%%")
-        : // On *nix we can use single quotes to escape $VARS.
-          // Note single quotes cannot be escaped inside single quotes.
-          (str) => `'${str.replace(/'/g, "'\\''")}'`
-
-    // Add headers from the header command.
-    let headerArg = ""
-    const headerCommand = getHeaderCommand(vscode.workspace.getConfiguration())
-    if (typeof headerCommand === "string" && headerCommand.trim().length > 0) {
-      headerArg = ` --header-command ${escapeSubcommand(headerCommand)}`
-    }
+    const headerArgs = getHeaderArgs(vscode.workspace.getConfiguration())
+    const headerArgList = headerArgs.length > 0 ? ` ${headerArgs.join(" ")}` : ""
 
     const hostPrefix = label ? `${AuthorityPrefix}.${label}--` : `${AuthorityPrefix}--`
 
     const proxyCommand = featureSet.wildcardSSH
-      ? `${escape(binaryPath)}${headerArg} --global-config ${escape(
+      ? `${escapeCommandArg(binaryPath)}${headerArgList} --global-config ${escapeCommandArg(
           path.dirname(this.storage.getSessionTokenPath(label)),
-        )} ssh --stdio --usage-app=vscode --disable-autostart --network-info-dir ${escape(this.storage.getNetworkInfoPath())}${await this.formatLogArg(logDir)} --ssh-host-prefix ${hostPrefix} %h`
-      : `${escape(binaryPath)}${headerArg} vscodessh --network-info-dir ${escape(
+        )} ssh --stdio --usage-app=vscode --disable-autostart --network-info-dir ${escapeCommandArg(this.storage.getNetworkInfoPath())}${await this.formatLogArg(logDir)} --ssh-host-prefix ${hostPrefix} %h`
+      : `${escapeCommandArg(binaryPath)}${headerArgList} vscodessh --network-info-dir ${escapeCommandArg(
           this.storage.getNetworkInfoPath(),
-        )}${await this.formatLogArg(logDir)} --session-token-file ${escape(this.storage.getSessionTokenPath(label))} --url-file ${escape(
+        )}${await this.formatLogArg(logDir)} --session-token-file ${escapeCommandArg(this.storage.getSessionTokenPath(label))} --url-file ${escapeCommandArg(
           this.storage.getUrlPath(label),
         )} %h`
 
diff --git a/src/util.ts b/src/util.ts
@@ -136,3 +136,7 @@ export function countSubstring(needle: string, haystack: string): number {
   }
   return count
 }
+
+export function escapeCommandArg(arg: string): string {
+  return `"${arg.replace(/"/g, '\\"')}"`
+}

Original file line number	Diff line number	Diff line change
`@@ -136,3 +136,7 @@ export function countSubstring(needle: string, haystack: string): number {`
`136`	`136`	`}`
`137`	`137`	`return count`
`138`	`138`	`}`
	`139`	`+`
	`140`	`+export function escapeCommandArg(arg: string): string {`
	`141`	+ return `"${arg.replace(/"/g, '\\"')}"`
	`142`	`+}`