Make header-command newline-delimited

code-asher · code-asher · commit 86e702beac86 · 2023-08-14T11:24:35.000-08:00
Instead of JSON.
diff --git a/cli/root.go b/cli/root.go
@@ -1,6 +1,8 @@
 package cli
 
 import (
+	"bufio"
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
@@ -361,7 +363,7 @@ func (r *RootCmd) Command(subcommands []*clibase.Cmd) (*clibase.Cmd, error) {
 		{
 			Flag:        varHeaderCommand,
 			Env:         "CODER_HEADER_COMMAND",
-			Description: "An external process that outputs JSON-encoded key-value pairs to be used as additional HTTP headers added to all requests.",
+			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,
 		},
@@ -605,36 +607,39 @@ func (r *RootCmd) setClient(ctx context.Context, client *codersdk.Client, server
 		transport: http.DefaultTransport,
 		header:    http.Header{},
 	}
-	for _, header := range r.header {
-		parts := strings.SplitN(header, "=", 2)
-		if len(parts) < 2 {
-			return xerrors.Errorf("split header %q had less than two parts", header)
-		}
-		transport.header.Add(parts[0], parts[1])
-	}
+	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())
-		out, err := cmd.Output()
+		cmd.Stdout = &outBuf
+		cmd.Stderr = io.Discard
+		err := cmd.Run()
 		if err != nil {
-			return xerrors.Errorf("failed to run %v (out: %q): %w", cmd.Args, out, err)
+			return xerrors.Errorf("failed to run %v: %w", cmd.Args, err)
 		}
-		var headers map[string]string
-		err = json.Unmarshal(out, &headers)
-		if err != nil {
-			return xerrors.Errorf("failed to parse json from %v (out: %q): %w", cmd.Args, out, err)
+		scanner := bufio.NewScanner(&outBuf)
+		for scanner.Scan() {
+			headers = append(headers, scanner.Text())
 		}
-		for key, value := range headers {
-			transport.header.Add(key, value)
+		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)
+		}
+		transport.header.Add(parts[0], parts[1])
+	}
 	client.URL = serverURL
 	client.HTTPClient = &http.Client{
 		Transport: transport,
diff --git a/cli/root_test.go b/cli/root_test.go
@@ -80,6 +80,7 @@ func TestRoot(t *testing.T) {
 			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()
@@ -94,7 +95,7 @@ func TestRoot(t *testing.T) {
 			"--no-version-warning",
 			"--header", "X-Testing=wow",
 			"--header", "Cool-Header=Dean was Here!",
-			"--header-command", "printf '{\"X-Process-Testing\": \"very-wow-'"+coderURLEnv+"'\"}'",
+			"--header-command", "printf X-Process-Testing=very-wow-"+coderURLEnv+"'\\r\\n'X-Process-Testing2=more-wow",
 			"login", srv.URL,
 		)
 		inv.Stdout = buf
@@ -169,7 +170,7 @@ func TestDERPHeaders(t *testing.T) {
 		"--no-version-warning",
 		"ping", workspace.Name,
 		"-n", "1",
-		"--header-command", "printf '{\"X-Process-Testing\": \"very-wow\"}'",
+		"--header-command", "printf X-Process-Testing=very-wow",
 	}
 	for k, v := range expectedHeaders {
 		if k != "X-Process-Testing" {
diff --git a/cli/testdata/coder_--help.golden b/cli/testdata/coder_--help.golden
@@ -63,8 +63,9 @@ variables or flags.
           Can be specified multiple times.
 
       --header-command string, $CODER_HEADER_COMMAND
-          An external command that outputs JSON-encoded key-value pairs to be
-          used as additional HTTP headers added to all requests.
+          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/docs/cli.md b/docs/cli.md
@@ -103,7 +103,7 @@ Additional HTTP headers added to all requests. Provide as key=value. Can be spec
 | Type        | <code>string</code>                |
 | Environment | <code>$CODER_HEADER_COMMAND</code> |
 
-An external command that outputs JSON-encoded key-value pairs to be used as additional HTTP headers added to all requests.
+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
@@ -34,8 +34,9 @@ variables or flags.
           Can be specified multiple times.
 
       --header-command string, $CODER_HEADER_COMMAND
-          An external command that outputs JSON-encoded key-value pairs to be
-          used as additional HTTP headers added to all requests.
+          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.
