coder · code-asher · Aug 14, 2023 · Aug 11, 2023 · Aug 14, 2023 · Aug 14, 2023
diff --git a/cli/login.go b/cli/login.go
@@ -76,7 +76,7 @@ func (r *RootCmd) login() *clibase.Cmd {
 				serverURL.Scheme = "https"
 			}
 
-			client, err := r.createUnauthenticatedClient(serverURL)
+			client, err := r.createUnauthenticatedClient(ctx, serverURL)
 			if err != nil {
 				return err
 			}

diff --git a/cli/root.go b/cli/root.go
@@ -1,6 +1,8 @@
 package cli
 
 import (
+	"bufio"
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -13,6 +15,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"os/exec"
 	"os/signal"
 	"path/filepath"
 	"runtime"
@@ -55,6 +58,7 @@ const (
 	varAgentToken       = "agent-token"
 	varAgentURL         = "agent-url"
 	varHeader           = "header"
+	varHeaderCommand    = "header-command"
 	varNoOpen           = "no-open"
 	varNoVersionCheck   = "no-version-warning"
 	varNoFeatureWarning = "no-feature-warning"
@@ -356,6 +360,13 @@ func (r *RootCmd) Command(subcommands []*clibase.Cmd) (*clibase.Cmd, error) {
 			Value:       clibase.StringArrayOf(&r.header),
 			Group:       globalGroup,
 		},
+		{
+			Flag:        varHeaderCommand,
+			Env:         "CODER_HEADER_COMMAND",
+			Description: "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.",
+			Value:       clibase.StringOf(&r.headerCommand),
+			Group:       globalGroup,
+		},
 		{
 			Flag:        varNoOpen,
 			Env:         "CODER_NO_OPEN",
@@ -437,6 +448,7 @@ type RootCmd struct {
 	token         string
 	globalConfig  string
 	header        []string
+	headerCommand string
 	agentToken    string
 	agentURL      *url.URL
 	forceTTY      bool
@@ -540,9 +552,7 @@ func (r *RootCmd) initClientInternal(client *codersdk.Client, allowTokenMissing
 					return err
 				}
 			}
-			err = r.setClient(
-				client, r.clientURL,
-			)
+			err = r.setClient(inv.Context(), client, r.clientURL)
 			if err != nil {
 				return err
 			}
@@ -592,12 +602,38 @@ func (r *RootCmd) initClientInternal(client *codersdk.Client, allowTokenMissing
 	}
 }
 
-func (r *RootCmd) setClient(client *codersdk.Client, serverURL *url.URL) error {
+func (r *RootCmd) setClient(ctx context.Context, client *codersdk.Client, serverURL *url.URL) error {
 	transport := &headerTransport{
 		transport: http.DefaultTransport,
 		header:    http.Header{},
 	}
-	for _, header := range r.header {
+	headers := r.header
+	if r.headerCommand != "" {
+		shell := "sh"
+		caller := "-c"
+		if runtime.GOOS == "windows" {
+			shell = "cmd.exe"
+			caller = "/c"
+		}
+		var outBuf bytes.Buffer
+		// #nosec
+		cmd := exec.CommandContext(ctx, shell, caller, r.headerCommand)
+		cmd.Env = append(os.Environ(), "CODER_URL="+serverURL.String())
+		cmd.Stdout = &outBuf
+		cmd.Stderr = io.Discard
+		err := cmd.Run()
+		if err != nil {
+			return xerrors.Errorf("failed to run %v: %w", cmd.Args, err)
+		}
+		scanner := bufio.NewScanner(&outBuf)
+		for scanner.Scan() {
+			headers = append(headers, scanner.Text())
+		}
+		if err := scanner.Err(); err != nil {
+			return xerrors.Errorf("scan %v: %w", cmd.Args, err)
+		}
+	}
+	for _, header := range headers {
 		parts := strings.SplitN(header, "=", 2)
 		if len(parts) < 2 {
 			return xerrors.Errorf("split header %q had less than two parts", header)
@@ -611,9 +647,9 @@ func (r *RootCmd) setClient(client *codersdk.Client, serverURL *url.URL) error {
 	return nil
 }
 
-func (r *RootCmd) createUnauthenticatedClient(serverURL *url.URL) (*codersdk.Client, error) {
+func (r *RootCmd) createUnauthenticatedClient(ctx context.Context, serverURL *url.URL) (*codersdk.Client, error) {
 	var client codersdk.Client
-	err := r.setClient(&client, serverURL)
+	err := r.setClient(ctx, &client, serverURL)
 	return &client, err
 }
 

diff --git a/cli/root_test.go b/cli/root_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"runtime"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -72,20 +73,29 @@ func TestRoot(t *testing.T) {
 	t.Run("Header", func(t *testing.T) {
 		t.Parallel()
 
+		var url string
 		var called int64
 		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			atomic.AddInt64(&called, 1)
 			assert.Equal(t, "wow", r.Header.Get("X-Testing"))
 			assert.Equal(t, "Dean was Here!", r.Header.Get("Cool-Header"))
+			assert.Equal(t, "very-wow-"+url, r.Header.Get("X-Process-Testing"))
+			assert.Equal(t, "more-wow", r.Header.Get("X-Process-Testing2"))
 			w.WriteHeader(http.StatusGone)
 		}))
 		defer srv.Close()
+		url = srv.URL
 		buf := new(bytes.Buffer)
+		coderURLEnv := "$CODER_URL"
+		if runtime.GOOS == "windows" {
+			coderURLEnv = "%CODER_URL%"
+		}
 		inv, _ := clitest.New(t,
 			"--no-feature-warning",
 			"--no-version-warning",
 			"--header", "X-Testing=wow",
 			"--header", "Cool-Header=Dean was Here!",
+			"--header-command", "printf X-Process-Testing=very-wow-"+coderURLEnv+"'\\r\\n'X-Process-Testing2=more-wow",
 			"login", srv.URL,
 		)
 		inv.Stdout = buf
@@ -97,8 +107,8 @@ func TestRoot(t *testing.T) {
 	})
 }
 
-// TestDERPHeaders ensures that the client sends the global `--header`s to the
-// DERP server when connecting.
+// TestDERPHeaders ensures that the client sends the global `--header`s and
+// `--header-command` to the DERP server when connecting.
 func TestDERPHeaders(t *testing.T) {
 	t.Parallel()
 
@@ -129,8 +139,9 @@ func TestDERPHeaders(t *testing.T) {
 	// Inject custom /derp handler so we can inspect the headers.
 	var (
 		expectedHeaders = map[string]string{
-			"X-Test-Header": "test-value",
-			"Cool-Header":   "Dean was Here!",
+			"X-Test-Header":     "test-value",
+			"Cool-Header":       "Dean was Here!",
+			"X-Process-Testing": "very-wow",
 		}
 		derpCalled int64
 	)
@@ -159,9 +170,12 @@ func TestDERPHeaders(t *testing.T) {
 		"--no-version-warning",
 		"ping", workspace.Name,
 		"-n", "1",
+		"--header-command", "printf X-Process-Testing=very-wow",
 	}
 	for k, v := range expectedHeaders {
-		args = append(args, "--header", fmt.Sprintf("%s=%s", k, v))
+		if k != "X-Process-Testing" {
+			args = append(args, "--header", fmt.Sprintf("%s=%s", k, v))
+		}
 	}
 	inv, root := clitest.New(t, args...)
 	clitest.SetupConfig(t, client, root)

diff --git a/cli/testdata/coder_--help.golden b/cli/testdata/coder_--help.golden
@@ -62,6 +62,11 @@ variables or flags.
           Additional HTTP headers added to all requests. Provide as key=value.
           Can be specified multiple times.
 
+      --header-command string, $CODER_HEADER_COMMAND
+          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.
+
       --no-feature-warning bool, $CODER_NO_FEATURE_WARNING
           Suppress warnings about unlicensed features.
 

diff --git a/cli/vscodessh.go b/cli/vscodessh.go
@@ -86,7 +86,7 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
 			client.SetSessionToken(string(sessionToken))
 
 			// This adds custom headers to the request!
-			err = r.setClient(client, serverURL)
+			err = r.setClient(ctx, client, serverURL)
 			if err != nil {
 				return xerrors.Errorf("set client: %w", err)
 			}

diff --git a/docs/cli.md b/docs/cli.md
@@ -96,6 +96,15 @@ Path to the global `coder` config directory.
 
 Additional HTTP headers added to all requests. Provide as key=value. Can be specified multiple times.
 
+### --header-command
+
+|             |                                    |
+| ----------- | ---------------------------------- |
+| Type        | <code>string</code>                |
+| Environment | <code>$CODER_HEADER_COMMAND</code> |
+
+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.
+
 ### --no-feature-warning
 
 |             |                                        |

diff --git a/enterprise/cli/testdata/coder_--help.golden b/enterprise/cli/testdata/coder_--help.golden
@@ -33,6 +33,11 @@ variables or flags.
           Additional HTTP headers added to all requests. Provide as key=value.
           Can be specified multiple times.
 
+      --header-command string, $CODER_HEADER_COMMAND
+          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.
+
       --no-feature-warning bool, $CODER_NO_FEATURE_WARNING
           Suppress warnings about unlicensed features.