Skip to content

Commit 3ab3aad

Browse files
code-ashercode-asher
authored
Add customizable header command (coder#119)
1 parent 7ddc95e commit 3ab3aad

File tree

10 files changed

+189
-21
lines changed

10 files changed

+189
-21
lines changed

‎.eslintrc.json

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,13 @@
3535
}],
3636
"import/no-unresolved": ["error", {
3737
"ignore": ["vscode"]
38-
}]
38+
}],
39+
"@typescript-eslint/no-unused-vars": [
40+
"error",
41+
{
42+
"varsIgnorePattern": "^_"
43+
}
44+
]
3945
},
4046
"ignorePatterns": [
4147
"out",

‎package.json

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,11 @@
5757
"markdownDescription": "The full path of the directory into which the Coder CLI will be downloaded. Defaults to the extension's global storage directory.",
5858
"type": "string",
5959
"default": ""
60+
},
61+
"coder.headerCommand": {
62+
"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`.",
63+
"type": "string",
64+
"default": ""
6065
}
6166
}
6267
},

‎src/commands.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,10 @@ export class Commands {
7070
severity: vscode.InputBoxValidationSeverity.Error,
7171
}
7272
}
73+
// This could be something like the header command erroring or an
74+
// invalid session token.
7375
return {
74-
message: "Invalid session token! (" + message + ")",
76+
message: "Failed to authenticate: " + message,
7577
severity: vscode.InputBoxValidationSeverity.Error,
7678
}
7779
})

‎src/extension.ts

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
4646
axios.interceptors.response.use(
4747
(r) => r,
4848
async (err) => {
49-
throw await CertificateError.maybeWrap(err, err.config.baseURL, storage)
49+
throw await CertificateError.maybeWrap(err, axios.getUri(err.config), storage)
5050
},
5151
)
5252

@@ -59,27 +59,42 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
5959
const storage = new Storage(output, ctx.globalState, ctx.secrets, ctx.globalStorageUri, ctx.logUri)
6060
await storage.init()
6161

62+
// Add headers from the header command.
63+
axios.interceptors.request.use(async (config) => {
64+
Object.entries(await storage.getHeaders(config.baseURL || axios.getUri(config))).forEach(([key, value]) => {
65+
config.headers[key] = value
66+
})
67+
return config
68+
})
69+
6270
const myWorkspacesProvider = new WorkspaceProvider(WorkspaceQuery.Mine, storage)
6371
const allWorkspacesProvider = new WorkspaceProvider(WorkspaceQuery.All, storage)
6472

6573
vscode.window.registerTreeDataProvider("myWorkspaces", myWorkspacesProvider)
6674
vscode.window.registerTreeDataProvider("allWorkspaces", allWorkspacesProvider)
6775

68-
getAuthenticatedUser()
69-
.then(async (user) => {
70-
if (user) {
71-
vscode.commands.executeCommand("setContext", "coder.authenticated", true)
72-
if (user.roles.find((role) => role.name === "owner")) {
73-
await vscode.commands.executeCommand("setContext", "coder.isOwner", true)
76+
const url = storage.getURL()
77+
if (url) {
78+
getAuthenticatedUser()
79+
.then(async (user) => {
80+
if (user) {
81+
vscode.commands.executeCommand("setContext", "coder.authenticated", true)
82+
if (user.roles.find((role) => role.name === "owner")) {
83+
await vscode.commands.executeCommand("setContext", "coder.isOwner", true)
84+
}
7485
}
75-
}
76-
})
77-
.catch(() => {
78-
// Not authenticated!
79-
})
80-
.finally(() => {
81-
vscode.commands.executeCommand("setContext", "coder.loaded", true)
82-
})
86+
})
87+
.catch((error) => {
88+
// This should be a failure to make the request, like the header command
89+
// errored.
90+
vscodeProposed.window.showErrorMessage("Failed to check user authentication: " + error.message)
91+
})
92+
.finally(() => {
93+
vscode.commands.executeCommand("setContext", "coder.loaded", true)
94+
})
95+
} else {
96+
vscode.commands.executeCommand("setContext", "coder.loaded", true)
97+
}
8398

8499
vscode.window.registerUriHandler({
85100
handleUri: async (uri) => {

‎src/headers.test.ts

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import * as os from "os"
2+
import { it, expect } from "vitest"
3+
import { getHeaders } from "./headers"
4+
5+
const logger = {
6+
writeToCoderOutputChannel() {
7+
// no-op
8+
},
9+
}
10+
11+
it("should return no headers", async () => {
12+
await expect(getHeaders(undefined, undefined, logger)).resolves.toStrictEqual({})
13+
await expect(getHeaders("localhost", undefined, logger)).resolves.toStrictEqual({})
14+
await expect(getHeaders(undefined, "command", logger)).resolves.toStrictEqual({})
15+
await expect(getHeaders("localhost", "", logger)).resolves.toStrictEqual({})
16+
await expect(getHeaders("", "command", logger)).resolves.toStrictEqual({})
17+
await expect(getHeaders("localhost", " ", logger)).resolves.toStrictEqual({})
18+
await expect(getHeaders(" ", "command", logger)).resolves.toStrictEqual({})
19+
})
20+
21+
it("should return headers", async () => {
22+
await expect(getHeaders("localhost", "printf foo=bar'\n'baz=qux", logger)).resolves.toStrictEqual({
23+
foo: "bar",
24+
baz: "qux",
25+
})
26+
await expect(getHeaders("localhost", "printf foo=bar'\r\n'baz=qux", logger)).resolves.toStrictEqual({
27+
foo: "bar",
28+
baz: "qux",
29+
})
30+
await expect(getHeaders("localhost", "printf foo=bar'\r\n'", logger)).resolves.toStrictEqual({ foo: "bar" })
31+
await expect(getHeaders("localhost", "printf foo=bar", logger)).resolves.toStrictEqual({ foo: "bar" })
32+
await expect(getHeaders("localhost", "printf foo=bar=", logger)).resolves.toStrictEqual({ foo: "bar=" })
33+
await expect(getHeaders("localhost", "printf foo=bar=baz", logger)).resolves.toStrictEqual({ foo: "bar=baz" })
34+
await expect(getHeaders("localhost", "printf foo=", logger)).resolves.toStrictEqual({ foo: "" })
35+
})
36+
37+
it("should error on malformed or empty lines", async () => {
38+
await expect(getHeaders("localhost", "printf foo=bar'\r\n\r\n'", logger)).rejects.toMatch(/Malformed/)
39+
await expect(getHeaders("localhost", "printf '\r\n'foo=bar", logger)).rejects.toMatch(/Malformed/)
40+
await expect(getHeaders("localhost", "printf =foo", logger)).rejects.toMatch(/Malformed/)
41+
await expect(getHeaders("localhost", "printf foo", logger)).rejects.toMatch(/Malformed/)
42+
await expect(getHeaders("localhost", "printf ' =foo'", logger)).rejects.toMatch(/Malformed/)
43+
await expect(getHeaders("localhost", "printf 'foo =bar'", logger)).rejects.toMatch(/Malformed/)
44+
await expect(getHeaders("localhost", "printf 'foo foo=bar'", logger)).rejects.toMatch(/Malformed/)
45+
await expect(getHeaders("localhost", "printf ''", logger)).rejects.toMatch(/Malformed/)
46+
})
47+
48+
it("should have access to environment variables", async () => {
49+
const coderUrl = "dev.coder.com"
50+
await expect(
51+
getHeaders(coderUrl, os.platform() === "win32" ? "printf url=%CODER_URL" : "printf url=$CODER_URL", logger),
52+
).resolves.toStrictEqual({ url: coderUrl })
53+
})
54+
55+
it("should error on non-zero exit", async () => {
56+
await expect(getHeaders("localhost", "exit 10", logger)).rejects.toMatch(/exited unexpectedly with code 10/)
57+
})

‎src/headers.ts

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
import * as cp from "child_process"
2+
import * as util from "util"
3+
4+
export interface Logger {
5+
writeToCoderOutputChannel(message: string): void
6+
}
7+
8+
interface ExecException {
9+
code?: number
10+
stderr?: string
11+
stdout?: string
12+
}
13+
14+
function isExecException(err: unknown): err is ExecException {
15+
return typeof (err as ExecException).code !== "undefined"
16+
}
17+
18+
// TODO: getHeaders might make more sense to directly implement on Storage
19+
// but it is difficult to test Storage right now since we use vitest instead of
20+
// the standard extension testing framework which would give us access to vscode
21+
// APIs. We should revert the testing framework then consider moving this.
22+
23+
// getHeaders executes the header command and parses the headers from stdout.
24+
// Both stdout and stderr are logged on error but stderr is otherwise ignored.
25+
// Throws an error if the process exits with non-zero or the JSON is invalid.
26+
// Returns undefined if there is no header command set. No effort is made to
27+
// validate the JSON other than making sure it can be parsed.
28+
export async function getHeaders(
29+
url: string | undefined,
30+
command: string | undefined,
31+
logger: Logger,
32+
): Promise<Record<string, string>> {
33+
const headers: Record<string, string> = {}
34+
if (typeof url === "string" && url.trim().length > 0 && typeof command === "string" && command.trim().length > 0) {
35+
let result: { stdout: string; stderr: string }
36+
try {
37+
result = await util.promisify(cp.exec)(command, {
38+
env: {
39+
...process.env,
40+
CODER_URL: url,
41+
},
42+
})
43+
} catch (error) {
44+
if (isExecException(error)) {
45+
logger.writeToCoderOutputChannel(`Header command exited unexpectedly with code ${error.code}`)
46+
logger.writeToCoderOutputChannel(`stdout: ${error.stdout}`)
47+
logger.writeToCoderOutputChannel(`stderr: ${error.stderr}`)
48+
throw new Error(`Header command exited unexpectedly with code ${error.code}`)
49+
}
50+
throw new Error(`Header command exited unexpectedly: ${error}`)
51+
}
52+
const lines = result.stdout.replace(/\r?\n$/, "").split(/\r?\n/)
53+
for (let i = 0; i < lines.length; ++i) {
54+
const [key, value] = lines[i].split(/=(.*)/)
55+
// Header names cannot be blank or contain whitespace and the Coder CLI
56+
// requires that there be an equals sign (the value can be blank though).
57+
if (key.length === 0 || key.indexOf(" ") !== -1 || typeof value === "undefined") {
58+
throw new Error(`Malformed line from header command: [${lines[i]}] (out: ${result.stdout})`)
59+
}
60+
headers[key] = value
61+
}
62+
}
63+
return headers
64+
}

‎src/remote.ts

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -514,9 +514,17 @@ export class Remote {
514514
}
515515

516516
const escape = (str: string): string => `"${str.replace(/"/g, '\\"')}"`
517+
518+
// Add headers from the header command.
519+
let headerArg = ""
520+
const headerCommand = vscode.workspace.getConfiguration().get("coder.headerCommand")
521+
if (typeof headerCommand === "string" && headerCommand.trim().length > 0) {
522+
headerArg = ` --header-command ${escape(headerCommand)}`
523+
}
524+
517525
const sshValues: SSHValues = {
518526
Host: `${Remote.Prefix}*`,
519-
ProxyCommand: `${escape(binaryPath)} vscodessh --network-info-dir ${escape(
527+
ProxyCommand: `${escape(binaryPath)}${headerArg} vscodessh --network-info-dir ${escape(
520528
this.storage.getNetworkInfoPath(),
521529
)} --session-token-file ${escape(this.storage.getSessionTokenPath())} --url-file ${escape(
522530
this.storage.getURLPath(),

‎src/sshSupport.test.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,14 @@ it("computes the config for a host", () => {
2929
Host coder-vscode--*
3030
StrictHostKeyChecking no
3131
Another=true
32+
ProxyCommand=/tmp/coder --header="X-FOO=bar" coder.dev
3233
# --- END CODER VSCODE ---
3334
`,
3435
)
3536

3637
expect(properties).toEqual({
3738
Another: "true",
3839
StrictHostKeyChecking: "yes",
40+
ProxyCommand: '/tmp/coder --header="X-FOO=bar" coder.dev',
3941
})
4042
})

‎src/sshSupport.ts

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,11 @@ export function computeSSHProperties(host: string, config: string): Record<strin
5252
if (line === "") {
5353
return
5454
}
55-
const [key, ...valueParts] = line.split(/\s+|=/)
55+
// The capture group here will include the captured portion in the array
56+
// which we need to join them back up with their original values. The first
57+
// separate is ignored since it splits the key and value but is not part of
58+
// the value itself.
59+
const [key, _, ...valueParts] = line.split(/(\s+|=)/)
5660
if (key.startsWith("#")) {
5761
// Ignore comments!
5862
return
@@ -62,15 +66,15 @@ export function computeSSHProperties(host: string, config: string): Record<strin
6266
configs.push(currentConfig)
6367
}
6468
currentConfig = {
65-
Host: valueParts.join(" "),
69+
Host: valueParts.join(""),
6670
properties: {},
6771
}
6872
return
6973
}
7074
if (!currentConfig) {
7175
return
7276
}
73-
currentConfig.properties[key] = valueParts.join(" ")
77+
currentConfig.properties[key] = valueParts.join("")
7478
})
7579
if (currentConfig) {
7680
configs.push(currentConfig)

‎src/storage.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import os from "os"
1111
import path from "path"
1212
import prettyBytes from "pretty-bytes"
1313
import * as vscode from "vscode"
14+
import { getHeaders } from "./headers"
1415

1516
export class Storage {
1617
public workspace?: Workspace
@@ -391,6 +392,10 @@ export class Storage {
391392
await fs.rm(this.getSessionTokenPath(), { force: true })
392393
}
393394
}
395+
396+
public async getHeaders(url = this.getURL()): Promise<Record<string, string>> {
397+
return getHeaders(url, vscode.workspace.getConfiguration().get("coder.headerCommand"), this)
398+
}
394399
}
395400

396401
// goos returns the Go format for the current platform.

0 commit comments

Comments
 (0)