feat: add --header-command flag (#9059)

code-asher · web-flow · commit 37f9d4b78329 · 2023-08-14T12:12:17.000-08:00
This allows specifying a command to run that can output headers for
cases where users require dynamic headers (like to authenticate to their
VPN).

The primary use case is to add this flag in SSH configs created by the
VS Code plugin, although maybe config-ssh should do the same.
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.
 

Original file line number	Diff line number	Diff line change
`@@ -76,7 +76,7 @@ func (r RootCmd) login() clibase.Cmd {`
`76`	`76`	`serverURL.Scheme = "https"`
`77`	`77`	`}`
`78`	`78`
`79`		`- client, err := r.createUnauthenticatedClient(serverURL)`
	`79`	`+ client, err := r.createUnauthenticatedClient(ctx, serverURL)`
`80`	`80`	`if err != nil {`
`81`	`81`	`return err`
`82`	`82`	`}`
Original file line number	Diff line number	Diff line change
`@@ -86,7 +86,7 @@ func (r RootCmd) vscodeSSH() clibase.Cmd {`
`86`	`86`	`client.SetSessionToken(string(sessionToken))`
`87`	`87`
`88`	`88`	`// This adds custom headers to the request!`
`89`		`- err = r.setClient(client, serverURL)`
	`89`	`+ err = r.setClient(ctx, client, serverURL)`
`90`	`90`	`if err != nil {`
`91`	`91`	`return xerrors.Errorf("set client: %w", err)`
`92`	`92`	`}`