From 37f9d4b783296ecc6498963bf0006a83298bb6b1 Mon Sep 17 00:00:00 2001 From: Asher Date: Mon, 14 Aug 2023 12:12:17 -0800 Subject: [PATCH] feat: add --header-command flag (#9059) 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. --- cli/login.go | 2 +- cli/root.go | 50 ++++++++++++++++++--- cli/root_test.go | 24 +++++++--- cli/testdata/coder_--help.golden | 5 +++ cli/vscodessh.go | 2 +- docs/cli.md | 9 ++++ enterprise/cli/testdata/coder_--help.golden | 5 +++ 7 files changed, 83 insertions(+), 14 deletions(-) diff --git a/cli/login.go b/cli/login.go index e16118dfec..2fa7eb231c 100644 --- 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 index 036be18a01..3197a3e3fc 100644 --- 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 index c892701a3a..fb19aae688 100644 --- 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 index e074756680..6e988a9f56 100644 --- 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 index 136a0d727c..7a576581fa 100644 --- 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 index 1a0c40a22f..96acc1cf94 100644 --- 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 | string | +| Environment | $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 | | | diff --git a/enterprise/cli/testdata/coder_--help.golden b/enterprise/cli/testdata/coder_--help.golden index 7fc6962ded..ae24592079 100644 --- 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.