coder · deansheather · Jun 14, 2023 · Jun 14, 2023 · mafredri · Jun 14, 2023
diff --git a/cli/root.go b/cli/root.go
@@ -21,16 +21,13 @@ import (
 	"text/tabwriter"
 	"time"
 
-	"golang.org/x/exp/slices"
-	"golang.org/x/xerrors"
-
-	"cdr.dev/slog"
-
 	"github.com/charmbracelet/lipgloss"
-	"github.com/gobwas/httphead"
 	"github.com/mattn/go-isatty"
 	"github.com/mitchellh/go-wordwrap"
+	"golang.org/x/exp/slices"
+	"golang.org/x/xerrors"
 
+	"cdr.dev/slog"
 	"github.com/coder/coder/buildinfo"
 	"github.com/coder/coder/cli/clibase"
 	"github.com/coder/coder/cli/cliui"
@@ -430,6 +427,15 @@ type RootCmd struct {
 }
 
 func addTelemetryHeader(client *codersdk.Client, inv *clibase.Invocation) {
+	transport, ok := client.HTTPClient.Transport.(*headerTransport)
+	if !ok {
+		transport = &headerTransport{
+			transport: client.HTTPClient.Transport,
+			header:    http.Header{},
+		}
+		client.HTTPClient.Transport = transport
+	}
+
 	var topts []telemetry.CLIOption
 	for _, opt := range inv.Command.FullOptions() {
 		if opt.ValueSource == clibase.ValueSourceNone || opt.ValueSource == clibase.ValueSourceDefault {
@@ -459,10 +465,7 @@ func addTelemetryHeader(client *codersdk.Client, inv *clibase.Invocation) {
 		return
 	}
 
-	client.ExtraHeaders.Set(
-		codersdk.CLITelemetryHeader,
-		s,
-	)
+	transport.header.Add(codersdk.CLITelemetryHeader, s)
 }
 
 // InitClient sets client to a new client.
@@ -560,18 +563,17 @@ func (r *RootCmd) setClient(client *codersdk.Client, serverURL *url.URL) error {
 		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])
+	}
 	client.URL = serverURL
 	client.HTTPClient = &http.Client{
 		Transport: transport,
 	}
-	client.ExtraHeaders = make(http.Header)
-	for _, hd := range r.header {
-		k, v, ok := httphead.ParseHeaderLine([]byte(hd))
-		if !ok {
-			return xerrors.Errorf("invalid header: %s", hd)
-		}
-		client.ExtraHeaders.Add(string(k), string(v))
-	}
 	return nil
 }
 

diff --git a/cli/root_test.go b/cli/root_test.go
@@ -2,11 +2,19 @@ package cli_test
 
 import (
 	"bytes"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/coder/coder/cli/clibase"
+	"github.com/coder/coder/coderd"
+	"github.com/coder/coder/coderd/coderdtest"
+	"github.com/coder/coder/codersdk"
+	"github.com/coder/coder/pty/ptytest"
+	"github.com/coder/coder/testutil"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -64,21 +72,112 @@ func TestRoot(t *testing.T) {
 	t.Run("Header", func(t *testing.T) {
 		t.Parallel()
 
-		done := make(chan struct{})
+		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"))
 			w.WriteHeader(http.StatusGone)
-			select {
-			case <-done:
-				close(done)
-			default:
-			}
 		}))
 		defer srv.Close()
 		buf := new(bytes.Buffer)
-		inv, _ := clitest.New(t, "--header", "X-Testing=wow", "login", srv.URL)
+		inv, _ := clitest.New(t,
+			"--no-feature-warning",
+			"--no-version-warning",
+			"--header", "X-Testing=wow",
+			"--header", "Cool-Header=Dean was Here!",
+			"login", srv.URL,
+		)
 		inv.Stdout = buf
-		// This won't succeed, because we're using the login cmd to assert requests.
-		_ = inv.Run()
+
+		err := inv.Run()
+		require.Error(t, err)
+		require.ErrorContains(t, err, "unexpected status code 410")
+		require.EqualValues(t, 1, atomic.LoadInt64(&called), "called exactly once")
+	})
+}
+
+// TestDERPHeaders ensures that the client sends the global `--header`s to the
+// DERP server when connecting.
+func TestDERPHeaders(t *testing.T) {
+	t.Parallel()
+
+	// Create a coderd API instance the hard way since we need to change the
+	// handler to inject our custom /derp handler.
+	setHandler, cancelFunc, serverURL, newOptions := coderdtest.NewOptions(t, nil)
+
+	// We set the handler after server creation for the access URL.
+	coderAPI := coderd.New(newOptions)
+	setHandler(coderAPI.RootHandler)
+	provisionerCloser := coderdtest.NewProvisionerDaemon(t, coderAPI)
+	t.Cleanup(func() {
+		_ = provisionerCloser.Close()
+	})
+	client := codersdk.New(serverURL)
+	t.Cleanup(func() {
+		cancelFunc()
+		_ = provisionerCloser.Close()
+		_ = coderAPI.Close()
+		client.HTTPClient.CloseIdleConnections()
+	})
+
+	var (
+		user      = coderdtest.CreateFirstUser(t, client)
+		workspace = runAgent(t, client, user.UserID)
+	)
+
+	// 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!",
+		}
+		derpCalled int64
+	)
+	setHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if strings.HasPrefix(r.URL.Path, "/derp") {
+			ok := true
+			for k, v := range expectedHeaders {
+				if r.Header.Get(k) != v {
+					ok = false
+					break
+				}
+			}
+			if ok {
+				// Only increment if all the headers are set, because the agent
+				// calls derp also.
+				atomic.AddInt64(&derpCalled, 1)
+			}
+		}
+
+		coderAPI.RootHandler.ServeHTTP(w, r)
+	}))
+
+	// Connect with the headers set as args.
+	args := []string{
+		"--no-feature-warning",
+		"--no-version-warning",
+		"ping", workspace.Name,
+		"-n", "1",
+	}
+	for k, v := range expectedHeaders {
+		args = append(args, "--header", fmt.Sprintf("%s=%s", k, v))
+	}
+	inv, root := clitest.New(t, args...)
+	clitest.SetupConfig(t, client, root)
+	pty := ptytest.New(t)
+	inv.Stdin = pty.Input()
+	inv.Stderr = pty.Output()
+	inv.Stdout = pty.Output()
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	cmdDone := tGo(t, func() {
+		err := inv.WithContext(ctx).Run()
+		assert.NoError(t, err)
 	})
+
+	pty.ExpectMatch("pong from " + workspace.Name)
+	<-cmdDone
+
+	require.Greater(t, atomic.LoadInt64(&derpCalled), int64(0), "expected /derp to be called at least once")
 }
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -256,7 +256,8 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 	var handler http.Handler
 	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		mutex.RLock()
-		defer mutex.RUnlock()
+		handler := handler
+		mutex.RUnlock()
 		if handler != nil {
 			handler.ServeHTTP(w, r)
 		}

diff --git a/codersdk/client.go b/codersdk/client.go
@@ -85,9 +85,8 @@ var loggableMimeTypes = map[string]struct{}{
 // New creates a Coder client for the provided URL.
 func New(serverURL *url.URL) *Client {
 	return &Client{
-		URL:          serverURL,
-		HTTPClient:   &http.Client{},
-		ExtraHeaders: make(http.Header),
+		URL:        serverURL,
+		HTTPClient: &http.Client{},
 	}
 }
 
@@ -97,9 +96,6 @@ type Client struct {
 	mu           sync.RWMutex // Protects following.
 	sessionToken string
 
-	// ExtraHeaders are headers to add to every request.
-	ExtraHeaders http.Header
-
 	HTTPClient *http.Client
 	URL        *url.URL
 
@@ -193,8 +189,6 @@ func (c *Client) Request(ctx context.Context, method, path string, body interfac
 		return nil, xerrors.Errorf("create request: %w", err)
 	}
 
-	req.Header = c.ExtraHeaders.Clone()
-
 	tokenHeader := c.SessionTokenHeader
 	if tokenHeader == "" {
 		tokenHeader = SessionTokenHeader

diff --git a/go.mod b/go.mod
@@ -80,6 +80,7 @@ require (
 	github.com/coreos/go-oidc/v3 v3.6.0
 	github.com/coreos/go-systemd v0.0.0-20191104093116-d3cd4ed1dbcf
 	github.com/creack/pty v1.1.18
+	github.com/dave/dst v0.27.2
 	github.com/elastic/go-sysinfo v1.11.0
 	github.com/fatih/color v1.15.0
 	github.com/fatih/structs v1.1.0
@@ -189,6 +190,7 @@ require (
 	github.com/Microsoft/go-winio v0.6.1 // indirect
 	github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5 // indirect
 	github.com/OneOfOne/xxhash v1.2.8 // indirect
+	github.com/ProtonMail/go-crypto v0.0.0-20230217124315-7d5c6f04bbb8 // indirect
 	github.com/agext/levenshtein v1.2.3 // indirect
 	github.com/agnivade/levenshtein v1.1.1 // indirect
 	github.com/akutz/memconn v0.1.0 // indirect
@@ -206,6 +208,7 @@ require (
 	github.com/charmbracelet/bubbles v0.15.0 // indirect
 	github.com/charmbracelet/bubbletea v0.23.2 // indirect
 	github.com/clbanning/mxj/v2 v2.5.7 // indirect
+	github.com/cloudflare/circl v1.3.3 // indirect
 	github.com/containerd/console v1.0.3 // indirect
 	github.com/containerd/continuity v0.3.0 // indirect
 	github.com/coreos/go-iptables v0.6.0 // indirect
@@ -217,6 +220,7 @@ require (
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/elastic/go-windows v1.0.0 // indirect
 	github.com/fxamacker/cbor/v2 v2.4.0 // indirect
+	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
 	github.com/ghodss/yaml v1.0.0 // indirect
 	github.com/gin-gonic/gin v1.9.1 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
@@ -227,6 +231,7 @@ require (
 	github.com/go-openapi/swag v0.19.15 // indirect
 	github.com/go-playground/locales v0.14.1 // indirect
 	github.com/go-playground/universal-translator v0.18.1 // indirect
+	github.com/go-test/deep v1.0.8 // indirect
 	github.com/go-toast/toast v0.0.0-20190211030409-01e6764cf0a4 // indirect
 	github.com/gobwas/glob v0.2.3 // indirect
 	github.com/gobwas/ws v1.1.0 // indirect
@@ -317,6 +322,7 @@ require (
 	github.com/tchap/go-patricia/v2 v2.3.1 // indirect
 	github.com/tcnksm/go-httpstat v0.2.0 // indirect
 	github.com/tdewolff/parse/v2 v2.6.6 // indirect
+	github.com/tdewolff/test v1.0.9 // indirect
 	github.com/u-root/uio v0.0.0-20221213070652-c3537552635f // indirect
 	github.com/vishvananda/netlink v1.1.1-0.20211118161826-650dca95af54 // indirect
 	github.com/vishvananda/netns v0.0.0-20211101163701-50045581ed74 // indirect
@@ -346,22 +352,9 @@ require (
 	golang.zx2c4.com/wireguard/windows v0.5.3 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto v0.0.0-20230530153820-e85fd2cbaebc // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20230530153820-e85fd2cbaebc // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20230530153820-e85fd2cbaebc // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	howett.net/plist v1.0.0 // indirect
 	inet.af/peercred v0.0.0-20210906144145-0893ea02156a // indirect
 )
-
-require (
-	github.com/dave/dst v0.27.2
-	github.com/gobwas/httphead v0.1.0
-)
-
-require (
-	github.com/ProtonMail/go-crypto v0.0.0-20230217124315-7d5c6f04bbb8 // indirect
-	github.com/cloudflare/circl v1.3.3 // indirect
-	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
-	github.com/go-test/deep v1.0.8 // indirect
-	github.com/tdewolff/test v1.0.9 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20230530153820-e85fd2cbaebc // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20230530153820-e85fd2cbaebc // indirect
-)