fix: fix --header flag in CLI (#8023)

This commit is contained in:
Dean Sheather 2023-06-14 21:52:01 +10:00 committed by GitHub
parent df842b31e8
commit 2c843f4011
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 140 additions and 51 deletions

View File

@ -21,16 +21,13 @@ import (
"text/tabwriter"
"time"
"github.com/charmbracelet/lipgloss"
"github.com/mattn/go-isatty"
"github.com/mitchellh/go-wordwrap"
"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"
"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
}

View File

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

View File

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

View File

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

23
go.mod
View File

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