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.
This commit is contained in:
Asher 2023-08-14 12:12:17 -08:00 committed by GitHub
parent b993cab49a
commit 37f9d4b783
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 83 additions and 14 deletions

View File

@ -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
}

View File

@ -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
}

View File

@ -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)

View File

@ -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.

View File

@ -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)
}

View File

@ -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
| | |

View File

@ -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.