Skip to content

Commit 9e3c749

Browse files
ammarioammario
committed
chore(cli): correctly report telemetry even when transport replaced (#7670)
By introducing the "ExtraHeaders" map, we can apply headers even when handlers replace the transport, as in the case of our scaletests. Also, only send telemetry header when it's small.
1 parent 3fb2073 commit 9e3c749

File tree

4 files changed

+54
-34
lines changed

4 files changed

+54
-34
lines changed

‎cli/root.go

+43-31
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727
"cdr.dev/slog"
2828

2929
"github.com/charmbracelet/lipgloss"
30+
"github.com/gobwas/httphead"
3031
"github.com/mattn/go-isatty"
3132

3233
"github.com/coder/coder/buildinfo"
@@ -428,9 +429,9 @@ type RootCmd struct {
428429
noFeatureWarning bool
429430
}
430431

431-
func telemetryInvocation(i *clibase.Invocation) telemetry.CLIInvocation {
432+
func addTelemetryHeader(client *codersdk.Client, inv *clibase.Invocation) {
432433
var topts []telemetry.CLIOption
433-
for _, opt := range i.Command.FullOptions() {
434+
for _, opt := range inv.Command.FullOptions() {
434435
if opt.ValueSource == clibase.ValueSourceNone || opt.ValueSource == clibase.ValueSourceDefault {
435436
continue
436437
}
@@ -439,11 +440,29 @@ func telemetryInvocation(i *clibase.Invocation) telemetry.CLIInvocation {
439440
ValueSource: string(opt.ValueSource),
440441
})
441442
}
442-
return telemetry.CLIInvocation{
443-
Command: i.Command.FullName(),
443+
ti := telemetry.CLIInvocation{
444+
Command: inv.Command.FullName(),
444445
Options: topts,
445446
InvokedAt: time.Now(),
446447
}
448+
449+
byt, err := json.Marshal(ti)
450+
if err != nil {
451+
// Should be impossible
452+
panic(err)
453+
}
454+
455+
// Per https://stackoverflow.com/questions/686217/maximum-on-http-header-values,
456+
// we don't want to send headers that are too long.
457+
s := base64.StdEncoding.EncodeToString(byt)
458+
if len(s) > 4096 {
459+
return
460+
}
461+
462+
client.ExtraHeaders.Set(
463+
codersdk.CLITelemetryHeader,
464+
s,
465+
)
447466
}
448467

449468
// InitClient sets client to a new client.
@@ -456,7 +475,7 @@ func (r *RootCmd) InitClient(client *codersdk.Client) clibase.MiddlewareFunc {
456475
panic("root is nil")
457476
}
458477
return func(next clibase.HandlerFunc) clibase.HandlerFunc {
459-
return func(i *clibase.Invocation) error {
478+
return func(inv *clibase.Invocation) error {
460479
conf := r.createConfig()
461480
var err error
462481
if r.clientURL == nil || r.clientURL.String() == "" {
@@ -485,23 +504,15 @@ func (r *RootCmd) InitClient(client *codersdk.Client) clibase.MiddlewareFunc {
485504
return err
486505
}
487506
}
488-
489-
telemInv := telemetryInvocation(i)
490-
byt, err := json.Marshal(telemInv)
491-
if err != nil {
492-
// Should be impossible
493-
panic(err)
494-
}
495507
err = r.setClient(
496508
client, r.clientURL,
497-
append(r.header, codersdk.CLITelemetryHeader+"="+
498-
base64.StdEncoding.EncodeToString(byt),
499-
),
500509
)
501510
if err != nil {
502511
return err
503512
}
504513

514+
addTelemetryHeader(client, inv)
515+
505516
client.SetSessionToken(r.token)
506517

507518
if r.debugHTTP {
@@ -515,57 +526,58 @@ func (r *RootCmd) InitClient(client *codersdk.Client) clibase.MiddlewareFunc {
515526
warningErr = make(chan error)
516527
)
517528
go func() {
518-
versionErr <- r.checkVersions(i, client)
529+
versionErr <- r.checkVersions(inv, client)
519530
close(versionErr)
520531
}()
521532

522533
go func() {
523-
warningErr <- r.checkWarnings(i, client)
534+
warningErr <- r.checkWarnings(inv, client)
524535
close(warningErr)
525536
}()
526537

527538
if err = <-versionErr; err != nil {
528539
// Just log the error here. We never want to fail a command
529540
// due to a pre-run.
530-
_, _ = fmt.Fprintf(i.Stderr,
541+
_, _ = fmt.Fprintf(inv.Stderr,
531542
cliui.Styles.Warn.Render("check versions error: %s"), err)
532-
_, _ = fmt.Fprintln(i.Stderr)
543+
_, _ = fmt.Fprintln(inv.Stderr)
533544
}
534545

535546
if err = <-warningErr; err != nil {
536547
// Same as above
537-
_, _ = fmt.Fprintf(i.Stderr,
548+
_, _ = fmt.Fprintf(inv.Stderr,
538549
cliui.Styles.Warn.Render("check entitlement warnings error: %s"), err)
539-
_, _ = fmt.Fprintln(i.Stderr)
550+
_, _ = fmt.Fprintln(inv.Stderr)
540551
}
541552

542-
return next(i)
553+
return next(inv)
543554
}
544555
}
545556
}
546557

547-
func (*RootCmd) setClient(client *codersdk.Client, serverURL *url.URL, headers []string) error {
558+
func (r *RootCmd) setClient(client *codersdk.Client, serverURL *url.URL) error {
548559
transport := &headerTransport{
549560
transport: http.DefaultTransport,
550561
header: http.Header{},
551562
}
552-
for _, header := range headers {
553-
parts := strings.SplitN(header, "=", 2)
554-
if len(parts) < 2 {
555-
return xerrors.Errorf("split header %q had less than two parts", header)
556-
}
557-
transport.header.Add(parts[0], parts[1])
558-
}
559563
client.URL = serverURL
560564
client.HTTPClient = &http.Client{
561565
Transport: transport,
562566
}
567+
client.ExtraHeaders = make(http.Header)
568+
for _, hd := range r.header {
569+
k, v, ok := httphead.ParseHeaderLine([]byte(hd))
570+
if !ok {
571+
return xerrors.Errorf("invalid header: %s", hd)
572+
}
573+
client.ExtraHeaders.Add(string(k), string(v))
574+
}
563575
return nil
564576
}
565577

566578
func (r *RootCmd) createUnauthenticatedClient(serverURL *url.URL) (*codersdk.Client, error) {
567579
var client codersdk.Client
568-
err := r.setClient(&client, serverURL, r.header)
580+
err := r.setClient(&client, serverURL)
569581
return &client, err
570582
}
571583

‎cli/vscodessh.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ func (r *RootCmd) vscodeSSH() *clibase.Cmd {
8383
client.SetSessionToken(string(sessionToken))
8484

8585
// This adds custom headers to the request!
86-
err = r.setClient(client, serverURL, r.header)
86+
err = r.setClient(client, serverURL)
8787
if err != nil {
8888
return xerrors.Errorf("set client: %w", err)
8989
}

‎codersdk/client.go

+8-2
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,9 @@ var loggableMimeTypes = map[string]struct{}{
8585
// New creates a Coder client for the provided URL.
8686
func New(serverURL *url.URL) *Client {
8787
return &Client{
88-
URL: serverURL,
89-
HTTPClient: &http.Client{},
88+
URL: serverURL,
89+
HTTPClient: &http.Client{},
90+
ExtraHeaders: make(http.Header),
9091
}
9192
}
9293

@@ -96,6 +97,9 @@ type Client struct {
9697
mu sync.RWMutex // Protects following.
9798
sessionToken string
9899

100+
// ExtraHeaders are headers to add to every request.
101+
ExtraHeaders http.Header
102+
99103
HTTPClient *http.Client
100104
URL *url.URL
101105

@@ -189,6 +193,8 @@ func (c *Client) Request(ctx context.Context, method, path string, body interfac
189193
return nil, xerrors.Errorf("create request: %w", err)
190194
}
191195

196+
req.Header = c.ExtraHeaders.Clone()
197+
192198
tokenHeader := c.SessionTokenHeader
193199
if tokenHeader == "" {
194200
tokenHeader = SessionTokenHeader

‎go.mod

+2
Original file line numberDiff line numberDiff line change
@@ -352,3 +352,5 @@ require (
352352
howett.net/plist v1.0.0 // indirect
353353
inet.af/peercred v0.0.0-20210906144145-0893ea02156a // indirect
354354
)
355+
356+
require github.com/gobwas/httphead v0.1.0

0 commit comments

Comments
 (0)