diff --git a/cli/root.go b/cli/root.go index cec87f09d114e..e8ee13b2476d6 100644 --- 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 index 34fbdf2dad7c8..163c9590f89ba 100644 --- 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 index b6623fbd6f942..26786b88df1e4 100644 --- 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 index 3c1bf39c192a8..19f765c097c5d 100644 --- 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 index edc5732a54cda..cfccde07c06af 100644 --- 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 -)